SqlAlchemy tweaks

README.md

@@ -89,6 +89,28 @@ nodes = Table(
 rows = connection.execute(select(nodes)).fetchall()
 ```
 
+In order to pass additional connection attributes use [connect_args](https://docs.sqlalchemy.org/en/14/core/engines.html#sqlalchemy.create_engine.params.connect_args) method.
+Attributes can be also passed in connection string.
+
+```python
+from sqlalchemy import create_engine
+
+engine = create_engine(
+    'trino://user@localhost:8080/system',
+    connect_args={
+      "session_properties": {'query_max_run_time': '1d'},
+      "client_tags": ["tag1", "tag2"]
+    }
+)
+
+# or in connection string
+engine = create_engine(
+    'trino://user@localhost:8080/system?'
+    'session_properties={"query_max_run_time": "1d"}'
+    '&client_tags=["tag1", "tag2"]',
+)
+```
+
 ## Authentications
 
 ### Basic Authentication
@@ -227,6 +249,9 @@ The OAuth2 token will be cached either per `trino.auth.OAuth2Authentication` ins
     from sqlalchemy import create_engine
     from trino.auth import CertificateAuthentication
 
+    engine = create_engine("trino://<username>@<host>:<port>/<catalog>/<schema>?cert=<cert>&key=<key>")
+
+    # or
     engine = create_engine(
     "trino://<username>@<host>:<port>/<catalog>",
         connect_args={

tests/unit/sqlalchemy/test_dialect.py

@@ -7,30 +7,29 @@
 
 from trino.auth import BasicAuthentication
 from trino.dbapi import Connection
-from trino.sqlalchemy.dialect import TrinoDialect
+from trino.sqlalchemy.dialect import CertificateAuthentication, JWTAuthentication, TrinoDialect
 from trino.transaction import IsolationLevel
 
 
 class TestTrinoDialect:
     def setup(self):
         self.dialect = TrinoDialect()
 
-    # TODO: Test more authentication methods and URL params (https://github.com/trinodb/trino-python-client/issues/106)
     @pytest.mark.parametrize(
         "url, expected_args, expected_kwargs",
         [
             (
                 make_url("trino://user@localhost"),
                 list(),
-                dict(host="localhost", catalog="system", user="user"),
+                dict(host="localhost", catalog="system", user="user", source="trino-sqlalchemy"),
             ),
             (
                 make_url("trino://user@localhost:8080"),
                 list(),
-                dict(host="localhost", port=8080, catalog="system", user="user"),
+                dict(host="localhost", port=8080, catalog="system", user="user", source="trino-sqlalchemy"),
             ),
             (
-                make_url("trino://user:pass@localhost:8080"),
+                make_url("trino://user:pass@localhost:8080?source=trino-rulez"),
                 list(),
                 dict(
                     host="localhost",
@@ -39,6 +38,27 @@ def setup(self):
                     user="user",
                     auth=BasicAuthentication("user", "pass"),
                     http_scheme="https",
+                    source="trino-rulez"
+                ),
+            ),
+            (
+                make_url(
+                    'trino://user@localhost:8080?'
+                    'session_properties={"query_max_run_time": "1d"}'
+                    '&http_headers={"trino": 1}'
+                    '&extra_credential=[("a", "b"), ("c", "d")]'
+                    '&client_tags=[1, "sql"]'),
+                list(),
+                dict(
+                    host="localhost",
+                    port=8080,
+                    catalog="system",
+                    user="user",
+                    source="trino-sqlalchemy",
+                    session_properties={"query_max_run_time": "1d"},
+                    http_headers={"trino": 1},
+                    extra_credential=[("a", "b"), ("c", "d")],
+                    client_tags=[1, "sql"]
                 ),
             ),
         ],
@@ -71,3 +91,40 @@ def test_isolation_level(self):
 
         isolation_level = self.dialect.get_isolation_level(dbapi_conn)
         assert isolation_level == "SERIALIZABLE"
+
+
+def test_trino_connection_basic_auth():
+    dialect = TrinoDialect()
+    username = 'trino-user'
+    password = 'trino-bunny'
+    url = make_url(f'trino://{username}:{password}@host')
+    _, cparams = dialect.create_connect_args(url)
+
+    assert cparams['http_scheme'] == "https"
+    assert isinstance(cparams['auth'], BasicAuthentication)
+    assert cparams['auth']._username == username
+    assert cparams['auth']._password == password
+
+
+def test_trino_connection_jwt_auth():
+    dialect = TrinoDialect()
+    access_token = 'sample-token'
+    url = make_url(f'trino://host/?access_token={access_token}')
+    _, cparams = dialect.create_connect_args(url)
+
+    assert cparams['http_scheme'] == "https"
+    assert isinstance(cparams['auth'], JWTAuthentication)
+    assert cparams['auth'].token == access_token
+
+
+def test_trino_connection_certificate_auth():
+    dialect = TrinoDialect()
+    cert = '/path/to/cert.pem'
+    key = '/path/to/key.pem'
+    url = make_url(f'trino://host/?cert={cert}&key={key}')
+    _, cparams = dialect.create_connect_args(url)
+
+    assert cparams['http_scheme'] == "https"
+    assert isinstance(cparams['auth'], CertificateAuthentication)
+    assert cparams['auth']._cert == cert
+    assert cparams['auth']._key == key

trino/sqlalchemy/dialect.py

@@ -9,6 +9,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import json
+from ast import literal_eval
 from textwrap import dedent
 from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple
 
@@ -18,7 +20,7 @@
 from sqlalchemy.engine.url import URL
 
 from trino import dbapi as trino_dbapi, logging
-from trino.auth import BasicAuthentication, JWTAuthentication
+from trino.auth import BasicAuthentication, CertificateAuthentication, JWTAuthentication
 from trino.dbapi import Cursor
 from trino.exceptions import TrinoUserError
 from trino.sqlalchemy import compiler, datatype, error
@@ -96,6 +98,27 @@ def create_connect_args(self, url: URL) -> Tuple[Sequence[Any], Mapping[str, Any
             kwargs["http_scheme"] = "https"
             kwargs["auth"] = JWTAuthentication(url.query["access_token"])
 
+        if "cert" and "key" in url.query:
+            kwargs["http_scheme"] = "https"
+            kwargs["auth"] = CertificateAuthentication(url.query['cert'], url.query['key'])
+
+        if "source" in url.query:
+            kwargs["source"] = url.query["source"]
+        else:
+            kwargs["source"] = "trino-sqlalchemy"
+
+        if "session_properties" in url.query:
+            kwargs["session_properties"] = json.loads(url.query["session_properties"])
+
+        if "http_headers" in url.query:
+            kwargs["http_headers"] = json.loads(url.query["http_headers"])
+
+        if "extra_credential" in url.query:
+            kwargs["extra_credential"] = literal_eval(url.query["extra_credential"])
+
+        if "client_tags" in url.query:
+            kwargs["client_tags"] = json.loads(url.query["client_tags"])
+
         return args, kwargs
 
     def get_columns(self, connection: Connection, table_name: str, schema: str = None, **kw) -> List[Dict[str, Any]]: