# cognite-databricks
A helper SDK for Databricks that provides Unity Catalog SQL UDTF registration utilities, Secret Manager integration, and Databricks-specific tooling for scalar UDTFs.
**Latest Release:**
- **Version 0.2.3**: Aligns Unity Catalog view registration with **cognite-pygen-spark** 0.2.3+ for reserved-word safe `UDTFField` naming.
- **Dependencies**: Now target **cognite-pygen-spark** 0.4.2+ for:
- Grouped aggregate rows that copy `space` and `externalId` onto the result.
- Shared pushdown parameter registry and a trailing `base_url` read at query time.
- WHERE / LIMIT / COUNT / MIN / MAX pushdown in generated data-model UDTFs.
- CDF audit headers in generated UDTFs.
- TypeConverter-based UDTF field typing.
- Timestamp normalization.
- Security-patched transitive dependencies (lockfile; includes `pyasn1>=0.6.3`).
- **Version 0.2.1**: Added SQL-native time series UDTF support with predicate pushdown hints and a SQL query analyzer for pushdown hints.
Full release notes are published on [GitHub Releases](https://github.com/cognitedata/cognite-databricks/releases). Patch releases may update locked dependencies and minimum constraints without API changes.
**Upgrade:** views and UDTFs registered before the trailing `base_url` argument must be regenerated and re-registered with `if_exists="replace"`. See [Upgrading views registered before query-time base_url](docs/catalog_based/troubleshooting.md#upgrading-views-registered-before-query-time-base_url).
**Note**: This package provides Databricks-specific utilities for Unity Catalog UDTF registration and Secret Manager integration.
## Overview
`cognite-databricks` is a **Databricks-specific helper SDK** that extends `pygen-spark` with Unity Catalog SQL registration, Secret Manager integration, and Databricks-specific utilities. It focuses on serverless-compatible scalar UDTF execution for SQL Warehouses.
**Package Purpose:**
- **Databricks-Specific Features**: Unity Catalog SQL registration, Secret Manager integration, and Databricks-specific utilities
- **Uses pygen-spark for Code Generation**: All UDTF code generation (both Data Model and Time Series UDTFs) is done by `pygen-spark` using template-based generation
- **Generic Components**: Generic utilities (`TypeConverter`, `CDFConnectionConfig`, `to_udtf_function_name`) are provided by `pygen-spark` and re-exported from `cognite.databricks` for backward compatibility
- **Notebook-Friendly API**: Aligned with `cognite.pygen`'s notebook workflow
It provides high-level APIs for:
- **UDTF Registration**: Register persistent UDTFs in Unity Catalog via SQL
- **Secret Manager Integration**: Manage OAuth2 credentials securely
- **SQL Usage**: Use UDTFs directly in SQL after registration
## Features
- **Unity Catalog SQL Registration**: Serverless-compatible UDTFs registered via `CREATE FUNCTION` statements
- **One-Line Registration**: Generate and register UDTFs in a single call
- **Secret Manager Integration**: Automatic credential management from TOML files, including `base_url` for Private Link and dedicated clusters
- **Scalar-Only Execution**: Compatible with SQL Warehouses and serverless execution
- **Catalog SQL rewrite helper**: `generator.rewrite_query()` binds WHERE / LIMIT / GROUP BY / aggregates and passes every UDTF parameter (`NULL` when unbound). `OR`, subqueries, and quoted identifiers are left unchanged.
- **SQL functions**: `generator.create_sql_function()` turns one view query into a Unity Catalog function for Power BI. Literals become arguments. Credentials stay `SECRET()` references. A statement that cannot be pushed raises and creates nothing.
- **Type Safety**: Full type hints and IDE support
- **Generic Components**: Uses template-generated UDTFs and generic utilities (`TypeConverter`, `CDFConnectionConfig`, `to_udtf_function_name`) from `cognite-pygen-spark` for generic Spark compatibility. These components are re-exported from `cognite.databricks` for backward compatibility, but the source is `cognite.pygen_spark`.
## Installation
```bash
pip install cognite-databricks
```
## Start here (recommended)
**Use the catalog-based quickstart** to get productive: Unity Catalog + Secret Manager + your CDF data model.
| Resource | What it is |
|----------|------------|
| **[docs/catalog_based/quickstart.md](docs/catalog_based/quickstart.md)** | Step-by-step guide with explained code blocks |
| **[quickstart.ipynb](https://github.com/cognitedata/cognite-databricks/blob/main/examples/catalog_based/quickstart.ipynb)** | **Notebook version** โ same flow with markdown introductions and **inline comments** in every code cell |
**Prerequisites:** [docs/catalog_based/prerequisites.md](docs/catalog_based/prerequisites.md) (Unity Catalog, Secret Manager, CDF TOML).
**Full documentation index:** [docs/index.md](docs/index.md).
The quickstart walks through: **install** โ **CDF + Databricks clients** โ **pick a SQL warehouse** โ **`generate_udtf_notebook`** โ **write CDF secrets to Secret Manager** โ **`register_udtfs`** โ **`register_views`**.
## Quick start (short reference)
Catalog-based registration (Unity Catalog) in a few lines โ for the full narrative, use the quickstart links above.
```python
from cognite.client.data_classes.data_modeling.ids import DataModelId
from cognite.databricks import generate_udtf_notebook
from cognite.pygen import load_cognite_client_from_toml
from databricks.sdk import WorkspaceClient
client = load_cognite_client_from_toml("/Workspace/Users/<you>/config.toml")
workspace_client = WorkspaceClient()
warehouses = list(workspace_client.warehouses.list())
warehouse = warehouses[0] # pick the warehouse you use for SQL
data_model_id = DataModelId(space="my_space", external_id="MyModel", version="v1")
generator = generate_udtf_notebook(
data_model_id,
client,
workspace_client=workspace_client,
output_dir="/Workspace/Users/<you>/udtf_generated",
catalog="my_catalog",
schema="my_schema",
warehouse_id=warehouse.id,
)
# Then: secret scope + set_cdf_credentials + register_udtfs + register_views
# (see quickstart โ do not skip warehouse selection or Secret Manager)
```
### Low-Level API
```python
from cognite.client.data_classes.data_modeling.ids import DataModelId
from cognite.databricks import UDTFGenerator, SecretManagerHelper
from cognite.pygen import load_cognite_client_from_toml
# Load client from TOML file
client = load_cognite_client_from_toml("config.toml")
# Create generator
generator = UDTFGenerator(
cognite_client=client,
catalog="main",
schema="cdf_models",
)
# Set up Secret Manager (one-time setup per data model)
data_model_id = DataModelId(space="sp_pygen_power", external_id="WindTurbine", version="1")
secret_scope = f"cdf_{data_model_id.space}_{data_model_id.external_id.lower()}"
generator.secret_helper.set_cdf_credentials(
scope_name=secret_scope,
project="my-project", # from config.toml
cdf_cluster="api.cognitedata.com", # from config.toml
client_id="...", # from config.toml
client_secret="...", # from config.toml
tenant_id="...", # from config.toml
base_url=None, # from config.toml when set (Private Link / dedicated); defaults to the public cluster URL
)
# Register UDTFs for catalog-based use (scalar-only)
registered = generator.register_session_scoped_udtfs()
```
## Unity Catalog UDTF Registration
Session-scoped registration is the primary mode for scalar-only UDTFs. Register functions in the current Spark session before running SQL queries.
### Register All UDTFs (Recommended)
```python
from cognite.databricks import generate_udtf_notebook
from cognite.pygen import load_cognite_client_from_toml
from cognite.client.data_classes.data_modeling.ids import DataModelId
# Load client and generate UDTFs
client = load_cognite_client_from_toml("config.toml")
data_model_id = DataModelId(space="sailboat", external_id="sailboat", version="v1")
generator = generate_udtf_notebook(
data_model_id,
client,
output_dir="/Workspace/Users/user@example.com/udtf",
)
# Install dependencies (run in separate cell first)
# %pip install cognite-sdk
# (Restart kernel after installation)
# Register all UDTFs for catalog-based use (includes time series UDTFs automatically)
registered = generator.register_session_scoped_udtfs()
# Returns: {"SmallBoat": "small_boat_udtf", "LargeBoat": "large_boat_udtf",
# "time_series_datapoints": "time_series_datapoints_udtf", ...}
# Use in SQL (always use SECRET() for credentials)
# SELECT * FROM small_boat_udtf(
# client_id => SECRET('cdf_sailboat_sailboat', 'client_id'),
# client_secret => SECRET('cdf_sailboat_sailboat', 'client_secret'),
# tenant_id => SECRET('cdf_sailboat_sailboat', 'tenant_id'),
# cdf_cluster => SECRET('cdf_sailboat_sailboat', 'cdf_cluster'),
# project => SECRET('cdf_sailboat_sailboat', 'project'),
# name => 'MyBoat',
# description => NULL
# ) LIMIT 10;
```
### Register Single UDTF
```python
from cognite.databricks import generate_udtf_notebook, register_udtf_from_file
# Generate UDTFs
generator = generate_udtf_notebook(data_model_id, client, ...)
# Register a single UDTF from generated file
register_udtf_from_file(
"/Workspace/Users/user@example.com/udtf/sailboat_sailboat_v1/SmallBoat_udtf.py",
function_name="small_boat_udtf"
)
```
### Session Scope Notes
Session-scoped registration is the supported mode for scalar-only UDTFs. Functions are temporary and must be registered at the start of each notebook/job before running SQL queries.
## Requirements
**Note:** This document uses PyPI package names for references:
- **PyPI:** `cognite-pygen` (repository: `pygen`; import: `cognite.pygen`)
- **PyPI:** `cognite-pygen-spark` (repository: `pygen-spark`; import: `cognite.pygen_spark`)
- Python 3.9+
- `cognite-pygen-spark` (PyPI package name; import: `cognite.pygen_spark`)
- `cognite-sdk-python` (dependency)
- `databricks-sdk` (dependency)
## Package Structure
```
cognite-databricks/
โโโ cognite/
โ โโโ databricks/
โ โโโ __init__.py # Exports generate_udtf_notebook, UDTFGenerator, etc.
โ โโโ udtf_registry.py # UDTF registration helpers
โ โโโ secret_manager.py # Secret Manager helpers
โ โโโ view_generator.py # View generation and registration
โ โโโ generator.py # generate_udtf_notebook helper function
โ โโโ utils.py # Utility functions
โโโ pyproject.toml
โโโ README.md
```
## Core Components
### `generate_udtf_notebook`
High-level function for notebook workflows, aligned with `pygen.generate_sdk_notebook`:
```python
from cognite.databricks import generate_udtf_notebook
generator = generate_udtf_notebook(
data_model_id,
client,
catalog="main",
schema="cdf_models",
)
```
### `UDTFGenerator`
Main class for orchestrating UDTF generation and registration:
```python
from cognite.databricks import UDTFGenerator
generator = UDTFGenerator(
workspace_client=workspace_client,
cognite_client=client,
catalog="main",
schema="cdf_models",
)
```
**Key Methods:**
- `register_session_scoped_udtfs()`: Register UDTFs for catalog-based use (scalar-only)
- `register_udtf_from_file()`: Register a single generated UDTF file in the current session
### `register_udtf_from_file`
Standalone function for registering a single UDTF from a generated Python file for catalog-based use:
```python
from cognite.databricks import register_udtf_from_file
register_udtf_from_file(
"/path/to/SmallBoat_udtf.py",
function_name="small_boat_udtf"
)
```
### `register_udtf_from_file`
Standalone function for registering a single UDTF from a generated Python file for catalog-based use. Useful when you only need to register one UDTF or want more control over the registration process.
```python
from cognite.databricks import register_udtf_from_file
register_udtf_from_file(
"/path/to/SmallBoat_udtf.py",
function_name="small_boat_udtf"
)
```
### `SecretManagerHelper`
Helper for managing OAuth2 credentials in Databricks Secret Manager:
```python
from cognite.databricks import SecretManagerHelper
secret_helper = SecretManagerHelper(workspace_client)
secret_helper.set_cdf_credentials(
scope_name="cdf_sp_pygen_power_windturbine",
project="my-project",
cdf_cluster="api.cognitedata.com",
client_id="...",
client_secret="...",
tenant_id="...",
base_url=None, # Private Link / dedicated API URL; defaults to the public cluster URL
)
```
## Development
### Setup
```bash
git clone <repository-url>
cd cognite-databricks
pip install -e ".[dev]"
```
### Running Tests
```bash
pytest tests/
```
## Related Packages
- **[cognite-pygen-spark](https://github.com/cognitedata/pygen-spark)** (PyPI: `cognite-pygen-spark`): Generic Spark UDTF code generation library that works with any Spark cluster. Provides template-based UDTF generation, type conversion utilities (`TypeConverter`), connection configuration (`CDFConnectionConfig`), and utility functions. `cognite-databricks` uses `pygen-spark` for all code generation.
- **[cognite-pygen](https://github.com/cognitedata/pygen)** (PyPI: `cognite-pygen`): Base code generation library for CDF Data Models
- **[cognite-sdk-python](https://github.com/cognitedata/cognite-sdk-python)**: Python SDK for CDF APIs
### Import Paths for Generic Components
Generic components (`TypeConverter`, `CDFConnectionConfig`, `to_udtf_function_name`) are provided by `pygen-spark` and re-exported from `cognite-databricks` for backward compatibility:
```python
# Preferred: Import directly from pygen-spark (source)
from cognite.pygen_spark import TypeConverter, CDFConnectionConfig, to_udtf_function_name
# Backward compatible: Still works (re-exported from pygen-spark)
from cognite.databricks import TypeConverter, CDFConnectionConfig, to_udtf_function_name
```
**Note**: These components are generic Spark utilities and work with any Spark cluster, not just Databricks. They were moved from `cognite-databricks` to `pygen-spark` to make them available for standalone Spark clusters.
## Documentation
- **[Start here โ catalog quickstart](docs/catalog_based/quickstart.md)** and **[quickstart notebook](https://github.com/cognitedata/cognite-databricks/blob/main/examples/catalog_based/quickstart.ipynb)**
- **[Documentation index](docs/index.md)**: full catalog-based and session-scoped guides
- **[Session-scoped (development)](docs/session_scoped/index.md)**: temporary session UDTFs
- **Technical plan** (if present in your workspace): CDF Databricks integration (UDTF-based) architecture
## License
[License information]
## Contributing
[Contributing guidelines]