diff --git a/pkg-py/src/querychat/_datasource.py b/pkg-py/src/querychat/_datasource.py index 56377502..466e0548 100644 --- a/pkg-py/src/querychat/_datasource.py +++ b/pkg-py/src/querychat/_datasource.py @@ -210,7 +210,7 @@ def get_schema(self, *, categorical_threshold: int) -> str: elif dtype == nw.Boolean: sql_type = "BOOLEAN" elif dtype == nw.Datetime: - sql_type = "TIME" + sql_type = "TIMESTAMP" elif dtype == nw.Date: sql_type = "DATE" else: @@ -224,7 +224,7 @@ def get_schema(self, *, categorical_threshold: int) -> str: categories = unique_values.to_list() categories_str = ", ".join([f"'{c}'" for c in categories]) column_info.append(f" Categorical values: {categories_str}") - elif sql_type in ["INTEGER", "FLOAT", "DATE", "TIME"]: + elif sql_type in ["INTEGER", "FLOAT", "DATE", "TIMESTAMP"]: rng = self._df[column].min(), self._df[column].max() if rng[0] is None and rng[1] is None: column_info.append(" Range: NULL to NULL") diff --git a/pkg-py/tests/test_datasource.py b/pkg-py/tests/test_datasource.py index 8ba00fee..d3b04e06 100644 --- a/pkg-py/tests/test_datasource.py +++ b/pkg-py/tests/test_datasource.py @@ -498,3 +498,27 @@ def test_check_query_escape_hatch_accepts_various_values(monkeypatch): for value in ["true", "TRUE", "1", "yes", "YES"]: monkeypatch.setenv("QUERYCHAT_ENABLE_UPDATE_QUERIES", value) check_query("INSERT INTO table VALUES (1)") # Should not raise + + +def test_dataframe_source_datetime_type_mapping(): + """Test that datetime columns are mapped to TIMESTAMP, not TIME.""" + # Create DataFrame with datetime column + test_df = pd.DataFrame( + { + "id": [1, 2, 3], + "name": ["a", "b", "c"], + "created_at": pd.to_datetime(["2023-01-01 10:00:00", "2023-06-15 14:30:00", "2023-12-31 23:59:59"]), + } + ) + + source = DataFrameSource(test_df, "test_table") + schema = source.get_schema(categorical_threshold=5) + + # Datetime columns should be TIMESTAMP, not TIME + assert "- created_at (TIMESTAMP)" in schema + assert "- created_at (TIME)" not in schema + + # Should also show range for datetime columns + assert "Range:" in schema + + source.cleanup()