vdk-trino: add oauth connection (#3417)

Add the ability to authenticate to trino using oAuth

Signed-off-by: Dako Dakov <[email protected]>

Author: dakodakov

projects/vdk-plugins/vdk-trino/src/vdk/plugin/trino/trino_config.py

@@ -3,6 +3,9 @@
 from typing import cast
 from typing import Optional
 
+from vdk.internal.builtin_plugins.config.vdk_config import TEAM_CLIENT_ID
+from vdk.internal.builtin_plugins.config.vdk_config import TEAM_CLIENT_SECRET
+from vdk.internal.builtin_plugins.config.vdk_config import TEAM_OAUTH_AUTHORIZE_URL
 from vdk.internal.core.config import Configuration
 
 TRINO_HOST = "TRINO_HOST"
@@ -15,6 +18,9 @@
 TRINO_SSL_VERIFY = "TRINO_SSL_VERIFY"
 TRINO_TIMEOUT_SECONDS = "TRINO_TIMEOUT_SECONDS"
 TRINO_TEMPLATES_DATA_TO_TARGET_STRATEGY = "TRINO_TEMPLATES_DATA_TO_TARGET_STRATEGY"
+TRINO_USE_TEAM_OAUTH = "TRINO_USE_TEAM_OAUTH"
+TRINO_RETRIES_ON_ERROR = "TRINO_RETRIES_ON_ERROR"
+TRINO_RETRIES_BACKOFF_SECONDS = "TRINO_RETRIES_BACKOFF_SECONDS"
 
 trino_templates_data_to_target_strategy: str = ""
 
@@ -121,6 +127,50 @@ def templates_data_to_target_strategy(self, section: Optional[str]) -> str:
             else "INSERT_SELECT",
         )
 
+    def use_team_oauth(self, section: Optional[str]) -> bool:
+        return (
+            parse_boolean(
+                self.__config.get_value(key=TRINO_USE_TEAM_OAUTH, section=section)
+            )
+            if (
+                self.__config.get_value(key=TRINO_USE_TEAM_OAUTH, section=section)
+                is not None
+            )
+            else False
+        )
+
+    def retries(self, section: Optional[str]) -> int:
+        return cast(
+            int, self.__config.get_value(key=TRINO_RETRIES_ON_ERROR, section=section)
+        )
+
+    def backoff_interval_seconds(self, section: Optional[str]) -> int:
+        return cast(
+            int,
+            self.__config.get_value(key=TRINO_RETRIES_BACKOFF_SECONDS, section=section),
+        )
+
+    def team_client_id(self) -> str:
+        return (
+            cast(str, self.__config.get_value(key=TEAM_CLIENT_ID))
+            if self.__config.get_value(key=TEAM_CLIENT_ID) is not None
+            else None
+        )
+
+    def team_client_secret(self) -> str:
+        return (
+            cast(str, self.__config.get_value(key=TEAM_CLIENT_SECRET))
+            if self.__config.get_value(key=TEAM_CLIENT_SECRET) is not None
+            else None
+        )
+
+    def team_oauth_url(self) -> str:
+        return (
+            cast(str, self.__config.get_value(key=TEAM_OAUTH_AUTHORIZE_URL))
+            if self.__config.get_value(key=TEAM_OAUTH_AUTHORIZE_URL) is not None
+            else None
+        )
+
     @staticmethod
     def add_definitions(config_builder):
         """
@@ -172,3 +222,18 @@ def add_definitions(config_builder):
             default_value=None,
             description="The trino query timeout in seconds.",
         )
+        config_builder.add(
+            key=TRINO_USE_TEAM_OAUTH,
+            default_value=False,
+            description="Should the connection use the team's oAuth credentials to connect to the DBs",
+        )
+        config_builder.add(
+            key=TRINO_RETRIES_ON_ERROR,
+            default_value=3,
+            description="The number of times the Trino plugin is going to retry a failing operation",
+        )
+        config_builder.add(
+            key=TRINO_RETRIES_BACKOFF_SECONDS,
+            default_value=30,
+            description="The backoff time in seconds between retries of a failed operation",
+        )

projects/vdk-plugins/vdk-trino/src/vdk/plugin/trino/trino_connection.py

@@ -1,19 +1,27 @@
 # Copyright 2023-2024 Broadcom
 # SPDX-License-Identifier: Apache-2.0
+import base64
+import json
 import logging
+from typing import Optional
 
+import requests
 from tenacity import before_sleep_log
 from tenacity import retry
 from tenacity import retry_if_exception_type
 from tenacity import stop_after_attempt
 from tenacity import wait_exponential
+from trino import constants
+from trino import dbapi
+from trino.auth import BasicAuthentication
 from trino.exceptions import TrinoExternalError
 from trino.exceptions import TrinoInternalError
 from vdk.internal.builtin_plugins.connection.managed_connection_base import (
     ManagedConnectionBase,
 )
 from vdk.internal.builtin_plugins.connection.recovery_cursor import RecoveryCursor
 from vdk.internal.util.decorators import closing_noexcept_on_close
+from vdk.plugin.trino.trino_config import TrinoConfiguration
 from vdk.plugin.trino.trino_error_handler import TrinoErrorHandler
 
 log = logging.getLogger(__name__)
@@ -22,18 +30,9 @@
 class TrinoConnection(ManagedConnectionBase):
     def __init__(
         self,
-        host,
-        port,
-        catalog,
-        schema,
-        user,
-        password,
-        use_ssl=True,
-        ssl_verify=True,
-        timeout_seconds=120,
+        configuration: TrinoConfiguration,
+        section: Optional[str],
         lineage_logger=None,
-        retries_on_error=3,
-        error_backoff_seconds=30,
     ):
         """
         Create a new database connection. Connection parameters are:
@@ -42,36 +41,93 @@ def __init__(
         - *port*: connection port number (defaults to 8080 if not provided)
         - *catalog*: the catalog name (only as keyword argument)
         - *schema*: the schema name (only as keyword argument)
-        - *user*: user name used to authenticate
+        - *user*: username used to authenticate
         """
         super().__init__(logging.getLogger(__name__))
 
-        self._host = host
-        self._port = port
-        self._catalog = catalog
-        self._schema = schema
-        self._user = user
-        self._password = password
-        self._use_ssl = use_ssl
-        self._ssl_verify = ssl_verify
-        self._timeout_seconds = timeout_seconds
+        self._host = configuration.host(section)
+        self._port = configuration.port(section)
+        self._catalog = configuration.catalog(section)
+        self._schema = configuration.schema(section)
+        self._user = configuration.user(section)
+        self._password = configuration.password(section)
+        self._use_ssl = configuration.use_ssl(section)
+        self._ssl_verify = configuration.ssl_verify(section)
+        self._timeout_seconds = configuration.timeout_seconds(section)
+        self._retries_on_error = configuration.retries(section)
+        self._error_backoff_seconds = configuration.backoff_interval_seconds(section)
+
         self._lineage_logger = lineage_logger
-        self._retries_on_error = retries_on_error
-        self._error_backoff_seconds = error_backoff_seconds
+
+        self._use_team_oauth = configuration.use_team_oauth(section)
+
+        if self._use_team_oauth:
+            self._team_client_id = configuration.team_client_id()
+            self._team_client_secret = configuration.team_client_secret()
+            self._team_oauth_url = configuration.team_oauth_url()
+            log.debug(
+                f"Creating new trino connection for oAuth ClientID: {self._team_client_id} to host: {self._host}:{self._port}"
+            )
+        else:
+            log.debug(
+                f"Creating new trino connection for user: {self._user} to host: {self._host}:{self._port}"
+            )
+
+    def _connect(self):
+        if self._use_team_oauth:
+            return self._team_oauth_connection()
+        else:
+            return self._basic_authentication_connection()
+
+    def _team_oauth_connection(self):
         log.debug(
-            f"Creating new trino connection for user: {user} to host: {host}:{port}"
+            f"Open Trino Connection: host: {self._host}:{self._port} with oAuth ClientID: {self._team_client_id}; "
+            f"catalog: {self._catalog}; schema: {self._schema}; timeout: {self._timeout_seconds}"
         )
 
-    def _connect(self):
-        from trino import dbapi
-        from trino import constants
+        oauth_token = self._get_access_token()
+
+        # Create an OAuth session
+        session = OAuthSession(oauth_token)
 
+        connection = dbapi.connect(
+            host=self._host,
+            port=self._port,
+            catalog=self._catalog,
+            schema=self._schema,
+            http_scheme=constants.HTTPS if self._use_ssl else constants.HTTP,
+            verify=self._ssl_verify,
+            request_timeout=self._timeout_seconds,
+            http_session=session,
+        )
+        return connection
+
+    def _get_access_token(self):
+        # Exchange client ID & Secret for an access token
+        # Original basic auth string
+        original_string = self._team_client_id + ":" + self._team_client_secret
+        # Encode
+        encoded_bytes = base64.b64encode(original_string.encode("utf-8"))
+        encoded_string = encoded_bytes.decode("utf-8")
+
+        headers = {
+            "Authorization": "Basic " + encoded_string,
+            "Content-Type": "application/x-www-form-urlencoded",
+        }
+        data = {"grant_type": "client_credentials"}
+        response = requests.post(self._team_oauth_url, headers=headers, data=data)
+        # If this call fails then, we better raise it as early as possible
+        response.raise_for_status()
+
+        response_json = json.loads(response.text)
+        oauth_token = response_json["access_token"]
+        return oauth_token
+
+    def _basic_authentication_connection(self):
         log.debug(
             f"Open Trino Connection: host: {self._host}:{self._port} with user: {self._user}; "
             f"catalog: {self._catalog}; schema: {self._schema}; timeout: {self._timeout_seconds}"
         )
-        from trino.auth import BasicAuthentication
-
         auth = (
             BasicAuthentication(self._user, self._password) if self._password else None
         )
@@ -164,3 +220,10 @@ def _get_lineage_data(self, query):
             return None
 
         return None
+
+
+# Define a custom requests session to add the OAuth token to the headers
+class OAuthSession(requests.Session):
+    def __init__(self, token):
+        super().__init__()
+        self.headers.update({"Authorization": f"Bearer {token}"})

projects/vdk-plugins/vdk-trino/src/vdk/plugin/trino/trino_plugin.py

@@ -112,29 +112,15 @@ def initialize_job(self, context: JobContext):
                 port = trino_conf.port(section)
 
                 if host and port:
-                    schema = trino_conf.schema(section)
-                    catalog = trino_conf.catalog(section)
-                    user = trino_conf.user(section)
-                    password = trino_conf.password(section)
-                    use_ssl = trino_conf.use_ssl(section)
-                    ssl_verify = trino_conf.ssl_verify(section)
-                    timeout_seconds = trino_conf.timeout_seconds(section)
                     lineage_logger = context.core_context.state.get(LINEAGE_LOGGER_KEY)
                     log.info(
                         f"Creating new Trino connection with name {connection_name} and host {host}"
                     )
                     context.connections.add_open_connection_factory_method(
                         connection_name.lower(),
-                        lambda t_host=host, t_port=port, t_schema=schema, t_catalog=catalog, t_user=user, t_password=password, t_use_ssl=use_ssl, t_ssl_verify=ssl_verify, t_timeout=timeout_seconds, t_lineage_logger=lineage_logger: TrinoConnection(
-                            host=t_host,
-                            port=t_port,
-                            schema=t_schema,
-                            catalog=t_catalog,
-                            user=t_user,
-                            password=t_password,
-                            use_ssl=t_use_ssl,
-                            ssl_verify=t_ssl_verify,
-                            timeout_seconds=t_timeout,
+                        lambda t_configuration=trino_conf, t_section=section, t_lineage_logger=lineage_logger: TrinoConnection(
+                            configuration=t_configuration,
+                            section=t_section,
                             lineage_logger=t_lineage_logger,
                         ),
                     )
@@ -157,7 +143,7 @@ def initialize_job(self, context: JobContext):
                     )
             except Exception as e:
                 raise Exception(
-                    "An error occurred while trying to create new  Trino connections and ingesters."
+                    f"An error occurred while trying to create new Trino connections and ingesters for connection:{connection_name}."
                     f"ERROR: {e}"
                 )
 
@@ -205,17 +191,10 @@ def vdk_start(plugin_registry: IPluginRegistry, command_line_args: List):
 @click.option("-q", "--query", type=click.STRING, required=True)
 @click.pass_context
 def trino_query(ctx: click.Context, query):
-    conf = ctx.obj.configuration
+    trino_conf = TrinoConfiguration(ctx.obj.configuration)
     conn = TrinoConnection(
-        host=conf.get_value(TRINO_HOST),
-        port=conf.get_value(TRINO_PORT),
-        schema=conf.get_value(TRINO_SCHEMA),
-        catalog=conf.get_value(TRINO_CATALOG),
-        user=conf.get_value(TRINO_USER),
-        password=conf.get_value(TRINO_PASSWORD),
-        use_ssl=conf.get_value(TRINO_USE_SSL),
-        ssl_verify=conf.get_value(TRINO_SSL_VERIFY),
-        timeout_seconds=conf.get_value(TRINO_TIMEOUT_SECONDS),
+        configuration=trino_conf,
+        section=None,
         lineage_logger=ctx.obj.state.get(LINEAGE_LOGGER_KEY),
     )
 

projects/vdk-plugins/vdk-trino/tests/test_trino_multiple_db.py

@@ -7,10 +7,12 @@
 from unittest import TestCase
 
 import pytest
+from vdk.internal.core.config import ConfigurationBuilder
 from vdk.plugin.test_utils.util_funcs import cli_assert_equal
 from vdk.plugin.test_utils.util_funcs import CliEntryBasedTestRunner
 from vdk.plugin.test_utils.util_funcs import get_test_job_path
 from vdk.plugin.trino import trino_plugin
+from vdk.plugin.trino.trino_config import TrinoConfiguration
 from vdk.plugin.trino.trino_connection import TrinoConnection
 
 VDK_DB_DEFAULT_TYPE = "VDK_DB_DEFAULT_TYPE"
@@ -61,16 +63,22 @@ def test_ingest_to_multiple_trino(self):
         )
 
         # check secondary db
+        builder = ConfigurationBuilder()
+        builder.add("TRINO_HOST", "localhost")
+        builder.add("TRINO_PORT", 8081)
+        builder.add("TRINO_SCHEMA", "default")
+        builder.add("TRINO_CATALOG", "memory")
+        builder.add("TRINO_USER", "unknown")
+        builder.add("TRINO_PASSWORD", None)
+        builder.add("TRINO_USE_SSL", False)
+        builder.add("TRINO_SSL_VERIFY", True)
+        builder.add("TRINO_TIMEOUT_SECONDS", None)
+        cfg = builder.build()
+
+        trino_conf = TrinoConfiguration(cfg)
         conn = TrinoConnection(
-            host="localhost",
-            port=8081,
-            schema="default",
-            catalog="memory",  # default
-            user="unknown",  # default
-            password=None,
-            use_ssl=False,
-            ssl_verify=True,  # default
-            timeout_seconds=None,
+            configuration=trino_conf,
+            section=None,
             lineage_logger=None,
         )