Migrate to the dbt catalog framework (#1012)

Author: mikealfare

CHANGELOG.md

@@ -3,6 +3,7 @@
 ### Features
 
 - Support constraint updates on incremental runs (with Materialization V2) ([1013](https://github.com/databricks/dbt-databricks/pull/1013))
+- Add catalog integration support - set table formats, file formats, and locations in `catalogs.yml` ([1012](https://github.com/databricks/dbt-databricks/pull/1012))
 
 ### Fixes
 

dbt/adapters/databricks/catalogs/__init__.py

@@ -0,0 +1,9 @@
+from dbt.adapters.databricks.catalogs._hive_metastore import HiveMetastoreCatalogIntegration
+from dbt.adapters.databricks.catalogs._relation import DatabricksCatalogRelation
+from dbt.adapters.databricks.catalogs._unity import UnityCatalogIntegration
+
+__all__ = [
+    "DatabricksCatalogRelation",
+    "HiveMetastoreCatalogIntegration",
+    "UnityCatalogIntegration",
+]

dbt/adapters/databricks/catalogs/_hive_metastore.py

@@ -0,0 +1,42 @@
+from typing import Optional
+
+from dbt.adapters.catalogs import CatalogIntegration, CatalogIntegrationConfig
+from dbt.adapters.contracts.relation import RelationConfig
+from dbt.adapters.databricks import constants, parse_model
+from dbt.adapters.databricks.catalogs._relation import DatabricksCatalogRelation
+
+
+class HiveMetastoreCatalogIntegration(CatalogIntegration):
+    catalog_type = constants.HIVE_METASTORE_CATALOG_TYPE
+    allows_writes = True
+
+    def __init__(self, config: CatalogIntegrationConfig) -> None:
+        super().__init__(config)
+        self.file_format: str = config.adapter_properties.get("file_format")
+
+    @property
+    def location_root(self) -> Optional[str]:
+        """
+        Volumes in Databricks are non-tabular datasets, which is something
+        different from what we mean by "external_volume" in dbt.
+        However, the protocol expects "external_volume" to be set.
+        """
+        return self.external_volume
+
+    @location_root.setter
+    def location_root(self, value: Optional[str]) -> None:
+        self.external_volume = value
+
+    def build_relation(self, model: RelationConfig) -> DatabricksCatalogRelation:
+        """
+        Args:
+            model: `config.model` (not `model`) from the jinja context
+        """
+        return DatabricksCatalogRelation(
+            catalog_type=self.catalog_type,
+            catalog_name=self.catalog_name,
+            table_format=parse_model.table_format(model) or self.table_format,
+            file_format=parse_model.file_format(model) or self.file_format,
+            external_volume=parse_model.location_root(model) or self.external_volume,
+            location_path=parse_model.location_path(model),
+        )

dbt/adapters/databricks/catalogs/_relation.py

@@ -0,0 +1,49 @@
+import posixpath
+from dataclasses import dataclass
+from typing import Optional
+
+from dbt_common.exceptions import DbtConfigError
+
+from dbt.adapters.databricks import constants
+
+
+@dataclass
+class DatabricksCatalogRelation:
+    catalog_type: str = constants.DEFAULT_CATALOG.catalog_type
+    catalog_name: Optional[str] = constants.DEFAULT_CATALOG.name
+    table_format: Optional[str] = constants.DEFAULT_CATALOG.adapter_properties.get("table_format")
+    file_format: Optional[str] = constants.DEFAULT_CATALOG.adapter_properties.get("file_format")
+    external_volume: Optional[str] = constants.DEFAULT_CATALOG.external_volume
+    location_path: Optional[str] = None
+
+    @property
+    def location_root(self) -> Optional[str]:
+        """
+        Volumes in Databricks are non-tabular datasets, which is something
+        different from what we mean by "external_volume" in dbt.
+        However, the protocol expects "external_volume" to be set.
+        """
+        return self.external_volume
+
+    @location_root.setter
+    def location_root(self, value: Optional[str]) -> None:
+        self.external_volume = value
+
+    @property
+    def location(self) -> Optional[str]:
+        if self.location_root and self.location_path:
+            return posixpath.join(self.location_root, self.location_path)
+        return None
+
+    @property
+    def iceberg_table_properties(self) -> dict[str, str]:
+        if self.table_format == constants.ICEBERG_TABLE_FORMAT:
+            if self.file_format != constants.DELTA_FILE_FORMAT:
+                raise DbtConfigError(
+                    "When table_format is 'iceberg', cannot set file_format to other than delta."
+                )
+            return {
+                "delta.enableIcebergCompatV2": "true",
+                "delta.universalFormat.enabledFormats": constants.ICEBERG_TABLE_FORMAT,
+            }
+        return {}

dbt/adapters/databricks/catalogs/_unity.py

@@ -0,0 +1,44 @@
+from typing import Optional
+
+from dbt.adapters.catalogs import CatalogIntegration, CatalogIntegrationConfig
+from dbt.adapters.contracts.relation import RelationConfig
+from dbt.adapters.databricks import constants, parse_model
+from dbt.adapters.databricks.catalogs._relation import DatabricksCatalogRelation
+
+
+class UnityCatalogIntegration(CatalogIntegration):
+    catalog_type = constants.UNITY_CATALOG_TYPE
+    allows_writes = True
+
+    def __init__(self, config: CatalogIntegrationConfig) -> None:
+        super().__init__(config)
+        if location_root := config.adapter_properties.get("location_root"):
+            self.external_volume: Optional[str] = location_root
+        self.file_format: str = config.adapter_properties.get("file_format")
+
+    @property
+    def location_root(self) -> Optional[str]:
+        """
+        Volumes in Databricks are non-tabular datasets, which is something
+        different from what we mean by "external_volume" in dbt.
+        However, the protocol expects "external_volume" to be set.
+        """
+        return self.external_volume
+
+    @location_root.setter
+    def location_root(self, value: Optional[str]) -> None:
+        self.external_volume = value
+
+    def build_relation(self, model: RelationConfig) -> DatabricksCatalogRelation:
+        """
+        Args:
+            model: `config.model` (not `model`) from the jinja context
+        """
+        return DatabricksCatalogRelation(
+            catalog_type=self.catalog_type,
+            catalog_name=self.catalog_name,
+            table_format=parse_model.table_format(model) or self.table_format,
+            file_format=parse_model.file_format(model) or self.file_format,
+            external_volume=parse_model.location_root(model) or self.external_volume,
+            location_path=parse_model.location_path(model),
+        )

dbt/adapters/databricks/constants.py

@@ -0,0 +1,38 @@
+from types import SimpleNamespace
+
+DEFAULT_TABLE_FORMAT = "default"
+ICEBERG_TABLE_FORMAT = "iceberg"
+
+
+DELTA_FILE_FORMAT = "delta"
+PARQUET_FILE_FORMAT = "parquet"
+HUDI_FILE_FORMAT = "hudi"
+
+
+UNITY_CATALOG_TYPE = "unity"
+HIVE_METASTORE_CATALOG_TYPE = "hive_metastore"
+
+
+DEFAULT_UNITY_CATALOG = SimpleNamespace(
+    name="unity",
+    catalog_name="default_unity",
+    catalog_type=UNITY_CATALOG_TYPE,
+    # requires the model to specify the external_volume (location_root) property
+    external_volume=None,
+    table_format=DEFAULT_TABLE_FORMAT,
+    adapter_properties={
+        "file_format": DELTA_FILE_FORMAT,
+    },
+)
+DEFAULT_HIVE_METASTORE_CATALOG = SimpleNamespace(
+    name="hive_metastore",
+    catalog_name="default_hive_metastore",
+    catalog_type=HIVE_METASTORE_CATALOG_TYPE,
+    # requires the model to specify the external_volume (location_root) property
+    external_volume=None,
+    table_format=DEFAULT_TABLE_FORMAT,
+    adapter_properties={
+        "file_format": DELTA_FILE_FORMAT,
+    },
+)
+DEFAULT_CATALOG = DEFAULT_UNITY_CATALOG

dbt/adapters/databricks/impl.py

@@ -23,14 +23,20 @@
 from dbt.adapters.base.meta import available
 from dbt.adapters.base.relation import BaseRelation
 from dbt.adapters.capability import Capability, CapabilityDict, CapabilitySupport, Support
+from dbt.adapters.catalogs import CatalogRelation
 from dbt.adapters.contracts.connection import AdapterResponse, Connection
 from dbt.adapters.contracts.relation import RelationConfig, RelationType
-from dbt.adapters.databricks import constraints
+from dbt.adapters.databricks import constants, constraints, parse_model
 from dbt.adapters.databricks.behaviors.columns import (
     GetColumnsBehavior,
     GetColumnsByDescribe,
     GetColumnsByInformationSchema,
 )
+from dbt.adapters.databricks.catalogs import (
+    DatabricksCatalogRelation,
+    HiveMetastoreCatalogIntegration,
+    UnityCatalogIntegration,
+)
 from dbt.adapters.databricks.column import DatabricksColumn
 from dbt.adapters.databricks.connections import DatabricksConnectionManager
 from dbt.adapters.databricks.global_state import GlobalState
@@ -180,6 +186,10 @@ class DatabricksAdapter(SparkAdapter):
         }
     )
 
+    CATALOG_INTEGRATIONS = [
+        HiveMetastoreCatalogIntegration,
+        UnityCatalogIntegration,
+    ]
     CONSTRAINT_SUPPORT = constraints.CONSTRAINT_SUPPORT
 
     get_column_behavior: GetColumnsBehavior
@@ -196,6 +206,9 @@ def __init__(self, config: Any, mp_context: SpawnContext) -> None:
         except CompilationError:
             pass
 
+        self.add_catalog_integration(constants.DEFAULT_UNITY_CATALOG)
+        self.add_catalog_integration(constants.DEFAULT_HIVE_METASTORE_CATALOG)
+
     @property
     def _behavior_flags(self) -> list[BehaviorFlag]:
         return [USE_INFO_SCHEMA_FOR_COLUMNS, USE_USER_FOLDER_FOR_PYTHON, USE_MATERIALIZATION_V2]
@@ -205,20 +218,18 @@ def update_tblproperties_for_iceberg(
         self, config: BaseConfig, tblproperties: Optional[dict[str, str]] = None
     ) -> dict[str, str]:
         result = tblproperties or config.get("tblproperties", {})
-        if config.get("table_format") == TableFormat.ICEBERG:
+
+        catalog_relation: DatabricksCatalogRelation = self.build_catalog_relation(config.model)  # type:ignore
+        if catalog_relation.table_format == constants.ICEBERG_TABLE_FORMAT:
             if self.compare_dbr_version(14, 3) < 0:
                 raise DbtConfigError("Iceberg support requires Databricks Runtime 14.3 or later.")
-            if config.get("file_format", "delta") != "delta":
-                raise DbtConfigError(
-                    "When table_format is 'iceberg', cannot set file_format to other than delta."
-                )
             if config.get("materialized") not in ("incremental", "table", "snapshot"):
                 raise DbtConfigError(
                     "When table_format is 'iceberg', materialized must be 'incremental'"
                     ", 'table', or 'snapshot'."
                 )
-            result["delta.enableIcebergCompatV2"] = "true"
-            result["delta.universalFormat.enabledFormats"] = "iceberg"
+            result.update(catalog_relation.iceberg_table_properties)
+
         return result
 
     @available.parse(lambda *a, **k: 0)
@@ -820,6 +831,26 @@ def is_cluster(self) -> bool:
     def clean_sql(self, sql: str) -> str:
         return SqlUtils.clean_sql(sql)
 
+    @available
+    def build_catalog_relation(self, model: RelationConfig) -> Optional[CatalogRelation]:
+        """
+        Builds a relation for a given configuration.
+
+        This method uses the provided configuration to determine the appropriate catalog
+        integration and config parser for building the relation. It defaults to the built-in Delta
+        catalog if none is provided in the configuration for backward compatibility.
+
+        Args:
+            model (RelationConfig): `config.model` (not `model`) from the jinja context
+
+        Returns:
+            Any: The relation object generated through the catalog integration and parser
+        """
+        if catalog := parse_model.catalog_name(model):
+            catalog_integration = self.get_catalog_integration(catalog)
+            return catalog_integration.build_relation(model)
+        return None
+
 
 @dataclass(frozen=True)
 class RelationAPIBase(ABC, Generic[DatabricksRelationConfig]):

dbt/adapters/databricks/parse_model.py

@@ -0,0 +1,48 @@
+import posixpath
+from typing import Optional
+
+from dbt.adapters.contracts.relation import RelationConfig
+from dbt.adapters.databricks import constants
+
+
+def catalog_name(model: RelationConfig) -> Optional[str]:
+    """
+    The user has not set the catalog, which is common as it's the legacy behavior.
+    We need to infer a catalog based on the model config since macros now expect it.
+    Luckily, all catalog attribution in the legacy behavior is on the model.
+    This means we can take a default catalog and rely on the model overrides to supply the rest.
+    """
+    return _get(model, "catalog") or constants.DEFAULT_CATALOG.name
+
+
+def file_format(model: RelationConfig) -> Optional[str]:
+    return _get(model, "file_format")
+
+
+def location_path(model: RelationConfig) -> Optional[str]:
+    if not model.config:
+        return None
+
+    if model.config.get("include_full_name_in_path"):
+        return posixpath.join(model.database, model.schema, model.identifier)
+    return model.identifier
+
+
+def location_root(model: RelationConfig) -> Optional[str]:
+    return _get(model, "location_root")
+
+
+def table_format(model: RelationConfig) -> Optional[str]:
+    return _get(model, "table_format")
+
+
+def _get(
+    model: RelationConfig, setting: str, case_sensitive: Optional[bool] = False
+) -> Optional[str]:
+    if not model.config:
+        return None
+
+    if value := model.config.get(setting):
+        return value if case_sensitive else value.casefold()
+
+    return None

dbt/adapters/databricks/relation_configs/table_format.py

@@ -1,5 +1,7 @@
 from dbt_common.dataclass_schema import StrEnum
 
+from dbt.adapters.databricks import constants
+
 
 class TableFormat(StrEnum):
     """
@@ -8,8 +10,8 @@ class TableFormat(StrEnum):
     simplify things for users.
     """
 
-    DEFAULT = "default"
-    ICEBERG = "iceberg"
+    DEFAULT = constants.DEFAULT_TABLE_FORMAT
+    ICEBERG = constants.ICEBERG_TABLE_FORMAT
 
     def __str__(self) -> str:
         return self.value

dbt/include/databricks/macros/materializations/incremental/incremental.sql

@@ -1,16 +1,17 @@
 {% materialization incremental, adapter='databricks', supported_languages=['sql', 'python'] -%}
   {{ log("MATERIALIZING INCREMENTAL") }}
-  
+
+  {%- set catalog_relation = adapter.build_catalog_relation(config.model) -%}
+
   {% set existing_relation = load_relation_with_metadata(this) %}
   {% set target_relation = this.incorporate(type='table') %}
-  {% set file_format = get_file_format() %}
-  {% set incremental_strategy = get_incremental_strategy(file_format) %}
+  {% set incremental_strategy = get_incremental_strategy(catalog_relation.file_format) %}
   {% set grant_config = config.get('grants') %}
   {% set full_refresh = should_full_refresh() %}
   {% set partition_by = config.get('partition_by') %}
   {% set language = model['language'] %}
   {% set on_schema_change = incremental_validate_on_schema_change(config.get('on_schema_change'), default='ignore') %}
-  {% set is_delta = (file_format == 'delta' and existing_relation.is_delta) %}
+  {% set is_delta = (catalog_relation.file_format == 'delta' and existing_relation.is_delta) %}
   {% set compiled_code = adapter.clean_sql(model['compiled_code']) %}
 
   {% if adapter.behavior.use_materialization_v2 %}