Fix connection issue in dbt 1.0.0 (#2230)

* Fix connection issue in dbt 1.0.0 * Set tests to use default Postgres database * Add postgres service to GH action ci tests * Fix test__templater_dbt_handle_database_connection_failure * Fixing * Correct database connection test * Add real incremental model to test database connection * Use mock to simulate database failure in "test__templater_dbt_handle_database_connection_failure" * Attempt to ignore templator tests for Windows * Move more tests * Linting fixes * Revert "Linting fixes" This reverts commit 68b1bfb. * Revert "Move more tests" This reverts commit 2a5e662. * Revert "Attempt to ignore templator tests for Windows" This reverts commit ca71069. * Attempt 2 to skip the tests for windows * Linting * Skip test on dbt versions below 1.0 * CONTRIBUTING.md: Document the dbt Docker Compose environment * jy_cleanup (#2257) * jy_cleanup * Remove unnecessary dbt_connection mark Co-authored-by: Barry Hart <barry.hart@mailchimp.com> Co-authored-by: Barry <barry@tunetheweb.com> Co-authored-by: Joseph Young <80432516+jpy-git@users.noreply.github.com>
sqlfluff · Jan 8, 2022 · 286a526 · 286a526
1 parent 3861582
commit 286a526
Show file tree

Hide file tree

Showing 9 changed files with 108 additions and 32 deletions.
diff --git a/.github/workflows/ci-tests.yml b/.github/workflows/ci-tests.yml
@@ -111,6 +111,23 @@ jobs:
         done
   dbt-tests:
     runs-on: ubuntu-latest
+    services:
+      # Label used to access the service container
+      postgres:
+        # Docker Hub image
+        image: postgres
+        # Provide the password for postgres
+        env:
+          POSTGRES_PASSWORD: password
+        # Set health checks to wait until postgres has started
+        options: >-
+          --health-cmd pg_isready
+          --health-interval 10s
+          --health-timeout 5s
+          --health-retries 5
+        ports:
+          # Maps tcp port 5432 on service container to the host
+          - 5432:5432
     strategy:
       matrix:
         dbt-version: [ 'dbt018', 'dbt019', 'dbt020', 'dbt021', 'dbt100' ]
@@ -186,7 +203,7 @@ jobs:
       # Do not set explicitly temp dir for dbt as causes problems
       # None of these test need temp dir set
       run: |
-        python -m tox -e dbt018-winpy -- plugins/sqlfluff-templater-dbt
+        python -m tox -e dbt018-winpy -- plugins/sqlfluff-templater-dbt -m "not dbt_connection_failure and not needs_dbt_connection"
     - name: Upload Coverage Report
       uses: codecov/codecov-action@v1
       with:

diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
@@ -168,6 +168,10 @@ tox -e py38 -- test/cli
 tox -e py38 -- test/cli/commands_test.py
 ```
 
+#### dbt templater tests
+
+The dbt templater tests require a locally running Postgres instance. See the required connection parameters in `plugins/sqlfluff-templater-dbt/test/fixtures/dbt/profiles.yml`. We recommend using https://postgresapp.com/.
+
 To run the dbt-related tests you will have to explicitly include these tests:
 
 ```shell
@@ -176,6 +180,26 @@ tox -e cov-init,dbt018-py38,cov-report-dbt -- plugins/sqlfluff-templater-dbt
 
 For more information on adding and running test cases see the [Parser Test README](test/fixtures/dialects/README.md) and the [Rules Test README](test/fixtures/rules/std_rule_cases/README.md).
 
+#### Running dbt templater tests in Docker Compose
+
+NOTE: If you prefer, you can develop and debug the dbt templater using a
+Docker Compose environment. It's a simple two-container configuration:
+* `app`: Hosts the SQLFluff development environment. The host's source
+  directory is mounted into the container, so you can iterate on code
+  changes without having to constantly rebuild and restart the container.
+* `postgres`: Hosts a transient Postgres database instance.
+
+Steps to use the Docker Compose environment:
+* Install Docker on your machine.
+* Run `plugins/sqlfluff-templater-dbt/docker/startup` to create the containers.
+* Run `plugins/sqlfluff-templater-dbt/docker/shell` to start a bash session in the `app` container.
+* Manually edit the `host` value in `plugins/sqlfluff-templater-dbt/test/fixtures/dbt/profiles_yml/profiles.yml`, changing it to `postgres`. (This will likely be automated somehow in the future.)
+
+Inside the container, run:
+```
+py.test -v plugins/sqlfluff-templater-dbt/test/
+```
+
 ### Pre-Commit Config
 
 For development convenience we also provide a `.pre-commit-config.yaml` file to allow the user to install a selection of pre-commit hooks via `tox -e pre-commit -- install`. These hooks can help the user identify and fix potential linting/typing violations prior to committing their code and therefore reduce having to deal with these sort of issues during code review.

diff --git a/plugins/sqlfluff-templater-dbt/sqlfluff_templater_dbt/templater.py b/plugins/sqlfluff-templater-dbt/sqlfluff_templater_dbt/templater.py
@@ -11,7 +11,7 @@
 
 from dbt.version import get_installed_version
 from dbt.config.runtime import RuntimeConfig as DbtRuntimeConfig
-from dbt.adapters.factory import register_adapter
+from dbt.adapters.factory import register_adapter, get_adapter
 from dbt.compilation import Compiler as DbtCompiler
 from dbt.exceptions import (
     CompilationException as DbtCompilationException,
@@ -54,6 +54,7 @@ class DbtConfigArgs:
     project_dir: Optional[str] = None
     profiles_dir: Optional[str] = None
     profile: Optional[str] = None
+    single_threaded: bool = False
 
 
 class DbtTemplater(JinjaTemplater):
@@ -446,6 +447,13 @@ def make_template(in_str):
 
         node = self._find_node(fname, config)
 
+        # We have to register the connection in dbt >= 1.0.0 ourselves
+        # In previous versions, we relied on the functionality removed in https://github.com/dbt-labs/dbt-core/pull/4062
+        if DBT_VERSION_TUPLE >= (1, 0):
+            adapter = get_adapter(self.dbt_config)
+            with adapter.connection_named("master"):
+                adapter.set_relations_cache(self.dbt_manifest)
+
         node = self.dbt_compiler.compile_node(
             node=node,
             manifest=self.dbt_manifest,

diff --git a/...qlfluff-templater-dbt/test/fixtures/dbt/dbt_project/models/my_new_project/incremental.sql b/...qlfluff-templater-dbt/test/fixtures/dbt/dbt_project/models/my_new_project/incremental.sql
@@ -1,15 +1,21 @@
 -- https://github.com/sqlfluff/sqlfluff/issues/780
--- Note we don't call is_incremental directly in the test
--- because that requires a database connection.
+
+{{
+  config(
+    materialized = 'incremental',
+    unique_key='product_id'
+    )
+}}
 
 select
     {#- Attributes #}
     products.product_id,
+    products.valid_date_local,
     products._fivetran_deleted
 from products
 inner join dispensaries
 where not products._fivetran_deleted
-    {% if true -%}
+{% if is_incremental() -%}
     and products.valid_date_local >= (
         select max(valid_date_local) from {{ this }})
-    {% endif %}
+{% endif %}
diff --git a/...plater-dbt/test/fixtures/dbt/profiles.yml → ...st/fixtures/dbt/profiles_yml/profiles.yml b/...plater-dbt/test/fixtures/dbt/profiles.yml → ...st/fixtures/dbt/profiles_yml/profiles.yml
@@ -4,10 +4,10 @@ default:
     dev:
       type: postgres
       host: localhost
-      user: alice
-      pass: <password>
+      user: postgres
+      pass: password
       port: 5432
-      dbname: jaffle_shop
+      dbname: postgres
       schema: dbt_alice
       threads: 4
 

diff --git a/plugins/sqlfluff-templater-dbt/test/fixtures/dbt/profiles_yml_fail/profiles.yml b/plugins/sqlfluff-templater-dbt/test/fixtures/dbt/profiles_yml_fail/profiles.yml
@@ -0,0 +1,15 @@
+default:
+  target: dev
+  outputs:
+    dev:
+      type: postgres
+      host: localhost
+      user: postgres
+      pass: password
+      port: 2345
+      dbname: postgres
+      schema: dbt_alice
+      threads: 4
+
+config:
+  send_anonymous_usage_stats: false
diff --git a/plugins/sqlfluff-templater-dbt/test/fixtures/dbt/templated_output/use_dbt_utils.sql b/plugins/sqlfluff-templater-dbt/test/fixtures/dbt/templated_output/use_dbt_utils.sql
@@ -2,7 +2,7 @@ with
 
 orders as (
     select *
-    from "jaffle_shop"."jaffle_shop"."orders"
+    from "postgres"."jaffle_shop"."orders"
 )
 
 select

diff --git a/plugins/sqlfluff-templater-dbt/test/fixtures/dbt/templater.py b/plugins/sqlfluff-templater-dbt/test/fixtures/dbt/templater.py
@@ -11,7 +11,7 @@
     },
     "templater": {
         "dbt": {
-            "profiles_dir": "plugins/sqlfluff-templater-dbt/test/fixtures/dbt",
+            "profiles_dir": "plugins/sqlfluff-templater-dbt/test/fixtures/dbt/profiles_yml",
             "project_dir": "plugins/sqlfluff-templater-dbt/test/fixtures/dbt/dbt_project",
         },
     },

diff --git a/plugins/sqlfluff-templater-dbt/test/templater_test.py b/plugins/sqlfluff-templater-dbt/test/templater_test.py
@@ -2,10 +2,11 @@
 
 import glob
 import os
-import pytest
 import logging
 from pathlib import Path
+from unittest import mock
 
+import pytest
 
 from sqlfluff.core import FluffConfig, Lexer, Linter
 from sqlfluff.core.errors import SQLTemplaterSkipFile
@@ -15,6 +16,7 @@
     dbt_templater,
     project_dir,
 )
+from sqlfluff_templater_dbt.templater import DbtFailedToConnectException
 
 
 def test__templater_dbt_missing(dbt_templater, project_dir):  # noqa: F811
@@ -224,8 +226,10 @@ def test__templater_dbt_skips_disabled_model(dbt_templater, project_dir):  # noq
     "fname",
     [
         "use_var.sql",
-        "incremental.sql",
-        "single_trailing_newline.sql",
+        pytest.param("incremental.sql", marks=pytest.mark.needs_dbt_connection),
+        pytest.param(
+            "single_trailing_newline.sql", marks=pytest.mark.needs_dbt_connection
+        ),
         "L034_test.sql",
     ],
 )
@@ -317,18 +321,29 @@ def test__templater_dbt_handle_exceptions(
     assert violations[0].desc().replace("\\", "/").startswith(exception_msg)
 
 
+@pytest.mark.skipif(
+    DBT_VERSION_TUPLE < (1, 0), reason="mocks a function that's only used in dbt >= 1.0"
+)
+@mock.patch("dbt.adapters.postgres.impl.PostgresAdapter.set_relations_cache")
+@pytest.mark.dbt_connection_failure
 def test__templater_dbt_handle_database_connection_failure(
-    project_dir, dbt_templater  # noqa: F811
+    set_relations_cache, project_dir, dbt_templater  # noqa: F811
 ):
     """Test the result of a failed database connection."""
     from dbt.adapters.factory import get_adapter
 
+    set_relations_cache.side_effect = DbtFailedToConnectException("dummy error")
+
     src_fpath = "plugins/sqlfluff-templater-dbt/test/fixtures/dbt/error_models/exception_connect_database.sql"
     target_fpath = os.path.abspath(
         os.path.join(
             project_dir, "models/my_new_project/exception_connect_database.sql"
         )
     )
+    dbt_fluff_config_fail = DBT_FLUFF_CONFIG.copy()
+    dbt_fluff_config_fail["templater"]["dbt"][
+        "profiles_dir"
+    ] = "plugins/sqlfluff-templater-dbt/test/fixtures/dbt/profiles_yml_fail"
     # We move the file that throws an error in and out of the project directory
     # as dbt throws an error if a node fails to parse while computing the DAG
     os.rename(src_fpath, target_fpath)
@@ -338,26 +353,17 @@ def test__templater_dbt_handle_database_connection_failure(
             fname=target_fpath,
             config=FluffConfig(configs=DBT_FLUFF_CONFIG),
         )
-    except Exception as e:
-        if DBT_VERSION_TUPLE >= (1, 0):
-            # In dbt 1.0.0, connection failures raise an exception
-            assert str(e).startswith(
-                "Runtime Error\n  connection never acquired for thread"
-            )
-        else:
-            raise (e)
     finally:
         get_adapter(dbt_templater.dbt_config).connections.release()
         os.rename(target_fpath, src_fpath)
-    if DBT_VERSION_TUPLE < (1, 0):
-        assert violations
-        # NB: Replace slashes to deal with different plaform paths being returned.
-        assert (
-            violations[0]
-            .desc()
-            .replace("\\", "/")
-            .startswith("dbt tried to connect to the database")
-        )
+    assert violations
+    # NB: Replace slashes to deal with different plaform paths being returned.
+    assert (
+        violations[0]
+        .desc()
+        .replace("\\", "/")
+        .startswith("dbt tried to connect to the database")
+    )
 
 
 def test__project_dir_does_not_exist_error(dbt_templater, caplog):  # noqa: F811