pt2

MetaphorData · Jul 23, 2024 · 5952fc6 · 5952fc6
1 parent c09837d
commit 5952fc6
Show file tree

Hide file tree

Showing 5 changed files with 35 additions and 15 deletions.
diff --git a/metaphor/dbt/cloud/discovery_api/graphql_client/client.py b/metaphor/dbt/cloud/discovery_api/graphql_client/client.py
@@ -227,22 +227,24 @@ def get_environment_adapter_type(
     def get_macro_arguments(
         self,
         environment_id: Any,
+        first: int,
         filter: Union[Optional[MacroDefinitionFilter], UnsetType] = UNSET,
         **kwargs: Any
     ) -> GetMacroArguments:
         query = gql(
             """
-            query GetMacroArguments($environmentId: BigInt!, $filter: MacroDefinitionFilter) {
+            query GetMacroArguments($environmentId: BigInt!, $filter: MacroDefinitionFilter, $first: Int!) {
               environment(id: $environmentId) {
                 definition {
-                  macros(filter: $filter, first: 1) {
+                  macros(filter: $filter, first: $first) {
                     edges {
                       node {
                         arguments {
                           description
                           name
                           type
                         }
+                        uniqueId
                       }
                     }
                   }
@@ -254,6 +256,7 @@ def get_macro_arguments(
         variables: Dict[str, object] = {
             "environmentId": environment_id,
             "filter": filter,
+            "first": first,
         }
         response = self.execute(
             query=query,

diff --git a/metaphor/dbt/cloud/discovery_api/graphql_client/get_macro_arguments.py b/metaphor/dbt/cloud/discovery_api/graphql_client/get_macro_arguments.py
@@ -3,6 +3,8 @@
 
 from typing import List, Optional
 
+from pydantic import Field
+
 from .base_model import BaseModel
 
 
@@ -28,6 +30,7 @@ class GetMacroArgumentsEnvironmentDefinitionMacrosEdges(BaseModel):
 
 class GetMacroArgumentsEnvironmentDefinitionMacrosEdgesNode(BaseModel):
     arguments: List["GetMacroArgumentsEnvironmentDefinitionMacrosEdgesNodeArguments"]
+    unique_id: str = Field(alias="uniqueId")
 
 
 class GetMacroArgumentsEnvironmentDefinitionMacrosEdgesNodeArguments(BaseModel):

diff --git a/metaphor/dbt/cloud/discovery_api/queries.graphql b/metaphor/dbt/cloud/discovery_api/queries.graphql
@@ -289,17 +289,18 @@ query GetEnvironmentAdapterType($environmentId: BigInt!) {
   }
 }
 
-query GetMacroArguments($environmentId: BigInt!, $filter: MacroDefinitionFilter) {
+query GetMacroArguments($environmentId: BigInt!, $filter: MacroDefinitionFilter, $first: Int!) {
   environment(id: $environmentId) {
     definition {
-      macros(filter: $filter, first: 1) {
+      macros(filter: $filter, first: $first) {
         edges {
           node {
             arguments {
               description
               name
               type
             }
+            uniqueId
           }
         }
       }

diff --git a/metaphor/dbt/cloud/parser/macro_parser.py b/metaphor/dbt/cloud/parser/macro_parser.py
@@ -1,3 +1,4 @@
+from collections import defaultdict
 from typing import Dict, List
 
 from metaphor.dbt.cloud.discovery_api import DiscoveryAPIClient
@@ -17,34 +18,46 @@ def __init__(
     ) -> None:
         self._discovery_api = discovery_api
 
-    def parse(self, macros: List[GetJobRunMacrosJobMacros]) -> Dict[str, DbtMacro]:
-        macro_map: Dict[str, DbtMacro] = {}
+    def _parse_macro_arguments(self, macros: List[GetJobRunMacrosJobMacros]):
+        macros_by_environment = defaultdict(list)
         for macro in macros:
+            macros_by_environment[macro.environment_id].append(macro.unique_id)
+
+        arguments: Dict[str, List[DbtMacroArgument]] = dict()
+        for environment_id, unique_ids in macros_by_environment.items():
             definition = self._discovery_api.get_macro_arguments(
-                macro.environment_id, MacroDefinitionFilter(uniqueIds=[macro.unique_id])
+                environment_id,
+                filter=MacroDefinitionFilter(
+                    uniqueIds=unique_ids,
+                ),
+                first=len(unique_ids),
             ).environment.definition
             if not definition:
                 continue
-
-            arguments = []
-            if len(definition.macros.edges) == 1:
-                macro_definition = definition.macros.edges[0].node
-                arguments = [
+            for edge in definition.macros.edges:
+                node = edge.node
+                arguments[node.unique_id] = [
                     DbtMacroArgument(
                         name=arg.name,
                         type=arg.type,
                         description=arg.description,
                     )
-                    for arg in macro_definition.arguments
+                    for arg in node.arguments
                 ]
+        return arguments
 
+    def parse(self, macros: List[GetJobRunMacrosJobMacros]) -> Dict[str, DbtMacro]:
+        arguments = self._parse_macro_arguments(macros)
+
+        macro_map: Dict[str, DbtMacro] = {}
+        for macro in macros:
             macro_depends_on = [x for x in macro.depends_on if x.startswith("macro.")]
             macro_map[macro.unique_id] = DbtMacro(
                 name=macro.name,
                 unique_id=macro.unique_id,
                 package_name=macro.package_name,
                 description=macro.description,
-                arguments=arguments,
+                arguments=arguments.get(macro.unique_id),
                 sql=macro.macro_sql,
                 depends_on_macros=macro_depends_on if macro_depends_on else None,
             )

diff --git a/metaphor/dbt/cloud/parser/node_parser.py b/metaphor/dbt/cloud/parser/node_parser.py
@@ -296,7 +296,7 @@ def _parse_depends_on(
 
         models = unique_list(
             [
-                get_virtual_view_id(self._virtual_views[n].logical_id)
+                get_virtual_view_id(self._virtual_views[n].logical_id)  # type: ignore
                 for n in depends_on
                 if n.startswith("model.") or n.startswith("snapshot.")
             ]