Merge pull request #355 from DataRecce/feature/drc-510-support-select-api

[Feature] Add the select nodes method

Author: popcornylu

recce/adapter/dbt_adapter/__init__.py

@@ -4,7 +4,7 @@
 import uuid
 from contextlib import contextmanager
 from dataclasses import dataclass, fields
-from typing import Callable, Dict, List, Optional, Tuple, Iterator, Any
+from typing import Callable, Dict, List, Optional, Tuple, Iterator, Any, Set
 
 import agate
 import dbt.adapters.factory
@@ -134,6 +134,7 @@ class DbtArgs:
     profile: Optional[str] = None,
     target_path: Optional[str] = None,
     project_only_flags: Optional[Dict[str, Any]] = None
+    which: Optional[str] = 'run'
 
 
 @dataclass
@@ -172,6 +173,7 @@ def load(cls, artifacts: ArtifactsRoot = None, no_artifacts=False, **kwargs):
             profiles_dir=profiles_dir,
             profile=profile_name,
             project_only_flags={},
+            which='recce'
         )
         set_from_args(args, args)
 
@@ -584,6 +586,16 @@ def create_relation(self, model, base=False):
 
         return self.adapter.Relation.create_from(self.runtime_config, node)
 
+    def select_nodes(self, select: str) -> Set[str]:
+        from dbt.graph import SelectionCriteria, NodeSelector
+        from dbt.compilation import Compiler
+
+        compiler = Compiler(self.runtime_config)
+        graph = compiler.compile(self.manifest)
+        selector = NodeSelector(graph, self.manifest)
+        spec = SelectionCriteria.from_single_spec(select)
+        return selector.get_selected(spec)
+
     def export_artifacts(self) -> ArtifactsRoot:
         artifacts = ArtifactsRoot()
         target_path = self.runtime_config.target_path

tests/adapter/dbt_adapter/test_dbt_adapter.py

@@ -0,0 +1,145 @@
+import os
+import textwrap
+import uuid
+
+import pytest
+from dbt.contracts.graph.nodes import ModelNode
+
+from recce.adapter.dbt_adapter import DbtAdapter, as_manifest, load_manifest
+from recce.core import RecceContext, set_default_context
+
+
+class DbtTestHelper:
+
+    def __init__(self):
+        schema_prefix = "schema_" + uuid.uuid4().hex
+        self.base_schema = f"{schema_prefix}_base"
+        self.curr_schema = f"{schema_prefix}_curr"
+
+        current_dir = os.path.dirname(os.path.abspath(__file__))
+        project_dir = os.path.join(current_dir, '.')
+        profiles_dir = project_dir
+        manifest_path = os.path.join(project_dir, 'manifest.json')
+
+        dbt_adapter = DbtAdapter.load(
+            no_artifacts=True,
+            project_dir=project_dir,
+            profiles_dir=profiles_dir,
+        )
+
+        context = RecceContext()
+        context.adapter_type = 'dbt'
+        context.adapter = dbt_adapter
+        context.schema_prefix = schema_prefix
+        self.adapter = dbt_adapter
+        self.context = context
+        self.curr_manifest = as_manifest(load_manifest(manifest_path))
+        self.base_manifest = as_manifest(load_manifest(manifest_path))
+        self.context = context
+
+        self.adapter.execute(f"CREATE schema IF NOT EXISTS {self.base_schema}")
+        self.adapter.execute(f"CREATE schema IF NOT EXISTS {self.curr_schema}")
+        self.adapter.set_artifacts(self.base_manifest, self.curr_manifest)
+
+    def create_model(self, model_name, base_csv, curr_csv):
+        package_name = "recce_test"
+
+        def _add_model_to_manifest(base):
+            if base:
+                schema = self.base_schema
+                manifest = self.base_manifest
+            else:
+                schema = self.curr_schema
+                manifest = self.curr_manifest
+
+            node = ModelNode.from_dict({
+                "resource_type": "model",
+                "name": model_name,
+                "package_name": package_name,
+                "path": "",
+                "original_file_path": "",
+                "unique_id": f"model.{package_name}.{model_name}",
+                "fqn": [
+                    package_name,
+                    model_name,
+                ],
+                "schema": schema,
+                "alias": model_name,
+                "checksum": {
+                    "name": "sha256",
+                    "checksum": ""
+                },
+                "raw_code": 'dummy',
+                "config": {
+                    "materialized": "table",
+                    "tags": ["test_tag"],
+                },
+                "tags": ["test_tag"],
+            })
+            manifest.add_node_nofile(node)
+
+        _add_model_to_manifest(base=True)
+        _add_model_to_manifest(base=False)
+
+        import pandas as pd
+        from io import StringIO
+        df_base = pd.read_csv(StringIO(textwrap.dedent(base_csv)))
+        df_curr = pd.read_csv(StringIO(textwrap.dedent(curr_csv)))
+        dbt_adapter = self.adapter
+        with dbt_adapter.connection_named('create model'):
+            dbt_adapter.execute(f"CREATE TABLE {self.base_schema}.{model_name} AS SELECT * FROM df_base")
+            dbt_adapter.execute(f"CREATE TABLE {self.curr_schema}.{model_name} AS SELECT * FROM df_curr")
+        self.adapter.set_artifacts(self.base_manifest, self.curr_manifest)
+
+    def remove_model(self, model_name):
+        dbt_adapter = self.adapter
+        with dbt_adapter.connection_named('cleanup'):
+            dbt_adapter.execute(f"DROP TABLE IF EXISTS {self.base_schema}.{model_name}")
+            dbt_adapter.execute(f"DROP TABLE IF EXISTS {self.curr_schema}.{model_name} ")
+
+    def cleanup(self):
+        dbt_adapter = self.adapter
+        with dbt_adapter.connection_named('cleanup'):
+            dbt_adapter.execute(f"DROP SCHEMA IF EXISTS {self.base_schema} CASCADE")
+            dbt_adapter.execute(f"DROP SCHEMA IF EXISTS {self.curr_schema} CASCADE")
+
+
+@pytest.fixture
+def helper():
+    helper = DbtTestHelper()
+    context = helper.context
+    set_default_context(context)
+    yield helper
+    helper.cleanup()
+
+
+def test_select(helper):
+    csv_data_curr = """
+        customer_id,name,age
+        1,Alice,30
+        2,Bob,25
+        3,Charlie,35
+        """
+
+    csv_data_base = """
+        customer_id,name,age
+        1,Alice,35
+        2,Bob,25
+        3,Charlie,35
+        """
+
+    helper.create_model("customers_1", csv_data_base, csv_data_curr)
+    helper.create_model("customers_2", csv_data_base, csv_data_curr)
+    adapter: DbtAdapter = helper.context.adapter
+    node_ids = adapter.select_nodes('resource_type:model')
+    assert len(node_ids) == 2
+    node_ids = adapter.select_nodes('customers_1')
+    assert len(node_ids) == 1
+    node_ids = adapter.select_nodes('tag:test_tag')
+    assert len(node_ids) == 2
+    node_ids = adapter.select_nodes('tag:test_tag2')
+    assert len(node_ids) == 0
+    node_ids = adapter.select_nodes("config.materialized:incremental")
+    assert len(node_ids) == 0
+    node_ids = adapter.select_nodes("config.materialized:table")
+    assert len(node_ids) == 2