restrict python submission

.pre-commit-config.yaml

@@ -6,7 +6,7 @@ exclude: ^test/
 
 # Force all unspecified python hooks to run python 3.8
 default_language_version:
-  python: python3.8
+  python: python3
 
 repos:
 - repo: https://github.com/pre-commit/pre-commit-hooks

core/dbt/adapters/base/impl.py

@@ -1220,7 +1220,6 @@ def python_submission_helpers(self) -> Dict[str, Type[PythonJobHelper]]:
     def default_python_submission_method(self) -> str:
         raise NotImplementedError("default_python_submission_method is not specified")
 
-    @available.parse_none
     @log_code_execution
     def submit_python_job(self, parsed_model: dict, compiled_code: str) -> AdapterResponse:
         submission_method = parsed_model["config"].get(

core/dbt/clients/jinja.py

@@ -27,6 +27,7 @@
 from dbt.clients._jinja_blocks import BlockIterator, BlockData, BlockTag
 from dbt.contracts.graph.compiled import CompiledGenericTestNode
 from dbt.contracts.graph.parsed import ParsedGenericTestNode
+
 from dbt.exceptions import (
     InternalException,
     raise_compiler_error,
@@ -305,13 +306,13 @@ def exception_handler(self) -> Iterator[None]:
     @contextmanager
     def track_call(self):
         # This is only called from __call__
-        if self.stack is None or self.node is None:
+        if self.stack is None:
             yield
         else:
             unique_id = self.macro.unique_id
             depth = self.stack.depth
-            # only mark depth=0 as a dependency
-            if depth == 0:
+            # only mark depth=0 as a dependency, when creating this dependency we don't pass in stack
+            if depth == 0 and self.node:
                 self.node.depends_on.add_macro(unique_id)
             self.stack.push(unique_id)
             try:

core/dbt/compilation.py

@@ -183,6 +183,7 @@ def _create_node_context(
 
         context = generate_runtime_model_context(node, self.config, manifest)
         context.update(extra_context)
+
         if isinstance(node, CompiledGenericTestNode):
             # for test nodes, add a special keyword args value to the context
             jinja.add_rendered_test_kwargs(context, node)

core/dbt/context/manifest.py

@@ -4,6 +4,7 @@
 from dbt.contracts.connection import AdapterRequiredConfig
 from dbt.contracts.graph.manifest import Manifest
 from dbt.context.macro_resolver import TestMacroNamespace
+from .base import contextproperty
 
 
 from .configured import ConfiguredContext
@@ -66,6 +67,10 @@ def to_dict(self):
             dct.update(self.namespace)
         return dct
 
+    @contextproperty
+    def context_macro_stack(self):
+        return self.macro_stack
+
 
 class QueryHeaderContext(ManifestContext):
     def __init__(self, config: AdapterRequiredConfig, manifest: Manifest) -> None:

core/dbt/context/providers.py

@@ -1244,6 +1244,19 @@ def selected_resources(self) -> List[str]:
         """
         return selected_resources.SELECTED_RESOURCES
 
+    @contextmember
+    def submit_python_job(self, parsed_model: Dict, compiled_code: str) -> AdapterResponse:
+        # Check macro_stack and that the unique id is for a materialization macro
+        if not (
+            self.context_macro_stack.depth == 2
+            and self.context_macro_stack.call_stack[1] == "macro.dbt.statement"
+            and "materialization" in self.context_macro_stack.call_stack[0]
+        ):
+            raise RuntimeException(
+                f"submit_python_job is not intended to be called here, at model {parsed_model['alias']}, with macro call_stack {self.context_macro_stack.call_stack}."
+            )
+        return self.adapter.submit_python_job(parsed_model, compiled_code)
+
 
 class MacroContext(ProviderContext):
     """Internally, macros can be executed like nodes, with some restrictions:

core/dbt/include/global_project/macros/etc/statement.sql

@@ -12,7 +12,7 @@ The macro override naming method (spark__statement) only works for macros which
     {%- if language == 'sql'-%}
       {%- set res, table = adapter.execute(compiled_code, auto_begin=auto_begin, fetch=fetch_result) -%}
     {%- elif language == 'python' -%}
-      {%- set res = adapter.submit_python_job(model, compiled_code) -%}
+      {%- set res = submit_python_job(model, compiled_code) -%}
       {#-- TODO: What should table be for python models? --#}
       {%- set table = None -%}
     {%- else -%}

core/dbt/task/run.py

@@ -278,7 +278,9 @@ def execute(self, model, manifest):
 
         hook_ctx = self.adapter.pre_model_hook(context_config)
         try:
-            result = MacroGenerator(materialization_macro, context)()
+            result = MacroGenerator(
+                materialization_macro, context, stack=context["context_macro_stack"]
+            )()
         finally:
             self.adapter.post_model_hook(context_config, hook_ctx)
 

test/unit/test_context.py

@@ -207,7 +207,7 @@ def assert_has_keys(required_keys: Set[str], maybe_keys: Set[str], ctx: Dict[str
 REQUIRED_TARGET_KEYS = REQUIRED_BASE_KEYS | {"target"}
 REQUIRED_DOCS_KEYS = REQUIRED_TARGET_KEYS | {"project_name"} | {"doc"}
 MACROS = frozenset({"macro_a", "macro_b", "root", "dbt"})
-REQUIRED_QUERY_HEADER_KEYS = REQUIRED_TARGET_KEYS | {"project_name"} | MACROS
+REQUIRED_QUERY_HEADER_KEYS = REQUIRED_TARGET_KEYS | {"project_name", "context_macro_stack"} | MACROS
 REQUIRED_MACRO_KEYS = REQUIRED_QUERY_HEADER_KEYS | {
     "_sql_results",
     "load_result",
@@ -239,8 +239,9 @@ def assert_has_keys(required_keys: Set[str], maybe_keys: Set[str], ctx: Dict[str
     "adapter_macro",
     "selected_resources",
     "invocation_args_dict",
+    "submit_python_job"
 }
-REQUIRED_MODEL_KEYS = REQUIRED_MACRO_KEYS | {"this", "compiled_code",}
+REQUIRED_MODEL_KEYS = REQUIRED_MACRO_KEYS | {"this", "compiled_code"}
 MAYBE_KEYS = frozenset({"debug"})