WIP

Author: ezvz

api/python/ai/chronon/airflow_helpers.py

@@ -16,11 +16,21 @@ def create_airflow_dependency(table, partition_column):
     Returns:
         A dictionary with name and spec for the Airflow dependency
     """
-    # Default partition column to 'ds' if not specified
-    partition_col = partition_column or 'ds'
+    assert partition_column is not None, """Partition column must be provided via the spark.chronon.partition.column 
+        config. This can be set as a default in teams.py, or at the individual config level. For example:
+        ``` 
+        Team(
+            conf=ConfigProperties(
+                common={
+                    "spark.chronon.partition.column": "_test_column",
+                }
+            )
+        )
+        ```
+        """
     return {
         "name": f"wf_{utils.sanitize(table)}",
-        "spec": f"{table}/{partition_col}={{{{ ds }}}}",
+        "spec": f"{table}/{partition_column}={{{{ ds }}}}",
     }
 
 def _get_partition_col_from_query(query):
@@ -53,10 +63,8 @@ def _get_airflow_deps_from_source(source, partition_column=None):
             tables.append(source.entities.mutationTable)
         source_partition_column = _get_partition_col_from_query(source.entities.query) or partition_column
     elif source.joinSource:
-        namespace = source.joinSource.join.metaData.outputNamespace
-        table = utils.sanitize(source.joinSource.join.metaData.name)
-        tables = [f"{namespace}.{table}"]
-        source_partition_column = _get_partition_col_from_query(source.joinSource.query) or partition_column
+        # TODO: Handle joinSource -- it doesn't work right now because the metadata isn't set on joinSource at this point
+        return []
     else:
         # Unknown source type
         return []
@@ -65,7 +73,7 @@ def _get_airflow_deps_from_source(source, partition_column=None):
 
 
 def extract_default_partition_column(obj):
-    return obj.metaData.executionInfo.env.common.get("partitionColumn")
+    return obj.metaData.executionInfo.conf.common.get("spark.chronon.partition.column")
 
 
 def _set_join_deps(join):

api/python/test/sample/teams.py

@@ -1,6 +1,6 @@
 from ai.chronon.api.ttypes import Team
 from ai.chronon.repo.constants import RunMode
-from ai.chronon.types import EnvironmentVariables, ConfigProperties
+from ai.chronon.types import ConfigProperties, EnvironmentVariables
 
 default = Team(
     description="Default team",
@@ -76,6 +76,11 @@
 
 sample_team = Team(
     outputNamespace="test",
+    conf=ConfigProperties(
+        common={
+            "spark.chronon.partition.column": "_test_column_sample",
+        }
+    ),
     env=EnvironmentVariables(
         common={
             "GCP_BIGTABLE_INSTANCE_ID": "test-instance"  # example, custom bigtable instance