tentaclio-databricks 2.1.0__py3-none-any.whl

tentaclio_databricks/__init__.py

@@ -0,0 +1,16 @@
+"""This package implements the tentaclio databricks database client"""
+
+from importlib.metadata import PackageNotFoundError
+from importlib.metadata import version as _version
+
+from tentaclio import *  # noqa
+
+from .clients.databricks_client import DatabricksClient
+
+try:
+    __version__ = _version("tentaclio-databricks")
+except PackageNotFoundError:
+    __version__ = "0.0.0"
+
+# Add DB registry
+DB_REGISTRY.register("databricks+thrift", DatabricksClient)  # type: ignore

tentaclio_databricks/__main__.py

@@ -0,0 +1,6 @@
+# Helper cli to encode urls.
+
+from tentaclio.__main__ import main
+
+if __name__ == "__main__":
+    main(prog_name="python -m tentaclio-databricks")

tentaclio_databricks/clients/__init__.py

@@ -0,0 +1,9 @@
+"""The client module defines stream and query based clients.
+
+Stream clients allow to read and write from stream based sources with ease,
+and in an interchangable manner.
+
+Query based clients unify how to access databases leveraging from sqlalchemy.
+"""
+
+from .databricks_client import *  # noqa

tentaclio_databricks/clients/databricks_client.py

@@ -0,0 +1,167 @@
+"""Databricks query client."""
+
+import importlib
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
+
+import pandas as pd
+from databricks import sql
+from tentaclio import URL
+
+if TYPE_CHECKING:
+    from polars import DataFrame as PolarsDataFrame  # pyright: ignore[reportMissingImports]
+else:
+    PolarsDataFrame = Any
+
+
+class DatabricksClientException(Exception):
+    """Databricks client specific exception."""
+
+
+class DatabricksClient:
+    """Databricks client, backed by an Apache Thrift connection."""
+
+    def __init__(
+        self,
+        url: URL,
+        arraysize: int = 1000000,
+        use_arrow: bool = False,
+        query_annotations: Optional[Dict[str, str]] = None,
+        **kwargs,
+    ):
+        # This is a very common issue reported by the users
+        if url.query is None or "HTTPPath" not in url.query:
+            raise DatabricksClientException(
+                "Missing the HTTPPath element in the http query. \n\n"
+                "The url should look like: "
+                "databricks+thrift://workspaceurl.databricks.com/?HTTPPath=value\n"
+                "Check the connection details of your databricks warehouse"
+            )
+        if url.username is None or url.username == "":
+            raise DatabricksClientException(
+                "Missing the token in the url:\n\n"
+                "The url should look like: "
+                "databricks+thrift://token@workspaceurl.databricks.com/?HTTPPath=value"
+            )
+
+        self.server_hostname = url.hostname
+        self.http_path = url.query["HTTPPath"]
+        self.access_token = url.username
+        self.arraysize = arraysize
+        self.use_arrow = use_arrow
+        self.query_annotations = query_annotations or {}
+
+    def __enter__(self):
+        conn_kwargs = {
+            "server_hostname": self.server_hostname,
+            "http_path": self.http_path,
+            "access_token": self.access_token,
+        }
+
+        if self.use_arrow:
+            conn_kwargs.update(
+                {
+                    "use_arrow_native_complex_types": True,
+                    "use_arrow_native_decimals": True,
+                    "use_arrow_native_timestamps": True,
+                }
+            )
+
+        self.conn = sql.connect(**conn_kwargs)
+
+        # Only set arraysize if not using Arrow
+        if self.use_arrow:
+            self.cursor = self.conn.cursor()
+        else:
+            self.cursor = self.conn.cursor(arraysize=self.arraysize)
+
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.cursor.close()
+        self.conn.close()
+
+    def _build_query_comment(self) -> str:
+        """Build a SQL comment from query annotations.
+
+        Returns:
+            Formatted SQL comment string, or empty string if no annotations.
+        """
+        if not self.query_annotations:
+            return ""
+
+        parts = []
+        for key, value in self.query_annotations.items():
+            # Escape single quotes in values
+            escaped_value = value.replace("'", "''")
+            parts.append(f"{key}='{escaped_value}'")
+
+        if parts:
+            return f"/* {', '.join(parts)} */\n"
+        return ""
+
+    def _prepend_comment(self, sql_query: str) -> str:
+        """Prepend query comment to SQL query if configured.
+
+        Args:
+            sql_query: The SQL query string.
+
+        Returns:
+            SQL query with prepended comment if annotations configured, else original query.
+        """
+        comment = self._build_query_comment()
+        if comment:
+            return f"{comment}{sql_query}"
+        return sql_query
+
+    def query(self, sql_query: str, **kwargs) -> List[Any]:
+        """Execute a SQL query, and return results."""
+        sql_query = self._prepend_comment(sql_query)
+        self.cursor.execute(sql_query, **kwargs)
+        return self.cursor.fetchall()
+
+    def execute(self, sql_query: str, **kwargs) -> None:
+        """Execute a raw SQL query command."""
+        sql_query = self._prepend_comment(sql_query)
+        self.cursor.execute(sql_query, **kwargs)
+
+    def get_df(self, sql_query: str, **kwargs) -> pd.DataFrame:
+        """Run a raw SQL query and return a data frame."""
+        sql_query = self._prepend_comment(sql_query)
+        self.cursor.execute(sql_query, **kwargs)
+
+        if self.use_arrow:
+            # Default to Arrow for fast data transfer
+            try:
+                arrow_table = self.cursor.fetchall_arrow()
+                return arrow_table.to_pandas()
+            except AttributeError:
+                data = self.cursor.fetchall()
+        else:
+            # Original row-by-row method
+            data = self.cursor.fetchall()
+
+        # Build DataFrame from rows (fallback path)
+        columns = (
+            [col_desc[0] for col_desc in self.cursor.description]
+            if self.cursor.description
+            else []
+        )
+        return pd.DataFrame(data, columns=columns)
+
+    def get_pl(self, sql_query: str, **kwargs) -> PolarsDataFrame:
+        """Run a raw SQL query and return a polars DataFrame."""
+
+        try:
+            pl = importlib.import_module("polars")
+        except ModuleNotFoundError as exc:
+            raise ModuleNotFoundError(
+                "Polars is not installed. Install it with "
+                "`pip install tentaclio-databricks[polars]` or `pip install polars` to use "
+                "the `.get_pl()` method."
+            ) from exc
+
+        sql_query = self._prepend_comment(sql_query)
+        self.cursor.execute(sql_query, **kwargs)
+        arrow_table = self.cursor.fetchall_arrow()
+
+        return pl.DataFrame(arrow_table)

tentaclio_databricks-2.1.0.dist-info/METADATA

@@ -0,0 +1,136 @@
+Metadata-Version: 2.4
+Name: tentaclio-databricks
+Version: 2.1.0
+Summary: A python project containing all the dependencies for schema databricks for tentaclio.
+Author-email: Octopus Energy <nerds@octopus.energy>
+License-Expression: MIT
+Classifier: Intended Audience :: Developers
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE.txt
+Requires-Dist: tentaclio
+Requires-Dist: databricks-sql-connector
+Requires-Dist: pandas<3
+Requires-Dist: pyarrow
+Provides-Extra: polars
+Requires-Dist: polars>=0.16; extra == "polars"
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: pytest-cov; extra == "dev"
+Requires-Dist: pytest-mock; extra == "dev"
+Requires-Dist: black; extra == "dev"
+Requires-Dist: flake8; extra == "dev"
+Requires-Dist: Flake8-pyproject; extra == "dev"
+Requires-Dist: isort>=5.11.2; extra == "dev"
+Requires-Dist: mypy; extra == "dev"
+Dynamic: license-file
+
+# tentaclio-databricks
+
+A package containing all the dependencies for the `databricks+thrift` tentaclio schema .
+
+## Quick Start
+
+This project comes with a `Makefile` which is ready to do basic common tasks
+
+```
+$ make help
+install                       Initalise the virtual env installing deps
+clean                         Remove all the unwanted clutter
+lock                          Lock dependencies
+update                        Update dependencies (whole tree)
+sync                          Install dependencies as per the lock file
+lint                          Lint files with flake and mypy
+format                        Run black and isort
+test                          Run unit tests
+circleci                      Validate circleci configuration (needs circleci cli)
+```
+
+## Configuring access to Databricks
+
+Your connection url should be in the following format:
+
+```
+databricks+thrift://<token>@<host>?HTTPPath=<http_path>
+```
+
+Example values:
+
+- token: dapi1213456789abc
+- host: myhost.databricks.com
+- http_path: /sql/1.0/endpoints/123456789
+
+## Query Comments
+
+Queries can be annotated with comments for observability using the `query_annotations` parameter.
+
+### Example: Basic usage
+
+```python
+import os
+from tentaclio import URL
+from tentaclio_databricks.clients.databricks_client import DatabricksClient
+
+url = URL("databricks+thrift://token@host.databricks.com?HTTPPath=/sql/1.0/endpoints/123")
+client = DatabricksClient(
+    url,
+    query_annotations={
+        "app_name": "JupyterHub",
+        "user": os.environ.get("USER_NAME", "unknown"),
+        "pipeline_id": "456"
+    }
+)
+```
+
+### Result
+
+All queries executed by the client will have prepended comments:
+
+```sql
+/* app_name='JupyterHub', user='john', pipeline_id='456' */
+SELECT * FROM table
+```
+
+## Arrow Usage
+
+`DatabricksClient` enables Arrow's usage for faster data transfer and more faithful type handling.
+
+### Example: Default row-based fetch
+
+```python
+from tentaclio import URL
+from tentaclio_databricks.clients.databricks_client import DatabricksClient
+
+url = URL("databricks+thrift://token@host.databricks.com?HTTPPath=/sql/1.0/endpoints/123")
+with DatabricksClient(url) as client:
+    df = client.get_df("SELECT * FROM sample_table")
+```
+
+When Arrow is disabled, the client falls back to `fetchall()` and uses `arraysize` for the cursor.
+
+### Example: Arrow enabled
+
+```python
+from tentaclio import URL
+from tentaclio_databricks.clients.databricks_client import DatabricksClient
+
+url = URL("databricks+thrift://token@host.databricks.com?HTTPPath=/sql/1.0/endpoints/123")
+with DatabricksClient(url, use_arrow=True) as client:
+    df = client.get_df("SELECT * FROM sample_table")
+```
+
+When Arrow is enabled, the client configures the connection with:
+
+- `use_arrow_native_complex_types`
+- `use_arrow_native_decimals`
+- `use_arrow_native_timestamps`
+
+

tentaclio_databricks-2.1.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+tentaclio_databricks/__init__.py,sha256=Te6JwNyyV3Pjfii-n63bydjLcXLFaypwkKbwwYLxSKA,472
+tentaclio_databricks/__main__.py,sha256=ed3jPmp9bmfc5ekIcjPxVEAPqQ3mXYacga2oUfEJgPI,147
+tentaclio_databricks/clients/__init__.py,sha256=FRSm7eaquCUtXPiaba1oxaT2Lv6Zf4gRfqhcKIwPZKU,296
+tentaclio_databricks/clients/databricks_client.py,sha256=9FWtMCVTa-xCg8rzcur5rE9tyWLCD4T1RtUytuTHncI,5636
+tentaclio_databricks-2.1.0.dist-info/licenses/LICENSE.txt,sha256=Aqv2h5R53YlGaQXl7XgUDMZPUxJK1HftERN2ExXJ_u4,1080
+tentaclio_databricks-2.1.0.dist-info/METADATA,sha256=W0CMq9tGkwHUxTI0TdchW_lBpFLLq-KALh-v7DVE77w,4158
+tentaclio_databricks-2.1.0.dist-info/WHEEL,sha256=K260EYznzXsJYBQGqmI8VTxEdiZYNvDZwW9cBh9-_MA,91
+tentaclio_databricks-2.1.0.dist-info/top_level.txt,sha256=ZoGTS3HcLGAiSp9c638A_y5eZ9J4uzl91KQylxs9ptY,21
+tentaclio_databricks-2.1.0.dist-info/RECORD,,

tentaclio_databricks-2.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (83.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

tentaclio_databricks-2.1.0.dist-info/licenses/LICENSE.txt

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2021 Octopus Energy Limited.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

tentaclio_databricks-2.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+tentaclio_databricks