mirror of
https://github.com/semantica-agi/semantica.git
synced 2026-08-29 04:26:20 +00:00
Merge pull request #748 from semantica-agi/feat/747-databricks-connector
Add Databricks connector (Unity Catalog + Delta Lake ingestion)
This commit is contained in:
Mohd Kaif
@@ -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`; each call closes the SQL connection it opened unless one is already open (e.g. via the `with DatabricksIngestor(...)` context manager), which reuses and closes it exactly once instead of leaking a second connection per call
|
||||
- `get_table_schema()`, `list_catalogs()`, `list_schemas()`, and `list_tables()` introspect Unity Catalog via `databricks-sdk`'s `WorkspaceClient`, validating both catalog and schema are resolved before calling the SDK; `get_table_lineage()` calls Unity Catalog's table-lineage REST API for upstream/downstream `Table --DEPENDS_ON--> Table` dependencies, plus an opt-in `include_column_lineage=True` that resolves per-column lineage via the column-lineage API
|
||||
- `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
|
||||
- 35 unit tests in `tests/test_databricks_ingestor.py` covering both auth methods, table/query ingestion, connection lifecycle (including reuse under the context manager), pagination, unsafe `ORDER BY` rejection, catalog/schema validation, schema/catalog/table listing, table and column lineage, document export, 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,199 @@
|
||||
---
|
||||
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 and column 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.
|
||||
|
||||
<Tip>
|
||||
Pass `include_column_lineage=True` to also resolve per-column upstream/downstream references (one extra Unity Catalog request per column, so it's opt-in):
|
||||
|
||||
```python
|
||||
lineage = ingestor.get_table_lineage(
|
||||
"customers", catalog="main", schema="default", include_column_lineage=True,
|
||||
)
|
||||
print(lineage["columns"]["email"])
|
||||
# {"upstream": ["main.default.raw_customers.email_address"], "downstream": []}
|
||||
```
|
||||
|
||||
</Tip>
|
||||
|
||||
|
||||
## 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",
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,912 @@
|
||||
"""
|
||||
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.oauth_service_principal")
|
||||
@patch("semantica.ingest.databricks_ingestor.Config")
|
||||
@patch("semantica.ingest.databricks_ingestor.databricks_sql")
|
||||
def test_connector_connect_oauth_m2m(
|
||||
self, mock_sql, mock_config, mock_oauth_sp, mock_databricks_connection
|
||||
):
|
||||
"""Test connection with OAuth M2M authentication uses credentials_provider.
|
||||
|
||||
databricks-sql-connector does NOT accept client_id/client_secret as
|
||||
direct kwargs to sql.connect(); the correct mechanism is a
|
||||
credentials_provider callable wrapping oauth_service_principal().
|
||||
"""
|
||||
from semantica.ingest.databricks_ingestor import DatabricksConnector
|
||||
|
||||
mock_conn, _ = mock_databricks_connection
|
||||
mock_sql.connect = Mock(return_value=mock_conn)
|
||||
mock_oauth_sp.return_value = {"Authorization": "Bearer fake-token"}
|
||||
|
||||
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]
|
||||
|
||||
# Must use credentials_provider, not bare client_id/client_secret
|
||||
assert "credentials_provider" in call_kwargs
|
||||
assert callable(call_kwargs["credentials_provider"])
|
||||
assert "client_id" not in call_kwargs
|
||||
assert "client_secret" not in call_kwargs
|
||||
assert "access_token" not in call_kwargs
|
||||
|
||||
# Invoke the provider to verify it wires Config + oauth_service_principal
|
||||
call_kwargs["credentials_provider"]()
|
||||
mock_config.assert_called_once_with(
|
||||
host="https://adb-xxx.azuredatabricks.net",
|
||||
client_id="test_client_id",
|
||||
client_secret="test_client_secret",
|
||||
)
|
||||
mock_oauth_sp.assert_called_once_with(mock_config.return_value)
|
||||
|
||||
@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_get_table_schema_requires_schema(self, mock_ws_client_cls, mock_sql):
|
||||
"""Test that get_table_schema() raises instead of calling the SDK with a missing schema."""
|
||||
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",
|
||||
)
|
||||
ingestor.connector.schema = None
|
||||
|
||||
with pytest.raises(ProcessingError, match="Schema name is required"):
|
||||
ingestor.get_table_schema("customers")
|
||||
|
||||
mock_ws_client.tables.get.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_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"]
|
||||
assert "columns" not in lineage
|
||||
|
||||
@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_with_column_lineage(self, mock_ws_client_cls, mock_sql):
|
||||
"""Test that include_column_lineage=True fetches per-column lineage."""
|
||||
from semantica.ingest.databricks_ingestor import DatabricksIngestor
|
||||
|
||||
mock_column_1 = Mock()
|
||||
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()
|
||||
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]
|
||||
|
||||
def do_side_effect(method, path, query=None):
|
||||
if path.endswith("table-lineage"):
|
||||
return {
|
||||
"upstreams": [{"tableInfo": {"name": "main.default.raw_customers"}}],
|
||||
"downstreams": [],
|
||||
}
|
||||
if path.endswith("column-lineage"):
|
||||
if query["column_name"] == "id":
|
||||
return {
|
||||
"upstream_cols": [
|
||||
{
|
||||
"catalog_name": "main",
|
||||
"schema_name": "default",
|
||||
"table_name": "raw_customers",
|
||||
"name": "customer_id",
|
||||
}
|
||||
],
|
||||
"downstream_cols": [],
|
||||
}
|
||||
return {"upstream_cols": [], "downstream_cols": []}
|
||||
raise AssertionError(f"unexpected path: {path}")
|
||||
|
||||
mock_ws_client = Mock()
|
||||
mock_ws_client.tables.get = Mock(return_value=mock_table_info)
|
||||
mock_ws_client.api_client.do = Mock(side_effect=do_side_effect)
|
||||
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", include_column_lineage=True)
|
||||
|
||||
assert lineage["upstream"] == ["main.default.raw_customers"]
|
||||
assert lineage["columns"]["id"]["upstream"] == ["main.default.raw_customers.customer_id"]
|
||||
assert lineage["columns"]["id"]["downstream"] == []
|
||||
assert lineage["columns"]["name"]["upstream"] == []
|
||||
|
||||
@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_requires_schema(self, mock_ws_client_cls, mock_sql):
|
||||
"""Test that get_table_lineage() validates schema before calling the REST API."""
|
||||
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="main",
|
||||
)
|
||||
ingestor.connector.schema = None
|
||||
|
||||
with pytest.raises(ProcessingError, match="Schema name is required"):
|
||||
ingestor.get_table_lineage("customers")
|
||||
|
||||
mock_ws_client.api_client.do.assert_not_called()
|
||||
|
||||
@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_context_manager_reuses_connection_across_calls(
|
||||
self, mock_sql, mock_databricks_connection
|
||||
):
|
||||
"""Test that ingest_table()/ingest_query() reuse (not leak) the connection
|
||||
opened by __enter__, and only close it once on __exit__."""
|
||||
from semantica.ingest.databricks_ingestor import DatabricksIngestor
|
||||
|
||||
mock_conn, mock_cursor = 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:
|
||||
ingestor.ingest_table("customers")
|
||||
# The connection opened by __enter__ must still be the live one:
|
||||
# ingest_table() should not have closed it out from under the
|
||||
# context manager.
|
||||
assert ingestor.connector.connection == mock_conn
|
||||
mock_conn.close.assert_not_called()
|
||||
|
||||
ingestor.ingest_query("SELECT * FROM sales")
|
||||
assert ingestor.connector.connection == mock_conn
|
||||
mock_conn.close.assert_not_called()
|
||||
|
||||
# databricks_sql.connect() should only have been called once (by
|
||||
# __enter__): both ingestion calls reused that same connection.
|
||||
mock_sql.connect.assert_called_once()
|
||||
mock_conn.close.assert_called_once()
|
||||
|
||||
@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