databao-context-engine 0.1.6__py3-none-any.whl → 0.1.7__py3-none-any.whl

databao_context_engine/__init__.py

@@ -23,7 +23,12 @@ from databao_context_engine.pluginlib.build_plugin import (
     DatasourceType,
 )
 from databao_context_engine.pluginlib.config import ConfigPropertyDefinition
-from databao_context_engine.project.info import DceInfo, DceProjectInfo, get_databao_context_engine_info
+from databao_context_engine.project.info import (
+    DceInfo,
+    DceProjectInfo,
+    get_databao_context_engine_info,
+    get_databao_context_engine_project_info,
+)
 from databao_context_engine.project.init_project import InitErrorReason, InitProjectError
 from databao_context_engine.services.chunk_embedding_service import ChunkEmbeddingMode
 
@@ -41,6 +46,7 @@ __all__ = [
     "DatasourceConfigFile",
     "DatasourceType",
     "get_databao_context_engine_info",
+    "get_databao_context_engine_project_info",
     "DceInfo",
     "DceProjectInfo",
     "init_dce_project",

databao_context_engine/cli/info.py

@@ -4,18 +4,21 @@ from pathlib import Path
 
 import click
 
-from databao_context_engine import DceInfo, get_databao_context_engine_info
+from databao_context_engine import DceInfo, DceProjectInfo, get_databao_context_engine_info
+from databao_context_engine.project.info import get_databao_context_engine_project_info
 
 
 def echo_info(project_dir: Path) -> None:
-    click.echo(_generate_info_string(get_databao_context_engine_info(project_dir=project_dir)))
+    dce_info = get_databao_context_engine_info()
+    dce_project_info = get_databao_context_engine_project_info(project_dir=project_dir)
+    click.echo(_generate_info_string(dce_info, dce_project_info))
 
 
-def _generate_info_string(command_info: DceInfo) -> str:
+def _generate_info_string(dce_info: DceInfo, project_info: DceProjectInfo) -> str:
     info_lines = []
-    info_lines.append(f"Databao context engine version: {command_info.version}")
-    info_lines.append(f"Databao context engine storage dir: {command_info.dce_path}")
-    info_lines.append(f"Databao context engine plugins: {command_info.plugin_ids}")
+    info_lines.append(f"Databao context engine version: {dce_info.version}")
+    info_lines.append(f"Databao context engine storage dir: {dce_info.dce_path}")
+    info_lines.append(f"Databao context engine plugins: {dce_info.plugin_ids}")
 
     info_lines.append("")
 
@@ -24,10 +27,10 @@ def _generate_info_string(command_info: DceInfo) -> str:
 
     info_lines.append("")
 
-    if command_info.project_info.is_initialized:
-        info_lines.append(f"Project dir: {command_info.project_info.project_path.resolve()}")
-        info_lines.append(f"Project ID: {str(command_info.project_info.project_id)}")
+    if project_info.is_initialized:
+        info_lines.append(f"Project dir: {project_info.project_path.resolve()}")
+        info_lines.append(f"Project ID: {str(project_info.project_id)}")
     else:
-        info_lines.append(f"Project not initialized at {command_info.project_info.project_path.resolve()}")
+        info_lines.append(f"Project not initialized at {project_info.project_path.resolve()}")
 
     return os.linesep.join(info_lines)

databao_context_engine/plugins/dbt/dbt_context_extractor.py

@@ -1,11 +1,18 @@
+import re
+from collections import defaultdict
 from pathlib import Path
+from typing import Any
 
 from databao_context_engine.plugins.dbt.types import (
+    DbtAcceptedValuesConstraint,
     DbtColumn,
     DbtConfigFile,
+    DbtConstraint,
     DbtContext,
     DbtMaterialization,
     DbtModel,
+    DbtRelationshipConstraint,
+    DbtSimpleConstraint,
 )
 from databao_context_engine.plugins.dbt.types_artifacts import (
     DbtArtifacts,
@@ -15,6 +22,7 @@ from databao_context_engine.plugins.dbt.types_artifacts import (
     DbtManifest,
     DbtManifestColumn,
     DbtManifestModel,
+    DbtManifestTest,
 )
 
 
@@ -47,25 +55,47 @@ def _read_dbt_artifacts(dbt_target_folder_path: Path) -> DbtArtifacts:
 
 def _extract_context_from_artifacts(artifacts: DbtArtifacts) -> DbtContext:
     manifest_models = [
-        manifest_model
-        for manifest_model in artifacts.manifest.nodes.values()
-        if isinstance(manifest_model, DbtManifestModel)
+        manifest_node
+        for manifest_node in artifacts.manifest.nodes.values()
+        if isinstance(manifest_node, DbtManifestModel)
     ]
 
+    manifest_tests_by_model_and_column = _get_manifest_tests(artifacts.manifest)
+
     catalog_nodes = artifacts.catalog.nodes if artifacts.catalog else {}
 
     # TODO: Extract the stages? Or at least the "highest-level" models (= marts?)
-    # TODO: Extract the constraints
     # TODO: Organize the models by schemas? Or by stages?
     return DbtContext(
         models=[
-            _manifest_model_to_dbt_model(manifest_model, catalog_nodes.get(manifest_model.unique_id, None))
+            _manifest_model_to_dbt_model(
+                manifest_model,
+                catalog_nodes.get(manifest_model.unique_id, None),
+                manifest_tests_by_model_and_column.get(manifest_model.unique_id, {}),
+            )
             for manifest_model in manifest_models
         ],
     )
 
 
-def _manifest_model_to_dbt_model(manifest_model: DbtManifestModel, catalog_node: DbtCatalogNode | None) -> DbtModel:
+def _get_manifest_tests(manifest: DbtManifest) -> dict[Any, dict[Any, list]]:
+    """Extract all tests nodes in the manifest and groups them by model and column."""
+    manifest_tests_by_model_and_column: dict[str, dict[str, list[DbtManifestTest]]] = defaultdict(
+        lambda: defaultdict(list)
+    )
+    for manifest_node in manifest.nodes.values():
+        if isinstance(manifest_node, DbtManifestTest) and manifest_node.attached_node and manifest_node.column_name:
+            manifest_tests_by_model_and_column[manifest_node.attached_node][manifest_node.column_name].append(
+                manifest_node
+            )
+    return manifest_tests_by_model_and_column
+
+
+def _manifest_model_to_dbt_model(
+    manifest_model: DbtManifestModel,
+    catalog_node: DbtCatalogNode | None,
+    test_nodes_by_column_name: dict[str, list[DbtManifestTest]],
+) -> DbtModel:
     catalog_columns = catalog_node.columns if catalog_node else {}
 
     return DbtModel(
@@ -75,7 +105,11 @@ def _manifest_model_to_dbt_model(manifest_model: DbtManifestModel, catalog_node:
         schema=manifest_model.schema_,
         description=manifest_model.description,
         columns=[
-            _manifest_column_to_dbt_column(manifest_column, catalog_columns.get(manifest_column.name))
+            _manifest_column_to_dbt_column(
+                manifest_column,
+                catalog_columns.get(manifest_column.name),
+                test_nodes_by_column_name.get(manifest_column.name, []),
+            )
             for manifest_column in manifest_model.columns.values()
         ],
         materialization=_manifest_materialization_to_dbt_materializaton(
@@ -87,15 +121,87 @@ def _manifest_model_to_dbt_model(manifest_model: DbtManifestModel, catalog_node:
 
 
 def _manifest_column_to_dbt_column(
-    manifest_column: DbtManifestColumn, catalog_column: DbtCatalogColumn | None
+    manifest_column: DbtManifestColumn, catalog_column: DbtCatalogColumn | None, test_nodes: list[DbtManifestTest]
 ) -> DbtColumn:
+    constraints = _manifest_test_to_dbt_constraint(test_nodes)
+
     return DbtColumn(
         name=manifest_column.name,
         description=manifest_column.description,
         type=catalog_column.type if catalog_column else manifest_column.data_type,
+        constraints=constraints,
     )
 
 
+def _manifest_test_to_dbt_constraint(test_nodes: list[DbtManifestTest]) -> list[DbtConstraint]:
+    result: list[DbtConstraint] = []
+
+    for test_node in test_nodes:
+        is_enforced = test_node.config.severity == "ERROR" if test_node.config else False
+
+        if test_node.test_metadata is None:
+            continue
+
+        match test_node.test_metadata.name:
+            case "not_null":
+                result.append(
+                    DbtSimpleConstraint(type="not_null", is_enforced=is_enforced, description=test_node.description)
+                )
+            case "unique":
+                result.append(
+                    DbtSimpleConstraint(type="unique", is_enforced=is_enforced, description=test_node.description)
+                )
+            case "accepted_values":
+                if test_node.test_metadata.kwargs is None:
+                    continue
+
+                accepted_values = test_node.test_metadata.kwargs.get("values", None)
+                if accepted_values is None:
+                    continue
+                result.append(
+                    DbtAcceptedValuesConstraint(
+                        type="accepted_values",
+                        is_enforced=is_enforced,
+                        description=test_node.description,
+                        accepted_values=accepted_values,
+                    )
+                )
+            case "relationships":
+                if test_node.test_metadata.kwargs is None:
+                    continue
+
+                target_model = _extract_ref_model(test_node.test_metadata.kwargs.get("to", None))
+                if target_model is None:
+                    continue
+                target_column = test_node.test_metadata.kwargs.get("field", None)
+                if target_column is None:
+                    continue
+                result.append(
+                    DbtRelationshipConstraint(
+                        type="relationships",
+                        is_enforced=is_enforced,
+                        description=test_node.description,
+                        target_model=target_model,
+                        target_column=target_column,
+                    )
+                )
+            case _:
+                continue
+
+    return result
+
+
+def _extract_ref_model(target_model_with_ref: str | None) -> str | None:
+    if target_model_with_ref is None:
+        return None
+
+    match = re.fullmatch(r"ref\(['\"]([\w.]+)['\"]\)", target_model_with_ref)
+    if match:
+        return match.group(1)
+
+    return None
+
+
 def _manifest_materialization_to_dbt_materializaton(materialized: str | None) -> DbtMaterialization | None:
     if materialized is None:
         return None

databao_context_engine/plugins/dbt/types.py

@@ -1,6 +1,7 @@
 from dataclasses import dataclass
 from enum import Enum
 from pathlib import Path
+from typing import Literal
 
 from pydantic import BaseModel, Field
 
@@ -19,11 +20,39 @@ class DbtMaterialization(str, Enum):
         return self.value
 
 
+@dataclass(kw_only=True)
+class DbtSimpleConstraint:
+    type: Literal["unique", "not_null"]
+    is_enforced: bool
+    description: str | None = None
+
+
+@dataclass(kw_only=True)
+class DbtAcceptedValuesConstraint:
+    type: Literal["accepted_values"]
+    is_enforced: bool
+    description: str | None = None
+    accepted_values: list[str]
+
+
+@dataclass(kw_only=True)
+class DbtRelationshipConstraint:
+    type: Literal["relationships"]
+    is_enforced: bool
+    description: str | None = None
+    target_model: str
+    target_column: str
+
+
+DbtConstraint = DbtSimpleConstraint | DbtAcceptedValuesConstraint | DbtRelationshipConstraint
+
+
 @dataclass(kw_only=True)
 class DbtColumn:
     name: str
     type: str | None = None
     description: str | None = None
+    constraints: list[DbtConstraint] | None = None
 
 
 @dataclass(kw_only=True)

databao_context_engine/plugins/dbt/types_artifacts.py

@@ -1,10 +1,10 @@
 from dataclasses import dataclass
-from typing import Annotated, Literal
+from typing import Annotated, Any, Literal
 
 from pydantic import BaseModel, Discriminator, Field
 
 
-class DbtManifestNodeConfig(BaseModel):
+class DbtManifestModelConfig(BaseModel):
     materialized: str
 
 
@@ -21,17 +21,36 @@ class DbtManifestModel(BaseModel):
     database: str
     schema_: str = Field(alias="schema")
     description: str | None = None
-    config: DbtManifestNodeConfig | None = None
+    config: DbtManifestModelConfig | None = None
     columns: dict[str, DbtManifestColumn]
     depends_on: dict[str, list[str]] | None = None
     primary_key: list[str] | None = None
 
 
+class DbtManifestTestConfig(BaseModel):
+    severity: str
+
+
+class DbtManifestTestMetadata(BaseModel):
+    name: str | None = None
+    kwargs: dict[str, Any] | None = None
+
+
+class DbtManifestTest(BaseModel):
+    resource_type: Literal["test"]
+    unique_id: str
+    attached_node: str | None = None
+    column_name: str | None = None
+    description: str | None = None
+    test_metadata: DbtManifestTestMetadata | None = None
+    config: DbtManifestTestConfig | None = None
+
+
 class DbtManifestOtherNode(BaseModel):
-    resource_type: Literal["seed", "analysis", "test", "operation", "sql_operation", "snapshot"]
+    resource_type: Literal["seed", "analysis", "operation", "sql_operation", "snapshot"]
 
 
-DbtManifestNode = Annotated[DbtManifestModel | DbtManifestOtherNode, Discriminator("resource_type")]
+DbtManifestNode = Annotated[DbtManifestModel | DbtManifestTest | DbtManifestOtherNode, Discriminator("resource_type")]
 
 
 class DbtManifest(BaseModel):

databao_context_engine/project/info.py

@@ -30,33 +30,38 @@ class DceInfo:
     Attributes:
         version: The version of the databao_context_engine package installed on the system.
         dce_path: The path where databao_context_engine stores its global data.
-        project_info: Information about the Databao Context Engine project.
     """
 
     version: str
     dce_path: Path
     plugin_ids: list[str]
 
-    project_info: DceProjectInfo
 
-
-def get_databao_context_engine_info(project_dir: Path) -> DceInfo:
+def get_databao_context_engine_info() -> DceInfo:
     """Return information about the current Databao Context Engine installation and project.
 
-    Args:
-        project_dir: The root directory of the Databao Context Project.
-
     Returns:
-        A DceInfo instance containing information about the Databao Context Engine installation and project.
+        A DceInfo instance containing information about the Databao Context Engine installation.
     """
     return DceInfo(
         version=get_dce_version(),
         dce_path=get_dce_path(),
         plugin_ids=[plugin.id for plugin in load_plugins().values()],
-        project_info=_get_project_info(project_dir),
     )
 
 
+def get_databao_context_engine_project_info(project_dir: Path) -> DceProjectInfo:
+    """Return information about the current Databao Context Engine project.
+
+    Args:
+        project_dir: The root directory of the Databao Context Project.
+
+    Returns:
+        A DceProjectInfo instance containing information about the Databao Context Engine project.
+    """
+    return _get_project_info(project_dir)
+
+
 def _get_project_info(project_dir: Path) -> DceProjectInfo:
     project_layout = validate_project_dir(project_dir)
 

{databao_context_engine-0.1.6.dist-info → databao_context_engine-0.1.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: databao-context-engine
-Version: 0.1.6
+Version: 0.1.7
 Summary: Semantic context for your LLMs — generated automatically
 Requires-Dist: click>=8.3.0
 Requires-Dist: duckdb>=1.4.3

{databao_context_engine-0.1.6.dist-info → databao_context_engine-0.1.7.dist-info}/RECORD

@@ -1,4 +1,4 @@
-databao_context_engine/__init__.py,sha256=sPhGoLZhsFvf1n-WLt9gFUL1IZegCQJPkOZrcrJYn2c,2239
+databao_context_engine/__init__.py,sha256=j7vQ0tiaGGwTiy9i4PGAhYaV9NC6uHHb3URbJx3GJ8w,2348
 databao_context_engine/build_sources/__init__.py,sha256=FteGp3MvSuX_fh-Fx1pq-jOjFcjYKQSXFSll4tQffvk,224
 databao_context_engine/build_sources/build_runner.py,sha256=aSxp99R4tbnbHlEDFrg8iw4Bi1dGvO_5uTO-mIrB080,3986
 databao_context_engine/build_sources/build_service.py,sha256=eywJzpvZUVrj-JYxXii_kxhTBzVUT0xfbNtdAFcpGWQ,1624
@@ -9,7 +9,7 @@ databao_context_engine/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJ
 databao_context_engine/cli/add_datasource_config.py,sha256=fK1MIaTCoz1jurldJZD_BpPIgc6kr-NaAyVjVgjsOTU,5523
 databao_context_engine/cli/commands.py,sha256=VJOHB3H28wNjOUuFC8hhF8-rnlIFjsZMlNT3C2ZbIUY,7807
 databao_context_engine/cli/datasources.py,sha256=kHbXhEnDVaYXTRAE3WeaDGohH1Y9cwyYjjcycQ8eh1U,2376
-databao_context_engine/cli/info.py,sha256=2iegfvRAxMae3KkWMDu-BOnljeJMVQ5u8P5k_8nOKpc,1203
+databao_context_engine/cli/info.py,sha256=ziduzb1IzH5rMklDsxk4ZREMIluPJ7qoo_iwrfnz0hY,1376
 databao_context_engine/config/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 databao_context_engine/config/log_config.yaml,sha256=hH2NPumufcyig8L-7Z_DbLQ88_G2N6CAEl5fvHafeAw,360
 databao_context_engine/config/logging.py,sha256=66k7x3ftEdHx-CKqbOss4pqZfGn4nMFv7AzVPs9aVnY,1525
@@ -84,10 +84,10 @@ databao_context_engine/plugins/databases/sqlite/sqlite_db_plugin.py,sha256=MI896
 databao_context_engine/plugins/databases/sqlite/sqlite_introspector.py,sha256=kKFNzclp1NmQ6BF3ylGn86R5PuQh6aqevj6E7_zAKmQ,8361
 databao_context_engine/plugins/dbt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 databao_context_engine/plugins/dbt/dbt_chunker.py,sha256=5OQo0Y9ouelZV8KQBA0Kqksx72wRhTdRztUOSCKwLZc,1516
-databao_context_engine/plugins/dbt/dbt_context_extractor.py,sha256=8ac2Fx5vv0t_AGqNbx9BXj73l9ZqjDK8r-VAAQukcP0,3748
+databao_context_engine/plugins/dbt/dbt_context_extractor.py,sha256=fbX2EUiMV_5yvtfWrQaEvfepgqgssXCuLrvOCmizyNc,7653
 databao_context_engine/plugins/dbt/dbt_plugin.py,sha256=-oAVh4u9YZ7j1b6w2K2lwcZT_bOdkN-lL1RRBWgB20k,1026
-databao_context_engine/plugins/dbt/types.py,sha256=yJUVSZa5Ohhx6Znl_20gKwjvgxkBBcxlKrAFFV6tfgY,881
-databao_context_engine/plugins/dbt/types_artifacts.py,sha256=OC2_DL7srcniVqrj9uj5XxNV-7VcA974xqgU1Rzfd5A,1349
+databao_context_engine/plugins/dbt/types.py,sha256=S7bgZT3tDGLX7so_hlrU5Qs3R9mnpy-n_BUEQt8YfFc,1590
+databao_context_engine/plugins/dbt/types_artifacts.py,sha256=umESx3Pm1QxbVPPJX1jbj12bfmZOsEqfNb7O0MAvR1w,1841
 databao_context_engine/plugins/duckdb_tools.py,sha256=46rctnTxDPAhHtaiTp1DxMuuDuRKrtKWJFSSM2w7uUU,645
 databao_context_engine/plugins/files/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 databao_context_engine/plugins/files/unstructured_files_plugin.py,sha256=eqs1anQhYBZh7xu4CwhfkqXQjGE5gJnKEwyJbtUR78E,2384
@@ -97,7 +97,7 @@ databao_context_engine/plugins/resources/parquet_chunker.py,sha256=R9WCOBqpKRTVN
 databao_context_engine/plugins/resources/parquet_introspector.py,sha256=Cn_yh6E-dOTOZstlavEGAsV6ZRKZXJraVAl_pzJJuGs,5629
 databao_context_engine/plugins/resources/parquet_plugin.py,sha256=FVRdAqhZkEBvUsW1CkYH4yPPs3P1pRsofR8PdHN26AE,1128
 databao_context_engine/project/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-databao_context_engine/project/info.py,sha256=3chVSZ4pV2EsYJSeVfOYTK746POlJImciIBOvrQuYbc,2479
+databao_context_engine/project/info.py,sha256=q7NwMs2E11ZUEIWFLOqm0f6-7ViuU920SPM7a7DgEVU,2610
 databao_context_engine/project/init_project.py,sha256=QiLalH--BMQGB-c2WBbPAUKBinDbO0cJkmilUfqPAuI,3970
 databao_context_engine/project/layout.py,sha256=3ttCJDLMh4AS1tBRcSROgT7hG6uOqkI1QGsQEcfGsqg,4087
 databao_context_engine/project/project_config.py,sha256=rI-Wkll7lca7RlYIaFScs5CIKIZ8uujtKRxSjibGIt8,1132
@@ -135,7 +135,7 @@ databao_context_engine/system/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5N
 databao_context_engine/system/properties.py,sha256=mQ7-_PZeYSESYn1cMUQ0IK7rJEnbhc7t4WesFjAgo-Q,429
 databao_context_engine/templating/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 databao_context_engine/templating/renderer.py,sha256=W2-0IGStAp6oxANmsKs_Z-UoIR6Gt_c4ILYFa3Hruo4,662
-databao_context_engine-0.1.6.dist-info/WHEEL,sha256=5w2T7AS2mz1-rW9CNagNYWRCaB0iQqBMYLwKdlgiR4Q,78
-databao_context_engine-0.1.6.dist-info/entry_points.txt,sha256=5EeQJ1W8zEFh4HuF1bs2zBeoP408oiwuM9UrkJiurgI,138
-databao_context_engine-0.1.6.dist-info/METADATA,sha256=weN9iS4ZtRnnt3tTJY-epHUgNdpMo-dWD4E3F-8dXyk,7773
-databao_context_engine-0.1.6.dist-info/RECORD,,
+databao_context_engine-0.1.7.dist-info/WHEEL,sha256=5w2T7AS2mz1-rW9CNagNYWRCaB0iQqBMYLwKdlgiR4Q,78
+databao_context_engine-0.1.7.dist-info/entry_points.txt,sha256=5EeQJ1W8zEFh4HuF1bs2zBeoP408oiwuM9UrkJiurgI,138
+databao_context_engine-0.1.7.dist-info/METADATA,sha256=REehDNFCn_yEYfgeo0UH_k2W9eVZPcbiIe2xY6IxzUY,7773
+databao_context_engine-0.1.7.dist-info/RECORD,,