From ea0cc1537f484d785294179196c79f6fd299f4b5 Mon Sep 17 00:00:00 2001
From: George Sittas <giwrgos.sittas@gmail.com>
Date: Fri, 18 Oct 2024 00:49:51 +0300
Subject: [PATCH] fix(cursor, fetch_dataframe): use column names in cursor's
description as is
Previously, the Cursor.fetch_dataframe method lowercased column names preemptively,
thus not respecting the active case-sensitivity configuration value. This could lead
to issues, when e.g. two different columns resolved to the same name when lowercased.
This commit ensures the resulting DataFrame's columns are the same as the ones in the
cursor's description, thus reflecting Redshift's behavior and respecting the underlying
case-sensitivity configuration.
---
redshift_connector/cursor.py | 2 +-
test/unit/test_cursor.py | 17 +++++++++++++++++
2 files changed, 18 insertions(+), 1 deletion(-)
diff --git a/redshift_connector/cursor.py b/redshift_connector/cursor.py
index 03e7af4..20ba7e5 100644
--- a/redshift_connector/cursor.py
+++ b/redshift_connector/cursor.py
@@ -523,7 +523,7 @@ def fetch_dataframe(self: "Cursor", num: typing.Optional[int] = None) -> "pandas
columns: typing.Optional[typing.List[typing.Union[str, bytes]]] = None
try:
- columns = [column[0].lower() for column in self.description]
+ columns = [column[0] for column in self.description]
except:
warn("No row description was found. pandas dataframe will be missing column labels.", stacklevel=2)
diff --git a/test/unit/test_cursor.py b/test/unit/test_cursor.py
index ec95b06..d002658 100644
--- a/test/unit/test_cursor.py
+++ b/test/unit/test_cursor.py
@@ -51,6 +51,23 @@ def test_fetch_dataframe_no_results(mocker) -> None:
assert mock_cursor.fetch_dataframe(1).size == 0
+@pandas_only
+def test_fetch_dataframe_respects_case_sensitivity(mocker) -> None:
+ import pandas as pd
+
+ mock_cursor: Cursor = Cursor.__new__(Cursor)
+ mocker.patch(
+ "redshift_connector.Cursor._getDescription",
+ return_value=[("C", 23, None, None, None, None, None)],
+ )
+ mocker.patch("redshift_connector.Cursor.__next__", side_effect=StopIteration("mocked end"))
+
+ df = mock_cursor.fetch_dataframe()
+
+ assert df.size == 0
+ assert df.columns.to_list() == ["C"]
+
+
def test_raw_connection_property_warns() -> None:
mock_cursor: Cursor = Cursor.__new__(Cursor)
mock_cursor._c = Connection.__new__(Connection)