hyperlake 0.1.1__py3-none-any.whl

hyperlake/__init__.py

@@ -0,0 +1,40 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+from . import auth
+from . import client
+from . import constants
+from . import dbapi
+from . import exceptions
+from . import logging
+from ._version import __author__
+from ._version import __author_email__
+from ._version import __description__
+from ._version import __license__
+from ._version import __title__
+from ._version import __url__
+from ._version import __version__
+
+__all__ = [
+    "auth",
+    "client",
+    "constants",
+    "dbapi",
+    "exceptions",
+    "logging",
+    "__author__",
+    "__author_email__",
+    "__description__",
+    "__license__",
+    "__title__",
+    "__url__",
+    "__version__",
+]

hyperlake/_version.py

@@ -0,0 +1,7 @@
+__title__ = "hyperlake"
+__description__ = "Python Client for the Hyperlake distributed SQL Engine"
+__url__ = "https://github.com/ockhamlabs/hyperlake-python-client"
+__version__ = "0.1.1"
+__author__ = "Ockham Labs Team"
+__author_email__ = "pythonclient+hello@ockhamlabs.ai"
+__license__ = "Apache 2.0"

hyperlake/auth.py

@@ -0,0 +1,609 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+import abc
+import importlib
+import json
+import os
+import re
+import threading
+import webbrowser
+from collections.abc import Mapping
+from typing import Any
+from typing import Callable
+from typing import Dict
+from typing import List
+from typing import Optional
+from typing import Tuple
+from urllib.parse import urlparse
+
+from requests import PreparedRequest
+from requests import Request
+from requests import Response
+from requests import Session
+from requests.auth import AuthBase
+from requests.auth import extract_cookies_to_jar
+
+import hyperlake.logging
+from hyperlake import exceptions
+from hyperlake.constants import HEADER_USER
+from hyperlake.constants import MAX_NT_PASSWORD_SIZE
+
+logger = hyperlake.logging.get_logger(__name__)
+
+
+class Authentication(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    def set_http_session(self, http_session: Session) -> Session:
+        pass
+
+    def get_exceptions(self) -> Tuple[Any, ...]:
+        return tuple()
+
+
+class KerberosAuthentication(Authentication):
+    def __init__(
+        self,
+        config: Optional[str] = None,
+        service_name: Optional[str] = None,
+        mutual_authentication: bool = False,
+        force_preemptive: bool = False,
+        hostname_override: Optional[str] = None,
+        sanitize_mutual_error_response: bool = True,
+        principal: Optional[str] = None,
+        delegate: bool = False,
+        ca_bundle: Optional[str] = None,
+    ) -> None:
+        self._config = config
+        self._service_name = service_name
+        self._mutual_authentication = mutual_authentication
+        self._force_preemptive = force_preemptive
+        self._hostname_override = hostname_override
+        self._sanitize_mutual_error_response = sanitize_mutual_error_response
+        self._principal = principal
+        self._delegate = delegate
+        self._ca_bundle = ca_bundle
+
+    def set_http_session(self, http_session: Session) -> Session:
+        try:
+            import requests_kerberos
+        except ImportError:
+            raise RuntimeError("unable to import requests_kerberos")
+
+        if self._config:
+            os.environ["KRB5_CONFIG"] = self._config
+        http_session.trust_env = False
+        http_session.auth = requests_kerberos.HTTPKerberosAuth(
+            mutual_authentication=self._mutual_authentication,
+            force_preemptive=self._force_preemptive,
+            hostname_override=self._hostname_override,
+            sanitize_mutual_error_response=self._sanitize_mutual_error_response,
+            principal=self._principal,
+            delegate=self._delegate,
+            service=self._service_name,
+        )
+        if self._ca_bundle:
+            http_session.verify = self._ca_bundle
+        return http_session
+
+    def get_exceptions(self) -> Tuple[Any, ...]:
+        try:
+            from requests_kerberos.exceptions import KerberosExchangeError
+
+            return KerberosExchangeError,
+        except ImportError:
+            raise RuntimeError("unable to import requests_kerberos")
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, KerberosAuthentication):
+            return False
+        return (self._config == other._config
+                and self._service_name == other._service_name
+                and self._mutual_authentication == other._mutual_authentication
+                and self._force_preemptive == other._force_preemptive
+                and self._hostname_override == other._hostname_override
+                and self._sanitize_mutual_error_response == other._sanitize_mutual_error_response
+                and self._principal == other._principal
+                and self._delegate == other._delegate
+                and self._ca_bundle == other._ca_bundle)
+
+
+class GSSAPIAuthentication(Authentication):
+    def __init__(
+        self,
+        config: Optional[str] = None,
+        service_name: Optional[str] = None,
+        mutual_authentication: bool = False,
+        force_preemptive: bool = False,
+        hostname_override: Optional[str] = None,
+        sanitize_mutual_error_response: bool = True,
+        principal: Optional[str] = None,
+        delegate: bool = False,
+        ca_bundle: Optional[str] = None,
+    ) -> None:
+        self._config = config
+        self._service_name = service_name
+        self._mutual_authentication = mutual_authentication
+        self._force_preemptive = force_preemptive
+        self._hostname_override = hostname_override
+        self._sanitize_mutual_error_response = sanitize_mutual_error_response
+        self._principal = principal
+        self._delegate = delegate
+        self._ca_bundle = ca_bundle
+
+    def set_http_session(self, http_session: Session) -> Session:
+        try:
+            import requests_gssapi
+        except ImportError:
+            raise RuntimeError("unable to import requests_gssapi")
+
+        if self._config:
+            os.environ["KRB5_CONFIG"] = self._config
+        http_session.trust_env = False
+        http_session.auth = requests_gssapi.HTTPSPNEGOAuth(
+            mutual_authentication=self._mutual_authentication,
+            opportunistic_auth=self._force_preemptive,
+            target_name=self._get_target_name(self._hostname_override, self._service_name),
+            sanitize_mutual_error_response=self._sanitize_mutual_error_response,
+            creds=self._get_credentials(self._principal),
+            delegate=self._delegate,
+        )
+        if self._ca_bundle:
+            http_session.verify = self._ca_bundle
+        return http_session
+
+    def _get_credentials(self, principal: Optional[str] = None) -> Any:
+        if principal:
+            try:
+                import gssapi
+            except ImportError:
+                raise RuntimeError("unable to import gssapi")
+
+            name = gssapi.Name(principal, gssapi.NameType.user)
+            return gssapi.Credentials(name=name, usage="initiate")
+
+        return None
+
+    def _get_target_name(
+            self,
+            hostname_override: Optional[str] = None,
+            service_name: Optional[str] = None,
+    ) -> Any:
+        if service_name is not None:
+            try:
+                import gssapi
+            except ImportError:
+                raise RuntimeError("unable to import gssapi")
+
+            if hostname_override is None:
+                raise ValueError("service name must be used together with hostname_override")
+
+            kerb_spn = "{0}@{1}".format(service_name, hostname_override)
+            return gssapi.Name(kerb_spn, gssapi.NameType.hostbased_service)
+
+        return hostname_override
+
+    def get_exceptions(self) -> Tuple[Any, ...]:
+        try:
+            from requests_gssapi.exceptions import SPNEGOExchangeError
+
+            return SPNEGOExchangeError,
+        except ImportError:
+            raise RuntimeError("unable to import requests_kerberos")
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, GSSAPIAuthentication):
+            return False
+        return (self._config == other._config
+                and self._service_name == other._service_name
+                and self._mutual_authentication == other._mutual_authentication
+                and self._force_preemptive == other._force_preemptive
+                and self._hostname_override == other._hostname_override
+                and self._sanitize_mutual_error_response == other._sanitize_mutual_error_response
+                and self._principal == other._principal
+                and self._delegate == other._delegate
+                and self._ca_bundle == other._ca_bundle)
+
+
+class BasicAuthentication(Authentication):
+    def __init__(self, username: str, password: str):
+        self._username = username
+        self._password = password
+
+    def set_http_session(self, http_session: Session) -> Session:
+        try:
+            import requests.auth
+        except ImportError:
+            raise RuntimeError("unable to import requests.auth")
+
+        http_session.auth = requests.auth.HTTPBasicAuth(self._username, self._password)
+        return http_session
+
+    def get_exceptions(self) -> Tuple[Any, ...]:
+        return ()
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, BasicAuthentication):
+            return False
+        return self._username == other._username and self._password == other._password
+
+
+class _BearerAuth(AuthBase):
+    """
+    Custom implementation of Authentication class for bearer token
+    """
+
+    def __init__(self, token: str):
+        self.token = token
+
+    def __call__(self, r: PreparedRequest) -> PreparedRequest:
+        r.headers["Authorization"] = "Bearer " + self.token
+        return r
+
+
+class JWTAuthentication(Authentication):
+
+    def __init__(self, token: str):
+        self.token = token
+
+    def set_http_session(self, http_session: Session) -> Session:
+        http_session.auth = _BearerAuth(self.token)
+        return http_session
+
+    def get_exceptions(self) -> Tuple[Any, ...]:
+        return ()
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, JWTAuthentication):
+            return False
+        return self.token == other.token
+
+
+class RedirectHandler(metaclass=abc.ABCMeta):
+    """
+    Abstract class for OAuth redirect handlers, inherit from this class to implement your own redirect handler.
+    """
+
+    @abc.abstractmethod
+    def __call__(self, url: str) -> None:
+        raise NotImplementedError()
+
+
+class ConsoleRedirectHandler(RedirectHandler):
+    """
+    Handler for OAuth redirections to log to console.
+    """
+
+    def __call__(self, url: str) -> None:
+        print("Open the following URL in browser for the external authentication:")
+        print(url)
+
+
+class WebBrowserRedirectHandler(RedirectHandler):
+    """
+    Handler for OAuth redirections to open in web browser.
+    """
+
+    def __call__(self, url: str) -> None:
+        webbrowser.open_new(url)
+
+
+class CompositeRedirectHandler(RedirectHandler):
+    """
+    Composite handler for OAuth redirect handlers.
+    """
+
+    def __init__(self, handlers: List[Callable[[str], None]]):
+        self.handlers = handlers
+
+    def __call__(self, url: str) -> None:
+        for handler in self.handlers:
+            handler(url)
+
+
+class _OAuth2TokenCache(metaclass=abc.ABCMeta):
+    """
+    Abstract class for OAuth token cache, inherit from this class to implement your own token cache.
+    """
+
+    @abc.abstractmethod
+    def get_token_from_cache(self, key: Optional[str]) -> Optional[str]:
+        pass
+
+    @abc.abstractmethod
+    def store_token_to_cache(self, key: Optional[str], token: str) -> None:
+        pass
+
+
+class _OAuth2TokenInMemoryCache(_OAuth2TokenCache):
+    """
+    Multiple clients can share the same cache only if each connection explicitly specifies
+    a user otherwise the first cached token will be used to authenticate all other users.
+    """
+
+    def __init__(self) -> None:
+        self._cache: Dict[Optional[str], str] = {}
+
+    def get_token_from_cache(self, key: Optional[str]) -> Optional[str]:
+        return self._cache.get(key)
+
+    def store_token_to_cache(self, key: Optional[str], token: str) -> None:
+        self._cache[key] = token
+
+
+class _OAuth2KeyRingTokenCache(_OAuth2TokenCache):
+    """
+    Keyring token cache implementation
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        try:
+            self._keyring = importlib.import_module("keyring")
+        except ImportError:
+            self._keyring = None  # type: ignore
+            logger.info("keyring module not found. OAuth2 token will not be stored in keyring.")
+
+    def is_keyring_available(self) -> bool:
+        return self._keyring is not None \
+            and not isinstance(self._keyring.get_keyring(), self._keyring.backends.fail.Keyring)
+
+    def get_token_from_cache(self, key: Optional[str]) -> Optional[str]:
+        password = self._keyring.get_password(key, "token")
+
+        try:
+            password_as_dict = json.loads(str(password))
+            if password_as_dict.get("sharded_password"):
+                # if password was stored shared, reconstruct it
+                shard_count = int(password_as_dict.get("shard_count"))
+
+                password = ""
+                for i in range(shard_count):
+                    password += str(self._keyring.get_password(key, f"token__{i}"))
+
+        except self._keyring.errors.NoKeyringError as e:
+            raise hyperlake.exceptions.NotSupportedError("Although keyring module is installed no backend has been "
+                                                     "detected, check https://pypi.org/project/keyring/ for more "
+                                                     "information.") from e
+        except ValueError:
+            pass
+
+        return password
+
+    def store_token_to_cache(self, key: Optional[str], token: str) -> None:
+        # keyring is installed, so we can store the token for reuse within multiple threads
+        try:
+            # if not Windows or "small" password, stick to the default
+            if os.name != "nt" or len(token) < MAX_NT_PASSWORD_SIZE:
+                self._keyring.set_password(key, "token", token)
+            else:
+                logger.debug(f"password is {len(token)} characters, sharding it.")
+
+                password_shards = [
+                    token[i: i + MAX_NT_PASSWORD_SIZE] for i in range(0, len(token), MAX_NT_PASSWORD_SIZE)
+                ]
+                shard_info = {
+                    "sharded_password": True,
+                    "shard_count": len(password_shards),
+                }
+
+                # store the "shard info" as the "base" password
+                self._keyring.set_password(key, "token", json.dumps(shard_info))
+                # then store all shards with the shard number as postfix
+                for i, s in enumerate(password_shards):
+                    self._keyring.set_password(key, f"token__{i}", s)
+        except self._keyring.errors.NoKeyringError as e:
+            raise hyperlake.exceptions.NotSupportedError("Although keyring module is installed no backend has been "
+                                                     "detected, check https://pypi.org/project/keyring/ for more "
+                                                     "information.") from e
+
+
+class _OAuth2TokenBearer(AuthBase):
+    """
+    Custom implementation of hyperlake Trino OAuth2 based authentication to get the token
+    """
+    MAX_OAUTH_ATTEMPTS = 5
+    _BEARER_PREFIX = re.compile(r"bearer", flags=re.IGNORECASE)
+
+    def __init__(self, redirect_auth_url_handler: Callable[[str], None]):
+        self._redirect_auth_url = redirect_auth_url_handler
+        keyring_cache = _OAuth2KeyRingTokenCache()
+        self._token_cache = keyring_cache if keyring_cache.is_keyring_available() else _OAuth2TokenInMemoryCache()
+        self._token_lock = threading.Lock()
+        self._inside_oauth_attempt_lock = threading.Lock()
+        self._inside_oauth_attempt_blocker = threading.Event()
+
+    def __call__(self, r: PreparedRequest) -> PreparedRequest:
+        host = self._determine_host(r.url)
+        user = self._determine_user(r.headers)
+        key = self._construct_cache_key(host, user)
+        token = self._get_token_from_cache(key)
+
+        if token is not None:
+            r.headers['Authorization'] = "Bearer " + token
+
+        r.register_hook('response', self._authenticate)
+
+        return r
+
+    def _authenticate(self, response: Response, **kwargs: Any) -> Optional[Response]:
+        if not 400 <= response.status_code < 500:
+            return response
+
+        acquired = self._inside_oauth_attempt_lock.acquire(blocking=False)
+        if acquired:
+            try:
+                # Lock is acquired, attempt the OAuth2 flow
+                self._attempt_oauth(response, **kwargs)
+                self._inside_oauth_attempt_blocker.set()
+            finally:
+                self._inside_oauth_attempt_lock.release()
+                self._inside_oauth_attempt_blocker.clear()
+        else:
+            # Lock is not acquired, we are already in the OAuth2 flow, so we block until OAuth2 flow is finished.
+            self._inside_oauth_attempt_blocker.wait()
+
+        return self._retry_request(response, **kwargs)
+
+    def _attempt_oauth(self, response: Response, **kwargs: Any) -> None:
+        # we have to handle the authentication, may be token the token expired, or it wasn't there at all
+        auth_info = response.headers.get('WWW-Authenticate')
+        if not auth_info:
+            raise exceptions.HyperlakeAuthError("Error: header WWW-Authenticate not available in the response.")
+
+        if not _OAuth2TokenBearer._BEARER_PREFIX.search(auth_info):
+            raise exceptions.HyperlakeAuthError(f"Error: header info didn't match {auth_info}")
+
+        # Example www-authenticate header value:
+        # 'Basic realm="Trino", Bearer realm="Trino", token_type="JWT",
+        # Bearer x_redirect_server="https://trino.com/oauth2/token/uuid4",
+        # x_token_server="https://trino.com/oauth2/token/uuid4"'
+        auth_info_headers = self._parse_authenticate_header(auth_info)
+
+        auth_server = auth_info_headers.get('bearer x_redirect_server', auth_info_headers.get('x_redirect_server'))
+        token_server = auth_info_headers.get('bearer x_token_server', auth_info_headers.get('x_token_server'))
+        if token_server is None:
+            raise exceptions.HyperlakeAuthError("Error: header info didn't have x_token_server")
+
+        if auth_server is not None:
+            # tell app that use this url to proceed with the authentication
+            self._redirect_auth_url(auth_server)
+
+        # Consume content and release the original connection
+        # to allow our new request to reuse the same one.
+        response.content
+        response.close()
+
+        token = self._get_token(token_server, response, **kwargs)
+
+        request = response.request
+        host = self._determine_host(request.url)
+        user = self._determine_user(request.headers)
+        key = self._construct_cache_key(host, user)
+        self._store_token_to_cache(key, token)
+
+    def _retry_request(self, response: Response, **kwargs: Any) -> Optional[Response]:
+        request = response.request.copy()
+        extract_cookies_to_jar(request._cookies, response.request, response.raw)
+        request.prepare_cookies(request._cookies)
+
+        host = self._determine_host(response.request.url)
+        user = self._determine_user(request.headers)
+        key = self._construct_cache_key(host, user)
+        token = self._get_token_from_cache(key)
+        if token is not None:
+            request.headers['Authorization'] = "Bearer " + token
+        retry_response = response.connection.send(request, **kwargs)
+        retry_response.history.append(response)
+        retry_response.request = request
+        return retry_response
+
+    def _get_token(self, token_server: str, response: Response, **kwargs: Any) -> str:
+        attempts = 0
+        while attempts < self.MAX_OAUTH_ATTEMPTS:
+            attempts += 1
+            with response.connection.send(Request(
+                    method='GET', url=token_server).prepare(), **kwargs) as response:
+                if response.status_code == 200:
+                    token_response = json.loads(response.text)
+                    token = token_response.get('token')
+                    if token:
+                        return token
+                    error = token_response.get('error')
+                    if error:
+                        raise exceptions.HyperlakeAuthError(f"Error while getting the token: {error}")
+                    else:
+                        token_server = token_response.get('nextUri')
+                        logger.debug(f"nextURi auth token server: {token_server}")
+                else:
+                    raise exceptions.HyperlakeAuthError(
+                        f"Error while getting the token response "
+                        f"status code: {response.status_code}, "
+                        f"body: {response.text}")
+
+        raise exceptions.HyperlakeAuthError("Exceeded max attempts while getting the token")
+
+    def _get_token_from_cache(self, key: Optional[str]) -> Optional[str]:
+        with self._token_lock:
+            return self._token_cache.get_token_from_cache(key)
+
+    def _store_token_to_cache(self, key: Optional[str], token: str) -> None:
+        with self._token_lock:
+            self._token_cache.store_token_to_cache(key, token)
+
+    @staticmethod
+    def _determine_host(url: Optional[str]) -> Any:
+        return urlparse(url).hostname
+
+    @staticmethod
+    def _determine_user(headers: Mapping[Any, Any]) -> Optional[Any]:
+        return headers.get(HEADER_USER)
+
+    @staticmethod
+    def _construct_cache_key(host: Optional[str], user: Optional[str]) -> Optional[str]:
+        if user is None:
+            return host
+        else:
+            return f"{host}@{user}"
+
+    @staticmethod
+    def _parse_authenticate_header(header: str) -> Dict[str, str]:
+        logger.debug(f"Authentication header: {header}")
+        components = header.split(",")
+        auth_info_headers = {}
+
+        for component in components:
+            component = component.strip()
+            if "=" in component:
+                key, value = component.split("=", 1)
+                if value[0] == '"' and value[-1] == '"':
+                    value = value[1:-1]
+                auth_info_headers[key.lower()] = value
+        return auth_info_headers
+
+
+class OAuth2Authentication(Authentication):
+    def __init__(self, redirect_auth_url_handler: CompositeRedirectHandler = CompositeRedirectHandler([
+        WebBrowserRedirectHandler(),
+        ConsoleRedirectHandler()
+    ])):
+        self._redirect_auth_url = redirect_auth_url_handler
+        self._bearer = _OAuth2TokenBearer(self._redirect_auth_url)
+
+    def set_http_session(self, http_session: Session) -> Session:
+        http_session.auth = self._bearer
+        return http_session
+
+    def get_exceptions(self) -> Tuple[Any, ...]:
+        return ()
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, OAuth2Authentication):
+            return False
+        return self._redirect_auth_url == other._redirect_auth_url
+
+
+class CertificateAuthentication(Authentication):
+    def __init__(self, cert: str, key: str):
+        self._cert = cert
+        self._key = key
+
+    def set_http_session(self, http_session: Session) -> Session:
+        http_session.cert = (self._cert, self._key)
+        return http_session
+
+    def get_exceptions(self) -> Tuple[Any, ...]:
+        return ()
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, CertificateAuthentication):
+            return False
+        return self._cert == other._cert and self._key == other._key