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

kumoai/connector/bigquery_connector.py

@@ -0,0 +1,200 @@
+import os
+from typing import Dict, List, Optional
+
+from kumoapi.data_source import (
+    BigQueryConnectorResourceConfig,
+    BigQueryCredentials,
+    CreateConnectorArgs,
+    DataSourceType,
+)
+from kumoapi.source_table import BigQuerySourceTableRequest
+from typing_extensions import Self, override
+
+from kumoai import global_state
+from kumoai.connector import Connector
+
+_ENV_BIGQUERY_PRIVATE_KEY_ID = 'BIGQUERY_PRIVATE_KEY_ID'
+_ENV_BIGQUERY_PRIVATE_KEY = 'BIGQUERY_PRIVATE_KEY'
+_ENV_BIGQUERY_CLIENT_ID = 'BIGQUERY_CLIENT_ID'
+_ENV_BIGQUERY_CLIENT_EMAIL = 'BIGQUERY_CLIENT_EMAIL'
+_ENV_BIGQUERY_TOKEN_URI = 'BIGQUERY_TOKEN_URI'
+_ENV_BIGQUERY_AUTH_URI = 'BIGQUERY_AUTH_URI'
+
+
+class BigQueryConnector(Connector):
+    r"""Establishes a connection to a
+    `BigQuery <https://cloud.google.com/bigquery>`_ database.
+
+    Authentication requires passing a private key ID, private key string,
+    client ID, client email, token URI, and authentication URI to the
+    connector, either via environment variables
+    (``BIGQUERY_PRIVATE_KEY_ID``, ``BIGQUERY_PRIVATE_KEY``,
+    ``BIGQUERY_CLIENT_ID``, ``BIGQUERY_CLIENT_EMAIL``, ``BIGQUERY_TOKEN_URI``,
+    ``BIGQUERY_AUTH_URI``), or via keys in the credentials dictionary
+    (:obj:`private_key_id`, :obj:`private_key`, :obj:`client_id`,
+    :obj:`client_email`, :obj:`token_uri`, :obj:`auth_uri`).
+
+    .. code-block:: python
+
+        import kumoai
+
+        # Either pass `credentials=dict(private_key_id=..., private_key=...,
+        # client_id=..., client_email=..., token_uri=..., auth_url=...)` or set
+        # the aforementioned environment variables:
+        connector = kumoai.BigQueryConnector(
+            name="<connector_name>",
+            project_id="<bigquery_project_id>",
+            dataset_id="<bigquery_dataset_id>",
+            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.
+        project_id: The project ID to connect to.
+        dataset_id: The dataset ID within the connected project.
+        credentials: The private key ID, private key, client ID, client email,
+            token URI, and auth URI that correspond to this BigQuery account.
+    """
+    def __init__(
+            self,
+            name: str,
+            project_id: str,
+            dataset_id: str,
+            credentials: Optional[Dict[str, str]] = None,
+            _bypass_creation: bool = False,  # INTERNAL ONLY.
+    ):
+        super().__init__()
+
+        self._name = name
+        self.project_id = project_id
+        self.dataset_id = dataset_id
+
+        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 {}
+        credentials_args = {
+            "private_key_id":
+            credentials.get("private_key_id",
+                            os.getenv(_ENV_BIGQUERY_PRIVATE_KEY_ID)),
+            "private_key":
+            credentials.get("private_key",
+                            os.getenv(_ENV_BIGQUERY_PRIVATE_KEY)),
+            "client_id":
+            credentials.get("client_id", os.getenv(_ENV_BIGQUERY_CLIENT_ID)),
+            "client_email":
+            credentials.get("client_email",
+                            os.getenv(_ENV_BIGQUERY_CLIENT_EMAIL)),
+            "token_uri":
+            credentials.get("token_uri", os.getenv(_ENV_BIGQUERY_TOKEN_URI)),
+            "auth_uri":
+            credentials.get("auth_uri", os.getenv(_ENV_BIGQUERY_AUTH_URI)),
+        }
+
+        self._create_connector(credentials_args)  # type: ignore
+
+    @classmethod
+    def get_by_name(cls, name: str) -> Self:
+        r"""Returns an instance of a named BigQuery 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, BigQueryConnectorResourceConfig)
+        return cls(
+            name=config.name,
+            project_id=config.project_id,
+            dataset_id=config.dataset_id,
+            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.BIGQUERY
+
+    @override
+    def _source_table_request(
+        self,
+        table_names: List[str],
+    ) -> BigQuerySourceTableRequest:
+        return BigQuerySourceTableRequest(
+            connector_id=self.name,
+            table_names=table_names,
+        )
+
+    def _create_connector(self, credentials: Dict[str, str]) -> None:
+        r"""Creates and persists a BigQuery connector in the REST DB.
+        Currently only intended for internal use.
+
+        Args:
+            credentials: Fully-specified credentials containing the username
+                and password for the BigQuery connector.
+
+        Raises:
+            RuntimeError: if connector creation failed
+        """
+        credentials = BigQueryCredentials(
+            private_key_id=credentials["private_key_id"] or '',
+            private_key=credentials["private_key"] or '',
+            client_id=credentials["client_id"] or '',
+            client_email=credentials["client_email"] or '',
+            token_uri=credentials["token_uri"] or '',
+            auth_uri=credentials["auth_uri"] or '',
+        )
+        args = CreateConnectorArgs(
+            config=BigQueryConnectorResourceConfig(
+                name=self.name,
+                project_id=self.project_id,
+                dataset_id=self.dataset_id,
+            ),
+            credentials=credentials,
+        )
+        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'(name=\"{self.name}\",'
+                f'project_id=\"{self.project_id}\", '
+                f'dataset_id=\"{self.dataset_id}\")')

kumoai/connector/databricks_connector.py

@@ -0,0 +1,213 @@
+import os
+from typing import Dict, List, Optional
+
+from kumoapi.data_source import (
+    CreateConnectorArgs,
+    DatabricksConnectorResourceConfig,
+    DatabricksCredentials,
+    DataSourceType,
+)
+from kumoapi.source_table import DatabricksSourceTableRequest
+from typing_extensions import Self, override
+
+from kumoai import global_state
+from kumoai.connector import Connector
+
+_ENV_DATABRICKS_CLIENT_ID = 'DATABRICKS_CLIENT_ID'
+_ENV_DATABRICKS_CLIENT_SECRET = 'DATABRICKS_CLIENT_SECRET'
+_ENV_DATABRICKS_TOKEN = 'DATABRICKS_TOKEN'
+
+
+class DatabricksConnector(Connector):
+    r"""Establishes a connection to a
+    `Databricks <https://www.databricks.com/>`_ database.
+
+    Authentication requires passing either a client ID and client secret, or a
+    personal access token, to the connector, either via environment variables
+    (``DATABRICKS_CLIENT_ID`` and ``DATABRICKS_CLIENT_SECRET``, or
+    ``DATABRICKS_TOKEN``), or via keys in the credentials dictionary
+    (``client_id`` and ``client_secret``, or ``token``).
+
+    .. code-block:: python
+
+        import kumoai
+
+        # Either pass `credentials=dict(client_id=..., client_secret=...,
+        # token=...) or set the 'DATABRICKS_CLIENT_ID' and
+        # 'DATABRICKS_CLIENT_SECRET' (or 'DATABRICKS_TOKEN') environment
+        # variables:
+        connector = kumoai.connector.DatabricksConnector(
+            name="<connector_name>",
+            host="<databricks_host_name>",
+            cluster_id="<databricks_cluster_id>",
+            warehouse_id="<databricks_warehouse_id>",
+            catalog="<databricks_catalog_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.
+        host: The host name.
+        cluster_id: The cluster ID of this warehouse.
+        warehouse_id: The warehouse ID of this warehous.
+        catalog: The name of the Databricks catalog.
+        credentials: The client ID, client secret, and personal access token
+            that correspond to this Databricks account.
+    """
+    def __init__(
+            self,
+            name: str,
+            host: str,
+            cluster_id: str,
+            warehouse_id: str,
+            catalog: str,
+            credentials: Optional[Dict[str, str]] = None,
+            _bypass_creation: bool = False,  # INTERNAL ONLY.
+    ):
+        super().__init__()
+
+        self._name = name
+        self.host = host
+        self.cluster_id = cluster_id
+        self.warehouse_id = warehouse_id
+        self.catalog = catalog
+
+        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 {}
+        credentials_args = {
+            "client_id":
+            credentials.get("client_id", os.getenv(_ENV_DATABRICKS_CLIENT_ID)),
+            "client_secret":
+            credentials.get("client_secret",
+                            os.getenv(_ENV_DATABRICKS_CLIENT_SECRET)),
+            "token":
+            credentials.get("token", os.getenv(_ENV_DATABRICKS_TOKEN))
+        }
+
+        has_pat = credentials_args["token"] is not None
+        has_client_id_secret = (credentials_args["client_id"] is not None and
+                                credentials_args["client_secret"] is not None)
+
+        if has_pat and has_client_id_secret:
+            raise ValueError(
+                "Please pass only one of a (Databricks client ID and "
+                "Databricks client secret) or a (Databricks PAT).")
+        elif not (has_pat or has_client_id_secret):
+            raise ValueError(
+                f"Please pass valid credentials to create a Databricks "
+                f"connector. You can do so either via the 'credentials' "
+                f"argument or the {_ENV_DATABRICKS_CLIENT_ID} and "
+                f"{_ENV_DATABRICKS_CLIENT_SECRET}, or "
+                f"{_ENV_DATABRICKS_TOKEN} environment variables.")
+
+        self._create_connector(credentials_args)  # type: ignore
+
+    @classmethod
+    def get_by_name(cls, name: str) -> Self:
+        r"""Returns an instance of a named Databricks Connector, including
+        those created in the Kumo UI.
+
+        Args:
+            name: The name of the existing connector.
+
+        Example:
+            >>> import kumoai
+            >>> connector = kumoai.DatabricksConnector.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, DatabricksConnectorResourceConfig)
+        return cls(
+            name=config.name,
+            host=config.host,
+            cluster_id=config.cluster_id,
+            warehouse_id=config.warehouse_id,
+            catalog=config.catalog,
+            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.DATABRICKS
+
+    @override
+    def _source_table_request(
+        self,
+        table_names: List[str],
+    ) -> DatabricksSourceTableRequest:
+        return DatabricksSourceTableRequest(
+            connector_id=self.name,
+            table_names=table_names,
+        )
+
+    def _create_connector(self, credentials: Dict[str, str]) -> None:
+        r"""Creates and persists a Databricks connector in the REST DB.
+        Currently only intended for internal use.
+
+        Args:
+            credentials: Fully-specified credentials containing the username
+                and password for the Databricks connector.
+
+        Raises:
+            RuntimeError: if connector creation failed
+        """
+        credentials = DatabricksCredentials(
+            client_id=credentials["client_id"] or '',
+            client_secret=credentials["client_secret"] or '',
+            pat=credentials["token"] or '',
+        )
+        args = CreateConnectorArgs(
+            config=DatabricksConnectorResourceConfig(
+                name=self.name,
+                host=self.host,
+                cluster_id=self.cluster_id,
+                warehouse_id=self.warehouse_id,
+                catalog=self.catalog,
+            ),
+            credentials=credentials,
+        )
+        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'(name=\"{self.name}\", host=\"{self.host}\", '
+                f'cluster_id=\"{self.cluster_id}\", '
+                f'warehouse_id=\"{self.warehouse_id}\", '
+                f'catalog=\"{self.catalog}\")')

kumoai/connector/file_upload_connector.py

@@ -0,0 +1,189 @@
+from typing import List
+
+from kumoapi.source_table import (
+    DataSourceType,
+    FileType,
+    S3SourceTableRequest,
+    SourceTableConfigRequest,
+    SourceTableConfigResponse,
+)
+from typing_extensions import override
+
+from kumoai import global_state
+from kumoai.connector.base import Connector
+from kumoai.connector.utils import delete_uploaded_table, upload_table
+
+
+class FileUploadConnector(Connector):
+    r"""Defines a connector to files directly uploaded to Kumo, either as
+    'parquet' or 'csv' (non-partitioned) data.
+
+    To get started with file upload, please first upload a table with
+    the :meth:`upload` method in the :class:`FileUploadConnector` class.
+    You can then access
+    this table behind the file upload connector as follows:
+
+    .. code-block:: python
+
+        import kumoai
+
+        # Create the file upload connector:
+        connector = kumoai.FileUploadConnector(file_type="parquet")
+
+        # Upload the table; assume it is stored at `/data/users.parquet`
+        connector.upload(name="users", path="/data/users.parquet")
+
+        # Check that the file upload connector has a `users` table:
+        assert connector.has_table("users")
+
+    Args:
+        file_type: The file type of uploaded data. Can be either ``"csv"``
+            or ``"parquet"``.
+    """
+    def __init__(self, file_type: str) -> None:
+        r"""Creates the connector to uploaded files of type
+        :obj:`file_type`.
+        """
+        assert file_type.lower() in {'parquet', 'csv'}
+        self._file_type = file_type.lower()
+
+    @property
+    def name(self) -> str:
+        return f'{self._file_type}_upload_connector'
+
+    @property
+    def source_type(self) -> DataSourceType:
+        return DataSourceType.S3
+
+    @property
+    def file_type(self) -> FileType:
+        return (FileType.PARQUET
+                if self._file_type == 'parquet' else FileType.CSV)
+
+    def _get_table_config(self, table_name: str) -> SourceTableConfigResponse:
+        req = SourceTableConfigRequest(connector_id=self.name,
+                                       table_name=table_name,
+                                       source_type=self.source_type,
+                                       file_type=None)
+        return global_state.client.source_table_api.get_table_config(req)
+
+    @override
+    def _source_table_request(self,
+                              table_names: List[str]) -> S3SourceTableRequest:
+        return S3SourceTableRequest(s3_root_dir="", connector_id=self.name,
+                                    table_names=table_names, file_type=None)
+
+    def upload(
+        self,
+        name: str,
+        path: str,
+        auto_partition: bool = True,
+        partition_size_mb: int = 250,
+    ) -> None:
+        r"""Upload a table to Kumo from a local or remote path.
+
+        Supports ``s3://``, ``gs://``, ``abfs://``, ``abfss://``, and ``az://``
+
+        Tables uploaded this way can be accessed from this
+        ``FileUploadConnector`` using the provided name, e.g.,
+        ``connector_obj["my_table"]``.
+
+        Local files
+        -----------
+        - Accepts one ``.parquet`` or ``.csv`` file (must match this
+          connector’s ``file_type``).
+        - If the file is > 1 GiB and ``auto_partition=True``, it is split
+          into ~``partition_size_mb`` MiB parts and uploaded under a common
+          prefix so the connector can read them as one table.
+
+        Remote paths
+        ------------
+        - **Single file** (``.parquet``/``.csv``): validated and uploaded via
+          multipart PUT. Files > 1 GiB are rejected — re-shard to ~200 MiB
+          and upload the directory instead.
+        - **Directory**: must contain only one format (all Parquet or all CSV)
+          matching this connector’s ``file_type``. Files are validated
+          (consistent schema; CSV headers sanitized) and uploaded in parallel
+          with memory-safe budgeting.
+
+        .. warning::
+           For local uploads, input must be a single CSV or Parquet file
+           (matching the connector type). For remote uploads, mixed
+           CSV/Parquet directories are not supported. Remote single files
+           larger than 1 GiB are not supported.
+
+        Examples:
+        ---------
+        .. code-block:: python
+
+            import kumoai
+            conn = kumoai.FileUploadConnector(file_type="parquet")
+
+            # Local: small file
+            conn.upload(name="users", path="/data/users.parquet")
+
+            # Local: large file (auto-partitions)
+            conn.upload(
+                name="txns",
+                path="/data/large_txns.parquet",
+            )
+
+            # Local: disable auto-partitioning (raises if > 1 GiB)
+            conn.upload(
+                name="users",
+                path="/data/users.parquet",
+                auto_partition=False,
+            )
+
+            # CSV connector
+            csv_conn = kumoai.FileUploadConnector(file_type="csv")
+            csv_conn.upload(name="sales", path="/data/sales.csv")
+
+            # Remote: single file (<= 1 GiB)
+            conn.upload(name="logs", path="s3://bkt/path/logs.parquet")
+
+            # Remote: directory of shards (uniform format)
+            csv_conn.upload(name="events", path="gs://mybkt/events_csv/")
+
+        Args:
+            name:
+                Table name to create in Kumo; access later via this connector.
+            path:
+                Local path or remote URL to a ``.parquet``/``.csv`` file or a
+                directory (uniform format). The format must match this
+                connector’s ``file_type``.
+            auto_partition:
+                Local-only. If ``True`` and the local file is > 1 GiB, split
+                into ~``partition_size_mb`` MiB parts.
+            partition_size_mb:
+                Local-only. Target partition size (100–1000 MiB) when
+                ``auto_partition`` is ``True``.
+        """
+        upload_table(name=name, path=path, auto_partition=auto_partition,
+                     partition_size_mb=partition_size_mb,
+                     file_type=self._file_type)
+
+    def delete(
+        self,
+        name: str,
+    ) -> None:
+        r"""Synchronously deletes a previously uploaded table from the Kumo
+        data plane.
+
+        .. code-block:: python
+
+            # Assume we have uploaded a `.parquet` table named `users`, and a
+            # `FileUploadConnector` has been created called `connector`, and
+            # we want to delete this table from Kumo:
+            connector.delete(name="users")
+
+        Args:
+            name: The name of the table to be deleted. This table must have
+                previously been uploaded with a call to
+                :meth:`~kumoai.connector.FileUploadConnector.upload`.
+        """
+        if not self.has_table(name):
+            raise ValueError(f"The table '{name}' does not exist in {self}. "
+                             f"Please check the existence of the source data.")
+
+        delete_uploaded_table(name, self._file_type)