kumoai 2.14.0.dev202601011731__cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl

kumoai/connector/glue_connector.py

@@ -0,0 +1,150 @@
+import logging
+from typing import List
+
+from kumoapi.data_source import (
+    CreateConnectorArgs,
+    DataSourceType,
+    GlueConnectorResourceConfig,
+)
+from kumoapi.source_table import GlueSourceTableRequest
+from typing_extensions import Self, override
+
+from kumoai import global_state
+from kumoai.connector import Connector
+
+logger = logging.getLogger(__name__)
+
+_DEFAULT_NAME = 'glue_connector'
+
+
+class GlueConnector(Connector):
+    r"""Defines a connector to a table stored in AWS Glue catalog. Currently,
+    only supports tables in partitioned parquet format. Authenticated via IAM
+    permissions on Glue catalog and data location in S3.
+
+    .. code-block:: python
+
+        import kumoai
+        connector = kumoai.GlueConnector(database="...", region="...", account="...")
+
+        # List all tables:
+        print(connector.table_names())  # Returns: ['articles', 'customers', 'users']
+
+        # Check whether a table is present:
+        assert "articles" in connector
+
+        # Fetch a source table (both approaches are equivalent):
+        source_table = connector["articles"]
+        source_table = connector.table("articles")
+
+    Args:
+        account: The account of the Glue catalog.
+        region: The region of the Glue catalog.
+        database: The name of the database in the Glue Catalog
+
+    """  # noqa
+
+    def __init__(
+            self,
+            name: str,
+            account: str,
+            region: str,
+            database: str,
+            _bypass_creation: bool = False,  # INTERNAL ONLY.
+    ) -> None:
+        self._name = name
+        self.account = account
+        self.region = region
+        self.database = database
+        if global_state.is_spcs and database is not None:
+            raise ValueError(
+                "Glue connectors are not supported when running Kumo in "
+                "Snowpark container services. Please use a Snowflake "
+                "connector instead.")
+
+        if _bypass_creation:
+            return
+
+        self._create_connector()
+
+    @override
+    @property
+    def name(self) -> str:
+        r"""Returns the name of this connector."""
+        return self._name
+
+    @override
+    @property
+    def source_type(self) -> DataSourceType:
+        return DataSourceType.GLUE
+
+    @override
+    def _source_table_request(
+        self,
+        table_names: List[str],
+    ) -> GlueSourceTableRequest:
+        return GlueSourceTableRequest(
+            connector_id=self.name,
+            table_names=table_names,
+        )
+
+    def _create_connector(self) -> None:
+        r"""Creates and persists a Glue connector in the REST DB.
+        Currently only intended for internal use.
+
+
+        Raises:
+            RuntimeError: if connector creation failed
+        """
+        args = CreateConnectorArgs(
+            config=GlueConnectorResourceConfig(
+                name=self.name,
+                account=self.account,
+                region=self.region,
+                database=self.database,
+            ), )
+        global_state.client.connector_api.create_if_not_exist(args)
+
+    def _delete_connector(self) -> None:
+        r"""Deletes a connector in the REST DB. Only intended for internal
+        use.
+        """
+        global_state.client.connector_api.delete_if_exists(self.name)
+
+    # Class properties ########################################################
+
+    @classmethod
+    def get_by_name(cls, name: str) -> Self:
+        r"""Returns an instance of a named Glue Connector, including those created
+        in the Kumo UI.
+
+        Args:
+            name: The name of the existing connector.
+
+        Example:
+            >>> import kumoai
+            >>> connector = kumoai.GlueConnector.get_by_name("name")  # doctest: +SKIP # noqa: E501
+        """
+        api = global_state.client.connector_api
+        resp = api.get(name)
+        if resp is None:
+            raise ValueError(
+                f"There does not exist an existing stored connector with name "
+                f"{name}.")
+        config = resp.config
+        assert isinstance(config, GlueConnectorResourceConfig)
+        return cls(
+            name=config.name,
+            account=config.account,
+            region=config.region,
+            database=config.database,
+            _bypass_creation=True,
+        )
+
+    @override
+    def __repr__(self) -> str:
+        account_name = f"\"{self.account}\"" if self.account else "None"
+        region_name = f"\"{self.region}\"" if self.region else "None"
+        database_name = f"\"{self.database}\"" if self.database else "None"
+        return (f'{self.__class__.__name__}(account={account_name}, '
+                f'region={region_name}, database={database_name})')

kumoai/connector/s3_connector.py

@@ -0,0 +1,278 @@
+import logging
+from typing import List, Optional
+
+from kumoapi.data_source import DataSourceType, FileConnectorResourceConfig
+from kumoapi.source_table import (
+    S3SourceTableRequest,
+    SourceTableConfigRequest,
+    SourceTableConfigResponse,
+    SourceTableListRequest,
+    SourceTableValidateRequest,
+    SourceTableValidateResponse,
+)
+from typing_extensions import Self, override
+
+from kumoai import global_state
+from kumoai.connector import Connector
+from kumoai.connector.source_table import SourceTable
+
+logger = logging.getLogger(__name__)
+
+_DEFAULT_NAME = 's3_connector'
+
+
+class S3Connector(Connector):
+    r"""Defines a connector to a table stored as a file (or partitioned
+    set of files) on the Amazon `S3 <https://aws.amazon.com/s3/>`__ object
+    store. Any table behind an S3 bucket accessible by the shared external IAM
+    role can be accessed through this connector.
+
+    .. code-block:: python
+
+        import kumoai
+        connector = kumoai.S3Connector(root_dir="s3://...")  # an S3 path.
+
+        # List all tables:
+        print(connector.table_names())  # Returns: ['articles', 'customers', 'users']
+
+        # Check whether a table is present:
+        assert "articles" in connector
+
+        # Fetch a source table (both approaches are equivalent):
+        source_table = connector["articles"]
+        source_table = connector.table("articles")
+
+    Args:
+        root_dir: The root directory of this connector. If provided, the root
+            directory is used as a prefix for tables in this connector. If not
+            provided, all tables must be specified by their full S3 paths.
+    """  # noqa
+
+    def __init__(
+        self,
+        root_dir: Optional[str] = None,
+        _connector_id: Optional[str] = None,
+    ) -> None:
+        if _connector_id is not None:
+            # UI S3Connector, named:
+            self._connector_id = _connector_id
+            self.root_dir = None
+            return
+
+        self._connector_id = _DEFAULT_NAME
+        if root_dir is not None:
+            # Remove trailing / to be consistent with boto s3
+            root_dir = root_dir.rstrip('/')
+        self.root_dir = root_dir
+        if global_state.is_spcs and root_dir is not None \
+                and root_dir.startswith('s3://'):
+            raise ValueError(
+                "S3 connectors are not supported when running Kumo in "
+                "Snowpark container services. Please use a Snowflake "
+                "connector instead.")
+
+    @override
+    @property
+    def name(self) -> str:
+        r"""Not supported by :class:`S3Connector`; returns an internal
+        specifier.
+        """
+        return self._connector_id
+
+    @override
+    @property
+    def source_type(self) -> DataSourceType:
+        return DataSourceType.S3
+
+    @override
+    def _source_table_request(
+        self,
+        table_names: List[str],
+    ) -> S3SourceTableRequest:
+        root_dir = self.root_dir
+        if not root_dir and self.name == _DEFAULT_NAME:
+            # Handle None root directories (table name is a path):
+            table_path = S3URI(table_names[0]).validate()
+            root_dir = table_path.root_dir
+            for i, v in enumerate(table_names):
+                uri = S3URI(v)
+                if uri.root_dir != root_dir:
+                    # TODO(manan): fix
+                    raise ValueError(
+                        f"Please ensure that all of your tables are behind "
+                        f"the same root directory ({root_dir}).")
+                table_names[i] = uri.object_name
+
+        connector_id = self.name if self.name != _DEFAULT_NAME else None
+        root_dir = root_dir if self.name == _DEFAULT_NAME else ""
+
+        # TODO(manan): file type?
+        return S3SourceTableRequest(
+            s3_root_dir=root_dir,
+            connector_id=connector_id,
+            table_names=table_names,
+        )
+
+    @override
+    def table(self, name: str) -> SourceTable:
+        r"""Returns a :class:`~kumoai.connector.SourceTable` object
+        corresponding to a source table on Amazon S3.
+
+        Args:
+            name: The name of the table on S3. If :obj:`root_dir` is provided,
+                the path will be specified as :obj:`root_dir/name`. If
+                :obj:`root_dir` is not provided, the name should be the full
+                path (e.g. starting with ``s3://``).
+
+        Raises:
+            :class:`ValueError`: if ``name`` does not exist in the backing
+                connector.
+        """
+        # NOTE only overridden for documentation purposes.
+        return super().table(name)
+
+    @override
+    def _list_tables(self) -> List[str]:
+        connector_id = self.name if self.name != _DEFAULT_NAME else None
+        root_dir = self.root_dir if self.name == _DEFAULT_NAME else None
+        if root_dir is None and connector_id is None:
+            raise ValueError(
+                "Listing tables without a specified root directory is not "
+                "supported. Please specify a root directory to continue; "
+                "alternatively, please access individual tables with their "
+                "full S3 paths.")
+
+        req = SourceTableListRequest(connector_id=connector_id,
+                                     root_dir=root_dir,
+                                     source_type=DataSourceType.S3)
+        return global_state.client.source_table_api.list_tables(req)
+
+    @override
+    def _get_table_config(self, table_name: str) -> SourceTableConfigResponse:
+        root_dir = self.root_dir
+        if not root_dir and self.name == _DEFAULT_NAME:
+            # Handle None root directories (table name is a path):
+            table_path = S3URI(table_name).validate()
+            root_dir = table_path.root_dir
+            table_name = table_path.object_name
+
+        connector_id = self.name if self.name != _DEFAULT_NAME else None
+        root_dir = root_dir if self.name == _DEFAULT_NAME else None
+
+        req = SourceTableConfigRequest(
+            connector_id=connector_id,
+            root_dir=root_dir,
+            table_name=table_name,
+            source_type=self.source_type,
+        )
+        return global_state.client.source_table_api.get_table_config(req)
+
+    @override
+    def _validate_table(self, table_name: str) -> SourceTableValidateResponse:
+        if table_name in self._validated_tables:
+            return SourceTableValidateResponse(is_valid=True, msg='')
+
+        if self.name == _DEFAULT_NAME:
+            # For S3 connector without name, pass root_dir to validate table -
+            # need to infer from table name if root_dir is not provided
+            if self.root_dir:
+                req = SourceTableValidateRequest(
+                    connector_id=None,
+                    table_name=table_name,
+                    source_type=self.source_type,
+                    root_dir=self.root_dir,
+                )
+            else:
+                req_root_dir = '/'.join(table_name.split('/')[:-1])
+                req_table_name = table_name.split('/')[-1]
+
+                req = SourceTableValidateRequest(
+                    connector_id=None,
+                    table_name=req_table_name,
+                    source_type=self.source_type,
+                    root_dir=req_root_dir,
+                )
+        else:
+            req = SourceTableValidateRequest(connector_id=self.name,
+                                             table_name=table_name,
+                                             source_type=self.source_type)
+
+        ret = global_state.client.source_table_api.validate_table(req)
+        # Cache the result for the whole session.
+        if ret.is_valid:
+            self._validated_tables.add(table_name)
+        return ret
+
+    # Class properties ########################################################
+
+    @classmethod
+    def get_by_name(cls, name: str) -> Self:
+        r"""Returns an instance of a named S3 Connector, created in the Kumo UI.
+
+        .. note::
+            Named S3 connectors are read-only: if you would like to modify the
+            root directory, please do so from the UI.
+
+        Args:
+            name: The name of the existing connector.
+
+        Example:
+            >>> import kumoai
+            >>> connector = kumoai.S3Connector.get_by_name("name")  # doctest: +SKIP # noqa: E501
+        """
+        api = global_state.client.connector_api
+        resp = api.get(name)
+        if resp is None:
+            raise ValueError(
+                f"There does not exist an existing stored connector with name "
+                f"{name}.")
+        config = resp.config
+        assert isinstance(config, FileConnectorResourceConfig)
+        return cls(
+            root_dir=None,
+            _connector_id=config.name,
+        )
+
+    def __repr__(self) -> str:
+        if self.name != _DEFAULT_NAME:
+            return f'{self.__class__.__name__}(name={self.name})'
+        root_dir_name = f"\"{self.root_dir}\"" if self.root_dir else "None"
+        return f'{self.__class__.__name__}(root_dir={root_dir_name})'
+
+
+class S3URI:
+    r"""A utility class to parse and navigate S3 URIs."""
+    def __init__(self, uri: str):
+        self.uri: str = uri
+        if uri.endswith('/'):  # remove trailing slash
+            self.uri = uri[:-1]
+
+    @property
+    def is_valid(self) -> bool:
+        # TODO(zeyuan): For SPCS, the path can be a local filesystem path
+        # For train/pred table.
+        if global_state.is_spcs:
+            return True
+        # TODO(manan): implement more checks...
+        return self.uri.startswith("s3://")
+
+    def validate(self) -> Self:
+        if not self.is_valid:
+            raise ValueError(f"Path {self.uri} is not a valid S3 URI.")
+        return self
+
+    @property
+    def root_dir(self) -> str:
+        self.validate()
+        return self.uri.rsplit('/', 1)[0]
+
+    @property
+    def object_name(self) -> str:
+        self.validate()
+        return self.uri.rsplit('/', 1)[1]
+
+    # Class properties ########################################################
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}('
+                f'uri={self.uri}, valid={self.is_valid})')

kumoai/connector/snowflake_connector.py

@@ -0,0 +1,252 @@
+import os
+from typing import Dict, List, Optional
+
+from kumoapi.data_source import (
+    CreateConnectorArgs,
+    DataSourceType,
+    KeyPair,
+    SnowflakeConnectorResourceConfig,
+    UsernamePassword,
+)
+from kumoapi.source_table import SnowflakeSourceTableRequest
+from typing_extensions import Self, override
+
+from kumoai import global_state
+from kumoai.connector import Connector
+
+_ENV_SNOWFLAKE_USER = 'SNOWFLAKE_USER'
+_ENV_SNOWFLAKE_PASSWORD = 'SNOWFLAKE_PASSWORD'
+_ENV_SNOWFLAKE_PRIVATE_KEY = 'SNOWFLAKE_PRIVATE_KEY'
+_ENV_SNOWFLAKE_PRIVATE_KEY_PASSPHRASE = 'SNOWFLAKE_PRIVATE_KEY_PASSPHRASE'
+
+
+class SnowflakeConnector(Connector):
+    r"""Establishes a connection to a `Snowflake <https://www.snowflake.com/>`_
+    database.
+
+    Multiple methods of authentication are available. Username/password
+    authentication is supported either via environment variables
+    (``SNOWFLAKE_USER`` and ``SNOWFLAKE_PASSWORD``) or via keys in the
+    credentials dictionary (:obj:`user` and :obj:`password`).
+
+    .. note::
+        Key-pair authentication is coming soon; please contact your Kumo POC if
+        you need access.
+
+    .. code-block:: python
+
+        import kumoai
+
+        # Either pass `credentials=dict(user=..., password=...)` or set the
+        # 'SNOWFLAKE_USER' and 'SNOWFLAKE_PASSWORD' environment variables:
+        connector = kumoai.SnowflakeConnector(
+            name="<connector_name>",
+            account="<snowflake_account_name>",
+            database="<snowflake_database_name>"
+            schema_name="<snowflake_schema_name>",
+            credentials=credentials,
+        )
+
+        # List all tables:
+        print(connector.table_names())
+
+        # Check whether a table is present:
+        assert "articles" in connector
+
+        # Fetch a source table (both approaches are equivalent):
+        source_table = connector["articles"]
+        source_table = connector.table("articles")
+
+    Args:
+        name: The name of the connector.
+        account: The account name.
+        warehouse: The name of the warehouse.
+        database: The name of the database.
+        schema_name: The name of the schema.
+        credentials: The username and password corresponding to this Snowflake
+            account, if not provided as environment variables.
+    """
+    def __init__(
+            self,
+            name: str,
+            account: str,
+            warehouse: str,
+            database: str,
+            schema_name: str,
+            credentials: Optional[Dict[str, str]] = None,
+            _bypass_creation: bool = False,  # INTERNAL ONLY.
+    ):
+        super().__init__()
+
+        self._name = name
+        self.account = account
+        self.warehouse = warehouse
+
+        # Snowflake DBs and schemas are all in upper-case:
+        self.database = database.upper()
+        self.schema_name = schema_name.upper()
+
+        if _bypass_creation:
+            # TODO(manan, siyang): validate that this connector actually exists
+            # in the REST DB:
+            return
+
+        # Fully specify credentials, create Kumo connector:
+        credentials = credentials or global_state._snowflake_credentials or {}
+
+        credentials_args = {
+            "user": credentials.get("user", os.getenv(_ENV_SNOWFLAKE_USER)),
+        }
+        password = credentials.get("password",
+                                   os.getenv(_ENV_SNOWFLAKE_PASSWORD))
+        private_key = credentials.get("private_key",
+                                      os.getenv(_ENV_SNOWFLAKE_PRIVATE_KEY))
+        private_key_passphrase = credentials.get(
+            "private_key_passphrase",
+            os.getenv(_ENV_SNOWFLAKE_PRIVATE_KEY_PASSPHRASE))
+
+        if not password and not private_key:
+            self._create_native_connector()
+            return
+
+        # Don't pass unused credential fields so that _create_connector can
+        # decide which credential class (KeyPair or UsernamePassword) to use
+        if private_key:
+            credentials_args["private_key"] = private_key
+            if private_key_passphrase:
+                credentials_args[
+                    "private_key_passphrase"] = private_key_passphrase
+        else:
+            credentials_args["password"] = password
+        error_name = None
+        error_var = None
+        if credentials_args["user"] is None:
+            error_name = "username"
+            error_var = _ENV_SNOWFLAKE_USER
+        elif password is None and private_key is None:
+            error_name = "password or private key"
+            error_var = f"{_ENV_SNOWFLAKE_PASSWORD} or " + \
+                        f"{_ENV_SNOWFLAKE_PRIVATE_KEY}"
+        if error_name is not None:
+            raise ValueError(
+                f"Please pass a valid {error_name} to create a Snowflake "
+                f"connector. You can do so either via the 'credentials' "
+                f"argument or the {error_var} environment variable.")
+
+        self._create_connector(credentials_args)  # type: ignore
+
+    @classmethod
+    def get_by_name(cls, name: str) -> Self:
+        r"""Returns an instance of a named Snowflake Connector, including
+        those created in the Kumo UI.
+
+        Args:
+            name: The name of the existing connector.
+
+        Example:
+            >>> import kumoai
+            >>> connector = kumoai.SnowflakeConnector.get_by_name("name")  # doctest: +SKIP # noqa: E501
+        """
+        api = global_state.client.connector_api
+        resp = api.get(name)
+        if resp is None:
+            raise ValueError(
+                f"There does not exist an existing stored connector with name "
+                f"{name}.")
+        config = resp.config
+        assert isinstance(config, SnowflakeConnectorResourceConfig)
+        return cls(
+            name=config.name,
+            account=config.account,
+            warehouse=config.warehouse,
+            database=config.database,
+            schema_name=config.schema_name,
+            credentials=None,
+            _bypass_creation=True,
+        )
+
+    @override
+    @property
+    def name(self) -> str:
+        r"""Returns the name of this connector."""
+        return self._name
+
+    @override
+    @property
+    def source_type(self) -> DataSourceType:
+        return DataSourceType.SNOWFLAKE
+
+    @override
+    def _source_table_request(
+        self,
+        table_names: List[str],
+    ) -> SnowflakeSourceTableRequest:
+        return SnowflakeSourceTableRequest(
+            connector_id=self.name,
+            table_names=table_names,
+        )
+
+    def _create_connector(self, credentials: Dict[str, str]) -> None:
+        r"""Creates and persists a Snowflake connector in the REST DB.
+        Currently only intended for internal use.
+
+        Args:
+            credentials: Fully-specified credentials containing the username
+                and password for the Snowflake connector.
+
+        Raises:
+            RuntimeError: if connector creation failed
+        """
+        # TODO(manan, siyang): consider avoiding connector persistence in the
+        # REST DB, instead moving towards global connectors. For now, to get
+        # a Snowflake experience working smoothly, using the old interface:
+        if credentials.get("password") is not None:
+            credentials = UsernamePassword(
+                username=credentials["user"],
+                password=credentials["password"],
+            )
+        else:
+            credentials = KeyPair(
+                user=credentials["user"],
+                private_key=credentials["private_key"],
+                private_key_passphrase=credentials.get(
+                    "private_key_passphrase"),
+            )
+        args = CreateConnectorArgs(
+            config=SnowflakeConnectorResourceConfig(
+                name=self.name,
+                account=self.account,
+                warehouse=self.warehouse,
+                database=self.database,
+                schema_name=self.schema_name,
+            ),
+            credentials=credentials,
+        )
+        global_state.client.connector_api.create_if_not_exist(args)
+
+    def _create_native_connector(self) -> None:
+        args = CreateConnectorArgs(
+            config=SnowflakeConnectorResourceConfig(
+                name=self.name,
+                account=self.account,
+                warehouse=self.warehouse,
+                database=self.database,
+                schema_name=self.schema_name,
+            ),
+            credentials=None,
+        )
+        global_state.client.connector_api.create_if_not_exist(args)
+
+    def _delete_connector(self) -> None:
+        r"""Deletes a connector in the REST DB. Only intended for internal
+        use.
+        """
+        global_state.client.connector_api.delete_if_exists(self.name)
+
+    # Class properties ########################################################
+
+    @override
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}'
+                f'(account=\"{self.account}\", database=\"{self.database}\", '
+                f'schema=\"{self.schema_name}\")')