Use threadpools for filling the cache and listing schemas

CHANGELOG.md

@@ -16,6 +16,9 @@
 - Fix an issue where dbt rendered source test args, fix issue where dbt ran an extra compile pass over the wrapped SQL. ([#2114](https://github.com/fishtown-analytics/dbt/issues/2114), [#2150](https://github.com/fishtown-analytics/dbt/pull/2150))
 - Set more upper bounds for jinja2,requests, and idna dependencies, upgrade snowflake-connector-python ([#2147](https://github.com/fishtown-analytics/dbt/issues/2147), [#2151](https://github.com/fishtown-analytics/dbt/pull/2151))
 
+### Under the hood
+- Parallelize filling the cache and listing schemas in each database during startup ([#2127](https://github.com/fishtown-analytics/dbt/issues/2127), [#2157](https://github.com/fishtown-analytics/dbt/pull/2157))
+
 Contributors:
  - [@bubbomb](https://github.com/bubbomb) ([#2080](https://github.com/fishtown-analytics/dbt/pull/2080))
  - [@sonac](https://github.com/sonac) ([#2078](https://github.com/fishtown-analytics/dbt/pull/2078))

core/dbt/adapters/base/impl.py

@@ -1,6 +1,5 @@
 import abc
-from concurrent.futures import ThreadPoolExecutor, as_completed
-from concurrent.futures import Future  # noqa - we use this for typing only
+from concurrent.futures import as_completed, Future
 from contextlib import contextmanager
 from datetime import datetime
 from typing import (
@@ -27,7 +26,7 @@
 from dbt.exceptions import warn_or_error
 from dbt.node_types import NodeType
 from dbt.logger import GLOBAL_LOGGER as logger
-from dbt.utils import filter_null_values
+from dbt.utils import filter_null_values, executor
 
 from dbt.adapters.base.connections import BaseConnectionManager, Connection
 from dbt.adapters.base.meta import AdapterMeta, available
@@ -358,23 +357,35 @@ def _get_cache_schemas(
         # databases
         return info_schema_name_map
 
+    def _list_relations_get_connection(
+        self, db: BaseRelation, schema: str
+    ) -> List[BaseRelation]:
+        with self.connection_named(f'list_{db.database}_{schema}'):
+            return self.list_relations_without_caching(db, schema)
+
     def _relations_cache_for_schemas(self, manifest: Manifest) -> None:
         """Populate the relations cache for the given schemas. Returns an
         iterable of the schemas populated, as strings.
         """
         if not dbt.flags.USE_CACHE:
             return
 
-        info_schema_name_map = self._get_cache_schemas(manifest,
-                                                       exec_only=True)
-        for db, schema in info_schema_name_map.search():
-            for relation in self.list_relations_without_caching(db, schema):
-                self.cache.add(relation)
+        schema_map = self._get_cache_schemas(manifest, exec_only=True)
+        with executor(self.config) as tpe:
+            futures: List[Future[List[BaseRelation]]] = [
+                tpe.submit(self._list_relations_get_connection, db, schema)
+                for db, schema in schema_map.search()
+            ]
+            for future in as_completed(futures):
+                # if we can't read the relations we need to just raise anyway,
+                # so just call future.result() and let that raise on failure
+                for relation in future.result():
+                    self.cache.add(relation)
 
         # it's possible that there were no relations in some schemas. We want
         # to insert the schemas we query into the cache's `.schemas` attribute
         # so we can check it later
-        self.cache.update_schemas(info_schema_name_map.schemas_searched())
+        self.cache.update_schemas(schema_map.schemas_searched())
 
     def set_relations_cache(
         self, manifest: Manifest, clear: bool = False
@@ -1047,13 +1058,11 @@ def _get_one_catalog(
     def get_catalog(
         self, manifest: Manifest
     ) -> Tuple[agate.Table, List[Exception]]:
-        # snowflake is super slow. split it out into the specified threads
-        num_threads = self.config.threads
         schema_map = self._get_cache_schemas(manifest)
 
-        with ThreadPoolExecutor(max_workers=num_threads) as executor:
-            futures = [
-                executor.submit(self._get_one_catalog, info, schemas, manifest)
+        with executor(self.config) as tpe:
+            futures: List[Future[agate.Table]] = [
+                tpe.submit(self._get_one_catalog, info, schemas, manifest)
                 for info, schemas in schema_map.items() if len(schemas) > 0
             ]
             catalogs, exceptions = catch_as_completed(futures)

core/dbt/task/runnable.py

@@ -1,8 +1,9 @@
 import os
 import time
+from concurrent.futures import as_completed, wait
 from datetime import datetime
 from multiprocessing.dummy import Pool as ThreadPool
-from typing import Optional, Dict, List, Set, Tuple
+from typing import Optional, Dict, List, Set, Tuple, Iterable
 
 from dbt.task.base import ConfiguredTask
 from dbt.adapters.factory import get_adapter
@@ -374,7 +375,9 @@ def interpret_results(self, results):
         failures = [r for r in results if r.error or r.fail]
         return len(failures) == 0
 
-    def get_model_schemas(self, selected_uids):
+    def get_model_schemas(
+        self, selected_uids: Iterable[str]
+    ) -> Set[Tuple[str, str]]:
         if self.manifest is None:
             raise InternalException('manifest was None in get_model_schemas')
 
@@ -387,19 +390,46 @@ def get_model_schemas(self, selected_uids):
 
         return schemas
 
-    def create_schemas(self, adapter, selected_uids):
+    def create_schemas(self, adapter, selected_uids: Iterable[str]):
         required_schemas = self.get_model_schemas(selected_uids)
         required_databases = set(db for db, _ in required_schemas)
 
         existing_schemas_lowered: Set[Tuple[str, str]] = set()
-        for db in required_databases:
-            existing_schemas_lowered.update(
-                (db.lower(), s.lower()) for s in adapter.list_schemas(db))
 
-        for db, schema in required_schemas:
-            if (db.lower(), schema.lower()) not in existing_schemas_lowered:
+        def list_schemas(db: str) -> List[Tuple[str, str]]:
+            with adapter.connection_named(f'list_{db}'):
+                return [
+                    (db.lower(), s.lower())
+                    for s in adapter.list_schemas(db)
+                ]
+
+        def create_schema(db: str, schema: str) -> None:
+            with adapter.connection_named(f'create_{db}_{schema}'):
                 adapter.create_schema(db, schema)
 
+        list_futures = []
+        create_futures = []
+
+        with dbt.utils.executor(self.config) as tpe:
+            list_futures = [
+                tpe.submit(list_schemas, db) for db in required_databases
+            ]
+
+            for ls_future in as_completed(list_futures):
+                existing_schemas_lowered.update(ls_future.result())
+
+            for db, schema in required_schemas:
+                db_schema = (db.lower(), schema.lower())
+                if db_schema not in existing_schemas_lowered:
+                    existing_schemas_lowered.add(db_schema)
+                    create_futures.append(
+                        tpe.submit(create_schema, db, schema)
+                    )
+
+            for create_future in as_completed(create_futures):
+                # trigger/re-raise any excceptions while creating schemas
+                create_future.result()
+
     def get_result(self, results, elapsed_time, generated_at):
         return ExecutionResult(
             results=results,

core/dbt/utils.py

@@ -1,4 +1,5 @@
 import collections
+import concurrent.futures
 import copy
 import datetime
 import decimal
@@ -8,6 +9,7 @@
 import json
 import os
 from enum import Enum
+from typing_extensions import Protocol
 from typing import (
     Tuple, Type, Any, Optional, TypeVar, Dict, Union, Callable
 )
@@ -489,3 +491,48 @@ def format_bytes(num_bytes):
         num_bytes /= 1024.0
 
     return "> 1024 TB"
+
+
+# a little concurrent.futures.Executor for single-threaded mode
+class SingleThreadedExecutor(concurrent.futures.Executor):
+    def submit(*args, **kwargs):
+        # this basic pattern comes from concurrent.futures.Executor itself,
+        # but without handling the `fn=` form.
+        if len(args) >= 2:
+            self, fn, *args = args
+        elif not args:
+            raise TypeError(
+                "descriptor 'submit' of 'SingleThreadedExecutor' object needs "
+                "an argument"
+            )
+        else:
+            raise TypeError(
+                'submit expected at least 1 positional argument, '
+                'got %d' % (len(args) - 1)
+            )
+        fut = concurrent.futures.Future()
+        try:
+            result = fn(*args, **kwargs)
+        except Exception as exc:
+            fut.set_exception(exc)
+        else:
+            fut.set_result(result)
+        return fut
+
+
+class ThreadedArgs(Protocol):
+    single_threaded: bool
+
+
+class HasThreadingConfig(Protocol):
+    args: ThreadedArgs
+    threads: Optional[int]
+
+
+def executor(config: HasThreadingConfig) -> concurrent.futures.Executor:
+    if config.args.single_threaded:
+        return SingleThreadedExecutor()
+    else:
+        return concurrent.futures.ThreadPoolExecutor(
+            max_workers=config.threads
+        )

test/integration/001_simple_copy_test/test_simple_copy.py

@@ -71,7 +71,6 @@ def test__postgres__simple_copy_with_materialized_views(self):
                 select * from {schema}.unrelated_materialized_view
             )
         '''.format(schema=self.unique_schema()))
-
         results = self.run_dbt(["seed"])
         self.assertEqual(len(results),  1)
         results = self.run_dbt()

test/unit/utils.py

@@ -25,6 +25,7 @@ def normalize(path):
 
 class Obj:
     which = 'blah'
+    single_threaded = False
 
 
 def mock_connection(name):