Adding oauth2 and 2.1 support to trino-python-client

Added documentaton for scopes and audiences.
moved set_session to a base class.

Author: dprophet

README.md

@@ -217,6 +217,10 @@ the [`JWT` authentication type](https://trino.io/docs/current/security/jwt.html)
 
 ### OAuth2 authentication
 
+Make sure that the OAuth2 support is installed using `pip install trino[oauth]`.
+
+#### Interactive Browser authentication
+
 The `OAuth2Authentication` class can be used to connect to a Trino cluster configured with
 the [OAuth2 authentication type](https://trino.io/docs/current/security/oauth2.html).
 
@@ -256,6 +260,152 @@ The OAuth2 token will be cached either per `trino.auth.OAuth2Authentication` ins
     )
     ```
 
+#### Client Credentials authentication
+
+The `ClientCredentials` class enables service-to-service authentication using standard OAuth2 Client Credentials flow.
+
+- DBAPI
+
+    ```python
+    from trino.dbapi import connect
+    from trino.auth import ClientCredentials, OidcConfig
+
+    conn = connect(
+        user="<username>",
+        auth=ClientCredentials(
+            client_id="<client_id>",
+            client_secret="<client_secret>",
+            url_config=OidcConfig(oidc_discovery_url="<oidc_discovery_url>")
+        ),
+        http_scheme="https",
+        ...
+    )
+    ```
+
+    Or using manual URLs:
+
+    ```python
+    from trino.dbapi import connect
+    from trino.auth import ClientCredentials, ManualUrlsConfig
+
+    conn = connect(
+        user="<username>",
+        auth=ClientCredentials(
+            client_id="<client_id>",
+            client_secret="<client_secret>",
+            url_config=ManualUrlsConfig(token_endpoint="<token_endpoint>")
+        ),
+        http_scheme="https",
+        ...
+    )
+    ```
+
+    With optional `scope` and `audience`:
+
+    ```python
+    from trino.dbapi import connect
+    from trino.auth import ClientCredentials, OidcConfig
+
+    conn = connect(
+        user="<username>",
+        auth=ClientCredentials(
+            client_id="<client_id>",
+            client_secret="<client_secret>",
+            scope="<scope>",
+            audience="<audience>",
+            url_config=OidcConfig(oidc_discovery_url="<oidc_discovery_url>")
+        ),
+        http_scheme="https",
+        ...
+    )
+    ```
+
+#### Device Code authentication
+
+The `DeviceCode` class enables authentication on devices with limited input capabilities using OAuth2 Device Code flow.
+This flow prints tries to open an browser to the URL but it will also fall back and print the verification code and
+a URL to visit on another device to authenticate.
+
+- DBAPI
+
+    ```python
+    from trino.dbapi import connect
+    from trino.auth import DeviceCode, OidcConfig
+
+    conn = connect(
+        user="<username>",
+        auth=DeviceCode(
+            client_id="<client_id>",
+            client_secret="<client_secret>",
+            url_config=OidcConfig(oidc_discovery_url="<oidc_discovery_url>")
+        ),
+        http_scheme="https",
+        ...
+    )
+    ```
+
+    With optional `scope` and `audience`:
+
+    ```python
+    from trino.dbapi import connect
+    from trino.auth import DeviceCode, OidcConfig
+
+    conn = connect(
+        user="<username>",
+        auth=DeviceCode(
+            client_id="<client_id>",
+            client_secret="<client_secret>",
+            scope="<scope>",
+            audience="<audience>",
+            url_config=OidcConfig(oidc_discovery_url="<oidc_discovery_url>")
+        ),
+        http_scheme="https",
+        ...
+    )
+    ```
+
+#### Authorization Code authentication
+
+The `AuthorizationCode` class enables the standard OAuth2 Authorization Code flow.
+
+- DBAPI
+
+    ```python
+    from trino.dbapi import connect
+    from trino.auth import AuthorizationCode, OidcConfig
+
+    conn = connect(
+        user="<username>",
+        auth=AuthorizationCode(
+            client_id="<client_id>",
+            client_secret="<client_secret>",
+            url_config=OidcConfig(oidc_discovery_url="<oidc_discovery_url>")
+        ),
+        http_scheme="https",
+        ...
+    )
+    ```
+
+    With optional `scope` and `audience`:
+
+    ```python
+    from trino.dbapi import connect
+    from trino.auth import AuthorizationCode, OidcConfig
+
+    conn = connect(
+        user="<username>",
+        auth=AuthorizationCode(
+            client_id="<client_id>",
+            client_secret="<client_secret>",
+            scope="<scope>",
+            audience="<audience>",
+            url_config=OidcConfig(oidc_discovery_url="<oidc_discovery_url>")
+        ),
+        http_scheme="https",
+        ...
+    )
+    ```
+
 ### Certificate authentication
 
 `CertificateAuthentication` class can be used to connect to Trino cluster configured with [certificate based authentication](https://trino.io/docs/current/security/certificate.html). `CertificateAuthentication` requires paths to a valid client certificate and private key.

setup.py

@@ -33,9 +33,10 @@
                   "krb5 == 0.5.1"]
 sqlalchemy_require = ["sqlalchemy >= 1.3"]
 external_authentication_token_cache_require = ["keyring"]
+oauth_require = ["trino.oauth2 @ git+https://github.com/dprophet/trino-python-oauth2"]
 
 # We don't add localstorage_require to all_require as users must explicitly opt in to use keyring.
-all_require = kerberos_require + sqlalchemy_require
+all_require = kerberos_require + sqlalchemy_require + oauth_require
 
 tests_require = all_require + gssapi_require + [
     # httpretty >= 1.1 duplicates requests in `httpretty.latest_requests`
@@ -96,6 +97,7 @@
         "all": all_require,
         "kerberos": kerberos_require,
         "gssapi": gssapi_require,
+        "oauth": oauth_require,
         "sqlalchemy": sqlalchemy_require,
         "tests": tests_require,
         "external-authentication-token-cache": external_authentication_token_cache_require,

trino/__init__.py

@@ -9,6 +9,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+__path__ = __import__("pkgutil").extend_path(__path__, __name__)
+
 from . import auth
 from . import client
 from . import constants

trino/auth.py

@@ -23,6 +23,7 @@
 from typing import List
 from typing import Optional
 from typing import Tuple
+from typing import Union
 from urllib.parse import urlparse
 
 from requests import PreparedRequest
@@ -37,6 +38,12 @@
 from trino.constants import HEADER_ORIGINAL_USER
 from trino.constants import HEADER_USER
 from trino.constants import MAX_NT_PASSWORD_SIZE
+from trino.oauth2 import OAuth2Client
+from trino.oauth2.models import AuthorizationCodeConfig
+from trino.oauth2.models import ClientCredentialsConfig
+from trino.oauth2.models import DeviceCodeConfig
+from trino.oauth2.models import ManualUrlsConfig
+from trino.oauth2.models import OidcConfig
 
 logger = trino.logging.get_logger(__name__)
 
@@ -50,6 +57,23 @@ def get_exceptions(self) -> Tuple[Any, ...]:
         return tuple()
 
 
+class OAuth2TokenAuthentication(Authentication):
+    """Shared base for OAuth2 strategies that authenticate with a bearer token."""
+
+    def __init__(self) -> None:
+        self._oauth2: Optional[OAuth2Client] = None
+
+    @property
+    def oauth2(self) -> OAuth2Client:
+        if self._oauth2 is None:
+            raise RuntimeError("OAuth2 client not initialized")
+        return self._oauth2
+
+    def set_http_session(self, http_session: Session) -> Session:
+        http_session.auth = _BearerAuth(self.oauth2.token())
+        return http_session
+
+
 class KerberosAuthentication(Authentication):
     MUTUAL_REQUIRED = 1
     MUTUAL_OPTIONAL = 2
@@ -276,6 +300,141 @@ def __eq__(self, other: object) -> bool:
         return self.token == other.token
 
 
+class ClientCredentials(OAuth2TokenAuthentication):
+    def __init__(self,
+                 client_id: str,
+                 client_secret: str,
+                 url_config: Union[OidcConfig, ManualUrlsConfig],
+                 scope: Optional[str] = None,
+                 audience: Optional[str] = None):
+        super().__init__()
+        self.client_id = client_id
+        self.client_secret = client_secret
+        self.url_config = url_config
+        self.scope = scope
+        self.audience = audience
+
+        config_args = {
+            "client_id": self.client_id,
+            "client_secret": self.client_secret,
+            "url_config": self.url_config,
+        }
+        if self.scope is not None:
+            config_args["scope"] = self.scope
+        if self.audience is not None:
+            config_args["audience"] = self.audience
+
+        self._oauth2 = OAuth2Client(
+            config=ClientCredentialsConfig(**config_args)
+        )
+
+    def get_exceptions(self) -> Tuple[Any, ...]:
+        return ()
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, ClientCredentials):
+            return False
+        return (
+            self.client_id == other.client_id
+            and self.client_secret == other.client_secret
+            and self.url_config == other.url_config
+        )
+
+
+class DeviceCode(OAuth2TokenAuthentication):
+    def __init__(self,
+                 client_id: str,
+                 url_config: Union[OidcConfig, ManualUrlsConfig],
+                 client_secret: Optional[str] = None,
+                 scope: Optional[str] = None,
+                 audience: Optional[str] = None,
+                 automation_callback: Optional[Callable[[str], None]] = None):
+
+        super().__init__()
+        self.client_id = client_id
+        self.client_secret = client_secret
+        self.url_config = url_config
+        self.scope = scope
+        self.audience = audience
+        self.automation_callback = automation_callback
+
+        config_args = {
+            "client_id": self.client_id,
+            "url_config": self.url_config,
+        }
+        if self.client_secret is not None:
+            config_args["client_secret"] = self.client_secret
+        if self.scope is not None:
+            config_args["scope"] = self.scope
+        if self.audience is not None:
+            config_args["audience"] = self.audience
+        if self.automation_callback is not None:
+            config_args["automation_callback"] = self.automation_callback
+
+        self._oauth2 = OAuth2Client(
+            config=DeviceCodeConfig(**config_args)
+        )
+
+    def get_exceptions(self) -> Tuple[Any, ...]:
+        return ()
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, DeviceCode):
+            return False
+        return (
+            self.client_id == other.client_id
+            and self.client_secret == other.client_secret
+            and self.url_config == other.url_config
+        )
+
+
+class AuthorizationCode(OAuth2TokenAuthentication):
+    def __init__(self,
+                 client_id: str,
+                 url_config: Union[OidcConfig, ManualUrlsConfig],
+                 client_secret: Optional[str] = None,
+                 scope: Optional[str] = None,
+                 audience: Optional[str] = None,
+                 automation_callback: Optional[Callable[[str], None]] = None):
+
+        super().__init__()
+        self.client_id = client_id
+        self.client_secret = client_secret
+        self.url_config = url_config
+        self.scope = scope
+        self.audience = audience
+        self.automation_callback = automation_callback
+
+        config_args = {
+            "client_id": self.client_id,
+            "url_config": self.url_config,
+        }
+        if self.client_secret is not None:
+            config_args["client_secret"] = self.client_secret
+        if self.scope is not None:
+            config_args["scope"] = self.scope
+        if self.audience is not None:
+            config_args["audience"] = self.audience
+        if self.automation_callback is not None:
+            config_args["automation_callback"] = self.automation_callback
+
+        self._oauth2 = OAuth2Client(
+            config=AuthorizationCodeConfig(**config_args)
+        )
+
+    def get_exceptions(self) -> Tuple[Any, ...]:
+        return ()
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, DeviceCode):
+            return False
+        return (
+            self.client_id == other.client_id
+            and self.client_secret == other.client_secret
+            and self.url_config == other.url_config
+        )
+
+
 class RedirectHandler(metaclass=abc.ABCMeta):
     """
     Abstract class for OAuth redirect handlers, inherit from this class to implement your own redirect handler.
@@ -292,7 +451,10 @@ class ConsoleRedirectHandler(RedirectHandler):
     """
 
     def __call__(self, url: str) -> None:
-        print(f"Open the following URL in browser for the external authentication:\n{url}", flush=True)
+        print(
+            f"Open the following URL in browser for the external authentication:\n{url}",
+            flush=True,
+        )
 
 
 class WebBrowserRedirectHandler(RedirectHandler):