Configure experimental_python_types through dbapi.Connection

Author: mdesmet

README.md

@@ -99,15 +99,17 @@ engine = create_engine(
     'trino://user@localhost:8080/system',
     connect_args={
       "session_properties": {'query_max_run_time': '1d'},
-      "client_tags": ["tag1", "tag2"]
+      "client_tags": ["tag1", "tag2"],
+      "experimental_python_types": True,
     }
 )
 
 # or in connection string
 engine = create_engine(
     'trino://user@localhost:8080/system?'
     'session_properties={"query_max_run_time": "1d"}'
-    '&client_tags=["tag1", "tag2"]',
+    '&client_tags=["tag1", "tag2"]'
+    '&experimental_python_types=true',
 )
 ```
 
@@ -361,10 +363,11 @@ import pytz
 from datetime import datetime
 
 conn = trino.dbapi.connect(
+    experimental_python_types=True,
     ...
 )
 
-cur = conn.cursor(experimental_python_types=True)
+cur = conn.cursor()
 
 params = datetime(2020, 1, 1, 16, 43, 22, 320000, tzinfo=pytz.timezone('America/Los_Angeles'))
 

tests/integration/test_dbapi_integration.py

@@ -166,8 +166,34 @@ def test_execute_many_select(trino_connection):
     assert "Query must return update type" in str(e.value)
 
 
-def test_python_types_not_used_when_experimental_python_types_is_not_set(trino_connection):
-    cur = trino_connection.cursor()
+@pytest.mark.parametrize("connection_experimental_python_types,cursor_experimental_python_types,expected",
+                         [
+                             (None, None, False),
+                             (None, False, False),
+                             (None, True, True),
+                             (False, None, False),
+                             (False, False, False),
+                             (False, True, True),
+                             (True, None, True),
+                             (True, False, False),
+                             (True, True, True),
+                         ])
+def test_experimental_python_types_with_connection_and_cursor(
+        connection_experimental_python_types,
+        cursor_experimental_python_types,
+        expected,
+        run_trino
+):
+    _, host, port = run_trino
+
+    connection = trino.dbapi.Connection(
+        host=host,
+        port=port,
+        user="test",
+        experimental_python_types=connection_experimental_python_types,
+    )
+
+    cur = connection.cursor(experimental_python_types=cursor_experimental_python_types)
 
     cur.execute("""
     SELECT
@@ -180,15 +206,23 @@ def test_python_types_not_used_when_experimental_python_types_is_not_set(trino_c
     """)
     rows = cur.fetchall()
 
-    for value in rows[0]:
-        assert isinstance(value, str)
-
-    assert rows[0][0] == '0.142857'
-    assert rows[0][1] == '2018-01-01'
-    assert rows[0][2] == '2019-01-01 00:00:00.000 +01:00'
-    assert rows[0][3] == '2019-01-01 00:00:00.000 UTC'
-    assert rows[0][4] == '2019-01-01 00:00:00.000'
-    assert rows[0][5] == '00:00:00.000'
+    if expected:
+        assert rows[0][0] == Decimal('0.142857')
+        assert rows[0][1] == date(2018, 1, 1)
+        assert rows[0][2] == datetime(2019, 1, 1, tzinfo=timezone(timedelta(hours=1)))
+        assert rows[0][3] == datetime(2019, 1, 1, tzinfo=pytz.timezone('UTC'))
+        assert rows[0][4] == datetime(2019, 1, 1)
+        assert rows[0][5] == time(0, 0, 0, 0)
+    else:
+        for value in rows[0]:
+            assert isinstance(value, str)
+
+        assert rows[0][0] == '0.142857'
+        assert rows[0][1] == '2018-01-01'
+        assert rows[0][2] == '2019-01-01 00:00:00.000 +01:00'
+        assert rows[0][3] == '2019-01-01 00:00:00.000 UTC'
+        assert rows[0][4] == '2019-01-01 00:00:00.000'
+        assert rows[0][5] == '00:00:00.000'
 
 
 def test_decimal_query_param(trino_connection):

tests/unit/sqlalchemy/test_dialect.py

@@ -47,7 +47,8 @@ def setup(self):
                     'session_properties={"query_max_run_time": "1d"}'
                     '&http_headers={"trino": 1}'
                     '&extra_credential=[("a", "b"), ("c", "d")]'
-                    '&client_tags=[1, "sql"]'),
+                    '&client_tags=[1, "sql"]'
+                    '&experimental_python_types=true'),
                 list(),
                 dict(
                     host="localhost",
@@ -58,7 +59,8 @@ def setup(self):
                     session_properties={"query_max_run_time": "1d"},
                     http_headers={"trino": 1},
                     extra_credential=[("a", "b"), ("c", "d")],
-                    client_tags=[1, "sql"]
+                    client_tags=[1, "sql"],
+                    experimental_python_types=True,
                 ),
             ),
         ],

trino/dbapi.py

@@ -109,7 +109,8 @@ def __init__(
         isolation_level=IsolationLevel.AUTOCOMMIT,
         verify=True,
         http_session=None,
-        client_tags=None
+        client_tags=None,
+        experimental_python_types=False,
     ):
         self.host = host
         self.port = port
@@ -136,6 +137,7 @@ def __init__(
         self._isolation_level = isolation_level
         self._request = None
         self._transaction = None
+        self.experimental_python_types = experimental_python_types
 
     @property
     def isolation_level(self):
@@ -199,17 +201,21 @@ def _create_request(self):
             client_tags=self.client_tags
         )
 
-    def cursor(self, experimental_python_types=False):
+    def cursor(self, experimental_python_types: bool = None):
         """Return a new :py:class:`Cursor` object using the connection."""
         if self.isolation_level != IsolationLevel.AUTOCOMMIT:
             if self.transaction is None:
                 self.start_transaction()
-            request = self.transaction._request
-        elif self.transaction is not None:
-            request = self.transaction._request
+        if self.transaction is not None:
+            request = self.transaction.request
         else:
             request = self._create_request()
-        return Cursor(self, request, experimental_python_types)
+        return Cursor(
+            self,
+            request,
+            # if experimental_python_types is not explicitly set in Cursor, take from Connection
+            experimental_python_types if experimental_python_types is not None else self.experimental_python_types
+        )
 
 
 class Cursor(object):

trino/sqlalchemy/dialect.py

@@ -122,6 +122,9 @@ def create_connect_args(self, url: URL) -> Tuple[Sequence[Any], Mapping[str, Any
         if "client_tags" in url.query:
             kwargs["client_tags"] = json.loads(url.query["client_tags"])
 
+        if "experimental_python_types" in url.query:
+            kwargs["experimental_python_types"] = json.loads(url.query["experimental_python_types"])
+
         return args, kwargs
 
     def get_columns(self, connection: Connection, table_name: str, schema: str = None, **kw) -> List[Dict[str, Any]]:

trino/transaction.py

@@ -59,6 +59,10 @@ def __init__(self, request):
     def id(self):
         return self._id
 
+    @property
+    def request(self):
+        return self._request
+
     def begin(self):
         response = self._request.post(START_TRANSACTION)
         if not response.ok: