mirror of
https://github.com/semantica-agi/semantica.git
synced 2026-08-29 04:26:20 +00:00
Add Databricks connector (Unity Catalog + Delta Lake ingestion)
Adds DatabricksIngestor to semantica/ingest/, mirroring SnowflakeIngestor's structure and public API shape: table/query ingestion via databricks-sql-connector, Unity Catalog metadata and lineage via databricks-sdk, and export-as-documents for KG construction. Closes #747
This commit is contained in:
KaifAhmad1
@@ -11,6 +11,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
||||
|
||||
### Added
|
||||
|
||||
- **Databricks Connector (Unity Catalog + Delta Lake ingestion)** (#747) by @KaifAhmad1
|
||||
- Added `DatabricksIngestor` (`semantica/ingest/databricks_ingestor.py`), mirroring `SnowflakeIngestor`'s structure and public API shape: a `DatabricksConnector` connection handler, a `DatabricksData` dataclass, and an optional-import guard for `databricks-sdk`/`databricks-sql-connector`
|
||||
- Supports personal access token and OAuth M2M (service principal `client_id`/`client_secret`) authentication, configurable via constructor args or `DATABRICKS_*` environment variables
|
||||
- `ingest_table()`/`ingest_query()` run against a SQL warehouse or cluster via `databricks-sql-connector`, with `where`/`order_by`/`limit`/`offset` support and the same identifier-escaping and unsafe-`ORDER BY` rejection as `SnowflakeIngestor`
|
||||
- `get_table_schema()`, `list_catalogs()`, `list_schemas()`, and `list_tables()` introspect Unity Catalog via `databricks-sdk`'s `WorkspaceClient`; `get_table_lineage()` calls Unity Catalog's table-lineage REST API to surface upstream/downstream dependencies for `Table --DEPENDS_ON--> Table` graph edges
|
||||
- `export_as_documents()` converts ingested rows into Semantica document dicts for KG construction, matching `SnowflakeIngestor.export_as_documents()`'s shape
|
||||
- Registered as a lazy export in `semantica.ingest` (`DatabricksIngestor`, `DatabricksData`, `DatabricksConnector`) and as the `db-databricks` optional extra (`pip install "semantica[db-databricks]"`) in `pyproject.toml`, included in `db-all`
|
||||
- New `docs/integrations/databricks.md` page modeled on `docs/integrations/snowflake.md`, plus a `DatabricksIngestor` section and table row in `docs/reference/ingest.md` and cross-links between the two integration pages
|
||||
- 27 unit tests in `tests/test_databricks_ingestor.py` covering both auth methods, table/query ingestion, pagination, unsafe `ORDER BY` rejection, schema/catalog/table listing, lineage, document export, the context manager, and the missing-dependency error path, closing #747
|
||||
|
||||
- **SQLite Vector Store Backend (`sqlite-vec`)** (#726) by @Luffy2208 and @KaifAhmad1
|
||||
- Added `SQLiteVecStore` (`semantica/vector_store/sqlite_vec_store.py`), a disk-backed local vector store using the `sqlite-vec` extension's `vec0` virtual tables, closing #240
|
||||
- Supports Cosine and L2 distance metrics, dynamic JSON metadata filtering, read-only mode, and an in-memory (`:memory:`) mode
|
||||
|
||||
+2
-1
@@ -103,7 +103,8 @@
|
||||
"pages": [
|
||||
"integrations/agno",
|
||||
"integrations/docling",
|
||||
"integrations/snowflake"
|
||||
"integrations/snowflake",
|
||||
"integrations/databricks"
|
||||
]
|
||||
},
|
||||
{
|
||||
|
||||
@@ -0,0 +1,186 @@
|
||||
---
|
||||
title: "Databricks Integration"
|
||||
description: "Ingest Unity Catalog metadata and Delta Lake tables from Databricks into Semantica's KG pipeline."
|
||||
icon: "cloud"
|
||||
---
|
||||
|
||||
> Extract Delta Lake tables and Unity Catalog metadata (schemas, lineage) from Databricks into Semantica with personal access token or OAuth M2M authentication.
|
||||
|
||||
|
||||
## Installation
|
||||
|
||||
```bash
|
||||
# Install with Databricks support
|
||||
pip install "semantica[db-databricks]"
|
||||
|
||||
# Or install the connectors separately
|
||||
pip install databricks-sdk databricks-sql-connector
|
||||
```
|
||||
|
||||
|
||||
## Basic Usage
|
||||
|
||||
```python
|
||||
from semantica.ingest import DatabricksIngestor
|
||||
import os
|
||||
|
||||
ingestor = DatabricksIngestor(
|
||||
host=os.getenv("DATABRICKS_HOST"), # e.g. https://adb-xxx.azuredatabricks.net
|
||||
token=os.getenv("DATABRICKS_TOKEN"),
|
||||
http_path=os.getenv("DATABRICKS_HTTP_PATH"), # SQL warehouse or cluster HTTP path
|
||||
catalog=os.getenv("DATABRICKS_CATALOG", "main"),
|
||||
schema=os.getenv("DATABRICKS_SCHEMA", "default"),
|
||||
)
|
||||
|
||||
data = ingestor.ingest_table("customers")
|
||||
print(f"Retrieved {data.row_count} rows: columns: {data.columns}")
|
||||
```
|
||||
|
||||
<Tip>
|
||||
Use environment variables (or a `.env` file with `python-dotenv`) to keep credentials out of source code. `DatabricksIngestor()` with no arguments reads from `DATABRICKS_*` environment variables automatically.
|
||||
</Tip>
|
||||
|
||||
|
||||
## Authentication Methods
|
||||
|
||||
<Tabs>
|
||||
<Tab title="Personal Access Token">
|
||||
```python
|
||||
ingestor = DatabricksIngestor(
|
||||
host="https://adb-xxx.azuredatabricks.net",
|
||||
token="dapi-xxxxxxxx",
|
||||
http_path="/sql/1.0/warehouses/xxxxxxxx",
|
||||
)
|
||||
```
|
||||
</Tab>
|
||||
<Tab title="OAuth M2M (Recommended)">
|
||||
```python
|
||||
ingestor = DatabricksIngestor(
|
||||
host="https://adb-xxx.azuredatabricks.net",
|
||||
client_id="your_service_principal_client_id",
|
||||
client_secret="your_service_principal_client_secret",
|
||||
http_path="/sql/1.0/warehouses/xxxxxxxx",
|
||||
)
|
||||
```
|
||||
Preferred for production: no long-lived personal token stored in config.
|
||||
</Tab>
|
||||
</Tabs>
|
||||
|
||||
<Note>
|
||||
`http_path` identifies the SQL warehouse or all-purpose cluster used for query execution. Find it in the Databricks UI under **SQL Warehouses → Connection details**. Unity Catalog metadata calls (`list_catalogs`, `get_table_schema`, `get_table_lineage`, …) only need `host` and credentials — `http_path` is not required for those.
|
||||
</Note>
|
||||
|
||||
|
||||
## Querying
|
||||
|
||||
### Ingest a table with filters
|
||||
|
||||
```python
|
||||
data = ingestor.ingest_table(
|
||||
"customers",
|
||||
catalog="main",
|
||||
schema="default",
|
||||
where="country = 'USA' AND created_date > '2024-01-01'",
|
||||
order_by="created_date DESC",
|
||||
limit=10000,
|
||||
)
|
||||
```
|
||||
|
||||
### Custom SQL
|
||||
|
||||
```python
|
||||
data = ingestor.ingest_query("""
|
||||
SELECT customer_id, SUM(amount) AS total_amount
|
||||
FROM main.default.sales
|
||||
WHERE date >= '2024-01-01'
|
||||
GROUP BY customer_id
|
||||
""")
|
||||
```
|
||||
|
||||
|
||||
## Unity Catalog Metadata
|
||||
|
||||
### Schema introspection
|
||||
|
||||
```python
|
||||
schema = ingestor.get_table_schema("customers")
|
||||
for column in schema["columns"]:
|
||||
print(f"{column['name']}: {column['type']}")
|
||||
```
|
||||
|
||||
### Catalogs, schemas, and tables
|
||||
|
||||
```python
|
||||
catalogs = ingestor.list_catalogs()
|
||||
schemas = ingestor.list_schemas(catalog="main")
|
||||
tables = ingestor.list_tables(catalog="main", schema="default")
|
||||
```
|
||||
|
||||
### Table lineage
|
||||
|
||||
```python
|
||||
lineage = ingestor.get_table_lineage("customers", catalog="main", schema="default")
|
||||
print(lineage["upstream"]) # tables that feed into `customers`
|
||||
print(lineage["downstream"]) # tables derived from `customers`
|
||||
```
|
||||
|
||||
Use `get_table_lineage` to build `Table --DEPENDS_ON--> Table` edges in the knowledge graph directly from Unity Catalog's lineage tracking, without re-deriving lineage from query logs.
|
||||
|
||||
|
||||
## Export as Semantica Documents
|
||||
|
||||
```python
|
||||
documents = ingestor.export_as_documents(
|
||||
data,
|
||||
id_field="customer_id",
|
||||
text_fields=["name", "email", "notes"],
|
||||
)
|
||||
print(f"Created {len(documents)} documents for processing")
|
||||
```
|
||||
|
||||
|
||||
## Batch Processing Large Tables
|
||||
|
||||
```python
|
||||
PAGE_SIZE = 5000
|
||||
for page in range(total_pages):
|
||||
data = ingestor.ingest_table(
|
||||
"large_table",
|
||||
limit=PAGE_SIZE,
|
||||
offset=page * PAGE_SIZE,
|
||||
)
|
||||
process_batch(data)
|
||||
```
|
||||
|
||||
Or use the built-in `batch_size` parameter:
|
||||
|
||||
```python
|
||||
data = ingestor.ingest_query(
|
||||
"SELECT * FROM main.default.large_table",
|
||||
batch_size=5000,
|
||||
)
|
||||
```
|
||||
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
```python
|
||||
from semantica.ingest import DatabricksConnector
|
||||
|
||||
connector = DatabricksConnector(
|
||||
host="https://adb-xxx.azuredatabricks.net",
|
||||
token="dapi-xxxxxxxx",
|
||||
http_path="/sql/1.0/warehouses/xxxxxxxx",
|
||||
)
|
||||
if not connector.test_connection():
|
||||
print("Connection failed: check host, http_path, and credentials")
|
||||
```
|
||||
|
||||
|
||||
## See Also
|
||||
|
||||
- [Ingest Module](../reference/ingest) — Full DatabricksIngestor and all other ingestors.
|
||||
- [Snowflake Integration](snowflake) — Companion connector for a Snowflake + Databricks hybrid estate.
|
||||
- [Pipeline](../reference/pipeline) — Use Databricks ingestion as a pipeline step.
|
||||
- [Installation](../installation) — All optional dependency extras.
|
||||
- [Knowledge Graph](../reference/kg) — Build a KG from ingested Databricks data.
|
||||
@@ -172,6 +172,7 @@ if not connector.test_connection():
|
||||
## See Also
|
||||
|
||||
- [Ingest Module](../reference/ingest) — Full SnowflakeIngestor and all other ingestors.
|
||||
- [Databricks Integration](databricks) — Companion connector for a Snowflake + Databricks hybrid estate.
|
||||
- [Pipeline](../reference/pipeline) — Use Snowflake ingestion as a pipeline step.
|
||||
- [Installation](../installation) — All optional dependency extras.
|
||||
- [Knowledge Graph](../reference/kg) — Build a KG from ingested Snowflake data.
|
||||
|
||||
@@ -27,6 +27,7 @@ icon: "database"
|
||||
| `RepoIngestor` | Git repositories: source files, commit history, and metadata |
|
||||
| `DBIngestor` | SQL databases via SQLAlchemy: tables, views, and custom queries |
|
||||
| `SnowflakeIngestor` | Snowflake data warehouse queries and table exports |
|
||||
| `DatabricksIngestor` | Databricks Unity Catalog metadata, Delta table queries, and lineage |
|
||||
| `ParquetIngestor` | Apache Parquet files and partitioned datasets with column selection |
|
||||
| `XMLIngestor` | XXE-safe XML parsing with optional XSD schema validation |
|
||||
| `EmailIngestor` | IMAP/POP3 email ingestion with attachment extraction |
|
||||
@@ -440,6 +441,24 @@ result = ingest("ontology.ttl") # -> {"ontology": OntologyData}
|
||||
result = ingestor.ingest_query("SELECT * FROM documents")
|
||||
result = ingestor.ingest_table("documents")
|
||||
```
|
||||
|
||||
### DatabricksIngestor
|
||||
|
||||
```python
|
||||
from semantica.ingest import DatabricksIngestor
|
||||
import os
|
||||
|
||||
ingestor = DatabricksIngestor(
|
||||
host=os.getenv("DATABRICKS_HOST"),
|
||||
token=os.getenv("DATABRICKS_TOKEN"),
|
||||
http_path=os.getenv("DATABRICKS_HTTP_PATH"),
|
||||
catalog="main",
|
||||
schema="default",
|
||||
)
|
||||
result = ingestor.ingest_query("SELECT * FROM documents")
|
||||
result = ingestor.ingest_table("documents")
|
||||
lineage = ingestor.get_table_lineage("documents")
|
||||
```
|
||||
</Tab>
|
||||
<Tab title="Stream">
|
||||
### StreamIngestor
|
||||
@@ -628,4 +647,5 @@ result = ingest_file("source_path", method="my_format")
|
||||
- [Parse](parse) — Parse raw sources into structured text and tables.
|
||||
- [Pipeline](pipeline) — Orchestrate ingest as the first pipeline step.
|
||||
- [Snowflake Integration](../integrations/snowflake) — Snowflake-specific setup and authentication guide.
|
||||
- [Databricks Integration](../integrations/databricks) — Databricks Unity Catalog setup, authentication, and lineage guide.
|
||||
- [Provenance](provenance) — Track lineage from ingest through to inference.
|
||||
|
||||
+2
-1
@@ -121,12 +121,13 @@ shacl = ["pyshacl>=0.25.0"]
|
||||
|
||||
# ---- Database Connectors ----
|
||||
db-snowflake = ["snowflake-connector-python>=4.6.0", "cryptography>=49.0.0"]
|
||||
db-databricks = ["databricks-sdk>=0.60.0", "databricks-sql-connector>=4.0.0"]
|
||||
db-arrow = ["pyarrow>=24.0.0"]
|
||||
ingest-parquet = ["pyarrow>=24.0.0"]
|
||||
ingest-arrow = ["pyarrow>=24.0.0"]
|
||||
|
||||
db-all = [
|
||||
"semantica[db-snowflake,db-arrow]"
|
||||
"semantica[db-snowflake,db-databricks,db-arrow]"
|
||||
]
|
||||
|
||||
# ---- Embedding / Models ----
|
||||
|
||||
@@ -218,6 +218,10 @@ _LAZY_EXPORTS: Dict[str, Tuple[str, str]] = {
|
||||
"SnowflakeIngestor": (".snowflake_ingestor", "SnowflakeIngestor"),
|
||||
"SnowflakeData": (".snowflake_ingestor", "SnowflakeData"),
|
||||
"SnowflakeConnector": (".snowflake_ingestor", "SnowflakeConnector"),
|
||||
# Databricks ingestion
|
||||
"DatabricksIngestor": (".databricks_ingestor", "DatabricksIngestor"),
|
||||
"DatabricksData": (".databricks_ingestor", "DatabricksData"),
|
||||
"DatabricksConnector": (".databricks_ingestor", "DatabricksConnector"),
|
||||
# Parquet ingestion
|
||||
"ParquetIngestor": (".parquet_ingestor", "ParquetIngestor"),
|
||||
"ParquetData": (".parquet_ingestor", "ParquetData"),
|
||||
@@ -341,6 +345,10 @@ __all__ = [
|
||||
"SnowflakeIngestor",
|
||||
"SnowflakeData",
|
||||
"SnowflakeConnector",
|
||||
# Databricks ingestion
|
||||
"DatabricksIngestor",
|
||||
"DatabricksData",
|
||||
"DatabricksConnector",
|
||||
# Parquet ingestion
|
||||
"ParquetIngestor",
|
||||
"ParquetData",
|
||||
|
||||
@@ -0,0 +1,935 @@
|
||||
"""
|
||||
Databricks Ingestion Module
|
||||
|
||||
This module provides comprehensive Databricks ingestion capabilities for the
|
||||
Semantica framework, enabling data extraction from Databricks lakehouses
|
||||
(Unity Catalog + Delta Lake) into the knowledge graph pipeline.
|
||||
|
||||
Key Features:
|
||||
- Native Databricks SQL connection using databricks-sql-connector
|
||||
- Unity Catalog metadata via databricks-sdk (catalogs, schemas, tables, columns)
|
||||
- Multiple authentication methods (personal access token, OAuth M2M)
|
||||
- Query execution with streaming and pagination
|
||||
- Table and schema introspection
|
||||
- Table-level lineage via Unity Catalog's lineage API
|
||||
- Progress tracking and error handling
|
||||
- Connection management with proper cleanup
|
||||
|
||||
Main Classes:
|
||||
- DatabricksIngestor: Main Databricks ingestion class
|
||||
- DatabricksConnector: Databricks connection handler
|
||||
- DatabricksData: Data representation for Databricks ingestion
|
||||
|
||||
Example Usage:
|
||||
>>> from semantica.ingest import DatabricksIngestor
|
||||
>>> ingestor = DatabricksIngestor(
|
||||
... host="https://adb-xxx.azuredatabricks.net",
|
||||
... token="dapi-xxxxxxxx",
|
||||
... http_path="/sql/1.0/warehouses/xxxxxxxx",
|
||||
... catalog="main",
|
||||
... schema="default",
|
||||
... )
|
||||
>>> data = ingestor.ingest_table("customers", limit=10000)
|
||||
>>> query_data = ingestor.ingest_query("SELECT * FROM sales WHERE date > '2024-01-01'")
|
||||
>>> documents = ingestor.export_as_documents(data)
|
||||
|
||||
Author: Semantica Contributors
|
||||
License: MIT
|
||||
"""
|
||||
|
||||
import os
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from ..utils.exceptions import ProcessingError, ValidationError
|
||||
from ..utils.logging import get_logger
|
||||
from ..utils.progress_tracker import get_progress_tracker
|
||||
|
||||
try:
|
||||
from databricks import sql as databricks_sql
|
||||
from databricks.sdk import WorkspaceClient
|
||||
|
||||
DATABRICKS_AVAILABLE = True
|
||||
except (ImportError, OSError):
|
||||
databricks_sql = None
|
||||
WorkspaceClient = None
|
||||
DATABRICKS_AVAILABLE = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class DatabricksData:
|
||||
"""Databricks data representation."""
|
||||
|
||||
data: List[Dict[str, Any]]
|
||||
row_count: int
|
||||
columns: List[str]
|
||||
table_name: Optional[str] = None
|
||||
query: Optional[str] = None
|
||||
catalog: Optional[str] = None
|
||||
schema: Optional[str] = None
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
ingested_at: datetime = field(default_factory=datetime.now)
|
||||
|
||||
|
||||
class DatabricksConnector:
|
||||
"""
|
||||
Databricks connection management.
|
||||
|
||||
This class manages connections to a Databricks workspace and provides
|
||||
support for multiple authentication methods.
|
||||
|
||||
Supported Authentication Methods:
|
||||
- Personal Access Token (PAT)
|
||||
- OAuth M2M (service principal client id / secret)
|
||||
|
||||
Example Usage:
|
||||
>>> connector = DatabricksConnector(
|
||||
... host="https://adb-xxx.azuredatabricks.net",
|
||||
... token="dapi-xxxxxxxx",
|
||||
... http_path="/sql/1.0/warehouses/xxxxxxxx",
|
||||
... )
|
||||
>>> conn = connector.connect()
|
||||
>>> connector.disconnect()
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
host: Optional[str] = None,
|
||||
token: Optional[str] = None,
|
||||
http_path: Optional[str] = None,
|
||||
catalog: Optional[str] = None,
|
||||
schema: Optional[str] = None,
|
||||
client_id: Optional[str] = None,
|
||||
client_secret: Optional[str] = None,
|
||||
**config,
|
||||
):
|
||||
"""
|
||||
Initialize Databricks connector.
|
||||
|
||||
Args:
|
||||
host: Databricks workspace URL (e.g. 'https://adb-xxx.azuredatabricks.net')
|
||||
token: Personal access token (for PAT authentication)
|
||||
http_path: HTTP path of a SQL warehouse or cluster
|
||||
(e.g. '/sql/1.0/warehouses/xxxxxxxx')
|
||||
catalog: Default Unity Catalog catalog to use
|
||||
schema: Default schema to use
|
||||
client_id: OAuth M2M service principal client ID
|
||||
client_secret: OAuth M2M service principal client secret
|
||||
**config: Additional Databricks connection configuration
|
||||
"""
|
||||
if not DATABRICKS_AVAILABLE:
|
||||
raise ImportError(
|
||||
"databricks-sdk and databricks-sql-connector are required for "
|
||||
"DatabricksConnector. Install them with: "
|
||||
"pip install databricks-sdk databricks-sql-connector"
|
||||
)
|
||||
|
||||
self.logger = get_logger("databricks_connector")
|
||||
|
||||
# Get configuration from environment variables if not provided
|
||||
self.host = host or os.getenv("DATABRICKS_HOST")
|
||||
self.token = token or os.getenv("DATABRICKS_TOKEN")
|
||||
self.http_path = http_path or os.getenv("DATABRICKS_HTTP_PATH")
|
||||
self.catalog = catalog or os.getenv("DATABRICKS_CATALOG", "main")
|
||||
self.schema = schema or os.getenv("DATABRICKS_SCHEMA", "default")
|
||||
self.client_id = client_id or os.getenv("DATABRICKS_CLIENT_ID")
|
||||
self.client_secret = client_secret or os.getenv("DATABRICKS_CLIENT_SECRET")
|
||||
|
||||
# Validate required parameters
|
||||
if not self.host:
|
||||
raise ValidationError(
|
||||
"Databricks host is required. "
|
||||
"Provide via 'host' parameter or DATABRICKS_HOST environment variable."
|
||||
)
|
||||
|
||||
if not (self.client_id and self.client_secret) and not self.token:
|
||||
raise ValidationError(
|
||||
"Databricks authentication is required. "
|
||||
"Provide either 'token' (or DATABRICKS_TOKEN) for personal access "
|
||||
"token authentication, or 'client_id'/'client_secret' (or "
|
||||
"DATABRICKS_CLIENT_ID/DATABRICKS_CLIENT_SECRET) for OAuth M2M "
|
||||
"authentication."
|
||||
)
|
||||
|
||||
self.config = config
|
||||
self.connection = None
|
||||
self._workspace_client: Optional["WorkspaceClient"] = None
|
||||
|
||||
self.logger.debug(
|
||||
f"Databricks connector initialized: host={self.host}, "
|
||||
f"catalog={self.catalog}, schema={self.schema}"
|
||||
)
|
||||
|
||||
def connect(self):
|
||||
"""
|
||||
Establish a Databricks SQL connection for table/query ingestion.
|
||||
|
||||
Returns:
|
||||
Connection: databricks-sql-connector connection object
|
||||
|
||||
Raises:
|
||||
ProcessingError: If connection fails
|
||||
ValidationError: If 'http_path' is not configured
|
||||
"""
|
||||
if not self.http_path:
|
||||
raise ValidationError(
|
||||
"Databricks 'http_path' is required for SQL ingestion. "
|
||||
"Provide via 'http_path' parameter or DATABRICKS_HTTP_PATH "
|
||||
"environment variable (the HTTP path of a SQL warehouse or cluster)."
|
||||
)
|
||||
|
||||
try:
|
||||
conn_params: Dict[str, Any] = {
|
||||
"server_hostname": self._hostname(),
|
||||
"http_path": self.http_path,
|
||||
}
|
||||
|
||||
if self.client_id and self.client_secret:
|
||||
conn_params["client_id"] = self.client_id
|
||||
conn_params["client_secret"] = self.client_secret
|
||||
else:
|
||||
conn_params["access_token"] = self.token
|
||||
|
||||
conn_params.update(self.config)
|
||||
|
||||
self.connection = databricks_sql.connect(**conn_params)
|
||||
|
||||
self.logger.info(f"Connected to Databricks: {self.host}")
|
||||
|
||||
return self.connection
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to connect to Databricks: {e}")
|
||||
raise ProcessingError(f"Failed to connect to Databricks: {e}") from e
|
||||
|
||||
def _hostname(self) -> str:
|
||||
"""Strip the scheme from the configured host for the SQL connector."""
|
||||
return re.sub(r"^https?://", "", self.host).rstrip("/")
|
||||
|
||||
def get_workspace_client(self) -> "WorkspaceClient":
|
||||
"""
|
||||
Get (and lazily create) a Unity Catalog WorkspaceClient.
|
||||
|
||||
Returns:
|
||||
WorkspaceClient: databricks-sdk workspace client
|
||||
|
||||
Raises:
|
||||
ProcessingError: If client creation fails
|
||||
"""
|
||||
if self._workspace_client is not None:
|
||||
return self._workspace_client
|
||||
|
||||
try:
|
||||
if self.client_id and self.client_secret:
|
||||
self._workspace_client = WorkspaceClient(
|
||||
host=self.host,
|
||||
client_id=self.client_id,
|
||||
client_secret=self.client_secret,
|
||||
)
|
||||
else:
|
||||
self._workspace_client = WorkspaceClient(
|
||||
host=self.host, token=self.token
|
||||
)
|
||||
return self._workspace_client
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to create Databricks workspace client: {e}")
|
||||
raise ProcessingError(
|
||||
f"Failed to create Databricks workspace client: {e}"
|
||||
) from e
|
||||
|
||||
def disconnect(self):
|
||||
"""Close Databricks SQL connection."""
|
||||
if self.connection:
|
||||
try:
|
||||
self.connection.close()
|
||||
self.connection = None
|
||||
self.logger.info("Disconnected from Databricks")
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Error during disconnect: {e}")
|
||||
|
||||
def test_connection(self) -> bool:
|
||||
"""
|
||||
Test the Databricks SQL connection without keeping it open.
|
||||
|
||||
Returns:
|
||||
bool: True if connection successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
conn = self.connect()
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("SELECT 1")
|
||||
cursor.fetchone()
|
||||
cursor.close()
|
||||
conn.close()
|
||||
self.connection = None
|
||||
return True
|
||||
except Exception as e:
|
||||
self.logger.debug(f"Connection test failed: {e}")
|
||||
return False
|
||||
|
||||
|
||||
class DatabricksIngestor:
|
||||
"""
|
||||
Databricks ingestion handler.
|
||||
|
||||
This class provides comprehensive Databricks ingestion capabilities,
|
||||
connecting to Databricks SQL warehouses/clusters for table and query
|
||||
ingestion, and to Unity Catalog for metadata and lineage introspection.
|
||||
|
||||
Features:
|
||||
- Table ingestion with pagination
|
||||
- Query execution with streaming
|
||||
- Unity Catalog schema introspection
|
||||
- Table-level lineage
|
||||
- Multiple authentication methods
|
||||
- Progress tracking and error handling
|
||||
- Connection management with proper cleanup
|
||||
|
||||
Example Usage:
|
||||
>>> ingestor = DatabricksIngestor(
|
||||
... host="https://adb-xxx.azuredatabricks.net",
|
||||
... token="dapi-xxxxxxxx",
|
||||
... http_path="/sql/1.0/warehouses/xxxxxxxx",
|
||||
... catalog="main",
|
||||
... )
|
||||
>>> data = ingestor.ingest_table("customers", limit=10000)
|
||||
>>> query_data = ingestor.ingest_query("SELECT * FROM sales")
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
host: Optional[str] = None,
|
||||
token: Optional[str] = None,
|
||||
http_path: Optional[str] = None,
|
||||
catalog: Optional[str] = None,
|
||||
schema: Optional[str] = None,
|
||||
client_id: Optional[str] = None,
|
||||
client_secret: Optional[str] = None,
|
||||
config: Optional[Dict[str, Any]] = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
Initialize Databricks ingestor.
|
||||
|
||||
Args:
|
||||
host: Databricks workspace URL
|
||||
token: Personal access token
|
||||
http_path: HTTP path of a SQL warehouse or cluster
|
||||
catalog: Default catalog (default: 'main')
|
||||
schema: Default schema (default: 'default')
|
||||
client_id: OAuth M2M service principal client ID
|
||||
client_secret: OAuth M2M service principal client secret
|
||||
config: Optional configuration dictionary
|
||||
**kwargs: Additional configuration parameters
|
||||
"""
|
||||
self.logger = get_logger("databricks_ingestor")
|
||||
self.config = config or {}
|
||||
self.config.update(kwargs)
|
||||
|
||||
# Initialize connector
|
||||
self.connector = DatabricksConnector(
|
||||
host=host,
|
||||
token=token,
|
||||
http_path=http_path,
|
||||
catalog=catalog,
|
||||
schema=schema,
|
||||
client_id=client_id,
|
||||
client_secret=client_secret,
|
||||
**self.config,
|
||||
)
|
||||
|
||||
# Initialize progress tracker
|
||||
self.progress_tracker = get_progress_tracker()
|
||||
# Ensure progress tracker is enabled
|
||||
if not self.progress_tracker.enabled:
|
||||
self.progress_tracker.enabled = True
|
||||
|
||||
self.logger.debug("Databricks ingestor initialized")
|
||||
|
||||
def _escape_identifier(self, identifier: str) -> str:
|
||||
"""Escape a SQL identifier by wrapping in backticks and doubling internal backticks.
|
||||
|
||||
Args:
|
||||
identifier: The identifier to escape
|
||||
|
||||
Returns:
|
||||
Properly escaped identifier safe for SQL interpolation
|
||||
"""
|
||||
escaped = identifier.replace("`", "``")
|
||||
return f"`{escaped}`"
|
||||
|
||||
def _full_table_name(
|
||||
self, table_name: str, catalog: Optional[str], schema: Optional[str]
|
||||
) -> str:
|
||||
"""Build a fully-qualified 'catalog.schema.table' identifier, omitting
|
||||
any component (catalog and/or schema) that wasn't provided."""
|
||||
parts = [part for part in (catalog, schema, table_name) if part]
|
||||
return ".".join(parts)
|
||||
|
||||
def _escaped_table_ref(
|
||||
self, table_name: str, catalog: Optional[str], schema: Optional[str]
|
||||
) -> str:
|
||||
"""Build a fully-qualified, identifier-escaped table reference for SQL,
|
||||
omitting any component (catalog and/or schema) that wasn't provided."""
|
||||
parts = [
|
||||
self._escape_identifier(part)
|
||||
for part in (catalog, schema, table_name)
|
||||
if part
|
||||
]
|
||||
return ".".join(parts)
|
||||
|
||||
def ingest_table(
|
||||
self,
|
||||
table_name: str,
|
||||
catalog: Optional[str] = None,
|
||||
schema: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
offset: Optional[int] = None,
|
||||
where: Optional[str] = None,
|
||||
order_by: Optional[str] = None,
|
||||
**options,
|
||||
) -> DatabricksData:
|
||||
"""
|
||||
Ingest data from a Databricks (Delta) table.
|
||||
|
||||
This method retrieves data from a Unity Catalog table with optional
|
||||
filtering, pagination, and ordering.
|
||||
|
||||
Args:
|
||||
table_name: Name of the table to ingest
|
||||
catalog: Catalog name (uses default if not provided)
|
||||
schema: Schema name (uses default if not provided)
|
||||
limit: Maximum number of rows to retrieve (optional)
|
||||
offset: Row offset for pagination (optional)
|
||||
where: WHERE clause for filtering (optional, must be trusted SQL)
|
||||
order_by: ORDER BY clause for sorting (optional, must be trusted SQL)
|
||||
**options: Additional query options
|
||||
|
||||
Warning:
|
||||
The 'where' and 'order_by' parameters accept raw SQL and must be
|
||||
trusted input from the caller. Do not pass untrusted user input.
|
||||
|
||||
Returns:
|
||||
DatabricksData: Ingested data object containing:
|
||||
- data: List of row dictionaries
|
||||
- row_count: Number of rows retrieved
|
||||
- columns: List of column names
|
||||
- table_name: Table name
|
||||
- catalog: Catalog name
|
||||
- schema: Schema name
|
||||
|
||||
Raises:
|
||||
ProcessingError: If table ingestion fails
|
||||
"""
|
||||
catalog = catalog or self.connector.catalog
|
||||
schema = schema or self.connector.schema
|
||||
|
||||
tracking_id = self.progress_tracker.start_tracking(
|
||||
file=f"{catalog}.{schema}.{table_name}",
|
||||
module="ingest",
|
||||
submodule="DatabricksIngestor",
|
||||
message=f"Table: {catalog}.{schema}.{table_name}",
|
||||
)
|
||||
|
||||
try:
|
||||
conn = self.connector.connect()
|
||||
try:
|
||||
table_ref = self._escaped_table_ref(table_name, catalog, schema)
|
||||
|
||||
query = f"SELECT * FROM {table_ref}"
|
||||
|
||||
if where:
|
||||
if ";" in where:
|
||||
raise ValueError(
|
||||
"Invalid WHERE clause: semicolons not permitted."
|
||||
)
|
||||
query += f" WHERE {where}"
|
||||
|
||||
if order_by:
|
||||
_SAFE_ORDER_RE = re.compile(
|
||||
r"^[A-Za-z_][A-Za-z0-9_]*(\s+(ASC|DESC))?"
|
||||
r"(\s*,\s*[A-Za-z_][A-Za-z0-9_]*(\s+(ASC|DESC))?)*$",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
if not _SAFE_ORDER_RE.match(order_by.strip()):
|
||||
raise ValueError(f"Invalid ORDER BY clause: '{order_by}'")
|
||||
query += f" ORDER BY {order_by}"
|
||||
|
||||
if limit is not None:
|
||||
query += f" LIMIT {int(limit)}"
|
||||
|
||||
if offset is not None:
|
||||
query += f" OFFSET {int(offset)}"
|
||||
|
||||
self.logger.debug(f"Executing query: {query}")
|
||||
|
||||
self.progress_tracker.update_tracking(
|
||||
tracking_id, message="Executing query..."
|
||||
)
|
||||
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(query)
|
||||
|
||||
self.progress_tracker.update_tracking(
|
||||
tracking_id, message="Fetching results..."
|
||||
)
|
||||
|
||||
columns = (
|
||||
[desc[0] for desc in cursor.description]
|
||||
if cursor.description
|
||||
else []
|
||||
)
|
||||
rows = cursor.fetchall()
|
||||
row_dicts = [dict(zip(columns, row)) for row in rows]
|
||||
|
||||
cursor.close()
|
||||
|
||||
data = self._convert_rows(row_dicts)
|
||||
|
||||
self.progress_tracker.stop_tracking(
|
||||
tracking_id,
|
||||
status="completed",
|
||||
message=f"Ingested {len(data)} rows",
|
||||
)
|
||||
|
||||
self.logger.info(f"Table ingestion completed: {len(data)} row(s)")
|
||||
|
||||
return DatabricksData(
|
||||
data=data,
|
||||
row_count=len(data),
|
||||
columns=columns,
|
||||
table_name=table_name,
|
||||
catalog=catalog,
|
||||
schema=schema,
|
||||
metadata={"query": query},
|
||||
)
|
||||
finally:
|
||||
self.connector.disconnect()
|
||||
|
||||
except Exception as e:
|
||||
self.progress_tracker.stop_tracking(
|
||||
tracking_id, status="failed", message=str(e)
|
||||
)
|
||||
self.logger.error(f"Failed to ingest table: {e}")
|
||||
raise ProcessingError(f"Failed to ingest table: {e}") from e
|
||||
|
||||
def ingest_query(
|
||||
self,
|
||||
query: str,
|
||||
parameters: Optional[Dict[str, Any]] = None,
|
||||
batch_size: Optional[int] = None,
|
||||
**options,
|
||||
) -> DatabricksData:
|
||||
"""
|
||||
Execute a Databricks SQL query and ingest results.
|
||||
|
||||
This method executes a SQL query and returns the results with
|
||||
optional batching for large result sets.
|
||||
|
||||
Args:
|
||||
query: SQL query to execute
|
||||
parameters: Query parameters for parameterized queries (optional)
|
||||
batch_size: Batch size for result fetching (optional)
|
||||
**options: Additional query options
|
||||
|
||||
Returns:
|
||||
DatabricksData: Query results as DatabricksData object
|
||||
|
||||
Raises:
|
||||
ProcessingError: If query execution fails
|
||||
|
||||
Example:
|
||||
>>> data = ingestor.ingest_query(
|
||||
... "SELECT * FROM sales WHERE date > :date",
|
||||
... parameters={"date": "2024-01-01"}
|
||||
... )
|
||||
"""
|
||||
tracking_id = self.progress_tracker.start_tracking(
|
||||
file="query",
|
||||
module="ingest",
|
||||
submodule="DatabricksIngestor",
|
||||
message="Executing query...",
|
||||
)
|
||||
|
||||
try:
|
||||
conn = self.connector.connect()
|
||||
try:
|
||||
self.logger.debug(f"Executing query: {query[:100]}...")
|
||||
|
||||
cursor = conn.cursor()
|
||||
|
||||
if parameters:
|
||||
cursor.execute(query, parameters)
|
||||
else:
|
||||
cursor.execute(query)
|
||||
|
||||
self.progress_tracker.update_tracking(
|
||||
tracking_id, message="Fetching results..."
|
||||
)
|
||||
|
||||
columns = (
|
||||
[desc[0] for desc in cursor.description]
|
||||
if cursor.description
|
||||
else []
|
||||
)
|
||||
|
||||
if batch_size:
|
||||
all_rows = []
|
||||
while True:
|
||||
rows = cursor.fetchmany(batch_size)
|
||||
if not rows:
|
||||
break
|
||||
all_rows.extend(rows)
|
||||
self.progress_tracker.update_tracking(
|
||||
tracking_id, message=f"Fetched {len(all_rows)} rows..."
|
||||
)
|
||||
else:
|
||||
all_rows = cursor.fetchall()
|
||||
|
||||
row_dicts = [dict(zip(columns, row)) for row in all_rows]
|
||||
|
||||
cursor.close()
|
||||
|
||||
data = self._convert_rows(row_dicts)
|
||||
|
||||
self.progress_tracker.stop_tracking(
|
||||
tracking_id,
|
||||
status="completed",
|
||||
message=f"Query completed: {len(data)} rows",
|
||||
)
|
||||
|
||||
self.logger.info(f"Query execution completed: {len(data)} row(s)")
|
||||
|
||||
return DatabricksData(
|
||||
data=data,
|
||||
row_count=len(data),
|
||||
columns=columns,
|
||||
query=query,
|
||||
metadata={"parameters": parameters} if parameters else {},
|
||||
)
|
||||
finally:
|
||||
self.connector.disconnect()
|
||||
|
||||
except Exception as e:
|
||||
self.progress_tracker.stop_tracking(
|
||||
tracking_id, status="failed", message=str(e)
|
||||
)
|
||||
self.logger.error(f"Failed to execute query: {e}")
|
||||
raise ProcessingError(f"Failed to execute query: {e}") from e
|
||||
|
||||
def get_table_schema(
|
||||
self,
|
||||
table_name: str,
|
||||
catalog: Optional[str] = None,
|
||||
schema: Optional[str] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Get schema information for a Unity Catalog table.
|
||||
|
||||
Args:
|
||||
table_name: Name of the table
|
||||
catalog: Catalog name (uses default if not provided)
|
||||
schema: Schema name (uses default if not provided)
|
||||
|
||||
Returns:
|
||||
dict: Table schema information containing:
|
||||
- columns: List of column dictionaries with name, type, nullable
|
||||
- primary_keys: List of primary key column names (if any)
|
||||
|
||||
Raises:
|
||||
ProcessingError: If schema retrieval fails
|
||||
"""
|
||||
try:
|
||||
catalog = catalog or self.connector.catalog
|
||||
schema = schema or self.connector.schema
|
||||
|
||||
if not catalog:
|
||||
raise ValidationError(
|
||||
"Catalog name is required for schema introspection. "
|
||||
"Provide via 'catalog' parameter or set default catalog in connector."
|
||||
)
|
||||
|
||||
client = self.connector.get_workspace_client()
|
||||
full_name = self._full_table_name(table_name, catalog, schema)
|
||||
|
||||
table_info = client.tables.get(full_name=full_name)
|
||||
|
||||
column_info = []
|
||||
primary_keys: List[str] = []
|
||||
for col in getattr(table_info, "columns", None) or []:
|
||||
column_info.append(
|
||||
{
|
||||
"name": col.name,
|
||||
"type": getattr(col, "type_text", None)
|
||||
or getattr(col, "type_name", None),
|
||||
"nullable": getattr(col, "nullable", True),
|
||||
"comment": getattr(col, "comment", None),
|
||||
}
|
||||
)
|
||||
|
||||
self.logger.debug(
|
||||
f"Retrieved schema for {full_name}: {len(column_info)} columns"
|
||||
)
|
||||
|
||||
return {"columns": column_info, "primary_keys": primary_keys}
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to get table schema: {e}")
|
||||
raise ProcessingError(f"Failed to get table schema: {e}") from e
|
||||
|
||||
def list_catalogs(self) -> List[str]:
|
||||
"""
|
||||
List all catalogs visible in Unity Catalog.
|
||||
|
||||
Returns:
|
||||
list: List of catalog names
|
||||
|
||||
Raises:
|
||||
ProcessingError: If listing fails
|
||||
"""
|
||||
try:
|
||||
client = self.connector.get_workspace_client()
|
||||
catalogs = [c.name for c in client.catalogs.list()]
|
||||
|
||||
self.logger.debug(f"Found {len(catalogs)} catalogs")
|
||||
|
||||
return catalogs
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to list catalogs: {e}")
|
||||
raise ProcessingError(f"Failed to list catalogs: {e}") from e
|
||||
|
||||
def list_schemas(self, catalog: Optional[str] = None) -> List[str]:
|
||||
"""
|
||||
List all schemas in a catalog.
|
||||
|
||||
Args:
|
||||
catalog: Catalog name (uses default if not provided)
|
||||
|
||||
Returns:
|
||||
list: List of schema names
|
||||
|
||||
Raises:
|
||||
ProcessingError: If listing fails
|
||||
"""
|
||||
try:
|
||||
catalog = catalog or self.connector.catalog
|
||||
|
||||
if not catalog:
|
||||
raise ValidationError(
|
||||
"Catalog name is required for listing schemas. "
|
||||
"Provide via 'catalog' parameter or set default catalog in connector."
|
||||
)
|
||||
|
||||
client = self.connector.get_workspace_client()
|
||||
schemas = [s.name for s in client.schemas.list(catalog_name=catalog)]
|
||||
|
||||
self.logger.debug(f"Found {len(schemas)} schemas in {catalog}")
|
||||
|
||||
return schemas
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to list schemas: {e}")
|
||||
raise ProcessingError(f"Failed to list schemas: {e}") from e
|
||||
|
||||
def list_tables(
|
||||
self, catalog: Optional[str] = None, schema: Optional[str] = None
|
||||
) -> List[str]:
|
||||
"""
|
||||
List all tables in a Unity Catalog catalog/schema.
|
||||
|
||||
Args:
|
||||
catalog: Catalog name (uses default if not provided)
|
||||
schema: Schema name (uses default if not provided)
|
||||
|
||||
Returns:
|
||||
list: List of table names
|
||||
|
||||
Raises:
|
||||
ProcessingError: If listing fails
|
||||
"""
|
||||
try:
|
||||
catalog = catalog or self.connector.catalog
|
||||
schema = schema or self.connector.schema
|
||||
|
||||
if not catalog:
|
||||
raise ValidationError(
|
||||
"Catalog name is required for listing tables. "
|
||||
"Provide via 'catalog' parameter or set default catalog in connector."
|
||||
)
|
||||
|
||||
if not schema:
|
||||
raise ValidationError(
|
||||
"Schema name is required for listing tables. "
|
||||
"Provide via 'schema' parameter or set default schema in connector."
|
||||
)
|
||||
|
||||
client = self.connector.get_workspace_client()
|
||||
tables = [
|
||||
t.name
|
||||
for t in client.tables.list(catalog_name=catalog, schema_name=schema)
|
||||
]
|
||||
|
||||
self.logger.debug(f"Found {len(tables)} tables in {catalog}.{schema}")
|
||||
|
||||
return tables
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to list tables: {e}")
|
||||
raise ProcessingError(f"Failed to list tables: {e}") from e
|
||||
|
||||
def get_table_lineage(
|
||||
self,
|
||||
table_name: str,
|
||||
catalog: Optional[str] = None,
|
||||
schema: Optional[str] = None,
|
||||
) -> Dict[str, List[str]]:
|
||||
"""
|
||||
Get table-level lineage from Unity Catalog.
|
||||
|
||||
Args:
|
||||
table_name: Name of the table
|
||||
catalog: Catalog name (uses default if not provided)
|
||||
schema: Schema name (uses default if not provided)
|
||||
|
||||
Returns:
|
||||
dict: Lineage information containing:
|
||||
- upstream: List of fully-qualified upstream table names
|
||||
- downstream: List of fully-qualified downstream table names
|
||||
|
||||
Raises:
|
||||
ProcessingError: If lineage retrieval fails
|
||||
"""
|
||||
try:
|
||||
catalog = catalog or self.connector.catalog
|
||||
schema = schema or self.connector.schema
|
||||
full_name = self._full_table_name(table_name, catalog, schema)
|
||||
|
||||
client = self.connector.get_workspace_client()
|
||||
|
||||
response = client.api_client.do(
|
||||
"GET",
|
||||
"/api/2.0/lineage-tracking/table-lineage",
|
||||
query={"table_name": full_name, "include_entity_lineage": False},
|
||||
)
|
||||
|
||||
upstream = [
|
||||
item.get("tableInfo", {}).get("name")
|
||||
for item in response.get("upstreams", []) or []
|
||||
if item.get("tableInfo")
|
||||
]
|
||||
downstream = [
|
||||
item.get("tableInfo", {}).get("name")
|
||||
for item in response.get("downstreams", []) or []
|
||||
if item.get("tableInfo")
|
||||
]
|
||||
|
||||
self.logger.debug(
|
||||
f"Retrieved lineage for {full_name}: "
|
||||
f"{len(upstream)} upstream, {len(downstream)} downstream"
|
||||
)
|
||||
|
||||
return {"upstream": upstream, "downstream": downstream}
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to get table lineage: {e}")
|
||||
raise ProcessingError(f"Failed to get table lineage: {e}") from e
|
||||
|
||||
def export_as_documents(
|
||||
self,
|
||||
data: DatabricksData,
|
||||
id_field: str = "id",
|
||||
text_fields: Optional[List[str]] = None,
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Convert Databricks data to document format for Semantica processing.
|
||||
|
||||
Args:
|
||||
data: DatabricksData object to convert
|
||||
id_field: Field to use as document ID (default: 'id')
|
||||
text_fields: List of fields to combine into document text (optional)
|
||||
|
||||
Returns:
|
||||
list: List of document dictionaries
|
||||
|
||||
Example:
|
||||
>>> data = ingestor.ingest_table("customers")
|
||||
>>> documents = ingestor.export_as_documents(data, text_fields=["name", "notes"])
|
||||
"""
|
||||
documents = []
|
||||
|
||||
for idx, row in enumerate(data.data):
|
||||
doc = {
|
||||
"id": str(row.get(id_field, idx)),
|
||||
"metadata": {
|
||||
"source": "databricks",
|
||||
"table": data.table_name,
|
||||
"catalog": data.catalog,
|
||||
"schema": data.schema,
|
||||
},
|
||||
}
|
||||
|
||||
if text_fields:
|
||||
text_parts = []
|
||||
for field_name in text_fields:
|
||||
if field_name in row and row[field_name] is not None:
|
||||
text_parts.append(str(row[field_name]))
|
||||
doc["text"] = " ".join(text_parts)
|
||||
else:
|
||||
text_parts = []
|
||||
for key, value in row.items():
|
||||
if isinstance(value, str):
|
||||
text_parts.append(value)
|
||||
doc["text"] = " ".join(text_parts)
|
||||
|
||||
doc["metadata"]["row_data"] = row
|
||||
|
||||
documents.append(doc)
|
||||
|
||||
self.logger.debug(f"Exported {len(documents)} documents")
|
||||
|
||||
return documents
|
||||
|
||||
def close(self):
|
||||
"""Close Databricks connection."""
|
||||
self.connector.disconnect()
|
||||
|
||||
def _convert_rows(self, rows: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Convert Databricks rows to JSON-serializable format.
|
||||
|
||||
Args:
|
||||
rows: List of row dictionaries from Databricks
|
||||
|
||||
Returns:
|
||||
list: Converted row dictionaries
|
||||
"""
|
||||
converted = []
|
||||
|
||||
for row in rows:
|
||||
converted_row = {}
|
||||
for key, value in row.items():
|
||||
if isinstance(value, datetime):
|
||||
converted_row[key] = value.isoformat()
|
||||
elif isinstance(value, bytes):
|
||||
try:
|
||||
converted_row[key] = value.decode("utf-8")
|
||||
except UnicodeDecodeError:
|
||||
converted_row[key] = str(value)
|
||||
else:
|
||||
converted_row[key] = value
|
||||
|
||||
converted.append(converted_row)
|
||||
|
||||
return converted
|
||||
|
||||
def __enter__(self):
|
||||
"""Context manager entry."""
|
||||
self.connector.connect()
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
"""Context manager exit."""
|
||||
self.close()
|
||||
@@ -0,0 +1,743 @@
|
||||
"""
|
||||
Unit tests for Databricks Ingestor
|
||||
|
||||
This test module uses mocks to test Databricks ingestion functionality
|
||||
without requiring a live Databricks workspace.
|
||||
"""
|
||||
|
||||
import os
|
||||
from datetime import datetime
|
||||
from unittest.mock import MagicMock, Mock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
# Test if databricks-sdk / databricks-sql-connector are available
|
||||
try:
|
||||
import databricks.sdk # noqa: F401
|
||||
from databricks import sql # noqa: F401
|
||||
|
||||
DATABRICKS_LIBS_AVAILABLE = True
|
||||
except ImportError:
|
||||
DATABRICKS_LIBS_AVAILABLE = False
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def mock_databricks_if_needed():
|
||||
"""Mock databricks modules if not installed."""
|
||||
if not DATABRICKS_LIBS_AVAILABLE:
|
||||
with patch.dict(
|
||||
"sys.modules",
|
||||
{
|
||||
"databricks": MagicMock(),
|
||||
"databricks.sql": MagicMock(),
|
||||
"databricks.sdk": MagicMock(),
|
||||
},
|
||||
):
|
||||
yield
|
||||
else:
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_databricks_connection():
|
||||
"""Create a mock Databricks SQL connection."""
|
||||
mock_conn = Mock()
|
||||
mock_cursor = Mock()
|
||||
|
||||
mock_cursor.execute = Mock()
|
||||
mock_cursor.fetchall = Mock(return_value=[])
|
||||
mock_cursor.fetchone = Mock(return_value=[1])
|
||||
mock_cursor.fetchmany = Mock(return_value=[])
|
||||
mock_cursor.description = [("id", None), ("name", None), ("value", None)]
|
||||
mock_cursor.close = Mock()
|
||||
|
||||
mock_conn.cursor = Mock(return_value=mock_cursor)
|
||||
mock_conn.close = Mock()
|
||||
|
||||
return mock_conn, mock_cursor
|
||||
|
||||
|
||||
class TestDatabricksConnector:
|
||||
"""Test DatabricksConnector class."""
|
||||
|
||||
@patch("semantica.ingest.databricks_ingestor.DATABRICKS_AVAILABLE", True)
|
||||
@patch("semantica.ingest.databricks_ingestor.databricks_sql")
|
||||
def test_connector_init_with_token(self, mock_sql):
|
||||
"""Test connector initialization with personal access token authentication."""
|
||||
from semantica.ingest.databricks_ingestor import DatabricksConnector
|
||||
|
||||
connector = DatabricksConnector(
|
||||
host="https://adb-xxx.azuredatabricks.net",
|
||||
token="test_token",
|
||||
http_path="/sql/1.0/warehouses/xxxx",
|
||||
catalog="TEST_CATALOG",
|
||||
)
|
||||
|
||||
assert connector.host == "https://adb-xxx.azuredatabricks.net"
|
||||
assert connector.token == "test_token"
|
||||
assert connector.http_path == "/sql/1.0/warehouses/xxxx"
|
||||
assert connector.catalog == "TEST_CATALOG"
|
||||
assert connector.schema == "default"
|
||||
|
||||
@patch("semantica.ingest.databricks_ingestor.DATABRICKS_AVAILABLE", True)
|
||||
@patch("semantica.ingest.databricks_ingestor.databricks_sql")
|
||||
def test_connector_init_from_env(self, mock_sql):
|
||||
"""Test connector initialization from environment variables."""
|
||||
from semantica.ingest.databricks_ingestor import DatabricksConnector
|
||||
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"DATABRICKS_HOST": "https://env-host.azuredatabricks.net",
|
||||
"DATABRICKS_TOKEN": "env_token",
|
||||
"DATABRICKS_HTTP_PATH": "/sql/1.0/warehouses/env",
|
||||
},
|
||||
):
|
||||
connector = DatabricksConnector()
|
||||
|
||||
assert connector.host == "https://env-host.azuredatabricks.net"
|
||||
assert connector.token == "env_token"
|
||||
assert connector.http_path == "/sql/1.0/warehouses/env"
|
||||
|
||||
@patch("semantica.ingest.databricks_ingestor.DATABRICKS_AVAILABLE", True)
|
||||
@patch("semantica.ingest.databricks_ingestor.databricks_sql")
|
||||
def test_connector_init_missing_host(self, mock_sql):
|
||||
"""Test connector initialization fails without host."""
|
||||
from semantica.ingest.databricks_ingestor import DatabricksConnector
|
||||
from semantica.utils.exceptions import ValidationError
|
||||
|
||||
with pytest.raises(ValidationError, match="Databricks host is required"):
|
||||
DatabricksConnector(token="test_token")
|
||||
|
||||
@patch("semantica.ingest.databricks_ingestor.DATABRICKS_AVAILABLE", True)
|
||||
@patch("semantica.ingest.databricks_ingestor.databricks_sql")
|
||||
def test_connector_init_missing_auth(self, mock_sql):
|
||||
"""Test connector initialization fails without any authentication method."""
|
||||
from semantica.ingest.databricks_ingestor import DatabricksConnector
|
||||
from semantica.utils.exceptions import ValidationError
|
||||
|
||||
with pytest.raises(ValidationError, match="Databricks authentication is required"):
|
||||
DatabricksConnector(host="https://adb-xxx.azuredatabricks.net")
|
||||
|
||||
@patch("semantica.ingest.databricks_ingestor.DATABRICKS_AVAILABLE", True)
|
||||
@patch("semantica.ingest.databricks_ingestor.databricks_sql")
|
||||
def test_connector_connect_token_auth(self, mock_sql, mock_databricks_connection):
|
||||
"""Test connection with personal access token authentication."""
|
||||
from semantica.ingest.databricks_ingestor import DatabricksConnector
|
||||
|
||||
mock_conn, _ = mock_databricks_connection
|
||||
mock_sql.connect = Mock(return_value=mock_conn)
|
||||
|
||||
connector = DatabricksConnector(
|
||||
host="https://adb-xxx.azuredatabricks.net",
|
||||
token="test_token",
|
||||
http_path="/sql/1.0/warehouses/xxxx",
|
||||
)
|
||||
|
||||
conn = connector.connect()
|
||||
|
||||
assert conn == mock_conn
|
||||
mock_sql.connect.assert_called_once()
|
||||
|
||||
call_kwargs = mock_sql.connect.call_args[1]
|
||||
assert call_kwargs["server_hostname"] == "adb-xxx.azuredatabricks.net"
|
||||
assert call_kwargs["http_path"] == "/sql/1.0/warehouses/xxxx"
|
||||
assert call_kwargs["access_token"] == "test_token"
|
||||
|
||||
@patch("semantica.ingest.databricks_ingestor.DATABRICKS_AVAILABLE", True)
|
||||
@patch("semantica.ingest.databricks_ingestor.databricks_sql")
|
||||
def test_connector_connect_oauth_m2m(self, mock_sql, mock_databricks_connection):
|
||||
"""Test connection with OAuth M2M authentication."""
|
||||
from semantica.ingest.databricks_ingestor import DatabricksConnector
|
||||
|
||||
mock_conn, _ = mock_databricks_connection
|
||||
mock_sql.connect = Mock(return_value=mock_conn)
|
||||
|
||||
connector = DatabricksConnector(
|
||||
host="https://adb-xxx.azuredatabricks.net",
|
||||
http_path="/sql/1.0/warehouses/xxxx",
|
||||
client_id="test_client_id",
|
||||
client_secret="test_client_secret",
|
||||
)
|
||||
|
||||
connector.connect()
|
||||
|
||||
call_kwargs = mock_sql.connect.call_args[1]
|
||||
assert call_kwargs["client_id"] == "test_client_id"
|
||||
assert call_kwargs["client_secret"] == "test_client_secret"
|
||||
assert "access_token" not in call_kwargs
|
||||
|
||||
@patch("semantica.ingest.databricks_ingestor.DATABRICKS_AVAILABLE", True)
|
||||
@patch("semantica.ingest.databricks_ingestor.databricks_sql")
|
||||
def test_connector_connect_missing_http_path(self, mock_sql):
|
||||
"""Test connection fails without http_path."""
|
||||
from semantica.ingest.databricks_ingestor import DatabricksConnector
|
||||
from semantica.utils.exceptions import ValidationError
|
||||
|
||||
connector = DatabricksConnector(
|
||||
host="https://adb-xxx.azuredatabricks.net",
|
||||
token="test_token",
|
||||
)
|
||||
|
||||
with pytest.raises(ValidationError, match="http_path"):
|
||||
connector.connect()
|
||||
|
||||
@patch("semantica.ingest.databricks_ingestor.DATABRICKS_AVAILABLE", True)
|
||||
@patch("semantica.ingest.databricks_ingestor.databricks_sql")
|
||||
def test_connector_disconnect(self, mock_sql, mock_databricks_connection):
|
||||
"""Test connection disconnect."""
|
||||
from semantica.ingest.databricks_ingestor import DatabricksConnector
|
||||
|
||||
mock_conn, _ = mock_databricks_connection
|
||||
mock_sql.connect = Mock(return_value=mock_conn)
|
||||
|
||||
connector = DatabricksConnector(
|
||||
host="https://adb-xxx.azuredatabricks.net",
|
||||
token="test_token",
|
||||
http_path="/sql/1.0/warehouses/xxxx",
|
||||
)
|
||||
|
||||
connector.connect()
|
||||
connector.disconnect()
|
||||
|
||||
mock_conn.close.assert_called_once()
|
||||
assert connector.connection is None
|
||||
|
||||
@patch("semantica.ingest.databricks_ingestor.DATABRICKS_AVAILABLE", True)
|
||||
@patch("semantica.ingest.databricks_ingestor.databricks_sql")
|
||||
def test_connector_test_connection_success(self, mock_sql, mock_databricks_connection):
|
||||
"""Test successful connection test."""
|
||||
from semantica.ingest.databricks_ingestor import DatabricksConnector
|
||||
|
||||
mock_conn, mock_cursor = mock_databricks_connection
|
||||
mock_sql.connect = Mock(return_value=mock_conn)
|
||||
|
||||
connector = DatabricksConnector(
|
||||
host="https://adb-xxx.azuredatabricks.net",
|
||||
token="test_token",
|
||||
http_path="/sql/1.0/warehouses/xxxx",
|
||||
)
|
||||
|
||||
result = connector.test_connection()
|
||||
|
||||
assert result is True
|
||||
mock_cursor.execute.assert_called_with("SELECT 1")
|
||||
|
||||
@patch("semantica.ingest.databricks_ingestor.DATABRICKS_AVAILABLE", True)
|
||||
@patch("semantica.ingest.databricks_ingestor.databricks_sql")
|
||||
def test_connector_test_connection_failure(self, mock_sql):
|
||||
"""Test connection test failure."""
|
||||
from semantica.ingest.databricks_ingestor import DatabricksConnector
|
||||
|
||||
mock_sql.connect = Mock(side_effect=Exception("Connection failed"))
|
||||
|
||||
connector = DatabricksConnector(
|
||||
host="https://adb-xxx.azuredatabricks.net",
|
||||
token="test_token",
|
||||
http_path="/sql/1.0/warehouses/xxxx",
|
||||
)
|
||||
|
||||
result = connector.test_connection()
|
||||
|
||||
assert result is False
|
||||
|
||||
|
||||
class TestDatabricksIngestor:
|
||||
"""Test DatabricksIngestor class."""
|
||||
|
||||
@patch("semantica.ingest.databricks_ingestor.DATABRICKS_AVAILABLE", True)
|
||||
@patch("semantica.ingest.databricks_ingestor.databricks_sql")
|
||||
def test_ingestor_init(self, mock_sql):
|
||||
"""Test ingestor initialization."""
|
||||
from semantica.ingest.databricks_ingestor import DatabricksIngestor
|
||||
|
||||
ingestor = DatabricksIngestor(
|
||||
host="https://adb-xxx.azuredatabricks.net",
|
||||
token="test_token",
|
||||
http_path="/sql/1.0/warehouses/xxxx",
|
||||
)
|
||||
|
||||
assert ingestor.connector is not None
|
||||
assert ingestor.connector.host == "https://adb-xxx.azuredatabricks.net"
|
||||
|
||||
@patch("semantica.ingest.databricks_ingestor.DATABRICKS_AVAILABLE", True)
|
||||
@patch("semantica.ingest.databricks_ingestor.databricks_sql")
|
||||
def test_ingest_table_basic(self, mock_sql, mock_databricks_connection):
|
||||
"""Test basic table ingestion."""
|
||||
from semantica.ingest.databricks_ingestor import DatabricksIngestor
|
||||
|
||||
mock_conn, mock_cursor = mock_databricks_connection
|
||||
mock_sql.connect = Mock(return_value=mock_conn)
|
||||
|
||||
mock_cursor.fetchall = Mock(
|
||||
return_value=[
|
||||
(1, "Alice", 100),
|
||||
(2, "Bob", 200),
|
||||
]
|
||||
)
|
||||
mock_cursor.description = [("id", None), ("name", None), ("value", None)]
|
||||
|
||||
ingestor = DatabricksIngestor(
|
||||
host="https://adb-xxx.azuredatabricks.net",
|
||||
token="test_token",
|
||||
http_path="/sql/1.0/warehouses/xxxx",
|
||||
catalog="TEST_CATALOG",
|
||||
schema="TEST_SCHEMA",
|
||||
)
|
||||
|
||||
data = ingestor.ingest_table("customers")
|
||||
|
||||
assert data.row_count == 2
|
||||
assert data.table_name == "customers"
|
||||
assert data.catalog == "TEST_CATALOG"
|
||||
assert data.schema == "TEST_SCHEMA"
|
||||
assert len(data.columns) == 3
|
||||
assert "id" in data.columns
|
||||
assert data.data[0]["name"] == "Alice"
|
||||
|
||||
@patch("semantica.ingest.databricks_ingestor.DATABRICKS_AVAILABLE", True)
|
||||
@patch("semantica.ingest.databricks_ingestor.databricks_sql")
|
||||
def test_ingest_table_closes_connection(self, mock_sql, mock_databricks_connection):
|
||||
"""Test that ingest_table() closes the SQL connection after use instead of leaking it."""
|
||||
from semantica.ingest.databricks_ingestor import DatabricksIngestor
|
||||
|
||||
mock_conn, _ = mock_databricks_connection
|
||||
mock_sql.connect = Mock(return_value=mock_conn)
|
||||
|
||||
ingestor = DatabricksIngestor(
|
||||
host="https://adb-xxx.azuredatabricks.net",
|
||||
token="test_token",
|
||||
http_path="/sql/1.0/warehouses/xxxx",
|
||||
)
|
||||
|
||||
ingestor.ingest_table("customers")
|
||||
|
||||
mock_conn.close.assert_called_once()
|
||||
assert ingestor.connector.connection is None
|
||||
|
||||
@patch("semantica.ingest.databricks_ingestor.DATABRICKS_AVAILABLE", True)
|
||||
@patch("semantica.ingest.databricks_ingestor.databricks_sql")
|
||||
def test_ingest_table_catalog_only(self, mock_sql, mock_databricks_connection):
|
||||
"""Test table ingestion still qualifies the reference when only catalog is provided."""
|
||||
from semantica.ingest.databricks_ingestor import DatabricksIngestor
|
||||
|
||||
mock_conn, mock_cursor = mock_databricks_connection
|
||||
mock_sql.connect = Mock(return_value=mock_conn)
|
||||
|
||||
ingestor = DatabricksIngestor(
|
||||
host="https://adb-xxx.azuredatabricks.net",
|
||||
token="test_token",
|
||||
http_path="/sql/1.0/warehouses/xxxx",
|
||||
)
|
||||
# Force a missing schema past the connector's "default" fallback.
|
||||
ingestor.connector.schema = None
|
||||
|
||||
ingestor.ingest_table("customers", catalog="main")
|
||||
|
||||
executed_query = mock_cursor.execute.call_args[0][0]
|
||||
assert "`main`.`customers`" in executed_query
|
||||
assert "None" not in executed_query
|
||||
|
||||
@patch("semantica.ingest.databricks_ingestor.DATABRICKS_AVAILABLE", True)
|
||||
@patch("semantica.ingest.databricks_ingestor.databricks_sql")
|
||||
def test_ingest_table_with_limit(self, mock_sql, mock_databricks_connection):
|
||||
"""Test table ingestion with limit."""
|
||||
from semantica.ingest.databricks_ingestor import DatabricksIngestor
|
||||
|
||||
mock_conn, mock_cursor = mock_databricks_connection
|
||||
mock_sql.connect = Mock(return_value=mock_conn)
|
||||
|
||||
ingestor = DatabricksIngestor(
|
||||
host="https://adb-xxx.azuredatabricks.net",
|
||||
token="test_token",
|
||||
http_path="/sql/1.0/warehouses/xxxx",
|
||||
)
|
||||
|
||||
ingestor.ingest_table("customers", limit=100)
|
||||
|
||||
executed_query = mock_cursor.execute.call_args[0][0]
|
||||
assert "LIMIT 100" in executed_query
|
||||
|
||||
@patch("semantica.ingest.databricks_ingestor.DATABRICKS_AVAILABLE", True)
|
||||
@patch("semantica.ingest.databricks_ingestor.databricks_sql")
|
||||
def test_ingest_table_with_where(self, mock_sql, mock_databricks_connection):
|
||||
"""Test table ingestion with WHERE clause."""
|
||||
from semantica.ingest.databricks_ingestor import DatabricksIngestor
|
||||
|
||||
mock_conn, mock_cursor = mock_databricks_connection
|
||||
mock_sql.connect = Mock(return_value=mock_conn)
|
||||
|
||||
ingestor = DatabricksIngestor(
|
||||
host="https://adb-xxx.azuredatabricks.net",
|
||||
token="test_token",
|
||||
http_path="/sql/1.0/warehouses/xxxx",
|
||||
)
|
||||
|
||||
ingestor.ingest_table("customers", where="value > 100")
|
||||
|
||||
executed_query = mock_cursor.execute.call_args[0][0]
|
||||
assert "WHERE value > 100" in executed_query
|
||||
|
||||
@patch("semantica.ingest.databricks_ingestor.DATABRICKS_AVAILABLE", True)
|
||||
@patch("semantica.ingest.databricks_ingestor.databricks_sql")
|
||||
def test_ingest_table_rejects_unsafe_order_by(self, mock_sql, mock_databricks_connection):
|
||||
"""Test table ingestion rejects unsafe ORDER BY clauses."""
|
||||
from semantica.ingest.databricks_ingestor import DatabricksIngestor
|
||||
|
||||
mock_conn, _ = mock_databricks_connection
|
||||
mock_sql.connect = Mock(return_value=mock_conn)
|
||||
|
||||
ingestor = DatabricksIngestor(
|
||||
host="https://adb-xxx.azuredatabricks.net",
|
||||
token="test_token",
|
||||
http_path="/sql/1.0/warehouses/xxxx",
|
||||
)
|
||||
|
||||
from semantica.utils.exceptions import ProcessingError
|
||||
|
||||
with pytest.raises(ProcessingError):
|
||||
ingestor.ingest_table("customers", order_by="value; DROP TABLE customers")
|
||||
|
||||
@patch("semantica.ingest.databricks_ingestor.DATABRICKS_AVAILABLE", True)
|
||||
@patch("semantica.ingest.databricks_ingestor.databricks_sql")
|
||||
def test_ingest_query_basic(self, mock_sql, mock_databricks_connection):
|
||||
"""Test basic query execution."""
|
||||
from semantica.ingest.databricks_ingestor import DatabricksIngestor
|
||||
|
||||
mock_conn, mock_cursor = mock_databricks_connection
|
||||
mock_sql.connect = Mock(return_value=mock_conn)
|
||||
|
||||
mock_cursor.fetchall = Mock(return_value=[(1000,)])
|
||||
mock_cursor.description = [("total", None)]
|
||||
|
||||
ingestor = DatabricksIngestor(
|
||||
host="https://adb-xxx.azuredatabricks.net",
|
||||
token="test_token",
|
||||
http_path="/sql/1.0/warehouses/xxxx",
|
||||
)
|
||||
|
||||
query = "SELECT SUM(value) AS total FROM sales"
|
||||
data = ingestor.ingest_query(query)
|
||||
|
||||
assert data.row_count == 1
|
||||
assert data.query == query
|
||||
assert data.data[0]["total"] == 1000
|
||||
|
||||
@patch("semantica.ingest.databricks_ingestor.DATABRICKS_AVAILABLE", True)
|
||||
@patch("semantica.ingest.databricks_ingestor.databricks_sql")
|
||||
def test_ingest_query_closes_connection(self, mock_sql, mock_databricks_connection):
|
||||
"""Test that ingest_query() closes the SQL connection after use instead of leaking it."""
|
||||
from semantica.ingest.databricks_ingestor import DatabricksIngestor
|
||||
|
||||
mock_conn, _ = mock_databricks_connection
|
||||
mock_sql.connect = Mock(return_value=mock_conn)
|
||||
|
||||
ingestor = DatabricksIngestor(
|
||||
host="https://adb-xxx.azuredatabricks.net",
|
||||
token="test_token",
|
||||
http_path="/sql/1.0/warehouses/xxxx",
|
||||
)
|
||||
|
||||
ingestor.ingest_query("SELECT * FROM sales")
|
||||
|
||||
mock_conn.close.assert_called_once()
|
||||
assert ingestor.connector.connection is None
|
||||
|
||||
@patch("semantica.ingest.databricks_ingestor.DATABRICKS_AVAILABLE", True)
|
||||
@patch("semantica.ingest.databricks_ingestor.databricks_sql")
|
||||
def test_ingest_query_with_batching(self, mock_sql, mock_databricks_connection):
|
||||
"""Test query execution with batch fetching."""
|
||||
from semantica.ingest.databricks_ingestor import DatabricksIngestor
|
||||
|
||||
mock_conn, mock_cursor = mock_databricks_connection
|
||||
mock_sql.connect = Mock(return_value=mock_conn)
|
||||
|
||||
batch1 = [(1,), (2,)]
|
||||
batch2 = [(3,)]
|
||||
mock_cursor.fetchmany = Mock(side_effect=[batch1, batch2, []])
|
||||
mock_cursor.description = [("id", None)]
|
||||
|
||||
ingestor = DatabricksIngestor(
|
||||
host="https://adb-xxx.azuredatabricks.net",
|
||||
token="test_token",
|
||||
http_path="/sql/1.0/warehouses/xxxx",
|
||||
)
|
||||
|
||||
data = ingestor.ingest_query("SELECT * FROM customers", batch_size=2)
|
||||
|
||||
assert data.row_count == 3
|
||||
assert len(data.data) == 3
|
||||
|
||||
@patch("semantica.ingest.databricks_ingestor.DATABRICKS_AVAILABLE", True)
|
||||
@patch("semantica.ingest.databricks_ingestor.databricks_sql")
|
||||
@patch("semantica.ingest.databricks_ingestor.WorkspaceClient")
|
||||
def test_get_table_schema(self, mock_ws_client_cls, mock_sql):
|
||||
"""Test getting table schema information."""
|
||||
from semantica.ingest.databricks_ingestor import DatabricksIngestor
|
||||
|
||||
mock_column_1 = Mock(name="id")
|
||||
mock_column_1.name = "id"
|
||||
mock_column_1.type_text = "BIGINT"
|
||||
mock_column_1.nullable = False
|
||||
mock_column_1.comment = None
|
||||
|
||||
mock_column_2 = Mock(name="name")
|
||||
mock_column_2.name = "name"
|
||||
mock_column_2.type_text = "STRING"
|
||||
mock_column_2.nullable = True
|
||||
mock_column_2.comment = None
|
||||
|
||||
mock_table_info = Mock()
|
||||
mock_table_info.columns = [mock_column_1, mock_column_2]
|
||||
|
||||
mock_ws_client = Mock()
|
||||
mock_ws_client.tables.get = Mock(return_value=mock_table_info)
|
||||
mock_ws_client_cls.return_value = mock_ws_client
|
||||
|
||||
ingestor = DatabricksIngestor(
|
||||
host="https://adb-xxx.azuredatabricks.net",
|
||||
token="test_token",
|
||||
http_path="/sql/1.0/warehouses/xxxx",
|
||||
catalog="TEST_CATALOG",
|
||||
schema="TEST_SCHEMA",
|
||||
)
|
||||
|
||||
schema = ingestor.get_table_schema("customers")
|
||||
|
||||
assert len(schema["columns"]) == 2
|
||||
assert schema["columns"][0]["name"] == "id"
|
||||
assert schema["columns"][0]["type"] == "BIGINT"
|
||||
assert schema["columns"][0]["nullable"] is False
|
||||
mock_ws_client.tables.get.assert_called_once_with(
|
||||
full_name="TEST_CATALOG.TEST_SCHEMA.customers"
|
||||
)
|
||||
|
||||
@patch("semantica.ingest.databricks_ingestor.DATABRICKS_AVAILABLE", True)
|
||||
@patch("semantica.ingest.databricks_ingestor.databricks_sql")
|
||||
@patch("semantica.ingest.databricks_ingestor.WorkspaceClient")
|
||||
def test_list_catalogs(self, mock_ws_client_cls, mock_sql):
|
||||
"""Test listing catalogs."""
|
||||
from semantica.ingest.databricks_ingestor import DatabricksIngestor
|
||||
|
||||
mock_catalog_1 = Mock()
|
||||
mock_catalog_1.name = "main"
|
||||
mock_catalog_2 = Mock()
|
||||
mock_catalog_2.name = "samples"
|
||||
|
||||
mock_ws_client = Mock()
|
||||
mock_ws_client.catalogs.list = Mock(return_value=[mock_catalog_1, mock_catalog_2])
|
||||
mock_ws_client_cls.return_value = mock_ws_client
|
||||
|
||||
ingestor = DatabricksIngestor(
|
||||
host="https://adb-xxx.azuredatabricks.net",
|
||||
token="test_token",
|
||||
http_path="/sql/1.0/warehouses/xxxx",
|
||||
)
|
||||
|
||||
catalogs = ingestor.list_catalogs()
|
||||
|
||||
assert len(catalogs) == 2
|
||||
assert "main" in catalogs
|
||||
|
||||
@patch("semantica.ingest.databricks_ingestor.DATABRICKS_AVAILABLE", True)
|
||||
@patch("semantica.ingest.databricks_ingestor.databricks_sql")
|
||||
@patch("semantica.ingest.databricks_ingestor.WorkspaceClient")
|
||||
def test_list_tables(self, mock_ws_client_cls, mock_sql):
|
||||
"""Test listing tables in a catalog/schema."""
|
||||
from semantica.ingest.databricks_ingestor import DatabricksIngestor
|
||||
|
||||
mock_table_1 = Mock()
|
||||
mock_table_1.name = "customers"
|
||||
mock_table_2 = Mock()
|
||||
mock_table_2.name = "orders"
|
||||
|
||||
mock_ws_client = Mock()
|
||||
mock_ws_client.tables.list = Mock(return_value=[mock_table_1, mock_table_2])
|
||||
mock_ws_client_cls.return_value = mock_ws_client
|
||||
|
||||
ingestor = DatabricksIngestor(
|
||||
host="https://adb-xxx.azuredatabricks.net",
|
||||
token="test_token",
|
||||
http_path="/sql/1.0/warehouses/xxxx",
|
||||
catalog="TEST_CATALOG",
|
||||
schema="TEST_SCHEMA",
|
||||
)
|
||||
|
||||
tables = ingestor.list_tables()
|
||||
|
||||
assert len(tables) == 2
|
||||
assert "customers" in tables
|
||||
mock_ws_client.tables.list.assert_called_once_with(
|
||||
catalog_name="TEST_CATALOG", schema_name="TEST_SCHEMA"
|
||||
)
|
||||
|
||||
@patch("semantica.ingest.databricks_ingestor.DATABRICKS_AVAILABLE", True)
|
||||
@patch("semantica.ingest.databricks_ingestor.databricks_sql")
|
||||
@patch("semantica.ingest.databricks_ingestor.WorkspaceClient")
|
||||
def test_list_tables_requires_schema(self, mock_ws_client_cls, mock_sql):
|
||||
"""Test that list_tables() raises instead of calling the SDK with schema_name=None."""
|
||||
from semantica.ingest.databricks_ingestor import DatabricksIngestor
|
||||
from semantica.utils.exceptions import ProcessingError
|
||||
|
||||
mock_ws_client = Mock()
|
||||
mock_ws_client_cls.return_value = mock_ws_client
|
||||
|
||||
ingestor = DatabricksIngestor(
|
||||
host="https://adb-xxx.azuredatabricks.net",
|
||||
token="test_token",
|
||||
http_path="/sql/1.0/warehouses/xxxx",
|
||||
catalog="TEST_CATALOG",
|
||||
)
|
||||
# Force a missing schema past the connector's "default" fallback.
|
||||
ingestor.connector.schema = None
|
||||
|
||||
with pytest.raises(ProcessingError, match="Schema name is required"):
|
||||
ingestor.list_tables()
|
||||
|
||||
mock_ws_client.tables.list.assert_not_called()
|
||||
|
||||
@patch("semantica.ingest.databricks_ingestor.DATABRICKS_AVAILABLE", True)
|
||||
@patch("semantica.ingest.databricks_ingestor.databricks_sql")
|
||||
@patch("semantica.ingest.databricks_ingestor.WorkspaceClient")
|
||||
def test_get_table_lineage(self, mock_ws_client_cls, mock_sql):
|
||||
"""Test getting table lineage."""
|
||||
from semantica.ingest.databricks_ingestor import DatabricksIngestor
|
||||
|
||||
mock_ws_client = Mock()
|
||||
mock_ws_client.api_client.do = Mock(
|
||||
return_value={
|
||||
"upstreams": [{"tableInfo": {"name": "main.default.raw_customers"}}],
|
||||
"downstreams": [{"tableInfo": {"name": "main.default.customer_summary"}}],
|
||||
}
|
||||
)
|
||||
mock_ws_client_cls.return_value = mock_ws_client
|
||||
|
||||
ingestor = DatabricksIngestor(
|
||||
host="https://adb-xxx.azuredatabricks.net",
|
||||
token="test_token",
|
||||
http_path="/sql/1.0/warehouses/xxxx",
|
||||
catalog="main",
|
||||
schema="default",
|
||||
)
|
||||
|
||||
lineage = ingestor.get_table_lineage("customers")
|
||||
|
||||
assert lineage["upstream"] == ["main.default.raw_customers"]
|
||||
assert lineage["downstream"] == ["main.default.customer_summary"]
|
||||
|
||||
@patch("semantica.ingest.databricks_ingestor.DATABRICKS_AVAILABLE", True)
|
||||
@patch("semantica.ingest.databricks_ingestor.databricks_sql")
|
||||
def test_export_as_documents(self, mock_sql):
|
||||
"""Test exporting data as documents."""
|
||||
from semantica.ingest.databricks_ingestor import DatabricksData, DatabricksIngestor
|
||||
|
||||
ingestor = DatabricksIngestor(
|
||||
host="https://adb-xxx.azuredatabricks.net",
|
||||
token="test_token",
|
||||
http_path="/sql/1.0/warehouses/xxxx",
|
||||
)
|
||||
|
||||
data = DatabricksData(
|
||||
data=[
|
||||
{"id": 1, "name": "Alice", "description": "Engineer"},
|
||||
{"id": 2, "name": "Bob", "description": "Designer"},
|
||||
],
|
||||
row_count=2,
|
||||
columns=["id", "name", "description"],
|
||||
table_name="employees",
|
||||
catalog="main",
|
||||
schema="default",
|
||||
)
|
||||
|
||||
documents = ingestor.export_as_documents(
|
||||
data, id_field="id", text_fields=["name", "description"]
|
||||
)
|
||||
|
||||
assert len(documents) == 2
|
||||
assert documents[0]["id"] == "1"
|
||||
assert documents[0]["text"] == "Alice Engineer"
|
||||
assert documents[0]["metadata"]["source"] == "databricks"
|
||||
assert documents[0]["metadata"]["table"] == "employees"
|
||||
|
||||
@patch("semantica.ingest.databricks_ingestor.DATABRICKS_AVAILABLE", True)
|
||||
@patch("semantica.ingest.databricks_ingestor.databricks_sql")
|
||||
def test_context_manager(self, mock_sql, mock_databricks_connection):
|
||||
"""Test using ingestor as context manager."""
|
||||
from semantica.ingest.databricks_ingestor import DatabricksIngestor
|
||||
|
||||
mock_conn, _ = mock_databricks_connection
|
||||
mock_sql.connect = Mock(return_value=mock_conn)
|
||||
|
||||
with DatabricksIngestor(
|
||||
host="https://adb-xxx.azuredatabricks.net",
|
||||
token="test_token",
|
||||
http_path="/sql/1.0/warehouses/xxxx",
|
||||
) as ingestor:
|
||||
assert ingestor.connector.connection == mock_conn
|
||||
|
||||
mock_conn.close.assert_called()
|
||||
|
||||
@patch("semantica.ingest.databricks_ingestor.DATABRICKS_AVAILABLE", True)
|
||||
@patch("semantica.ingest.databricks_ingestor.databricks_sql")
|
||||
def test_convert_datetime(self, mock_sql):
|
||||
"""Test datetime conversion in _convert_rows."""
|
||||
from semantica.ingest.databricks_ingestor import DatabricksIngestor
|
||||
|
||||
ingestor = DatabricksIngestor(
|
||||
host="https://adb-xxx.azuredatabricks.net",
|
||||
token="test_token",
|
||||
http_path="/sql/1.0/warehouses/xxxx",
|
||||
)
|
||||
|
||||
test_dt = datetime(2024, 1, 15, 10, 30, 0)
|
||||
rows = [{"timestamp": test_dt, "value": 100}]
|
||||
|
||||
converted = ingestor._convert_rows(rows)
|
||||
|
||||
assert converted[0]["timestamp"] == "2024-01-15T10:30:00"
|
||||
assert converted[0]["value"] == 100
|
||||
|
||||
def test_import_error_without_databricks(self):
|
||||
"""Test that proper error is raised when databricks libraries are not installed."""
|
||||
with patch("semantica.ingest.databricks_ingestor.DATABRICKS_AVAILABLE", False):
|
||||
from semantica.ingest.databricks_ingestor import DatabricksConnector
|
||||
|
||||
with pytest.raises(ImportError, match="databricks-sdk"):
|
||||
DatabricksConnector(
|
||||
host="https://adb-xxx.azuredatabricks.net",
|
||||
token="test_token",
|
||||
)
|
||||
|
||||
|
||||
class TestDatabricksData:
|
||||
"""Test DatabricksData dataclass."""
|
||||
|
||||
def test_databricks_data_creation(self):
|
||||
"""Test DatabricksData creation."""
|
||||
from semantica.ingest.databricks_ingestor import DatabricksData
|
||||
|
||||
data = DatabricksData(
|
||||
data=[{"col1": "val1"}],
|
||||
row_count=1,
|
||||
columns=["col1"],
|
||||
table_name="test_table",
|
||||
)
|
||||
|
||||
assert data.row_count == 1
|
||||
assert data.table_name == "test_table"
|
||||
assert len(data.data) == 1
|
||||
assert isinstance(data.ingested_at, datetime)
|
||||
|
||||
def test_databricks_data_with_metadata(self):
|
||||
"""Test DatabricksData with metadata."""
|
||||
from semantica.ingest.databricks_ingestor import DatabricksData
|
||||
|
||||
metadata = {"custom_field": "value"}
|
||||
data = DatabricksData(
|
||||
data=[],
|
||||
row_count=0,
|
||||
columns=[],
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
assert data.metadata["custom_field"] == "value"
|
||||
Reference in New Issue
Block a user