Add certificate auth support via connection string in SQLAlchemy

Author: hovaesco

README.md

@@ -241,6 +241,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,15 +7,14 @@
 
 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",
         [
@@ -72,3 +71,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

@@ -18,7 +18,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 +96,10 @@ 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: