Merge branch 'master' into devin/1772278180-cleanup-stale-ci-schemas

Author: haritamar

.github/workflows/test-warehouse.yml

@@ -277,3 +277,10 @@ jobs:
 
       - name: Run Python package e2e tests
         run: pytest -vv tests/e2e --warehouse-type ${{ inputs.warehouse-type }}
+
+      - name: Drop test schemas
+        if: always()
+        working-directory: ${{ env.E2E_DBT_PROJECT_DIR }}
+        continue-on-error: true
+        run: |
+          dbt run-operation elementary_integration_tests.drop_test_schemas --target "${{ inputs.warehouse-type }}"

elementary/clients/dbt/api_dbt_runner.py

@@ -49,15 +49,24 @@ def collect_dbt_command_logs(event):
             with with_chdir(self.project_dir):
                 res: dbtRunnerResult = dbt.invoke(dbt_command_args)
         output = "\n".join(dbt_logs) or None
+        # Surface the exception text so that transient-error detection in
+        # _inner_run_command_with_retries can match against it.  The dbt
+        # Python API doesn't write to stderr, so we repurpose that field
+        # for the exception string (analogous to how SubprocessDbtRunner
+        # captures subprocess stderr).
+        exception_text = str(res.exception) if res.exception else None
         if self.raise_on_failure and not res.success:
             raise DbtCommandError(
                 base_command_args=dbt_command_args,
-                err_msg=(str(res.exception) if res.exception else output),
+                err_msg=(exception_text or output),
                 logs=[DbtLog.from_log_line(log) for log in dbt_logs],
             )
 
         return APIDbtCommandResult(
-            success=res.success, output=output, stderr=None, result_obj=res
+            success=res.success,
+            output=output,
+            stderr=exception_text,
+            result_obj=res,
         )
 
     def _parse_ls_command_result(

tests/e2e_dbt_project/dbt_project.yml

@@ -20,6 +20,8 @@ vars:
   days_back: 30
   debug_logs: "{{ env_var('DBT_EDR_DEBUG', False) }}"
   custom_run_started_at: "{{ modules.datetime.datetime.utcfromtimestamp(0) }}"
+  clean_elementary_temp_tables: false
+  disable_dbt_artifacts_autoupload: true
 
 seeds:
   +schema: test_seeds

tests/e2e_dbt_project/macros/system/drop_test_schemas.sql

@@ -0,0 +1,36 @@
+{% macro drop_test_schemas() %}
+    {# Drop both the main test schema and the elementary schema used by the CLI.
+       The schema names are derived from the profile's target schema. #}
+    {% set main_schema = target.schema %}
+    {% set elementary_schema = main_schema ~ '_elementary' %}
+
+    {% do elementary_integration_tests.edr_drop_schema(elementary_schema) %}
+    {% do elementary_integration_tests.edr_drop_schema(main_schema) %}
+    {% do log("Dropped schemas: " ~ main_schema ~ ", " ~ elementary_schema, info=true) %}
+{% endmacro %}
+
+{% macro edr_drop_schema(schema_name) %}
+    {% do return(adapter.dispatch('edr_drop_schema', 'elementary_integration_tests')(schema_name)) %}
+{% endmacro %}
+
+{% macro default__edr_drop_schema(schema_name) %}
+    {% set schema_relation = api.Relation.create(database=target.database, schema=schema_name) %}
+    {% do dbt.drop_schema(schema_relation) %}
+    {% do adapter.commit() %}
+{% endmacro %}
+
+{% macro bigquery__edr_drop_schema(schema_name) %}
+    {% set schema_relation = api.Relation.create(database=target.database, schema=schema_name) %}
+    {% do dbt.drop_schema(schema_relation) %}
+{% endmacro %}
+
+{% macro clickhouse__edr_drop_schema(schema_name) %}
+    {% set quoted_schema = adapter.quote(schema_name) %}
+    {% do run_query("DROP DATABASE IF EXISTS " ~ quoted_schema) %}
+    {% do adapter.commit() %}
+{% endmacro %}
+
+{% macro athena__edr_drop_schema(schema_name) %}
+    {% set schema_relation = api.Relation.create(database=target.database, schema=schema_name) %}
+    {% do dbt.drop_schema(schema_relation) %}
+{% endmacro %}

tests/unit/clients/dbt_runner/test_retry_logic.py

@@ -201,3 +201,116 @@ def test_successful_command_not_retried(self, mock_subprocess_run):
 
         assert mock_subprocess_run.call_count == 1
         assert result is True
+
+
+def _make_api_runner(**kwargs):
+    """Create an APIDbtRunner with deps/packages stubbed out."""
+    defaults = dict(
+        project_dir="fake_project",
+        profiles_dir="fake_profiles",
+        target=None,
+        raise_on_failure=False,
+        run_deps_if_needed=False,
+    )
+    defaults.update(kwargs)
+    from elementary.clients.dbt.api_dbt_runner import APIDbtRunner
+
+    with mock.patch.object(APIDbtRunner, "_run_deps_if_needed"):
+        return APIDbtRunner(**defaults)
+
+
+@_ZERO_WAIT
+class TestAPIDbtRunnerTransientDetection:
+    """Test that APIDbtRunner surfaces exception text for transient error detection.
+
+    The dbt Python API (APIDbtRunner) only captures JinjaLogInfo and
+    RunningOperationCaughtError events into ``output``.  Transient errors
+    like RemoteDisconnected appear as ``res.exception`` — not in the
+    captured output.  Without surfacing this, the retry logic has nothing
+    to match against and never fires.
+    """
+
+    @mock.patch(
+        "elementary.clients.dbt.api_dbt_runner.with_chdir",
+        return_value=mock.MagicMock(
+            __enter__=mock.MagicMock(), __exit__=mock.MagicMock()
+        ),
+    )
+    @mock.patch("elementary.clients.dbt.api_dbt_runner.dbtRunner")
+    def test_transient_exception_triggers_retry(self, mock_dbt_runner_cls, _mock_chdir):
+        """A transient exception in res.exception should be retried."""
+        # Simulate dbtRunnerResult with a transient exception.
+        fail_result = mock.MagicMock()
+        fail_result.success = False
+        fail_result.exception = ConnectionError(
+            "('Connection aborted.', "
+            "RemoteDisconnected('Remote end closed connection without response'))"
+        )
+
+        success_result = mock.MagicMock()
+        success_result.success = True
+        success_result.exception = None
+
+        # dbtRunner().invoke returns fail first, then success.
+        mock_dbt_instance = mock.MagicMock()
+        mock_dbt_instance.invoke.side_effect = [fail_result, success_result]
+        mock_dbt_runner_cls.return_value = mock_dbt_instance
+
+        runner = _make_api_runner(raise_on_failure=False)
+        result = runner.seed()
+
+        assert mock_dbt_instance.invoke.call_count == 2
+        assert result is True
+
+    @mock.patch(
+        "elementary.clients.dbt.api_dbt_runner.with_chdir",
+        return_value=mock.MagicMock(
+            __enter__=mock.MagicMock(), __exit__=mock.MagicMock()
+        ),
+    )
+    @mock.patch("elementary.clients.dbt.api_dbt_runner.dbtRunner")
+    def test_non_transient_exception_not_retried(
+        self, mock_dbt_runner_cls, _mock_chdir
+    ):
+        """A non-transient exception should NOT be retried."""
+        fail_result = mock.MagicMock()
+        fail_result.success = False
+        fail_result.exception = Exception("Compilation Error in model foo")
+
+        mock_dbt_instance = mock.MagicMock()
+        mock_dbt_instance.invoke.return_value = fail_result
+        mock_dbt_runner_cls.return_value = mock_dbt_instance
+
+        runner = _make_api_runner(raise_on_failure=False)
+        result = runner.seed()
+
+        assert mock_dbt_instance.invoke.call_count == 1
+        assert result is False
+
+    @mock.patch(
+        "elementary.clients.dbt.api_dbt_runner.with_chdir",
+        return_value=mock.MagicMock(
+            __enter__=mock.MagicMock(), __exit__=mock.MagicMock()
+        ),
+    )
+    @mock.patch("elementary.clients.dbt.api_dbt_runner.dbtRunner")
+    def test_transient_exception_exhausts_retries(
+        self, mock_dbt_runner_cls, _mock_chdir
+    ):
+        """After exhausting retries, the last failed result is returned."""
+        fail_result = mock.MagicMock()
+        fail_result.success = False
+        fail_result.exception = ConnectionError(
+            "('Connection aborted.', "
+            "RemoteDisconnected('Remote end closed connection without response'))"
+        )
+
+        mock_dbt_instance = mock.MagicMock()
+        mock_dbt_instance.invoke.return_value = fail_result
+        mock_dbt_runner_cls.return_value = mock_dbt_instance
+
+        runner = _make_api_runner(raise_on_failure=False)
+        result = runner.seed()
+
+        assert mock_dbt_instance.invoke.call_count == _TRANSIENT_MAX_RETRIES
+        assert result is False