register + get catalog integration methods on BaseAdapter

dbt-labs · Dec 12, 2024 · 778554f · 778554f
1 parent 0f9bbbd
commit 778554f
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 3 deletions.
diff --git a/dbt/adapters/base/connections.py b/dbt/adapters/base/connections.py
@@ -77,7 +77,6 @@ def __init__(self, profile: AdapterRequiredConfig, mp_context: SpawnContext) ->
         self.thread_connections: Dict[Hashable, Connection] = {}
         self.lock: RLock = mp_context.RLock()
         self.query_header: Optional[MacroQueryStringSetter] = None
-        self.catalog_integrations = catalog_integrations or []
 
     def set_query_header(self, query_header_context: Dict[str, Any]) -> None:
         self.query_header = MacroQueryStringSetter(self.profile, query_header_context)

diff --git a/dbt/adapters/base/impl.py b/dbt/adapters/base/impl.py
@@ -65,6 +65,7 @@
 )
 from dbt.adapters.cache import RelationsCache, _make_ref_key_dict
 from dbt.adapters.capability import Capability, CapabilityDict
+from dbt.adapters.contracts.catalog_integration import CatalogIntegration
 from dbt.adapters.contracts.connection import Credentials
 from dbt.adapters.contracts.macros import MacroResolverProtocol
 from dbt.adapters.contracts.relation import RelationConfig
@@ -288,6 +289,7 @@ def __init__(self, config, mp_context: SpawnContext) -> None:
         self.connections = self.ConnectionManager(config, mp_context)
         self._macro_resolver: Optional[MacroResolverProtocol] = None
         self._macro_context_generator: Optional[MacroContextGeneratorCallable] = None
+        self._catalog_integrations: Dict[str, CatalogIntegration] = {}
         self.behavior = DEFAULT_BASE_BEHAVIOR_FLAGS  # type: ignore
 
     ###
@@ -308,6 +310,15 @@ def set_macro_context_generator(
         macro_context_generator: MacroContextGeneratorCallable,
     ) -> None:
         self._macro_context_generator = macro_context_generator
+
+    ###
+    # Methods to set / access catalog integrations
+    ###
+    def register_catalog_integration(self, catalog_name: str, catalog_integration: CatalogIntegration):
+        self._catalog_integrations[catalog_name] = catalog_integration
+
+    def get_catalog_integration(self, catalog_name: str) -> Optional[CatalogIntegration]:
+        return self._catalog_integrations.get(catalog_name)
 
     @available_property
     def behavior(self) -> Behavior:

diff --git a/dbt/adapters/contracts/connection.py b/dbt/adapters/contracts/connection.py
@@ -28,7 +28,6 @@
 from mashumaro.jsonschema.annotations import Pattern
 from typing_extensions import Protocol, Annotated
 
-from dbt.adapters.contracts.catalog_integration import CatalogIntegration
 from dbt.adapters.events.types import NewConnectionOpening
 from dbt.adapters.utils import translate_aliases
 
@@ -228,4 +227,4 @@ class AdapterRequiredConfig(HasCredentials, Protocol):
     cli_vars: Dict[str, Any]
     target_path: str
     log_cache_events: bool
-    catalog_integrations: Optional[List[CatalogIntegration]]
+