104 lines
3.8 KiB
Python
104 lines
3.8 KiB
Python
import json
|
|
from types import SimpleNamespace
|
|
|
|
from tht.mschema.catalog_snapshot import load_catalog_metadata_snapshot
|
|
from tht.mschema.context import load_schema_context
|
|
from tht.search import SearchResult, schema_tables
|
|
from tht.vectorstore.records import catalog_schema_records, qdrant_semantic_kind
|
|
|
|
|
|
def _snapshot():
|
|
return {
|
|
"schemaVersion": 1,
|
|
"workspaceId": "sales",
|
|
"databaseId": "db-1",
|
|
"databaseName": "warehouse",
|
|
"schemaName": "analytics",
|
|
"metadataContentRevision": 7,
|
|
"tables": [
|
|
{
|
|
"id": "table-orders",
|
|
"name": "orders",
|
|
"description": "Customer orders",
|
|
"descriptionSource": "curated",
|
|
"columns": [
|
|
{
|
|
"id": "column-customer",
|
|
"name": "customer_id",
|
|
"ordinalPosition": 1,
|
|
"dataType": "uuid",
|
|
"isNullable": False,
|
|
"defaultExpression": None,
|
|
"primaryKeyPosition": None,
|
|
"sensitive": False,
|
|
"description": "Ordering customer",
|
|
"descriptionSource": "generated",
|
|
}
|
|
],
|
|
},
|
|
{
|
|
"id": "table-customers",
|
|
"name": "customers",
|
|
"description": None,
|
|
"descriptionSource": None,
|
|
"columns": [
|
|
{
|
|
"id": "column-id",
|
|
"name": "id",
|
|
"ordinalPosition": 1,
|
|
"dataType": "uuid",
|
|
"isNullable": False,
|
|
"defaultExpression": None,
|
|
"primaryKeyPosition": 1,
|
|
"sensitive": True,
|
|
"description": "Customer identifier",
|
|
"descriptionSource": "source_comment",
|
|
}
|
|
],
|
|
},
|
|
],
|
|
"relationships": [
|
|
{
|
|
"id": "relationship-1",
|
|
"origin": "physical",
|
|
"sourceTable": "orders",
|
|
"sourceColumns": ["customer_id"],
|
|
"targetTable": "customers",
|
|
"targetColumns": ["id"],
|
|
}
|
|
],
|
|
}
|
|
|
|
|
|
def test_catalog_snapshot_is_the_schema_context_and_vector_source(tmp_path):
|
|
path = tmp_path / "catalog-metadata.json"
|
|
path.write_text(json.dumps(_snapshot()))
|
|
cfg = SimpleNamespace(
|
|
paths=SimpleNamespace(catalog_metadata_snapshot=path),
|
|
_workspace_id="sales",
|
|
)
|
|
|
|
context = load_schema_context(cfg)
|
|
assert context.physical.database == "warehouse"
|
|
assert context.physical.tables["orders"].comment == "Customer orders"
|
|
assert context.physical.tables["customers"].columns["id"].eligibility_reason == "sensitive"
|
|
assert context.effective_relationships["orders"][0].ref_table == "customers"
|
|
|
|
snapshot = load_catalog_metadata_snapshot(path, "sales")
|
|
records = catalog_schema_records(snapshot)
|
|
assert [record.kind for record in records].count("schema_relationship") == 1
|
|
relationship = next(record for record in records if record.kind == "schema_relationship")
|
|
assert relationship.metadata["tables"] == ["orders", "customers"]
|
|
assert qdrant_semantic_kind("schema_relationship") == "schema"
|
|
|
|
|
|
def test_relationship_hit_promotes_both_endpoint_tables():
|
|
result = SearchResult(
|
|
key="relationship:orders->customers",
|
|
label="orders->customers",
|
|
kind="schema_relationship",
|
|
signals={"vector": {"rank": 1, "score": 0.9}},
|
|
rrf=0.5,
|
|
)
|
|
assert schema_tables([result], 10) == [("customers", 0.5), ("orders", 0.5)]
|