moving to partitioned by

dataders · Oct 5, 2023 · d1df913 · d1df913
1 parent fc8b3fb
commit d1df913
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 29 deletions.
diff --git a/dbt/adapters/databricks/relation_configs/base.py b/dbt/adapters/databricks/relation_configs/base.py
@@ -31,7 +31,7 @@ def quote_policy(cls) -> Policy:
         return DatabricksQuotePolicy()
 
     @classmethod
-    def from_model_node(cls, model_node: ModelNode) -> "RelationConfigBase":
+    def from_model_node(cls, model_node: ModelNode) -> RelationConfigBase:
         relation_config = cls.parse_model_node(model_node)
         relation = cls.from_dict(relation_config)
         return relation
@@ -43,7 +43,7 @@ def parse_model_node(cls, model_node: ModelNode) -> dict:
         )
 
     @classmethod
-    def from_relation_results(cls, relation_results: RelationResults) -> "RelationConfigBase":
+    def from_relation_results(cls, relation_results: RelationResults) -> RelationConfigBase:
         relation_config = cls.parse_relation_results(relation_results)
         relation = cls.from_dict(relation_config)
         return relation

diff --git a/dbt/adapters/databricks/relation_configs/materialized_view.py b/dbt/adapters/databricks/relation_configs/materialized_view.py
@@ -13,9 +13,9 @@
 from dbt.exceptions import DbtRuntimeError
 
 from dbt.adapters.databricks.relation_configs.base import DatabricksRelationConfigBase
-from dbt.adapters.databricks.relation_configs.partition import (
-    DatabricksPartitionConfig,
-    DatabricksPartitionConfigChange,
+from dbt.adapters.databricks.relation_configs.partitioned_by import (
+    DatabricksPartitionedByConfig,
+    DatabricksPartitionedByConfigChange,
 )
 from dbt.adapters.databricks.relation_configs.schedule import DatabricksScheduleConfigChange
 
@@ -93,7 +93,7 @@ def parse_model_node(cls, model_node: ModelNode) -> dict:
 
         if model_node.config.get("partition"):
             config_dict.update(
-                {"partition": DatabricksPartitionConfig.parse_model_node(model_node)}
+                {"partition": DatabricksPartitionedByConfig.parse_model_node(model_node)}
             )
 
         return config_dict
@@ -142,7 +142,7 @@ def parse_relation_results(cls, relation_results: RelationResults) -> dict:
 
 @dataclass
 class DatabricksMaterializedViewConfigChangeset:
-    partition: Optional[DatabricksPartitionConfigChange] = None
+    partition: Optional[DatabricksPartitionedByConfigChange] = None
     schedule: Optional[DatabricksScheduleConfigChange] = None
 
     @property

diff --git a/.../databricks/relation_configs/partition.py → ...bricks/relation_configs/partitioned_by.py b/.../databricks/relation_configs/partition.py → ...bricks/relation_configs/partitioned_by.py
@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Optional, FrozenSet, Set
+from typing import Dict, FrozenSet, Optional, Set
 
 import agate
 from dbt.adapters.relation_configs import (
@@ -16,7 +16,7 @@
 
 
 @dataclass(frozen=True, eq=True, unsafe_hash=True)
-class DatabricksPartitionConfig(DatabricksRelationConfigBase, RelationConfigValidationMixin):
+class DatabricksPartitionedByConfig(DatabricksRelationConfigBase, RelationConfigValidationMixin):
     """
     This config fallows the specs found here:
     https://docs.databricks.com/en/sql/language-manual/sql-ref-partition.html#partitioned-by
@@ -26,32 +26,20 @@ class DatabricksPartitionConfig(DatabricksRelationConfigBase, RelationConfigVali
     - column_type: the type of the specified thing
     """
 
-    # sortkey: Optional[FrozenSet[str]] = None
-
-    def __post_init__(self):
-        # # maintains `frozen=True` while allowing for a variable default on `sort_type`
-        # if self.sortstyle is None and self.sortkey is None:
-        #     object.__setattr__(self, "sortstyle", DatabrickSortStyle.default())
-        # elif self.sortstyle is None:
-        #     object.__setattr__(self, "sortstyle", DatabrickSortStyle.default_with_columns())
-        super().__post_init__()
+    partitioned_by_columns: FrozenSet[str]
 
     @property
     def validation_rules(self) -> Set[RelationConfigValidationRule]:
-        # TODO
-
-        pass
-
-    @classmethod
-    def from_dict(cls, config_dict):
-        # TODO
-
+        # TODO: can we reference the schema here to ensure that the columns are valid?
         pass
 
     @classmethod
     def parse_model_node(cls, model_node: ModelNode) -> dict:
-        # TODO
-        pass
+        config = {}
+        partitioned_by = model_node.config.extra.get("partitioned_by", None)
+
+        if partitioned_by:
+            config["partitioned_by_columns"] = frozenset(partitioned_by.split(","))
 
     @classmethod
     def parse_relation_results(cls, relation_results_entry: agate.Row) -> dict:
@@ -78,7 +66,7 @@ def parse_relation_results(cls, relation_results_entry: agate.Row) -> dict:
 
 
 @dataclass(frozen=True, eq=True, unsafe_hash=True)
-class DatabricksPartitionConfigChange(RelationConfigChange):
+class DatabricksPartitionedByConfigChange(RelationConfigChange):
     context: Optional[bool] = None
 
     @property