Skip to content

Commit

Permalink
Better feature detection mechanism for adapters.
Browse files Browse the repository at this point in the history
  • Loading branch information
peterallenwebb committed Sep 28, 2023
1 parent 471d29d commit a9addc0
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 4 deletions.
18 changes: 15 additions & 3 deletions core/dbt/adapters/base/impl.py
Original file line number Diff line number Diff line change
Expand Up @@ -162,6 +162,14 @@ def submit(self, compiled_code: str) -> Any:
raise NotImplementedError("PythonJobHelper submit function is not implemented yet")


class AdapterFeature(str, Enum):
"""Enumeration of optional adapter features which can be probed using BaseAdapter.has_feature()"""

CatalogByRelations = "GetCatalogByRelations"
"""Flags support for retrieving catalog information using a list of relations, rather than always retrieving all
the relations in a schema """


class BaseAdapter(metaclass=AdapterMeta):
"""The BaseAdapter provides an abstract base class for adapters.
Expand Down Expand Up @@ -223,8 +231,6 @@ class BaseAdapter(metaclass=AdapterMeta):
ConstraintType.foreign_key: ConstraintSupport.ENFORCED,
}

CATALOG_BY_RELATION_SUPPORT = False

def __init__(self, config) -> None:
self.config = config
self.cache = RelationsCache()
Expand Down Expand Up @@ -1141,7 +1147,7 @@ def get_catalog(self, manifest: Manifest) -> Tuple[agate.Table, List[Exception]]
with executor(self.config) as tpe:
futures: List[Future[agate.Table]] = []
relation_count = len(self._get_catalog_relations(manifest))
if relation_count <= 100 and self.CATALOG_BY_RELATION_SUPPORT:
if relation_count <= 100 and self.has_feature(AdapterFeature.CatalogByRelations):
relations_by_schema = self._get_catalog_relations_by_info_schema(manifest)
for info_schema in relations_by_schema:
name = ".".join([str(info_schema.database), "information_schema"])
Expand Down Expand Up @@ -1495,6 +1501,12 @@ def render_model_constraint(cls, constraint: ModelLevelConstraint) -> Optional[s
else:
return None

def has_feature(self, feature: AdapterFeature) -> bool:
# The base adapter implementation does not implement any optional
# features, so always return false. Adapters which wish to provide
# optional features will have to override this function.
return False

Check warning on line 1508 in core/dbt/adapters/base/impl.py

View check run for this annotation

Codecov / codecov/patch

core/dbt/adapters/base/impl.py#L1508

Added line #L1508 was not covered by tests


COLUMNS_EQUAL_SQL = """
with diff_count as (
Expand Down
7 changes: 6 additions & 1 deletion plugins/postgres/dbt/adapters/postgres/impl.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
from typing import Optional, Set, List, Any

from dbt.adapters.base.meta import available
from dbt.adapters.base.impl import AdapterConfig, ConstraintSupport
from dbt.adapters.base.impl import AdapterConfig, AdapterFeature, ConstraintSupport
from dbt.adapters.sql import SQLAdapter
from dbt.adapters.postgres import PostgresConnectionManager
from dbt.adapters.postgres.column import PostgresColumn
Expand Down Expand Up @@ -75,6 +75,8 @@ class PostgresAdapter(SQLAdapter):

CATALOG_BY_RELATION_SUPPORT = True

SUPPORTED_FEATURES: Set[AdapterFeature] = frozenset(AdapterFeature.CatalogByRelations)

@classmethod
def date_function(cls):
return "now()"
Expand Down Expand Up @@ -145,3 +147,6 @@ def valid_incremental_strategies(self):

def debug_query(self):
self.execute("select 1 as id")

def has_feature(self, feature: AdapterFeature) -> bool:
return feature in self.SUPPORTED_FEATURES

0 comments on commit a9addc0

Please sign in to comment.