Bridging the Lakehouse: Ingesting Databricks Unity Catalog into Cognee Knowledge Graphs

In large enterprise lakehouses, finding the right dataset, understanding table schemas, and tracing ML model lineage is often a major hurdle for data scientists, product managers, and AI agents.

While catalog tools like Databricks Unity Catalog centralize table schemas, storage locations, and column-level comments, interacting with these assets via natural language requires a structured knowledge layer. Vector similarity alone fails when users ask questions that span organizational boundaries:

• "Which production tables contain customer billing information, and what are their primary keys?"
• "What AI/ML models are registered under champion status for credit risk prediction?"
• "Show me all schemas in the analytics catalog that partition by event_date."

To solve this, I designed and implemented the official Databricks Unity Catalog data-source connector (cognee-community-connector-unitycatalog) for Cognee as part of the WeMakeDevs × Cognee Mergetober Hackathon.

Here is how we turn Unity Catalog metadata into a connected GraphRAG layer.

---

Architectural Blueprint

The connector traverses Unity Catalog’s three-level namespace hierarchy (catalog.schema.table) along with the AI/ML Model Registry, converting tabular definitions into rich semantic markdown nodes:

flowchart TD
    subgraph UnityCatalog ["Databricks Unity Catalog API"]
        A["/api/2.1/unity-catalog/catalogs"]
        B["/api/2.1/unity-catalog/schemas"]
        C["/api/2.1/unity-catalog/tables (Managed, External, Views)"]
        D["/api/2.1/unity-catalog/models (Registered AI Models)"]
    end

    subgraph ConnectorEngine ["Connector Ingestion & Normalization"]
        E["Selective Scope Filter (catalogs / schemas)"]
        F["Column Schema Serializer (Types, Comments, Partitions)"]
        G["Model Registry Metadata (Aliases, Runs, Tags)"]
        H["Credential Scrubbing (Databricks PAT & Bearer Redaction)"]
    end

    subgraph CogneeGraph ["Cognee GraphRAG Ecosystem"]
        I["dlt Source (DOCUMENT_SOURCE_ATTR='unitycatalog')"]
        J["Cognify Entity Extractor"]
        K["Lakehouse Knowledge Graph (FalkorDB / NetworkX)"]
        L["Vector Index (Qdrant)"]
        M["Forget-on-Delete Orphan Pruner"]
    end

    A & B & C & D --> E --> F & G --> H --> I
    I --> J --> K & L
    I --> M

---

1. Modeling the Three-Level Namespace & AI Models

Databricks Unity Catalog organizes assets in a strict hierarchy: catalog $\rightarrow$ schema $\rightarrow$ table/view. The connector extracts four distinct resource tiers:

A. Catalogs (unitycatalog_catalogs)

Captures top-level governance containers, storage roots, catalog owners, and business descriptions.

B. Schemas (unitycatalog_schemas)

Captures functional namespaces (e.g. finance, marketing, ml_feature_store), properties, and regional metadata.

C. Tables & Views with Rich Column Types (unitycatalog_tables)

Tables are transformed into structured markdown documents detailing table types (MANAGED, EXTERNAL, VIEW), underlying data formats (DELTA, PARQUET), and complete column schemas:

# Table: production.finance.settled_transactions
- **Catalog:** production | **Schema:** finance
- **Table Type:** MANAGED (DELTA)
- **Owner:** finance-data-eng
- **Storage Location:** s3://company-lakehouse/finance/settled_transactions

### Columns
| Column Name | Type | Nullable | Comment |
| :--- | :--- | :--- | :--- |
| `transaction_id` | STRING | No | Unique payment transaction identifier |
| `account_id` | STRING | No | Foreign key to core accounts |
| `amount_cents` | BIGINT | No | Settled amount in cents |
| `settled_at` | TIMESTAMP | No | Monotonic settlement timestamp |

D. AI/ML Model Registry (unitycatalog_models)

Unity Catalog registers MLflow models directly within schemas. The connector ingests registered model names, current aliases (e.g., champion, challenger), model versions, run artifacts, and tags.

---

2. Dynamic Scoping & Forget-on-Delete Invariants

Enterprise data lakehouses often host thousands of tables. Ingesting an entire data lake at once can create unnecessary overhead. The connector supports granular scoping:

source = unitycatalog_source(
    catalogs=["production", "gold_analytics"],
    schemas=["finance", "risk_modeling"],
    include_tables=True,
    include_models=True,
)

Automatic Lifecycle Management ("Forget-on-Delete")

When data engineers decommission legacy tables or drop experimental schemas, stale nodes shouldn't persist in company knowledge graphs:

setattr(source, DOCUMENT_SOURCE_ATTR, "unitycatalog")

@dlt.resource(name="unitycatalog_tables", write_disposition="replace")
def tables_resource():
    ...

By tagging DOCUMENT_SOURCE_ATTR = "unitycatalog" and write_disposition="replace", Cognee's internal diff engine detects removed tables upon each resync. Removed tables, columns, and model entries are automatically pruned from FalkorDB/NetworkX and Qdrant, guaranteeing 100% data consistency.

---

3. End-to-End Implementation

Here is how easily you can connect Databricks Unity Catalog to Cognee to query enterprise governance:

import asyncio
import os
import cognee
from cognee_community_connector_unitycatalog import unitycatalog_source

async def main():
    # 1. Connect to Databricks Workspace or Unity Catalog OSS
    source = unitycatalog_source(
        base_url=os.environ.get("DATABRICKS_HOST", "https://adb-1234.azuredatabricks.net"),
        token=os.environ["DATABRICKS_TOKEN"],
        catalogs=["production"],
        include_tables=True,
        include_models=True,
    )

    # 2. Ingest into Cognee and build Knowledge Graph
    print("Ingesting Unity Catalog lakehouse metadata...")
    await cognee.add(source)

    print("Running Cognify to construct Governance Graph...")
    await cognee.cognify()

    # 3. Ask natural language questions across schemas & models
    query = "What columns exist in the production finance tables, and what is our champion credit model?"
    print(f"\nQuerying Cognee: '{query}'")

    results = await cognee.search(query)
    for res in results:
        print(f"\n[Governance Answer]:\n{res}")

if __name__ == "__main__":
    asyncio.run(main())

---

Verification & Standards

• 100% Offline Testing: 18 comprehensive unit tests pass offline using mocked API payloads (tests/test_unitycatalog.py).
• Security Guardrails: Pre-ingestion credential masking strips Databricks tokens from all logs and error traces.
• Maintainer Alignment: Fully compliant with Cognee’s DCO sign-off guidelines and Hatchling packaging conventions.

---

Pull Request & Resources

• Pull Request: topoteretes/cognee-community#472
• Repository: topoteretes/cognee

How is your organization approaching conversational governance across Databricks and Iceberg catalogs? Share your feedback below!