Added defined tags for data flow

Author: lu-ohai

ads/jobs/builders/infrastructure/dataflow.py

@@ -384,6 +384,8 @@ class DataFlow(Infrastructure):
     CONST_OCPUS = "ocpus"
     CONST_ID = "id"
     CONST_PRIVATE_ENDPOINT_ID = "private_endpoint_id"
+    CONST_FREEFORM_TAGS = "freeform_tags"
+    CONST_DEFINED_TAGS = "defined_tags"
 
     attribute_map = {
         CONST_COMPARTMENT_ID: "compartmentId",
@@ -402,6 +404,8 @@ class DataFlow(Infrastructure):
         CONST_OCPUS: CONST_OCPUS,
         CONST_ID: CONST_ID,
         CONST_PRIVATE_ENDPOINT_ID: "privateEndpointId",
+        CONST_FREEFORM_TAGS: "freeformTags",
+        CONST_DEFINED_TAGS: "definedTags"
     }
 
     def __init__(self, spec: dict = None, **kwargs):
@@ -414,7 +418,9 @@ def __init__(self, spec: dict = None, **kwargs):
             spec = {
                 k: v
                 for k, v in spec.items()
-                if f"with_{camel_to_snake(k)}" in self.__dir__() and v is not None
+                if (f"with_{camel_to_snake(k)}" in self.__dir__()
+                or (k == "defined_tags" or "freeform_tags"))
+                and v is not None
             }
             defaults.update(spec)
             super().__init__(defaults, **kwargs)
@@ -775,6 +781,36 @@ def with_private_endpoint_id(self, private_endpoint_id: str) -> "DataFlow":
             the Data Flow instance itself
         """
         return self.set_spec(self.CONST_PRIVATE_ENDPOINT_ID, private_endpoint_id)
+    
+    def with_freeform_tag(self, **kwargs) -> "DataFlow":
+        """Sets freeform tags
+
+        Returns
+        -------
+        DataFlow
+            The DataFlow instance (self)
+        """
+        return self.set_spec(self.CONST_FREEFORM_TAGS, kwargs)
+
+    def with_defined_tag(self, **kwargs) -> "DataFlow":
+        """Sets defined tags
+
+        Returns
+        -------
+        DataFlow
+            The DataFlow instance (self)
+        """
+        return self.set_spec(self.CONST_DEFINED_TAGS, kwargs)
+
+    @property
+    def freeform_tags(self) -> dict:
+        """Freeform tags"""
+        return self.get_spec(self.CONST_FREEFORM_TAGS, {})
+
+    @property
+    def defined_tags(self) -> dict:
+        """Defined tags"""
+        return self.get_spec(self.CONST_DEFINED_TAGS, {})
 
     def __getattr__(self, item):
         if f"with_{item}" in self.__dir__():
@@ -849,7 +885,8 @@ def create(self, runtime: DataFlowRuntime, **kwargs) -> "DataFlow":
             {
                 "display_name": self.name,
                 "file_uri": runtime.script_uri,
-                "freeform_tags": runtime.freeform_tags,
+                "freeform_tags": runtime.freeform_tags or self.freeform_tags,
+                "defined_tags": runtime.defined_tags or self.defined_tags,
                 "archive_uri": runtime.archive_uri,
                 "configuration": runtime.configuration,
             }
@@ -915,6 +952,7 @@ def run(
         args: List[str] = None,
         env_vars: Dict[str, str] = None,
         freeform_tags: Dict[str, str] = None,
+        defined_tags: Dict[str, Dict[str, object]] = None,
         wait: bool = False,
         **kwargs,
     ) -> DataFlowRun:
@@ -932,6 +970,8 @@ def run(
             dictionary of environment variables (not used for data flow)
         freeform_tags: Dict[str, str], optional
             freeform tags
+        defined_tags: Dict[str, Dict[str, object]], optional
+            defined tags
         wait: bool, optional
             whether to wait for a run to terminate
         kwargs
@@ -950,7 +990,8 @@ def run(
         # Set default display_name if not specified - randomly generated easy to remember name generated
         payload["display_name"] = name if name else utils.get_random_name_for_resource()
         payload["arguments"] = args if args and len(args) > 0 else None
-        payload["freeform_tags"] = freeform_tags
+        payload["freeform_tags"] = freeform_tags or self.freeform_tags
+        payload["defined_tags"] = defined_tags or self.defined_tags
         payload.pop("spark_version", None)
         logger.debug(f"Creating a DataFlow Run with payload {payload}")
         run = DataFlowRun(**payload).create()

tests/unitary/default_setup/jobs/test_jobs_dataflow.py

@@ -320,13 +320,27 @@ def df(self):
                 2
             ).with_private_endpoint_id(
                 "test_private_endpoint"
+            ).with_freeform_tag(
+                test_freeform_tags_key="test_freeform_tags_value",
+            ).with_defined_tag(
+                test_defined_tags_namespace={
+                    "test_defined_tags_key": "test_defined_tags_value"
+                }
             )
         return df
 
     def test_create_with_builder_pattern(self, mock_to_dict, mock_client, df):
         assert df.language == "PYTHON"
         assert df.spark_version == "3.2.1"
         assert df.num_executors == 2
+        assert df.freeform_tags == {
+            "test_freeform_tags_key": "test_freeform_tags_value"
+        }
+        assert df.defined_tags == {
+            "test_defined_tags_namespace": {
+                "test_defined_tags_key": "test_defined_tags_value"
+            }
+        }
 
         rt = (
             DataFlowRuntime()
@@ -335,9 +349,25 @@ def test_create_with_builder_pattern(self, mock_to_dict, mock_client, df):
             .with_custom_conda(
                 "oci://my_bucket@my_namespace/conda_environments/cpu/PySpark 3.0 and Data Flow/5.0/pyspark30_p37_cpu_v5"
             )
+            .with_freeform_tag(
+                test_freeform_tags_runtime_key="test_freeform_tags_runtime_value"
+            )
+            .with_defined_tag(
+                test_defined_tags_namespace={
+                    "test_defined_tags_runtime_key": "test_defined_tags_runtime_value"
+                }
+            )
             .with_overwrite(True)
         )
         assert rt.overwrite == True
+        assert rt.freeform_tags == {
+            "test_freeform_tags_runtime_key": "test_freeform_tags_runtime_value"
+        }
+        assert rt.defined_tags == {
+            "test_defined_tags_namespace": {
+                "test_defined_tags_runtime_key": "test_defined_tags_runtime_value"
+            }
+        }
 
         with patch.object(DataFlowApp, "client", mock_client):
             with patch.object(DataFlowApp, "to_dict", mock_to_dict):
@@ -429,6 +459,14 @@ def test_to_and_from_dict(self, df):
         assert df_dict["spec"]["privateEndpointId"] == "test_private_endpoint"
         assert df_dict["spec"]["driverShapeConfig"] == {"memoryInGBs": 1, "ocpus": 16}
         assert df_dict["spec"]["executorShapeConfig"] == {"memoryInGBs": 1, "ocpus": 16}
+        assert df_dict["spec"]["freeformTags"] == {
+            "test_freeform_tags_key": "test_freeform_tags_value"
+        }
+        assert df_dict["spec"]["definedTags"] == {
+            "test_defined_tags_namespace": {
+                "test_defined_tags_key": "test_defined_tags_value"
+            }
+        }
 
         df_dict["spec"].pop("language")
         df_dict["spec"].pop("numExecutors")