From 58e65ad5bb2178583dec0a33605580fc7af81d96 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C4=90=E1=BA=B7ng=20Minh=20D=C5=A9ng?= <dungdm93@live.com>
Date: Thu, 5 May 2022 23:43:31 +0700
Subject: [PATCH] feat: switch from `sqlalchemy-trino` to `trino` python client
 (#19957)

---
 docs/docs/databases/trino.mdx | 90 +++++++++++++++++++++++++++++++----
 setup.py                      |  4 +-
 2 files changed, 84 insertions(+), 10 deletions(-)

diff --git a/docs/docs/databases/trino.mdx b/docs/docs/databases/trino.mdx
index 7b75e4466..50ccf1f27 100644
--- a/docs/docs/databases/trino.mdx
+++ b/docs/docs/databases/trino.mdx
@@ -9,19 +9,93 @@ version: 1
 
 Supported trino version 352 and higher
 
-The [sqlalchemy-trino](https://pypi.org/project/sqlalchemy-trino/) library is the recommended way to connect to Trino through SQLAlchemy.
-
-The expected connection string is formatted as follows:
-
+### Connection String
+The connection string format is as follows:
 ```
 trino://{username}:{password}@{hostname}:{port}/{catalog}
 ```
 
-If you are running trino with docker on local machine please use the following connection URL
-
+If you are running Trino with docker on local machine, please use the following connection URL
 ```
 trino://trino@host.docker.internal:8080
 ```
 
-Reference:
-[Trino-Superset-Podcast](https://trino.io/episodes/12.html)
+### Authentications
+#### 1. Basic Authentication
+You can provide `username`/`password` in the connection string or in the `Secure Extra` field at `Advanced / Security`
+* In Connection String
+    ```
+    trino://{username}:{password}@{hostname}:{port}/{catalog}
+    ```
+
+* In `Secure Extra` field
+    ```json
+    {
+        "auth_method": "basic",
+        "auth_params": {
+            "username": "<username>",
+            "password": "<password>"
+        }
+    }
+    ```
+
+NOTE: if both are provided, `Secure Extra` always takes higher priority.
+
+#### 2. Kerberos Authentication
+In `Secure Extra` field, config as following example:
+```json
+{
+    "auth_method": "kerberos",
+    "auth_params": {
+        "service_name": "superset",
+        "config": "/path/to/krb5.config",
+        ...
+    }
+}
+```
+
+All fields in `auth_params` are passed directly to the [`KerberosAuthentication`](https://github.com/trinodb/trino-python-client/blob/0.306.0/trino/auth.py#L40) class.
+
+#### 3. JWT Authentication
+Config `auth_method` and provide token in `Secure Extra` field
+```json
+{
+    "auth_method": "jwt",
+    "auth_params": {
+        "token": "<your-jwt-token>"
+    }
+}
+```
+
+#### 4. Custom Authentication
+To use custom authentication, first you need to add it into
+`ALLOWED_EXTRA_AUTHENTICATIONS` allow list in Superset config file:
+```python
+from your.module import AuthClass
+from another.extra import auth_method
+
+ALLOWED_EXTRA_AUTHENTICATIONS: Dict[str, Dict[str, Callable[..., Any]]] = {
+    "trino": {
+        "custom_auth": AuthClass,
+        "another_auth_method": auth_method,
+    },
+}
+```
+
+Then in `Secure Extra` field:
+```json
+{
+    "auth_method": "custom_auth",
+    "auth_params": {
+        ...
+    }
+}
+```
+
+You can also use custom authentication by providing reference to your `trino.auth.Authentication` class
+or factory function (which returns an `Authentication` instance) to `auth_method`.
+
+All fields in `auth_params` are passed directly to your class/function.
+
+**Reference**:
+* [Trino-Superset-Podcast](https://trino.io/episodes/12.html)
diff --git a/setup.py b/setup.py
index 0f5a3cab4..e0a6f7556 100644
--- a/setup.py
+++ b/setup.py
@@ -64,7 +64,7 @@ setup(
     zip_safe=False,
     entry_points={
         "console_scripts": ["superset=superset.cli.main:superset"],
-        "sqlalchemy.dialects": ["trinonative = sqlalchemy_trino.dialect:TrinoDialect"],
+        "sqlalchemy.dialects": ["trinonative = trino.sqlalchemy.dialect:TrinoDialect"],
     },
     install_requires=[
         "backoff>=1.8.0",
@@ -152,7 +152,7 @@ setup(
         "pinot": ["pinotdb>=0.3.3, <0.4"],
         "postgres": ["psycopg2-binary==2.9.1"],
         "presto": ["pyhive[presto]>=0.4.0"],
-        "trino": ["sqlalchemy-trino>=0.2"],
+        "trino": ["trino>=0.313.0"],
         "prophet": ["prophet>=1.0.1, <1.1", "pystan<3.0"],
         "redshift": ["sqlalchemy-redshift>=0.8.1, < 0.9"],
         "rockset": ["rockset>=0.8.10, <0.9"],