ob-metaflow-extensions 1.1.110__tar.gz → 1.1.112__tar.gz

{ob-metaflow-extensions-1.1.110 → ob-metaflow-extensions-1.1.112}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ob-metaflow-extensions
-Version: 1.1.110
+Version: 1.1.112
 Summary: Outerbounds Platform Extensions for Metaflow
 Author: Outerbounds, Inc.
 License: Commercial

{ob-metaflow-extensions-1.1.110 → ob-metaflow-extensions-1.1.112}/metaflow_extensions/outerbounds/plugins/__init__.py

@@ -336,3 +336,5 @@ ENVIRONMENTS_DESC = [
 SECRETS_PROVIDERS_DESC = [
     ("outerbounds", ".secrets.secrets.OuterboundsSecretsProvider"),
 ]
+# Adding an override here so the library can be imported at the metaflow.plugins level
+__mf_promote_submodules__ = ["snowflake"]

{ob-metaflow-extensions-1.1.110 → ob-metaflow-extensions-1.1.112}/metaflow_extensions/outerbounds/plugins/secrets/secrets.py

@@ -37,15 +37,13 @@ class OuterboundsSecretsProvider(SecretsProvider):
         secrets manager on the core oss and returns the secrets.
         """
         headers = {"Content-Type": "application/json", "Connection": "keep-alive"}
-        perimeter, integrations_secrets_metadata_url = self._get_secret_configs()
+        perimeter, integrations_url = self._get_secret_configs()
         integration_name = secret_id
         request_payload = {
             "perimeter_name": perimeter,
             "integration_name": integration_name,
         }
-        response = self._make_request(
-            integrations_secrets_metadata_url, headers, request_payload
-        )
+        response = self._make_request(integrations_url, headers, request_payload)
         secret_resource_id = response.secret_resource_id
         secret_backend_type = response.secret_backend_type
 
@@ -105,26 +103,23 @@ class OuterboundsSecretsProvider(SecretsProvider):
             # if the perimeter is not in metaflow config, try to get it from the environment
             perimeter = environ.get("OBP_PERIMETER", "")
 
-        if "OBP_INTEGRATIONS_SECRETS_METADATA_URL" in conf:
-            integrations_secrets_metadata_url = conf[
-                "OBP_INTEGRATIONS_SECRETS_METADATA_URL"
-            ]
+        if "OBP_INTEGRATIONS_URL" in conf:
+            integrations_url = conf["OBP_INTEGRATIONS_URL"]
         else:
-            # if the integrations secrets metadata url is not in metaflow config, try to get it from the environment
-            integrations_secrets_metadata_url = environ.get(
-                "OBP_INTEGRATIONS_SECRETS_METADATA_URL", ""
-            )
+            # if the integrations is not in metaflow config, try to get it from the environment
+            integrations_url = environ.get("OBP_INTEGRATIONS_URL", "")
 
         if not perimeter:
             raise OuterboundsSecretsException(
                 "No perimeter set. Please make sure to run `outerbounds configure <...>` command which can be found on the Ourebounds UI or reach out to your Outerbounds support team."
             )
 
-        if not integrations_secrets_metadata_url:
+        if not integrations_url:
             raise OuterboundsSecretsException(
-                "No integrations secrets metadata url set. Please notify your Outerbounds support team about this issue."
+                "No integrations url set. Please notify your Outerbounds support team about this issue."
             )
 
+        integrations_secrets_metadata_url = f"{integrations_url}/secrets/metadata"
         return perimeter, integrations_secrets_metadata_url
 
     def _make_request(self, url, headers: Dict, payload: Dict):

ob-metaflow-extensions-1.1.112/metaflow_extensions/outerbounds/plugins/snowflake/__init__.py

@@ -0,0 +1,3 @@
+from .snowflake import connect, get_snowflake_token, Snowflake
+
+__mf_promote_submodules__ = ["plugins.snowflake"]

ob-metaflow-extensions-1.1.112/metaflow_extensions/outerbounds/plugins/snowflake/snowflake.py

@@ -0,0 +1,308 @@
+"""
+This library is an abstraction layer for connecting to snowflake using Outerbounds
+OIDC tokens. It expects that a security integration that authenticates tokens minted
+by Outerbounds has already been configured in the target snowflake account.
+"""
+from metaflow.metaflow_config import SERVICE_URL
+from metaflow.metaflow_config_funcs import init_config
+from typing import Dict
+from os import environ
+
+import json
+import requests
+import time
+
+
+class OuterboundsSnowflakeConnectorException(Exception):
+    pass
+
+
+class OuterboundsSnowflakeIntegrationSpecApiResponse:
+    def __init__(self, response):
+        self.response = response
+
+    @property
+    def account(self):
+        return self.response["account"]
+
+    @property
+    def user(self):
+        return self.response["user"]
+
+    @property
+    def default_role(self):
+        return self.response["default_role"]
+
+    @property
+    def warehouse(self):
+        return self.response["warehouse"]
+
+    @property
+    def database(self):
+        return self.response["database"]
+
+
+def get_snowflake_token(user: str = "", role: str = "", integration: str = "") -> str:
+    """
+    Uses the Outerbounds source token to request for a snowflake compatible OIDC
+    token. This token can then be used to connect to snowflake.
+    user: str
+        The user the token will be minted for
+    role: str
+        The role to which the token will be scoped to
+    integration: str
+        The name of the snowflake integration to use. If not set, an existing integration will be used provided that only one exists per perimeter. If integration is not set and more than one exists, then we raise an exception.
+    """
+    provisioner = SnowflakeIntegrationProvisioner(integration)
+    if not user or not role or not integration:
+        integration_spec = provisioner.get_snowflake_integration_spec()
+        if not user:
+            user = integration_spec.user
+
+        if not role:
+            role = integration_spec.default_role
+
+        if not integration:
+            integration = provisioner.get_integration_name()
+
+    snowflake_token_url = provisioner.get_snowflake_token_url()
+    perimeter = provisioner.get_perimeter()
+    payload = {
+        "perimeterName": perimeter,
+        "snowflakeUser": user,
+        "snowflakeRole": role,
+        "integrationName": integration,
+    }
+    json_payload = json.dumps(payload)
+    headers = provisioner.get_service_auth_header()
+    response = requests.get(snowflake_token_url, data=json_payload, headers=headers)
+    response.raise_for_status()
+    return response.json()["token"]
+
+
+def connect(user: str = "", role: str = "", integration: str = "", **kwargs):
+    """
+    Connect to snowflake using the token minted by Outerbounds
+    user: str
+        The user name used to authenticate with snowflake
+    role: str
+        The role to request when connect with snowflake
+    integration: str
+        The name of the snowflake integration to use. If not set, an existing integration will be used provided that only one exists in the current perimeter. If integration is not set and more than one exists in the current perimeter, then we raise an exception.
+    kwargs: dict
+        Additional arguments to pass to the python snowflake connector
+    """
+    # ensure password is not set
+    if "password" in kwargs:
+        raise OuterboundsSnowflakeConnectorException(
+            "Password should not be set when using Outerbounds snowflake connector."
+        )
+
+    provisioner = SnowflakeIntegrationProvisioner(integration)
+    get_defaults = any(
+        key not in kwargs for key in ["account", "warehouse", "database"]
+    )
+    if not user or not role or not integration or get_defaults:
+        integration_spec = provisioner.get_snowflake_integration_spec()
+        if not user:
+            user = integration_spec.user
+
+        if not role:
+            role = integration_spec.default_role
+
+        if not integration:
+            integration = provisioner.get_integration_name()
+
+        if "account" not in kwargs:
+            kwargs["account"] = integration_spec.account
+
+        if "warehouse" not in kwargs:
+            kwargs["warehouse"] = integration_spec.warehouse
+
+        # if the user is attempting to use a warehouse different from what is specified in the
+        # integration we will not set the database
+        if (
+            "database" not in kwargs
+            and kwargs["warehouse"] == integration_spec.warehouse
+        ):
+            kwargs["database"] = integration_spec.database
+
+    # get snowflake token
+    token = get_snowflake_token(user=user, role=role, integration=integration)
+    kwargs["token"] = token
+    kwargs["authenticator"] = "oauth"
+    kwargs["role"] = role
+    kwargs["user"] = user
+
+    # connect to snowflake
+    try:
+        from snowflake.connector import connect
+
+        cn = connect(**kwargs)
+        return cn
+    except ImportError as ie:
+        raise OuterboundsSnowflakeConnectorException(
+            f"Error importing snowflake connector: {ie}.\nPlease make sure the 'snowflake-connector-python' package has been installed by running 'pip install -U \"outerbounds[snowflake]\"' or using the Metaflow decorators @pypi or @conda."
+        )
+    except Exception as e:
+        raise OuterboundsSnowflakeConnectorException(
+            f"Error connecting to snowflake: {e}"
+        )
+
+
+class Snowflake:
+    def __init__(
+        self, user: str = "", role: str = "", integration: str = "", **kwargs
+    ) -> None:
+        self.cn = connect(user, role, integration, **kwargs)
+
+    def __enter__(self):
+        return self.cn
+
+    def __exit__(self, exception_type, exception_value, traceback):
+        self.cn.close()
+
+    def close(self):
+        self.cn.close()
+
+
+class SnowflakeIntegrationProvisioner:
+    def __init__(self, integration_name: str) -> None:
+        self.conf = init_config()
+        self.integration_name = integration_name
+
+    def get_snowflake_integration_spec(
+        self,
+    ) -> OuterboundsSnowflakeIntegrationSpecApiResponse:
+        integrations_url = self._get_integration_url()
+        perimeter = self.get_perimeter()
+        headers = {"Content-Type": "application/json", "Connection": "keep-alive"}
+        request_payload = {
+            "perimeter_name": perimeter,
+        }
+        # if integration is not set, list all integrations
+        if not self.integration_name:
+            list_snowflake_integrations_url = f"{integrations_url}/snowflake"
+            response = self._make_request(
+                list_snowflake_integrations_url, headers, request_payload
+            )
+            snowflake_integrations = response.get("integrations", [])
+            if not snowflake_integrations:
+                raise OuterboundsSnowflakeConnectorException(
+                    "No snowflake integrations found. Please make sure you have created a Snowflake integration on the Outerbounds UI first."
+                )
+
+            if len(snowflake_integrations) > 1:
+                raise OuterboundsSnowflakeConnectorException(
+                    f"Multiple snowflake integrations found. Please specify a specific integration name you would like to use."
+                )
+
+            self.integration_name = snowflake_integrations[0]["integration_name"]
+            return OuterboundsSnowflakeIntegrationSpecApiResponse(
+                snowflake_integrations[0]["integration_spec"]
+            )
+
+        get_snowflake_integration_url = (
+            f"{integrations_url}/snowflake/{self.integration_name}"
+        )
+        response = self._make_request(
+            get_snowflake_integration_url, headers, request_payload
+        )
+        self.integration_name = response["integration_name"]
+        return OuterboundsSnowflakeIntegrationSpecApiResponse(
+            response["integration_spec"]
+        )
+
+    def get_integration_name(self) -> str:
+        return self.integration_name
+
+    def get_perimeter(self) -> str:
+        if "OBP_PERIMETER" in self.conf:
+            perimeter = self.conf["OBP_PERIMETER"]
+        else:
+            # if the perimeter is not in metaflow config, try to get it from the environment
+            perimeter = environ.get("OBP_PERIMETER", "")
+        if not perimeter:
+            raise OuterboundsSnowflakeConnectorException(
+                "No perimeter set. Please make sure to run `outerbounds configure <...>` command which can be found on the Ourebounds UI or reach out to your Outerbounds support team."
+            )
+
+        return perimeter
+
+    def get_snowflake_token_url(self) -> str:
+        if "OBP_AUTH_SERVER" in self.conf:
+            auth_host = self.conf["OBP_AUTH_SERVER"]
+        else:
+            from urllib.parse import urlparse
+
+            auth_host = "auth." + urlparse(SERVICE_URL).hostname.split(".", 1)[1]
+
+        return "https://" + auth_host + "/generate/snowflake"
+
+    def get_service_auth_header(self) -> str:
+        if "METAFLOW_SERVICE_AUTH_KEY" in self.conf:
+            return {"x-api-key": self.conf["METAFLOW_SERVICE_AUTH_KEY"]}
+        else:
+            return json.loads(environ.get("METAFLOW_SERVICE_HEADERS"))
+
+    def _get_integration_url(self) -> str:
+        from metaflow_extensions.outerbounds.remote_config import init_config
+        from os import environ
+
+        if "OBP_INTEGRATIONS_URL" in self.conf:
+            integrations_url = self.conf["OBP_INTEGRATIONS_URL"]
+        else:
+            # if the integrations url is not in metaflow config, try to get it from the environment
+            integrations_url = environ.get("OBP_INTEGRATIONS_URL", "")
+
+        if not integrations_url:
+            raise OuterboundsSnowflakeConnectorException(
+                "No integrations url set. Please notify your Outerbounds support team about this issue."
+            )
+
+        return integrations_url
+
+    def _make_request(self, url, headers: Dict, payload: Dict) -> Dict:
+        try:
+            from metaflow.metaflow_config import SERVICE_HEADERS
+
+            request_headers = {**headers, **(SERVICE_HEADERS or {})}
+        except ImportError:
+            headers = headers
+
+        retryable_status_codes = [409]
+        json_payload = json.dumps(payload)
+        for attempt in range(2):  # 0 = initial attempt, 1-2 = retries
+            response = requests.get(url, data=json_payload, headers=request_headers)
+            if response.status_code not in retryable_status_codes:
+                break
+
+            if attempt < 2:  # Don't sleep after the last attempt
+                sleep_time = 0.5 * (attempt + 1)
+                time.sleep(sleep_time)
+
+        response = requests.get(url, data=json_payload, headers=request_headers)
+        self._handle_error_response(response)
+        return response.json()
+
+    @staticmethod
+    def _handle_error_response(response: requests.Response):
+        if response.status_code >= 500:
+            raise OuterboundsSnowflakeConnectorException(
+                f"Server error: {response.text}. Please reach out to your Outerbounds support team."
+            )
+
+        body = response.json()
+        status_code = body.get("error", {}).get("statusCode", response.status_code)
+        if status_code == 404:
+            raise OuterboundsSnowflakeConnectorException(f"Secret not found: {body}")
+
+        if status_code >= 400:
+            try:
+                raise OuterboundsSnowflakeConnectorException(
+                    f"status_code={status_code}\t*{body['error']['details']['kind']}*\n{body['error']['details']['message']}"
+                )
+            except KeyError:
+                raise OuterboundsSnowflakeConnectorException(
+                    f"status_code={status_code} Unexpected error: {body}"
+                )

{ob-metaflow-extensions-1.1.110 → ob-metaflow-extensions-1.1.112}/metaflow_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.py

@@ -51,3 +51,4 @@ def S3(*args, **kwargs):
 
 
 from .. import profilers
+from ..plugins.snowflake import Snowflake

ob-metaflow-extensions-1.1.112/metaflow_extensions/outerbounds/toplevel/plugins/snowflake/__init__.py

@@ -0,0 +1 @@
+__mf_promote_submodules__ = ["plugins.snowflake"]

{ob-metaflow-extensions-1.1.110 → ob-metaflow-extensions-1.1.112}/ob_metaflow_extensions.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ob-metaflow-extensions
-Version: 1.1.110
+Version: 1.1.112
 Summary: Outerbounds Platform Extensions for Metaflow
 Author: Outerbounds, Inc.
 License: Commercial

{ob-metaflow-extensions-1.1.110 → ob-metaflow-extensions-1.1.112}/ob_metaflow_extensions.egg-info/SOURCES.txt

@@ -24,6 +24,8 @@ metaflow_extensions/outerbounds/plugins/profilers/deco_injector.py
 metaflow_extensions/outerbounds/plugins/profilers/gpu_profile_decorator.py
 metaflow_extensions/outerbounds/plugins/secrets/__init__.py
 metaflow_extensions/outerbounds/plugins/secrets/secrets.py
+metaflow_extensions/outerbounds/plugins/snowflake/__init__.py
+metaflow_extensions/outerbounds/plugins/snowflake/snowflake.py
 metaflow_extensions/outerbounds/plugins/snowpark/__init__.py
 metaflow_extensions/outerbounds/plugins/snowpark/snowpark.py
 metaflow_extensions/outerbounds/plugins/snowpark/snowpark_cli.py
@@ -40,6 +42,7 @@ metaflow_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.py
 metaflow_extensions/outerbounds/toplevel/plugins/azure/__init__.py
 metaflow_extensions/outerbounds/toplevel/plugins/gcp/__init__.py
 metaflow_extensions/outerbounds/toplevel/plugins/kubernetes/__init__.py
+metaflow_extensions/outerbounds/toplevel/plugins/snowflake/__init__.py
 ob_metaflow_extensions.egg-info/PKG-INFO
 ob_metaflow_extensions.egg-info/SOURCES.txt
 ob_metaflow_extensions.egg-info/dependency_links.txt

ob-metaflow-extensions-1.1.112/ob_metaflow_extensions.egg-info/requires.txt

@@ -0,0 +1,3 @@
+boto3
+kubernetes
+ob-metaflow==2.12.36.3

{ob-metaflow-extensions-1.1.110 → ob-metaflow-extensions-1.1.112}/setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_namespace_packages
 from pathlib import Path
 
 
-version = "1.1.110"
+version = "1.1.112"
 this_directory = Path(__file__).parent
 long_description = (this_directory / "README.md").read_text()
 
@@ -18,6 +18,6 @@ setup(
     install_requires=[
         "boto3",
         "kubernetes",
-        "ob-metaflow == 2.12.36.1",
+        "ob-metaflow == 2.12.36.3",
     ],
 )

ob-metaflow-extensions-1.1.110/ob_metaflow_extensions.egg-info/requires.txt

@@ -1,3 +0,0 @@
-boto3
-kubernetes
-ob-metaflow==2.12.36.1