Repository: fal-ai/dbt-fal

Issue:
Support source status and test results in run context
**Context**
fal provides information about model transformation status and their tests, but not for source tests. That prevents a use case of, for example, notifying all the failures in a `dbt build` command.

**Describe the solution you'd like**
Ability to access source test results in the `context` of a `fal run`.

**Describe alternatives you've considered**
Directly parsing manifest files (not a good option considering all `fal` can do).

**Is there an existing feature request for this?**
 * [x]  I have searched the existing issues


Discussion:

Resolving patch:
diff --git a/src/faldbt/project.py b/src/faldbt/project.py
--- a/src/faldbt/project.py
+++ b/src/faldbt/project.py
@@ -42,12 +42,19 @@ class FalGeneralException(Exception):
 
 
 @dataclass
-class DbtTest:
-    node: Any
+class _DbtNode:
     name: str = field(init=False)
+    status: str = field(init=False)
+
+    def set_status(self, status: str):
+        self.status = status
+
+
+@dataclass
+class DbtTest(_DbtNode):
+    node: Any
     model: str = field(init=False)
     column: str = field(init=False)
-    status: str = field(init=False)
 
     def __post_init__(self):
         node = self.node
@@ -66,19 +73,26 @@ def __post_init__(self):
             else:
                 logger.debug(f"Non-generic test was not processed: {node.name}")
 
-    def set_status(self, status: str):
-        self.status = status
+
+@dataclass
+class DbtSource(_DbtNode):
+    name: str = field()
+    table_name: str = field()
+    unique_id: str = field()
+    tests: List[DbtTest] = field(init=False)
+
+    def __post_init__(self):
+        self.tests = []
 
 
 @dataclass
-class DbtModel:
+class DbtModel(_DbtNode):
     node: ParsedModelNode
-    name: str = field(init=False)
     meta: Dict[str, Any] = field(init=False)
-    status: NodeStatus = field(init=False)
     columns: Dict[str, Any] = field(init=False)
     tests: List[DbtTest] = field(init=False)
     python_model: Optional[Path] = field(init=False)
+    unique_id: str = field(init=False)
 
     def __post_init__(self):
         node = self.node
@@ -120,9 +134,6 @@ def get_scripts(self, keyword: str, before: bool) -> List[str]:
         else:
             return []
 
-    def set_status(self, status: str):
-        self.status = status
-
 
 @dataclass
 class DbtManifest:
@@ -267,7 +278,7 @@ def __init__(
 
         self._manifest = DbtManifest(self._compile_task.manifest)
 
-        self.models, self.tests = _map_nodes_to_models(
+        self.models, self.tests, self.sources = _map_nodes(
             self._run_results, self._manifest, generated_models
         )
 
@@ -310,13 +321,9 @@ def project_name(self):
     @telemetry.log_call("list_sources")
     def list_sources(self):
         """
-        List tables available for `source` usage, formatting like `[[source_name, table_name], ...]`
+        List tables available for `source` usage
         """
-        res = []
-        for source in self._manifest.get_sources():
-            res.append([source.source_name, source.name])
-
-        return res
+        return self.sources
 
     @telemetry.log_call("list_models_ids")
     def list_models_ids(self) -> Dict[str, str]:
@@ -634,32 +641,40 @@ def _firestore_dict_to_document(data: Dict, key_column: str):
     return output
 
 
-def _map_nodes_to_models(
+def _map_nodes(
     run_results: DbtRunResult, manifest: DbtManifest, generated_models: Dict[str, Path]
 ):
     models = manifest.get_models()
     tests = manifest.get_tests()
+    sources = [
+        DbtSource(
+            name=source.source_name, table_name=source.name, unique_id=source.unique_id
+        )
+        for source in manifest.get_sources()
+    ]
     status_map = dict(
         map(
             lambda res: [res.unique_id, res.status],
             run_results.results,
         )
     )
-    for model in models:
-        if model.name in generated_models:
-            model.python_model = generated_models[model.name]
+    for parent in models + sources:
+        if parent.name in generated_models:
+            parent.python_model = generated_models[parent.name]
 
-        model.set_status(status_map.get(model.unique_id, NodeStatus.Skipped))
+        parent.set_status(status_map.get(parent.unique_id, NodeStatus.Skipped))
         for test in tests:
-            if hasattr(test, "model") and test.model == model.name:
-                model.tests.append(test)
+            if (hasattr(test, "model") and test.model == parent.name) or (
+                hasattr(test, "source") and test.source == parent.name
+            ):
+                parent.tests.append(test)
                 test.set_status(status_map.get(test.unique_id, NodeStatus.Skipped))
                 if (
                     test.status != NodeStatus.Skipped
-                    and model.status == NodeStatus.Skipped
+                    and parent.status == NodeStatus.Skipped
                 ):
-                    model.set_status("tested")
-    return (models, tests)
+                    parent.set_status("tested")
+    return (models, tests, sources)
 
 
 def _get_custom_target(run_results: DbtRunResult):


Test patch:
diff --git a/integration_tests/features/dbt_test.feature b/integration_tests/features/dbt_test.feature
new file mode 100644
--- /dev/null
+++ b/integration_tests/features/dbt_test.feature
@@ -0,0 +1,33 @@
+Feature: `write_to_source` function
+  Background: Project Setup
+    Given the project 005_functions_and_variables
+
+  Scenario: Source tests are present in context
+    When the following shell command is invoked:
+      """
+      dbt run --profiles-dir $profilesDir --project-dir $baseDir
+      """
+    And the following command is invoked:
+      """
+      fal run --profiles-dir $profilesDir --project-dir $baseDir
+      """
+    Then the following scripts are ran:
+      | some_model.write_to_source_twice.py |
+    And the script some_model.write_to_source_twice.py output file has the lines:
+      | source results has 2 tests, source status is skipped   |
+      | model some_model has 2 tests, model status is success  |
+      | model other_model has 0 tests, model status is success |
+    When the following shell command is invoked:
+      """
+      dbt test --profiles-dir $profilesDir --project-dir $baseDir
+      """
+    And the following command is invoked:
+      """
+      fal run --profiles-dir $profilesDir --project-dir $baseDir
+      """
+    Then the following scripts are ran:
+      | some_model.write_to_source_twice.py |
+    And the script some_model.write_to_source_twice.py output file has the lines:
+      | source results has 2 tests, source status is tested    |
+      | model some_model has 2 tests, model status is tested   |
+      | model other_model has 0 tests, model status is skipped |
diff --git a/integration_tests/projects/005_functions_and_variables/fal_scripts/write_to_source_twice.py b/integration_tests/projects/005_functions_and_variables/fal_scripts/write_to_source_twice.py
--- a/integration_tests/projects/005_functions_and_variables/fal_scripts/write_to_source_twice.py
+++ b/integration_tests/projects/005_functions_and_variables/fal_scripts/write_to_source_twice.py
@@ -19,6 +19,18 @@
 source_size = len(source("results", "some_source"))
 output += f"source size {source_size}\n"
 
+for source in list_sources():
+    output += (
+        f"source {source.name} has {len(source.tests)} tests,"
+        f" source status is {source.status}\n"
+    )
+
+for model in list_models():
+    output += (
+        f"model {model.name} has {len(model.tests)} tests,"
+        f" model status is {model.status}\n"
+    )
+
 path = reduce(
     os.path.join, [os.environ["temp_dir"], model_name + ".write_to_source_twice.txt"]
 )
diff --git a/integration_tests/projects/005_functions_and_variables/models/schema.yml b/integration_tests/projects/005_functions_and_variables/models/schema.yml
--- a/integration_tests/projects/005_functions_and_variables/models/schema.yml
+++ b/integration_tests/projects/005_functions_and_variables/models/schema.yml
@@ -6,9 +6,24 @@ sources:
     schema: "{{ env_var('DBT_SCHEMA', 'dbt_fal') }}"
     tables:
       - name: some_source
+        columns:
+          - name: my_text
+            tests:
+              - not_null
+          - name: my_int
+            tests:
+              - not_null
 
 models:
   - name: some_model
+    columns:
+      - name: my_text
+        tests:
+          - not_null
+      - name: my_int
+        tests:
+          - not_null
+
     meta:
       fal:
         scripts:
diff --git a/tests/test_import.py b/tests/test_import.py
--- a/tests/test_import.py
+++ b/tests/test_import.py
@@ -27,7 +27,7 @@ def test_initialize():
     df = faldbt.ref("model_with_scripts")
 
     sources = faldbt.list_sources()
-    assert ["test_sources", "single_col"] in sources
+    assert "test_sources" in [source.name for source in sources]
 
     models = faldbt.list_models()
     assert "model_with_scripts" in [model.name for model in models]