Support USE catalog.schema and USE schema

Author: mdesmet

tests/integration/test_dbapi_integration.py

@@ -983,3 +983,54 @@ def retrieve_client_tags_from_query(run_trino, client_tags):
 
     query_client_tags = query_info['session']['clientTags']
     return query_client_tags
+
+
+@pytest.mark.skipif(trino_version() == '351', reason="current_catalog not supported in older Trino versions")
+def test_use_catalog_schema(trino_connection):
+    cur = trino_connection.cursor()
+    cur.execute('SELECT current_catalog, current_schema')
+    result = cur.fetchall()
+    assert result[0][0] is None
+    assert result[0][1] is None
+
+    cur.execute('USE tpch.tiny')
+    cur.fetchall()
+    cur.execute('SELECT current_catalog, current_schema')
+    result = cur.fetchall()
+    assert result[0][0] == 'tpch'
+    assert result[0][1] == 'tiny'
+
+    cur.execute('USE tpcds.sf1')
+    cur.fetchall()
+    cur.execute('SELECT current_catalog, current_schema')
+    result = cur.fetchall()
+    assert result[0][0] == 'tpcds'
+    assert result[0][1] == 'sf1'
+
+
+@pytest.mark.skipif(trino_version() == '351', reason="current_catalog not supported in older Trino versions")
+def test_use_catalog(run_trino):
+    _, host, port = run_trino
+
+    trino_connection = trino.dbapi.Connection(
+        host=host, port=port, user="test", source="test", catalog="tpch", max_attempts=1
+    )
+    cur = trino_connection.cursor()
+    cur.execute('SELECT current_catalog, current_schema')
+    result = cur.fetchall()
+    assert result[0][0] == 'tpch'
+    assert result[0][1] is None
+
+    cur.execute('USE tiny')
+    cur.fetchall()
+    cur.execute('SELECT current_catalog, current_schema')
+    result = cur.fetchall()
+    assert result[0][0] == 'tpch'
+    assert result[0][1] == 'tiny'
+
+    cur.execute('USE sf1')
+    cur.fetchall()
+    cur.execute('SELECT current_catalog, current_schema')
+    result = cur.fetchall()
+    assert result[0][0] == 'tpch'
+    assert result[0][1] == 'sf1'

trino/client.py

@@ -126,11 +126,23 @@ def user(self):
 
     @property
     def catalog(self):
-        return self._catalog
+        with self._object_lock:
+            return self._catalog
+
+    @catalog.setter
+    def catalog(self, catalog):
+        with self._object_lock:
+            self._catalog = catalog
 
     @property
     def schema(self):
-        return self._schema
+        with self._object_lock:
+            return self._schema
+
+    @schema.setter
+    def schema(self, schema):
+        with self._object_lock:
+            self._schema = schema
 
     @property
     def source(self):
@@ -489,6 +501,12 @@ def process(self, http_response) -> TrinoStatus:
             ):
                 self._client_session.properties[key] = value
 
+        if constants.HEADER_SET_CATALOG in http_response.headers:
+            self._client_session.catalog = http_response.headers[constants.HEADER_SET_CATALOG]
+
+        if constants.HEADER_SET_SCHEMA in http_response.headers:
+            self._client_session.schema = http_response.headers[constants.HEADER_SET_SCHEMA]
+
         self._next_uri = response.get("nextUri")
 
         return TrinoStatus(

trino/constants.py

@@ -45,3 +45,6 @@
 HEADER_PREPARED_STATEMENT = 'X-Trino-Prepared-Statement'
 HEADER_ADDED_PREPARE = 'X-Trino-Added-Prepare'
 HEADER_DEALLOCATED_PREPARE = 'X-Trino-Deallocated-Prepare'
+
+HEADER_SET_SCHEMA = "X-Trino-Set-Schema"
+HEADER_SET_CATALOG = "X-Trino-Set-Catalog"