Runtime: Prevent introspective queries at compile (SL only) (#5926) (#…

…5943)

* Preliminary changes to keep compile from connecting to the warehouse for runtime calls

* Adds option to lib to skip connecting to warehouse for compile; adds prelim tests

* Removes unused imports

* Simplifies test and renames to SqlCompileRunnerNoIntrospection

* Updates name in tests

* Spacing

* Updates test to check for adapter connection call instead of compile and execute

* Removes commented line

* Fixes test names

* Updates plugin to postgres type as snowflake isn't available

* Fixes docstring

* Fixes formatting

* Moves conditional logic out of class

* Fixes formatting

* Removes commented line

* Moves import

* Unmoves import

* Updates changelog

* Adds further info to method docstring

(cherry picked from commit f1326f5)

Co-authored-by: Rachel <[email protected]>

.changes/unreleased/Features-20220926-130627.yaml

@@ -0,0 +1,9 @@
+kind: Features
+body: This conditionally no-ops warehouse connection at compile depending on an env
+  var, disabling introspection/queries during compilation only. This is a temporary
+  solution to more complex permissions requirements for the semantic layer.
+time: 2022-09-26T13:06:27.591061-05:00
+custom:
+  Author: racheldaniel
+  Issue: "5936"
+  PR: "5926"

core/dbt/lib.py

@@ -1,6 +1,10 @@
 import os
+from dbt.contracts.results import RunningStatus, collect_timing_info
+from dbt.events.functions import fire_event
+from dbt.events.types import NodeCompiling, NodeExecuting
 from dbt.exceptions import RuntimeException
 from dbt import flags
+from dbt.task.sql import SqlCompileRunner
 from dataclasses import dataclass
 
 
@@ -13,6 +17,50 @@ class RuntimeArgs:
     target: str
 
 
+class SqlCompileRunnerNoIntrospection(SqlCompileRunner):
+    def compile_and_execute(self, manifest, ctx):
+        """
+        This version of this method does not connect to the data warehouse.
+        As a result, introspective queries at compilation will not be supported
+        and will throw an error.
+
+        TODO: This is a temporary solution to more complex permissions requirements
+        for the semantic layer, and thus largely duplicates the code in the parent class
+        method. Once conditional credential usage is enabled, this should be removed.
+        """
+        result = None
+        ctx.node._event_status["node_status"] = RunningStatus.Compiling
+        fire_event(
+            NodeCompiling(
+                node_info=ctx.node.node_info,
+                unique_id=ctx.node.unique_id,
+            )
+        )
+        with collect_timing_info("compile") as timing_info:
+            # if we fail here, we still have a compiled node to return
+            # this has the benefit of showing a build path for the errant
+            # model
+            ctx.node = self.compile(manifest)
+        ctx.timing.append(timing_info)
+
+        # for ephemeral nodes, we only want to compile, not run
+        if not ctx.node.is_ephemeral_model:
+            ctx.node._event_status["node_status"] = RunningStatus.Executing
+            fire_event(
+                NodeExecuting(
+                    node_info=ctx.node.node_info,
+                    unique_id=ctx.node.unique_id,
+                )
+            )
+            with collect_timing_info("execute") as timing_info:
+                result = self.run(ctx.node, manifest)
+                ctx.node = result.node
+
+            ctx.timing.append(timing_info)
+
+        return result
+
+
 def get_dbt_config(project_dir, args=None, single_threaded=False):
     from dbt.config.runtime import RuntimeConfig
     import dbt.adapters.factory
@@ -104,12 +152,17 @@ def _get_operation_node(manifest, project_path, sql, node_name):
 
 
 def compile_sql(manifest, project_path, sql, node_name="query"):
-    from dbt.task.sql import SqlCompileRunner
-
     config, node, adapter = _get_operation_node(manifest, project_path, sql, node_name)
+    allow_introspection = str(os.environ.get("__DBT_ALLOW_INTROSPECTION", "1")).lower() in (
+        "true",
+        "1",
+        "on",
+    )
 
-    runner = SqlCompileRunner(config, adapter, node, 1, 1)
-
+    if allow_introspection:
+        runner = SqlCompileRunner(config, adapter, node, 1, 1)
+    else:
+        runner = SqlCompileRunnerNoIntrospection(config, adapter, node, 1, 1)
     return runner.safe_run(manifest)
 
 

test/unit/test_lib.py

@@ -0,0 +1,63 @@
+import os
+import unittest
+from unittest import mock
+from dbt.contracts.results import RunningStatus
+from dbt.lib import compile_sql
+from dbt.adapters.postgres import Plugin
+
+from test.unit.utils import clear_plugin, inject_adapter
+
+
+class MockContext:
+    def __init__(self, node):
+        self.timing = []
+        self.node = mock.MagicMock()
+        self.node._event_status = {
+            "node_status": RunningStatus.Started
+        }
+        self.node.is_ephemeral_model = True
+
+def noop_ephemeral_result(*args):
+    return None
+
+class TestSqlCompileRunnerNoIntrospection(unittest.TestCase):
+    def setUp(self):
+            self.manifest = {'mock':'manifest'}
+            self.adapter = Plugin.adapter({})
+            self.adapter.connection_for = mock.MagicMock()
+            self.ephemeral_result = lambda: None
+            inject_adapter(self.adapter, Plugin)
+
+    def tearDown(self):
+        clear_plugin(Plugin)
+
+    @mock.patch('dbt.lib._get_operation_node')
+    @mock.patch('dbt.task.sql.GenericSqlRunner.compile')
+    @mock.patch('dbt.task.sql.GenericSqlRunner.ephemeral_result', noop_ephemeral_result)
+    @mock.patch('dbt.task.base.ExecutionContext', MockContext)
+    def test__compile_and_execute__with_connection(self, mock_compile, mock_get_node):
+        """
+        By default, env var for allowing introspection is true, and calling this
+        method should defer to the parent method.
+        """
+        mock_get_node.return_value = ({}, None, self.adapter)
+        compile_sql(self.manifest, 'some/path', None)
+
+        mock_compile.assert_called_once_with(self.manifest)
+        self.adapter.connection_for.assert_called_once()
+
+
+    @mock.patch('dbt.lib._get_operation_node')
+    @mock.patch('dbt.task.sql.GenericSqlRunner.compile')
+    @mock.patch('dbt.task.sql.GenericSqlRunner.ephemeral_result', noop_ephemeral_result)
+    @mock.patch('dbt.task.base.ExecutionContext', MockContext)
+    def test__compile_and_execute__without_connection(self, mock_compile, mock_get_node):
+        """
+        Ensure that compile is called but does not attempt warehouse connection
+        """
+        with mock.patch.dict(os.environ, {"__DBT_ALLOW_INTROSPECTION": "0"}):
+            mock_get_node.return_value = ({}, None, self.adapter)
+            compile_sql(self.manifest, 'some/path', None)
+
+            mock_compile.assert_called_once_with(self.manifest)
+            self.adapter.connection_for.assert_not_called()