chore: adjust time zone approach back to legacy (#178)

eakmanrq · Sep 22, 2024 · ed4c24a · ed4c24a
1 parent 880f0cb
commit ed4c24a
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 4 deletions.
diff --git a/sqlframe/duckdb/session.py b/sqlframe/duckdb/session.py
@@ -52,6 +52,10 @@ def __init__(self, conn: t.Optional[DuckDBPyConnection] = None, *args, **kwargs)
             super().__init__(conn, *args, **kwargs)
             self._last_result = None
 
+    @cached_property
+    def _cur(self) -> DuckDBPyConnection:  # type: ignore
+        return self._conn
+
     @classmethod
     def _try_get_map(cls, value: t.Any) -> t.Optional[t.Dict[str, t.Any]]:
         if value and isinstance(value, dict):

diff --git a/tests/common_fixtures.py b/tests/common_fixtures.py
@@ -117,8 +117,11 @@ def spark_session(pyspark_session: PySparkSession) -> SparkSession:
 
 @pytest.fixture(scope="function")
 def duckdb_session() -> DuckDBSession:
-    connection = duckdb.connect(config={"TimeZone": "UTC"})
-    return DuckDBSession(conn=connection)
+    connector = duckdb.connect()
+    connector.execute("set TimeZone = 'UTC'")
+    connector.execute("SELECT * FROM duckdb_settings() WHERE name = 'TimeZone'")
+    assert connector.fetchone()[1] == "UTC"  # type: ignore
+    return DuckDBSession(conn=connector)
 
 
 @pytest.fixture

diff --git a/tests/integration/engines/duck/test_duckdb_catalog.py b/tests/integration/engines/duck/test_duckdb_catalog.py
@@ -8,7 +8,10 @@
 def duckdb_session() -> DuckDBSession:
     import duckdb
 
-    connector = duckdb.connect(config={"TimeZone": "UTC"})
+    connector = duckdb.connect()
+    connector.execute("set TimeZone = 'UTC'")
+    connector.execute("SELECT * FROM duckdb_settings() WHERE name = 'TimeZone'")
+    assert connector.fetchone()[1] == "UTC"  # type: ignore
     session = DuckDBSession(connector)
     session._execute('''ATTACH ':memory:' AS "default"''')
     session._execute('''ATTACH ':memory:' AS "catalog1"''')

diff --git a/tests/integration/engines/test_int_functions.py b/tests/integration/engines/test_int_functions.py
@@ -18,7 +18,7 @@
     get_func_from_session as get_func_from_session_without_fallback,
 )
 from sqlframe.bigquery import BigQuerySession
-from sqlframe.duckdb import DuckDBSession
+from sqlframe.duckdb import DuckDBCatalog, DuckDBSession
 from sqlframe.postgres import PostgresDataFrame, PostgresSession
 from sqlframe.snowflake import SnowflakeSession
 from sqlframe.spark.session import SparkSession