fix(execute): fixed a bug which resulted in Pandas Timestamp datatype to be sent to Redshift as the DATE instead of TIMESTAMP/TIMESTAMPTZ when statements are executed using bind parameters. issue #206

Author: Brooke-white

redshift_connector/core.py

@@ -1609,11 +1609,19 @@ def make_params(self: "Connection", values) -> typing.Tuple[typing.Tuple[int, in
         for value in values:
             typ: typing.Type = type(value)
             try:
+                # 1) check if we have a direct match for this datatype in PY_TYPES mapping
                 params.append(self.py_types[typ])
             except KeyError:
                 try:
+                    # 2) if no match was found in 1) check if we have a match in inspect_funcs.
+                    # note that inspect_funcs inspect the data value to determine the type.
+                    # e.g. if the datatype is a Datetime and has a timezone, we want to map it
+                    # to TIMETSTAMPTZ rather than TIMESTAMP.
                     params.append(self.inspect_funcs[typ](value))
                 except KeyError as e:
+                    # 3) if no match was found in 1) nor 2), we again iterate through PY_TYPES but
+                    # check if our data is an instance of any datatypes found in PY_TYPES
+                    # rather than looking for an exact match as was performed in 1)
                     param: typing.Optional[typing.Tuple[int, int, typing.Callable]] = None
                     for k, v in self.py_types.items():
                         try:
@@ -1624,6 +1632,9 @@ def make_params(self: "Connection", values) -> typing.Tuple[typing.Tuple[int, in
                             pass
 
                     if param is None:
+                        # 4) if no match was found in 1) nor 2) nor 3), we again iterate through
+                        # inspect_funcs but check if our data is an instance of any datatype
+                        # found in inspect_funcs
                         for k, v in self.inspect_funcs.items():  # type: ignore
                             try:
                                 if isinstance(value, k):
@@ -1634,6 +1645,23 @@ def make_params(self: "Connection", values) -> typing.Tuple[typing.Tuple[int, in
                                 pass
                             except KeyError:
                                 pass
+                    elif param[0] == RedshiftOID.DATE:
+                        # 5) if we classified this data as DATE in 3), we perform a secondary check to
+                        # ensure this data was not misclassified in 3). Misclassification occurs in the case
+                        # where data having a type that is a subclass of datetime.date is also a subclass of
+                        # datetime.datetime. For this example, the misclassification leads to the loss of
+                        # time precision in transformed data sent to Redshift on the wire. The simplest
+                        # example of this edge case can be seen in
+                        # https://github.com/aws/amazon-redshift-python-driver/issues/206
+                        # where a pandas Timestamp is misclassified as Redshift DATE in 3).
+                        if isinstance(value, Datetime) and Datetime in self.inspect_funcs:
+                            try:
+                                v_func = typing.cast(typing.Callable, self.inspect_funcs[Datetime])
+                                param = v_func(value)
+                            except TypeError:
+                                pass
+                            except KeyError:
+                                pass
 
                     if param is None:
                         raise NotSupportedError("type " + str(e) + " not mapped to pg type")

test/unit/test_connection.py

@@ -2,6 +2,7 @@
 import typing
 from collections import deque
 from decimal import Decimal
+from test.utils import pandas_only
 from unittest import mock
 from unittest.mock import patch
 
@@ -520,3 +521,42 @@ def test_broken_pipe_timeout_on_connect(db_kwargs) -> None:
                     db_kwargs.pop("region")
                     db_kwargs.pop("cluster_identifier")
                     Connection(**db_kwargs)
+
+
+def make_mock_connection(db_kwargs):
+    db_kwargs["ssl"] = False
+    db_kwargs["timeout"] = 60
+
+    with mock.patch("socket.getaddrinfo") as mock_getaddrinfo:
+        addr_tuple = [(0, 1, 2, "", ("3.226.18.73", 5439)), (2, 1, 6, "", ("3.226.18.73", 5439))]
+        mock_getaddrinfo.return_value = addr_tuple
+        with mock.patch("socket.socket.connect") as mock_usock:
+            mock_usock.side_effect = lambda *args, **kwargs: None
+
+            with mock.patch("socket.socket.makefile") as mock_sock:
+                mock_file = mock_sock.return_value
+                mock_file.read.return_value = b"Zasej"
+                db_kwargs.pop("region")
+                db_kwargs.pop("cluster_identifier")
+                return Connection(**db_kwargs)
+
+
+@pandas_only
+def test_make_params_maps_pandas_timestamp_to_timestamp(db_kwargs):
+    import datetime
+
+    import pandas as pd
+
+    from redshift_connector.utils.oids import RedshiftOID
+    from redshift_connector.utils.type_utils import py_types, timestamptz_send_integer
+
+    columns = ["dw_inserted_at"]
+    values = pd.DataFrame({col: [datetime.datetime.now(datetime.timezone.utc)] * 1 for col in columns}).values.tolist()[
+        0
+    ]
+
+    mock_connection: Connection = make_mock_connection(db_kwargs)
+    res = mock_connection.make_params(values)
+    assert res[0][0] == RedshiftOID.TIMESTAMPTZ
+    assert res[0][1] == 1
+    assert res[0][2] == timestamptz_send_integer

test/unit/test_cursor.py

@@ -1,3 +1,4 @@
+import datetime
 import typing
 from io import StringIO
 from math import ceil
@@ -502,3 +503,75 @@ def test_write_dataframe_handles_npdtyes(mocker):
         assert len(spy.mock_calls[1].args[1]) == 1
         # bind parameter list should not contain numpy objects
         assert not isinstance(spy.mock_calls[1].args[1][0], np.generic)
+
+
+@pandas_only
+def test_write_dataframe_handles_pandas_types(mocker):
+    import pandas as pd
+
+    mocker.patch("redshift_connector.Cursor.execute", return_value=None)
+    mocker.patch("redshift_connector.Cursor.fetchone", return_value=[1])
+    mock_cursor: Cursor = Cursor.__new__(Cursor)
+    mock_connection: Connection = Connection.__new__(Connection)
+    mock_cursor._c = mock_connection
+
+    mock_cursor.paramstyle = "mocked_val"
+
+    for datatype, data, _type in (
+        ("int64", pd.Series([42]), int),
+        ("float64", pd.Series([3.14]), float),
+        ("object", pd.Series(["Hello, Pandas!"]), str),
+        ("bool", pd.Series([True]), bool),
+        ("datetime64", pd.Series([pd.Timestamp("2022-01-01")]), int),
+        ("timedelta64", pd.Series([pd.Timedelta(days=5)]), int),
+    ):
+        spy = mocker.spy(mock_cursor, "execute")
+        dataframe = pd.DataFrame(data)
+        mock_cursor.write_dataframe(df=dataframe, table=datatype)
+
+        assert spy.called
+        assert spy.call_count == 2  # once for __is_valid_table, once for write_dataframe
+        assert not isinstance(spy.mock_calls[1].args[1], pd.core.base.PandasObject)
+        assert isinstance(spy.mock_calls[1].args[1], list)
+        assert len(spy.mock_calls[1].args[1]) == 1
+        # bind parameter list should not contain numpy objects
+        assert isinstance(spy.mock_calls[1].args[1][0], _type)
+
+
+@pandas_only
+@pytest.mark.parametrize(
+    "datatype,data,_type",
+    (
+        ("int", 42, int),
+        ("float", 3.14, float),
+        ("str", "H", str),
+        ("bool", True, bool),
+        ("list", [1, 2, 3], list),
+        ("tuple", (4, 5, 6), tuple),
+        ("set", {1, 2, 3}, set),
+        ("datetime", datetime.datetime.now(datetime.timezone.utc), datetime.datetime),
+    ),
+)
+def test_write_dataframe_handles_python_types(mocker, datatype, data, _type):
+    import datetime
+
+    import pandas as pd
+
+    mocker.patch("redshift_connector.Cursor.execute", return_value=None)
+    mocker.patch("redshift_connector.Cursor.fetchone", return_value=[1])
+    mock_cursor: Cursor = Cursor.__new__(Cursor)
+    mock_connection: Connection = Connection.__new__(Connection)
+    mock_cursor._c = mock_connection
+
+    mock_cursor.paramstyle = "mocked_val"
+
+    spy = mocker.spy(mock_cursor, "execute")
+    dataframe = pd.DataFrame({col: [data] * 1 for col in (datatype,)})
+    mock_cursor.write_dataframe(df=dataframe, table=datatype)
+
+    assert spy.called
+    assert spy.call_count == 2  # once for __is_valid_table, once for write_dataframe
+    assert not isinstance(spy.mock_calls[1].args[1], pd.core.base.PandasObject)
+    assert isinstance(spy.mock_calls[1].args[1], list)
+    assert len(spy.mock_calls[1].args[1]) == 1
+    assert isinstance((spy.mock_calls[1].args[1][0]), _type)