parse test

metaphor/dbt/cloud/discovery_api/graphql_client/__init__.py

@@ -53,6 +53,7 @@
     GetJobRunModelsJob,
     GetJobRunModelsJobModels,
     GetJobRunModelsJobModelsColumns,
+    GetJobRunModelsJobModelsRunResults,
 )
 from .get_job_run_snapshots import (
     GetJobRunSnapshots,
@@ -66,6 +67,7 @@
     GetJobRunSourcesJobSources,
     GetJobRunSourcesJobSourcesColumns,
 )
+from .get_job_run_tests import GetJobRunTests, GetJobRunTestsJob, GetJobRunTestsJobTests
 from .get_job_tests import GetJobTests, GetJobTestsJob, GetJobTestsJobTests
 from .get_macro_arguments import (
     GetMacroArguments,
@@ -127,6 +129,7 @@
     "GetJobRunModelsJob",
     "GetJobRunModelsJobModels",
     "GetJobRunModelsJobModelsColumns",
+    "GetJobRunModelsJobModelsRunResults",
     "GetJobRunSnapshots",
     "GetJobRunSnapshotsJob",
     "GetJobRunSnapshotsJobSnapshots",
@@ -135,6 +138,9 @@
     "GetJobRunSourcesJob",
     "GetJobRunSourcesJobSources",
     "GetJobRunSourcesJobSourcesColumns",
+    "GetJobRunTests",
+    "GetJobRunTestsJob",
+    "GetJobRunTestsJobTests",
     "GetJobTests",
     "GetJobTestsJob",
     "GetJobTestsJobTests",

metaphor/dbt/cloud/discovery_api/graphql_client/client.py

@@ -12,6 +12,7 @@
 from .get_job_run_models import GetJobRunModels
 from .get_job_run_snapshots import GetJobRunSnapshots
 from .get_job_run_sources import GetJobRunSources
+from .get_job_run_tests import GetJobRunTests
 from .get_job_tests import GetJobTests
 from .get_macro_arguments import GetMacroArguments
 from .input_types import (
@@ -79,6 +80,10 @@ def get_job_run_models(
                   packageName
                   rawCode
                   rawSql
+                  runResults {
+                    status
+                    executeCompletedAt
+                  }
                   schema
                   tags
                   uniqueId
@@ -202,6 +207,33 @@ def get_job_run_sources(
         data = self.get_data(response)
         return GetJobRunSources.model_validate(data)
 
+    def get_job_run_tests(
+        self, job_id: Any, run_id: Any, **kwargs: Any
+    ) -> GetJobRunTests:
+        query = gql(
+            """
+            query GetJobRunTests($jobId: BigInt!, $runId: BigInt!) {
+              job(id: $jobId, runId: $runId) {
+                tests {
+                  columnName
+                  compiledSql
+                  compiledCode
+                  dependsOn
+                  name
+                  runId
+                  uniqueId
+                }
+              }
+            }
+            """
+        )
+        variables: Dict[str, object] = {"jobId": job_id, "runId": run_id}
+        response = self.execute(
+            query=query, operation_name="GetJobRunTests", variables=variables, **kwargs
+        )
+        data = self.get_data(response)
+        return GetJobRunTests.model_validate(data)
+
     def get_environment_adapter_type(
         self, environment_id: Any, **kwargs: Any
     ) -> GetEnvironmentAdapterType:

metaphor/dbt/cloud/discovery_api/graphql_client/get_job_run_models.py

@@ -31,6 +31,7 @@ class GetJobRunModelsJobModels(BaseModel):
     package_name: Optional[str] = Field(alias="packageName")
     raw_code: Optional[str] = Field(alias="rawCode")
     raw_sql: Optional[str] = Field(alias="rawSql")
+    run_results: List["GetJobRunModelsJobModelsRunResults"] = Field(alias="runResults")
     schema_: Optional[str] = Field(alias="schema")
     tags: Optional[List[str]]
     unique_id: str = Field(alias="uniqueId")
@@ -45,6 +46,11 @@ class GetJobRunModelsJobModelsColumns(BaseModel):
     type: Optional[str]
 
 
+class GetJobRunModelsJobModelsRunResults(BaseModel):
+    status: Optional[str]
+    execute_completed_at: Optional[Any] = Field(alias="executeCompletedAt")
+
+
 GetJobRunModels.model_rebuild()
 GetJobRunModelsJob.model_rebuild()
 GetJobRunModelsJobModels.model_rebuild()

metaphor/dbt/cloud/discovery_api/graphql_client/get_job_run_tests.py

@@ -0,0 +1,30 @@
+# Generated by ariadne-codegen
+# Source: queries.graphql
+
+from typing import Any, List, Optional
+
+from pydantic import Field
+
+from .base_model import BaseModel
+
+
+class GetJobRunTests(BaseModel):
+    job: Optional["GetJobRunTestsJob"]
+
+
+class GetJobRunTestsJob(BaseModel):
+    tests: List["GetJobRunTestsJobTests"]
+
+
+class GetJobRunTestsJobTests(BaseModel):
+    column_name: Optional[str] = Field(alias="columnName")
+    compiled_sql: Optional[str] = Field(alias="compiledSql")
+    compiled_code: Optional[str] = Field(alias="compiledCode")
+    depends_on: List[str] = Field(alias="dependsOn")
+    name: Optional[str]
+    run_id: Any = Field(alias="runId")
+    unique_id: str = Field(alias="uniqueId")
+
+
+GetJobRunTests.model_rebuild()
+GetJobRunTestsJob.model_rebuild()

metaphor/dbt/cloud/discovery_api/queries.graphql

@@ -35,6 +35,10 @@ query GetJobRunModels($jobId: BigInt!, $runId: BigInt!) {
       packageName
       rawCode
       rawSql
+      runResults {
+        status
+        executeCompletedAt
+      }
       schema
       tags
       uniqueId
@@ -102,6 +106,8 @@ query GetJobRunSources($jobId: BigInt!, $runId: BigInt!) {
       schema
       uniqueId
     }
+  }
+}
     # metrics {
     #   accountId
     #   calculation_method
@@ -244,42 +250,17 @@ query GetJobRunSources($jobId: BigInt!, $runId: BigInt!) {
     #   type
     #   uniqueId
     # }
-    # tests {
-    #   accountId
-    #   columnName
-    #   compileCompletedAt
-    #   compileStartedAt
-    #   compiledCode
-    #   compiledSql
-    #   dbtVersion
-    #   dependsOn
-    #   description
-    #   environmentId
-    #   error
-    #   executeCompletedAt
-    #   executeStartedAt
-    #   executionTime
-    #   fail
-    #   invocationId
-    #   jobId
-    #   language
-    #   meta
-    #   name
-    #   projectId
-    #   rawCode
-    #   rawSql
-    #   resourceType
-    #   runElapsedTime
-    #   runGeneratedAt
-    #   runId
-    #   skip
-    #   state
-    #   status
-    #   tags
-    #   threadId
-    #   uniqueId
-    #   warn
-    # }
+
+query GetJobRunTests($jobId: BigInt!, $runId: BigInt!) {
+  job(id: $jobId, runId: $runId) {
+    tests {
+      columnName
+      compiledSql
+      compiledCode
+      dependsOn
+      name
+      uniqueId
+    }
   }
 }
 

metaphor/dbt/cloud/parser/node_parser.py

@@ -69,7 +69,7 @@ def __init__(
         self._metrics = metrics
 
     @staticmethod
-    def _get_node_name(node: NODE_TYPE):
+    def get_node_name(node: NODE_TYPE):
         if isinstance(node, GetJobRunModelsJobModels):
             return get_model_name_from_unique_id
         return get_snapshot_name_from_unique_id
@@ -251,7 +251,7 @@ def parse(
         source_map: Dict[str, EntityId],
         macro_map: Dict[str, DbtMacro],
     ):
-        node_name_getter = self._get_node_name(node)
+        node_name_getter = self.get_node_name(node)
         virtual_view = init_virtual_view(
             self._virtual_views, node.unique_id, node_name_getter
         )

metaphor/dbt/cloud/parser/parser.py

@@ -4,9 +4,11 @@
 from metaphor.common.snowflake import normalize_snowflake_account
 from metaphor.dbt.cloud.client import DbtRun
 from metaphor.dbt.cloud.discovery_api import DiscoveryAPIClient
+from metaphor.dbt.cloud.discovery_api.graphql_client.get_job_run_models import GetJobRunModelsJobModels
 from metaphor.dbt.cloud.parser.macro_parser import MacroParser
 from metaphor.dbt.cloud.parser.node_parser import NodeParser
 from metaphor.dbt.cloud.parser.source_parser import SourceParser
+from metaphor.dbt.cloud.parser.test_parser import TestParser
 from metaphor.dbt.config import DbtRunConfig
 from metaphor.dbt.util import init_virtual_view
 from metaphor.models.metadata_change_event import (
@@ -52,6 +54,12 @@ def __init__(
             self._virtual_views,
             self._metrics,
         )
+        self._test_parser = TestParser(
+            self._platform,
+            self._account,
+            self._virtual_views,
+            self._datasets,
+        )
 
     def _get_source_map(self, run: DbtRun):
         job_run_sources = self._discovery_api.get_job_run_sources(
@@ -75,24 +83,37 @@ def _get_nodes(self, run: DbtRun):
         assert job_run_snapshots.job
         return job_run_models.job.models + job_run_snapshots.job.snapshots
 
+    def _get_tests(self, run: DbtRun):
+        job_run_tests = self._discovery_api.get_job_run_tests(run.job_id, run.run_id)
+        assert job_run_tests.job
+        return job_run_tests.job.tests
+
     def parse_run(self, run: DbtRun):
+        """
+        Parses a single job run.
+        """
         source_map = self._get_source_map(run)
         macro_map = self._get_macro_map(run)
 
         nodes = self._get_nodes(run)
 
         for node in nodes:
             init_virtual_view(
-                self._virtual_views, node.unique_id, NodeParser._get_node_name(node)
+                self._virtual_views, node.unique_id, NodeParser.get_node_name(node)
             )
             if self._project_name and node.package_name != self._project_name:
                 self._referenced_virtual_views.add(node.unique_id)
 
         for node in nodes:
             self._node_parser.parse(node, source_map, macro_map)
 
-        # for _, test in job.tests:
-        #     self._parse_test(test, run_results, models)
+        models = {
+            node.unique_id: node
+            for node in nodes
+            if isinstance(node, GetJobRunModelsJobModels)
+        }
+        for test in self._get_tests(run):
+            self._test_parser.parse(test, models)
 
         # for _, metric in job.metrics:
         #     self._parse_metric(metric, source_map, macro_map)