Merge pull request #1983 from fishtown-analytics/fix/restore-drop-schema

Restore drop_schema to adapter in jinja context (#1980)

core/dbt/adapters/base/impl.py

@@ -760,6 +760,7 @@ def create_schema(self, database: str, schema: str):
         )
 
     @abc.abstractmethod
+    @available.parse_none
     def drop_schema(self, database: str, schema: str):
         """Drop the given schema (and everything in it) if it exists."""
         raise dbt.exceptions.NotImplementedException(

core/dbt/adapters/cache.py

@@ -1,6 +1,8 @@
 from collections import namedtuple
-import threading
 from copy import deepcopy
+from typing import List, Iterable
+import threading
+
 from dbt.logger import CACHE_LOGGER as logger
 import dbt.exceptions
 
@@ -181,29 +183,36 @@ def __init__(self):
         self.lock = threading.RLock()
         self.schemas = set()
 
-    def add_schema(self, database, schema):
+    def add_schema(self, database: str, schema: str):
         """Add a schema to the set of known schemas (case-insensitive)
 
-        :param str database: The database name to add.
-        :param str schema: The schema name to add.
+        :param database: The database name to add.
+        :param schema: The schema name to add.
         """
         self.schemas.add((_lower(database), _lower(schema)))
 
-    def remove_schema(self, database, schema):
-        """Remove a schema from the set of known schemas (case-insensitive)
+    def drop_schema(self, database: str, schema: str):
+        """Drop the given schema and remove it from the set of known schemas.
 
-        If the schema does not exist, it will be ignored - it could just be a
-        temporary table.
-
-        :param str database: The database name to remove.
-        :param str schema: The schema name to remove.
+        Then remove all its contents (and their dependents, etc) as well.
         """
-        self.schemas.discard((_lower(database), _lower(schema)))
+        key = (_lower(database), _lower(schema))
+        if key not in self.schemas:
+            return
 
-    def update_schemas(self, schemas):
+        # avoid iterating over self.relations while removing things by
+        # collecting the list first.
+
+        with self.lock:
+            to_remove = self._list_relations_in_schema(database, schema)
+            self._remove_all(to_remove)
+            # handle a drop_schema race by using discard() over remove()
+            self.schemas.discard(key)
+
+    def update_schemas(self, schemas: Iterable[str]):
         """Add multiple schemas to the set of known schemas (case-insensitive)
 
-        :param Iterable[str] schemas: An iterable of the schema names to add.
+        :param schemas: An iterable of the schema names to add.
         """
         self.schemas.update((_lower(d), _lower(s)) for (d, s) in schemas)
 
@@ -402,7 +411,6 @@ def _rename_relation(self, old_key, new_relation):
 
         self.relations[new_key] = relation
         # also fixup the schemas!
-        self.remove_schema(old_key.database, old_key.schema)
         self.add_schema(new_key.database, new_key.schema)
 
         return True
@@ -489,3 +497,25 @@ def clear(self):
         with self.lock:
             self.relations.clear()
             self.schemas.clear()
+
+    def _list_relations_in_schema(
+        self, database: str, schema: str
+    ) -> List[_CachedRelation]:
+        """Get the relations in a schema. Callers should hold the lock."""
+        key = (_lower(database), _lower(schema))
+
+        to_remove: List[_CachedRelation] = []
+        for cachekey, relation in self.relations.items():
+            if (cachekey.database, cachekey.schema) == key:
+                to_remove.append(relation)
+        return to_remove
+
+    def _remove_all(self, to_remove: List[_CachedRelation]):
+        """Remove all the listed relations. Ignore relations that have been
+        cascaded out.
+        """
+        for relation in to_remove:
+            # it may have been cascaded out already
+            drop_key = _make_key(relation)
+            if drop_key in self.relations:
+                self.drop(drop_key)

core/dbt/adapters/sql/impl.py

@@ -172,15 +172,18 @@ def create_schema(self, database, schema):
         }
         self.execute_macro(CREATE_SCHEMA_MACRO_NAME, kwargs=kwargs)
         self.commit_if_has_connection()
+        # we can't update the cache here, as if the schema already existed we
+        # don't want to (incorrectly) say that it's empty
 
     def drop_schema(self, database, schema):
         logger.debug('Dropping schema "{}"."{}".', database, schema)
         kwargs = {
             'database_name': self.quote_as_configured(database, 'database'),
             'schema_name': self.quote_as_configured(schema, 'schema'),
         }
-        self.execute_macro(DROP_SCHEMA_MACRO_NAME,
-                           kwargs=kwargs)
+        self.execute_macro(DROP_SCHEMA_MACRO_NAME, kwargs=kwargs)
+        # we can update the cache here
+        self.cache.drop_schema(database, schema)
 
     def list_relations_without_caching(self, information_schema, schema):
         kwargs = {'information_schema': information_schema, 'schema': schema}

plugins/bigquery/dbt/adapters/bigquery/impl.py

@@ -201,6 +201,7 @@ def create_schema(self, database: str, schema: str) -> None:
     def drop_schema(self, database: str, schema: str) -> None:
         logger.debug('Dropping schema "{}.{}".', database, schema)
         self.connections.drop_dataset(database, schema)
+        self.cache.drop_schema(database, schema)
 
     @classmethod
     def quote(cls, identifier: str) -> str:

test/integration/054_adapter_methods_test/models/expected.sql

@@ -0,0 +1,3 @@
+-- make sure this runs after 'model'
+-- {{ ref('model') }}
+select 2 as id

test/integration/054_adapter_methods_test/models/model.sql

@@ -0,0 +1,19 @@
+
+{% set upstream = ref('upstream') %}
+
+{% if execute %}
+    {# don't ever do any of this #}
+    {%- do adapter.drop_schema(upstream.database, upstream.schema) -%}
+    {% set existing = adapter.get_relation(upstream.database, upstream.schema, upstream.identifier) %}
+    {% if existing is not none %}
+        {% do exceptions.raise_compiler_error('expected ' ~ ' to not exist, but it did') %}
+    {% endif %}
+
+    {%- do adapter.create_schema(upstream.database, upstream.schema) -%}
+
+    {% set sql = create_view_as(upstream, 'select 2 as id') %}
+    {% do run_query(sql) %}
+{% endif %}
+
+
+select * from {{ upstream }}

test/integration/054_adapter_methods_test/models/upstream.sql

@@ -0,0 +1 @@
+select 1 as id

test/integration/054_adapter_methods_test/test_adapter_methods.py

@@ -0,0 +1,31 @@
+from test.integration.base import DBTIntegrationTest, use_profile
+
+
+class TestBaseCaching(DBTIntegrationTest):
+    @property
+    def schema(self):
+        return "caching_038"
+
+    @property
+    def models(self):
+        return "models"
+
+    @use_profile('postgres')
+    def test_postgres_adapter_methods(self):
+        self.run_dbt()
+        self.assertTablesEqual('model', 'expected')
+
+    @use_profile('redshift')
+    def test_redshift_adapter_methods(self):
+        self.run_dbt()
+        self.assertTablesEqual('model', 'expected')
+
+    @use_profile('snowflake')
+    def test_snowflake_adapter_methods(self):
+        self.run_dbt()
+        self.assertTablesEqual('MODEL', 'EXPECTED')
+
+    @use_profile('bigquery')
+    def test_bigquery_adapter_methods(self):
+        self.run_dbt()
+        self.assertTablesEqual('model', 'expected')

test/unit/test_cache.py

@@ -131,7 +131,8 @@ def test_dest_different_db(self):
                           make_relation('DBT_2', 'schema', 'foo'))
         self.assert_relations_exist('DBT_2', 'schema', 'foo')
         self.assert_relations_do_not_exist('DBT', 'schema', 'foo')
-        self.assertEqual(self.cache.schemas, {('dbt_2', 'schema')})
+        # we know about both schemas: dbt has nothing, dbt_2 has something.
+        self.assertEqual(self.cache.schemas, {('dbt_2', 'schema'), ('dbt', 'schema')})
         self.assertEqual(len(self.cache.relations), 1)
 
     def test_rename_identifier(self):
@@ -156,7 +157,8 @@ def test_rename_db(self):
         self.assertEqual(len(self.cache.get_relations('DBT_2', 'schema')), 1)
         self.assert_relations_exist('DBT_2', 'schema', 'foo')
         self.assert_relations_do_not_exist('DBT', 'schema', 'foo')
-        self.assertEqual(self.cache.schemas, {('dbt_2', 'schema')})
+        # we know about both schemas: dbt has nothing, dbt_2 has something.
+        self.assertEqual(self.cache.schemas, {('dbt_2', 'schema'), ('dbt', 'schema')})
 
         relation = self.cache.relations[('dbt_2', 'schema', 'foo')]
         self.assertEqual(relation.inner.database, 'DBT_2')
@@ -174,7 +176,8 @@ def test_rename_schema(self):
         self.assertEqual(len(self.cache.get_relations('DBT', 'schema_2')), 1)
         self.assert_relations_exist('DBT', 'schema_2', 'foo')
         self.assert_relations_do_not_exist('DBT', 'schema', 'foo')
-        self.assertEqual(self.cache.schemas, {('dbt', 'schema_2')})
+        # we know about both schemas: schema has nothing, schema_2 has something.
+        self.assertEqual(self.cache.schemas, {('dbt', 'schema_2'), ('dbt', 'schema')})
 
         relation = self.cache.relations[('dbt', 'schema_2', 'foo')]
         self.assertEqual(relation.inner.database, 'DBT')