feat(trino): add support for user impersonation (#14843)

* trino impersonation feature

* Extra options label update

* Update superset/db_engine_specs/trino.py

Co-authored-by: Đặng Minh Dũng <dungdm93@live.com>

Co-authored-by: rijojoseph01 <rijo.joseph@myntra.com>
Co-authored-by: Đặng Minh Dũng <dungdm93@live.com>

superset-frontend/src/views/CRUD/data/database/DatabaseModal/ExtraOptions.tsx

@@ -292,12 +292,12 @@ const ExtraOptions = ({
               checked={!!db?.impersonate_user}
               onChange={onInputChange}
               labelText={t(
-                'Impersonate Logged In User (Presto, Hive, and GSheets)',
+                'Impersonate Logged In User (Presto, Trino, Hive, and GSheets)',
               )}
             />
             <InfoTooltip
               tooltip={t(
-                'If Presto, all the queries in SQL Lab are going to be executed as the ' +
+                'If Presto or Trino, all the queries in SQL Lab are going to be executed as the ' +
                   'currently logged on user who must have permission to run them. If Hive ' +
                   'and hive.server2.enable.doAs is enabled, will run the queries as ' +
                   'service account, but impersonate the currently logged on user via ' +

superset/db_engine_specs/trino.py

@@ -15,10 +15,10 @@
 # specific language governing permissions and limitations
 # under the License.
 from datetime import datetime
-from typing import Optional
+from typing import Any, Dict, Optional
 from urllib import parse
 
-from sqlalchemy.engine.url import URL
+from sqlalchemy.engine.url import make_url, URL
 
 from superset.db_engine_specs.base import BaseEngineSpec
 from superset.utils import core as utils
@@ -69,3 +69,37 @@ class TrinoEngineSpec(BaseEngineSpec):
             selected_schema = parse.quote(selected_schema, safe="")
             database = database.split("/")[0] + "/" + selected_schema
             uri.database = database
+
+    @classmethod
+    def update_impersonation_config(
+        cls, connect_args: Dict[str, Any], uri: str, username: Optional[str],
+    ) -> None:
+        """
+        Update a configuration dictionary
+        that can set the correct properties for impersonating users
+        :param connect_args: config to be updated
+        :param uri: URI string
+        :param impersonate_user: Flag indicating if impersonation is enabled
+        :param username: Effective username
+        :return: None
+        """
+        url = make_url(uri)
+        backend_name = url.get_backend_name()
+
+        # Must be Trino connection, enable impersonation, and set optional param
+        # auth=LDAP|KERBEROS
+        # Set principal_username=$effective_username
+        if backend_name == "trino" and username is not None:
+            connect_args["user"] = username
+
+    @classmethod
+    def modify_url_for_impersonation(
+        cls, url: URL, impersonate_user: bool, username: Optional[str]
+    ) -> None:
+        """
+        Modify the SQL Alchemy URL object with the user to impersonate if applicable.
+        :param url: SQLAlchemy URL object
+        :param impersonate_user: Flag indicating if impersonation is enabled
+        :param username: Effective username
+        """
+        # Do nothing and let update_impersonation_config take care of impersonation

tests/model_tests.py

@@ -157,6 +157,33 @@ class TestDatabaseModel(SupersetTestCase):
             "password": "original_user_password",
         }
 
+    @mock.patch("superset.models.core.create_engine")
+    def test_impersonate_user_trino(self, mocked_create_engine):
+        uri = "trino://localhost"
+        principal_user = "logged_in_user"
+
+        model = Database(database_name="test_database", sqlalchemy_uri=uri)
+
+        model.impersonate_user = True
+        model.get_sqla_engine(user_name=principal_user)
+        call_args = mocked_create_engine.call_args
+
+        assert str(call_args[0][0]) == "trino://localhost"
+
+        assert call_args[1]["connect_args"] == {
+            "user": "logged_in_user",
+        }
+
+        uri = "trino://original_user:original_user_password@localhost"
+        model = Database(database_name="test_database", sqlalchemy_uri=uri)
+        model.impersonate_user = True
+        model.get_sqla_engine(user_name=principal_user)
+        call_args = mocked_create_engine.call_args
+
+        assert str(call_args[0][0]) == "trino://original_user@localhost"
+
+        assert call_args[1]["connect_args"] == {"user": "logged_in_user"}
+
     @mock.patch("superset.models.core.create_engine")
     def test_impersonate_user_hive(self, mocked_create_engine):
         uri = "hive://localhost"