|
| 1 | +"""Tests that polars and pandas correctly infer schemas from cursor.description type codes.""" |
| 2 | + |
| 3 | +import datetime |
| 4 | +import inspect |
| 5 | +import pytest |
| 6 | + |
| 7 | +try: |
| 8 | + import polars as pl |
| 9 | + |
| 10 | + HAS_POLARS = True |
| 11 | +except ImportError: |
| 12 | + HAS_POLARS = False |
| 13 | + |
| 14 | +try: |
| 15 | + import pandas as pd |
| 16 | + |
| 17 | + HAS_PANDAS = True |
| 18 | +except ImportError: |
| 19 | + HAS_PANDAS = False |
| 20 | + |
| 21 | + |
| 22 | +# ── cursor.description type_code verification ───────────────────────────── |
| 23 | + |
| 24 | + |
| 25 | +class TestCursorDescriptionTypeCodes: |
| 26 | + """Verify cursor.description returns isclass-compatible Python types.""" |
| 27 | + |
| 28 | + def test_date_type_code_is_datetime_date(self, cursor): |
| 29 | + """DATE columns must report datetime.date, not str.""" |
| 30 | + cursor.execute("SELECT CAST('2024-01-15' AS DATE) AS d") |
| 31 | + type_code = cursor.description[0][1] |
| 32 | + assert type_code is datetime.date |
| 33 | + assert inspect.isclass(type_code) |
| 34 | + cursor.fetchall() |
| 35 | + |
| 36 | + def test_time_type_code_is_datetime_time(self, cursor): |
| 37 | + """TIME columns must report datetime.time.""" |
| 38 | + cursor.execute("SELECT CAST('13:45:30' AS TIME) AS t") |
| 39 | + type_code = cursor.description[0][1] |
| 40 | + assert type_code is datetime.time |
| 41 | + assert inspect.isclass(type_code) |
| 42 | + cursor.fetchall() |
| 43 | + |
| 44 | + def test_datetime_type_code_is_datetime_datetime(self, cursor): |
| 45 | + """DATETIME columns must report datetime.datetime.""" |
| 46 | + cursor.execute("SELECT CAST('2024-01-15 13:45:30' AS DATETIME) AS dt") |
| 47 | + type_code = cursor.description[0][1] |
| 48 | + assert type_code is datetime.datetime |
| 49 | + assert inspect.isclass(type_code) |
| 50 | + cursor.fetchall() |
| 51 | + |
| 52 | + def test_datetime2_type_code_is_datetime_datetime(self, cursor): |
| 53 | + """DATETIME2 columns must report datetime.datetime.""" |
| 54 | + cursor.execute("SELECT CAST('2024-01-15 13:45:30.1234567' AS DATETIME2) AS dt2") |
| 55 | + type_code = cursor.description[0][1] |
| 56 | + assert type_code is datetime.datetime |
| 57 | + assert inspect.isclass(type_code) |
| 58 | + cursor.fetchall() |
| 59 | + |
| 60 | + def test_smalldatetime_type_code_is_datetime_datetime(self, cursor): |
| 61 | + """SMALLDATETIME columns must report datetime.datetime.""" |
| 62 | + cursor.execute("SELECT CAST('2024-01-15 13:45:00' AS SMALLDATETIME) AS sdt") |
| 63 | + type_code = cursor.description[0][1] |
| 64 | + assert type_code is datetime.datetime |
| 65 | + assert inspect.isclass(type_code) |
| 66 | + cursor.fetchall() |
| 67 | + |
| 68 | + def test_datetimeoffset_type_code_is_datetime_datetime(self, cursor): |
| 69 | + """DATETIMEOFFSET columns must report datetime.datetime.""" |
| 70 | + cursor.execute("SELECT CAST('2024-01-15 13:45:30.123 +05:30' AS DATETIMEOFFSET) AS dto") |
| 71 | + type_code = cursor.description[0][1] |
| 72 | + assert type_code is datetime.datetime |
| 73 | + assert inspect.isclass(type_code) |
| 74 | + cursor.fetchall() |
| 75 | + |
| 76 | + def test_all_types_are_isclass(self, cursor): |
| 77 | + """Every type_code in cursor.description must pass inspect.isclass().""" |
| 78 | + cursor.execute(""" |
| 79 | + SELECT |
| 80 | + CAST(1 AS INT) AS i, |
| 81 | + CAST('x' AS NVARCHAR(10)) AS s, |
| 82 | + CAST('2024-01-15' AS DATE) AS d, |
| 83 | + CAST('13:45:30' AS TIME) AS t, |
| 84 | + CAST('2024-01-15 13:45:30' AS DATETIME2) AS dt2, |
| 85 | + CAST(1.5 AS DECIMAL(10,2)) AS dec, |
| 86 | + CAST(1 AS BIT) AS b, |
| 87 | + CAST(0x01 AS VARBINARY(10)) AS bin |
| 88 | + """) |
| 89 | + for desc in cursor.description: |
| 90 | + col_name = desc[0] |
| 91 | + type_code = desc[1] |
| 92 | + assert inspect.isclass( |
| 93 | + type_code |
| 94 | + ), f"Column '{col_name}': type_code={type_code!r} fails isclass()" |
| 95 | + cursor.fetchall() |
| 96 | + |
| 97 | + |
| 98 | +# ── Polars integration ──────────────────────────────────────────────────── |
| 99 | + |
| 100 | + |
| 101 | +@pytest.mark.skipif(not HAS_POLARS, reason="polars not installed") |
| 102 | +class TestPolarsIntegration: |
| 103 | + """Polars read_database must infer correct dtypes from cursor.description.""" |
| 104 | + |
| 105 | + def test_polars_date_column(self, db_connection): |
| 106 | + """Issue #352: DATE columns caused ComputeError in polars.""" |
| 107 | + df = pl.read_database( |
| 108 | + query="SELECT CAST('2024-01-15' AS DATE) AS d", |
| 109 | + connection=db_connection, |
| 110 | + ) |
| 111 | + assert df.schema["d"] == pl.Date |
| 112 | + assert df["d"][0] == datetime.date(2024, 1, 15) |
| 113 | + |
| 114 | + def test_polars_all_datetime_types(self, db_connection): |
| 115 | + """All date/time types must produce correct polars dtypes.""" |
| 116 | + df = pl.read_database( |
| 117 | + query=""" |
| 118 | + SELECT |
| 119 | + CAST('2024-01-15' AS DATE) AS d, |
| 120 | + CAST('2024-01-15 13:45:30' AS DATETIME) AS dt, |
| 121 | + CAST('2024-01-15 13:45:30.123' AS DATETIME2) AS dt2, |
| 122 | + CAST('2024-01-15 13:45:00' AS SMALLDATETIME) AS sdt |
| 123 | + """, |
| 124 | + connection=db_connection, |
| 125 | + ) |
| 126 | + assert df.schema["d"] == pl.Date |
| 127 | + assert df.schema["dt"] == pl.Datetime |
| 128 | + assert df.schema["dt2"] == pl.Datetime |
| 129 | + assert df.schema["sdt"] == pl.Datetime |
| 130 | + |
| 131 | + def test_polars_mixed_types(self, db_connection): |
| 132 | + """Mixed column types with DATE must not cause schema mismatch.""" |
| 133 | + df = pl.read_database( |
| 134 | + query=""" |
| 135 | + SELECT |
| 136 | + CAST(42 AS INT) AS i, |
| 137 | + CAST('hello' AS NVARCHAR(50)) AS s, |
| 138 | + CAST('2024-06-15' AS DATE) AS d, |
| 139 | + CAST(99.95 AS DECIMAL(10,2)) AS amount |
| 140 | + """, |
| 141 | + connection=db_connection, |
| 142 | + ) |
| 143 | + assert df["i"][0] == 42 |
| 144 | + assert df["s"][0] == "hello" |
| 145 | + assert df["d"][0] == datetime.date(2024, 6, 15) |
| 146 | + assert df.schema["d"] == pl.Date |
| 147 | + |
| 148 | + def test_polars_date_with_nulls(self, db_connection): |
| 149 | + """DATE columns with NULLs must still infer Date dtype.""" |
| 150 | + cursor = db_connection.cursor() |
| 151 | + try: |
| 152 | + cursor.execute(""" |
| 153 | + CREATE TABLE #polars_null_test ( |
| 154 | + id INT, |
| 155 | + d DATE |
| 156 | + ) |
| 157 | + """) |
| 158 | + cursor.execute(""" |
| 159 | + INSERT INTO #polars_null_test VALUES |
| 160 | + (1, '2024-01-15'), |
| 161 | + (2, NULL), |
| 162 | + (3, '2024-03-20') |
| 163 | + """) |
| 164 | + db_connection.commit() |
| 165 | + |
| 166 | + df = pl.read_database( |
| 167 | + query="SELECT * FROM #polars_null_test ORDER BY id", |
| 168 | + connection=db_connection, |
| 169 | + ) |
| 170 | + assert df.schema["d"] == pl.Date |
| 171 | + assert df["d"][0] == datetime.date(2024, 1, 15) |
| 172 | + assert df["d"][1] is None |
| 173 | + assert df["d"][2] == datetime.date(2024, 3, 20) |
| 174 | + finally: |
| 175 | + try: |
| 176 | + cursor.execute("DROP TABLE IF EXISTS #polars_null_test") |
| 177 | + db_connection.commit() |
| 178 | + except Exception: |
| 179 | + # Intentionally ignore cleanup errors: the temp table may not exist |
| 180 | + # or the connection may already be closed, and this should not fail the test. |
| 181 | + pass |
| 182 | + cursor.close() |
| 183 | + |
| 184 | + |
| 185 | +# ── Pandas integration ──────────────────────────────────────────────────── |
| 186 | + |
| 187 | + |
| 188 | +@pytest.mark.skipif(not HAS_PANDAS, reason="pandas not installed") |
| 189 | +@pytest.mark.filterwarnings( |
| 190 | + "ignore:pandas only supports SQLAlchemy connectable" ":UserWarning:pandas.io.sql" |
| 191 | +) |
| 192 | +class TestPandasIntegration: |
| 193 | + """Pandas read_sql must handle date/time columns correctly.""" |
| 194 | + |
| 195 | + def test_pandas_date_column(self, db_connection): |
| 196 | + """DATE columns must be readable by pandas without error.""" |
| 197 | + df = pd.read_sql( |
| 198 | + "SELECT CAST('2024-01-15' AS DATE) AS d", |
| 199 | + db_connection, |
| 200 | + ) |
| 201 | + assert len(df) == 1 |
| 202 | + val = df["d"].iloc[0] |
| 203 | + # pandas may return datetime or date depending on version |
| 204 | + if isinstance(val, datetime.datetime): |
| 205 | + assert val.date() == datetime.date(2024, 1, 15) |
| 206 | + else: |
| 207 | + assert val == datetime.date(2024, 1, 15) |
| 208 | + |
| 209 | + def test_pandas_all_datetime_types(self, db_connection): |
| 210 | + """All date/time types must be readable by pandas.""" |
| 211 | + df = pd.read_sql( |
| 212 | + """ |
| 213 | + SELECT |
| 214 | + CAST('2024-01-15' AS DATE) AS d, |
| 215 | + CAST('2024-01-15 13:45:30' AS DATETIME) AS dt, |
| 216 | + CAST('2024-01-15 13:45:30.123' AS DATETIME2) AS dt2, |
| 217 | + CAST('2024-01-15 13:45:00' AS SMALLDATETIME) AS sdt |
| 218 | + """, |
| 219 | + db_connection, |
| 220 | + ) |
| 221 | + assert len(df) == 1 |
| 222 | + assert len(df.columns) == 4 |
| 223 | + |
| 224 | + def test_pandas_mixed_types_with_date(self, db_connection): |
| 225 | + """Mixed column types including DATE must work correctly.""" |
| 226 | + df = pd.read_sql( |
| 227 | + """ |
| 228 | + SELECT |
| 229 | + CAST(42 AS INT) AS i, |
| 230 | + CAST('hello' AS NVARCHAR(50)) AS s, |
| 231 | + CAST('2024-06-15' AS DATE) AS d, |
| 232 | + CAST(99.95 AS DECIMAL(10,2)) AS amount |
| 233 | + """, |
| 234 | + db_connection, |
| 235 | + ) |
| 236 | + assert df["i"].iloc[0] == 42 |
| 237 | + assert df["s"].iloc[0] == "hello" |
| 238 | + val = df["d"].iloc[0] |
| 239 | + if isinstance(val, datetime.datetime): |
| 240 | + assert val.date() == datetime.date(2024, 6, 15) |
| 241 | + else: |
| 242 | + assert val == datetime.date(2024, 6, 15) |
0 commit comments