Fix limit bug #2 (#1329)

AltimateAI · Jul 25, 2024 · be8e8fd · be8e8fd
1 parent 8e359d0
commit be8e8fd
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/dbt_core_integration.py b/dbt_core_integration.py
@@ -630,16 +630,19 @@ def get_server_node(self, sql: str, node_name="name", original_node: Optional[Un
         return sql_node
 
     @lru_cache(maxsize=100)
-    def get_macro_function(self, macro_name: str) -> Callable[[Dict[str, Any]], Any]:
+    def get_macro_function(self, macro_name: str, compiled_code: Optional[str] = None) -> Callable[[Dict[str, Any]], Any]:
         """Get macro as a function which takes a dict via argument named `kwargs`,
         ie: `kwargs={"relation": ...}`
 
         make_schema_fn = get_macro_function('make_schema')\n
         make_schema_fn({'name': '__test_schema_1'})\n
         make_schema_fn({'name': '__test_schema_2'})"""
         if DBT_MAJOR_VER >= 1 and DBT_MINOR_VER >= 8:
+            model_context = {}
+            if compiled_code is not None:
+                model_context["compiled_code"] = compiled_code
             return partial(
-                self.adapter.execute_macro, macro_name=macro_name
+                self.adapter.execute_macro, macro_name=macro_name, context_override=model_context,
             )
         else:
             return partial(
@@ -656,9 +659,10 @@ def execute_macro(
         self,
         macro: str,
         kwargs: Optional[Dict[str, Any]] = None,
+        compiled_code: Optional[str] = None
     ) -> Any:
         """Wraps adapter execute_macro. Execute a macro like a function."""
-        return self.get_macro_function(macro)(kwargs=kwargs)
+        return self.get_macro_function(macro, compiled_code)(kwargs=kwargs)
 
     def execute_sql(self, raw_sql: str, original_node: Optional[Union["ManifestNode", str]] = None) -> DbtAdapterExecutionResult:
         """Execute dbt SQL statement against database"""

diff --git a/src/dbt_client/dbtCoreIntegration.ts b/src/dbt_client/dbtCoreIntegration.ts
@@ -327,7 +327,7 @@ export class DBTCoreProjectIntegration
         const args = { compiled_code: query, limit };
         const queryTemplateFromMacro = await this.python?.lock(
           (python) =>
-            python!`to_dict(project.execute_macro('get_show_sql', ${args}))`,
+            python!`to_dict(project.execute_macro('get_show_sql', ${args}, ${query}))`,
         );
 
         this.dbtTerminal.debug(