Add impersonation option in trino datasource

Author: kimsehwan96

redash/query_runner/trino.py

@@ -1,5 +1,6 @@
 import logging
 
+from redash.models.users import ApiUser, User
 from redash.query_runner import (
     TYPE_BOOLEAN,
     TYPE_DATE,
@@ -57,20 +58,34 @@ def configuration_schema(cls):
                 "port": {"type": "number"},
                 "username": {"type": "string"},
                 "password": {"type": "string"},
+                "source": {"type": "string", "default": "redash"},
                 "catalog": {"type": "string"},
                 "schema": {"type": "string"},
+                "impersonation": {"type": "boolean", "default": False},
+                "impersonationField": {
+                    "type": "string",
+                    "title": "Impersonation User Attribute",
+                    "default": "email",
+                    "extendedEnum": [{"value": "email", "name": "Email"}, {"value": "name", "name": "Name"}],
+                },
             },
             "order": [
                 "protocol",
                 "host",
                 "port",
                 "username",
                 "password",
+                "source",
                 "catalog",
                 "schema",
+                "impersonation",
             ],
             "required": ["host", "username"],
             "secret": ["password"],
+            "extra_options": [
+                "impersonation",
+                "impersonationField",
+            ],
         }
 
     @classmethod
@@ -127,20 +142,41 @@ def _get_catalogs(self):
             catalogs.append(catalog)
         return catalogs
 
+    def _get_trino_user(self, user):
+        """Determine the Trino user based on impersonation settings."""
+        default_user = self.configuration.get("username")
+
+        if not self.configuration.get("impersonation") or user is None:
+            return default_user
+
+        impersonation_field = self.configuration.get("impersonationField", "email")
+
+        if isinstance(user, User):
+            if impersonation_field == "email":
+                return user.email or default_user
+            elif impersonation_field == "name":
+                return user.name or default_user
+        elif isinstance(user, ApiUser):
+            return user.name or default_user
+
+        return default_user
+
     def run_query(self, query, user):
         if self.configuration.get("password"):
             auth = trino.auth.BasicAuthentication(
                 username=self.configuration.get("username"), password=self.configuration.get("password")
             )
         else:
             auth = trino.constants.DEFAULT_AUTH
+
         connection = trino.dbapi.connect(
             http_scheme=self.configuration.get("protocol", "http"),
             host=self.configuration.get("host", ""),
+            source=self.configuration.get("source", "redash"),
             port=self.configuration.get("port", 8080),
             catalog=self.configuration.get("catalog", ""),
             schema=self.configuration.get("schema", ""),
-            user=self.configuration.get("username"),
+            user=self._get_trino_user(user),
             auth=auth,
         )