Added defined tags for data flow (#210)

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,9 +781,31 @@ 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)
 
     def __getattr__(self, item):
-        if f"with_{item}" in self.__dir__():
+        if item == self.CONST_DEFINED_TAGS or item == self.CONST_FREEFORM_TAGS:
+            return self.get_spec(item)
+        elif f"with_{item}" in self.__dir__() and item != "defined_tag" and item != "freeform_tag":
             return self.get_spec(item)
         raise AttributeError(f"Attribute {item} not found.")
 
@@ -849,7 +877,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 +944,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 +962,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 +982,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()

docs/source/user_guide/apachespark/dataflow.rst

@@ -207,6 +207,8 @@ You can set them using the ``with_{property}`` functions:
 - ``with_spark_version``
 - ``with_warehouse_bucket_uri``
 - ``with_private_endpoint_id`` (`doc <https://docs.oracle.com/en-us/iaas/data-flow/using/pe-allowing.htm#pe-allowing>`__)
+- ``with_defined_tags``
+- ``with_freeform_tags``
 
 For more details, see `Data Flow class documentation <https://docs.oracle.com/en-us/iaas/tools/ads-sdk/latest/ads.jobs.html#module-ads.jobs.builders.infrastructure.dataflow>`__.
 
@@ -229,10 +231,10 @@ create applications.
 
 In the following "hello-world" example, ``DataFlow`` is populated with ``compartment_id``,
 ``driver_shape``, ``driver_shape_config``, ``executor_shape``, ``executor_shape_config``
-and ``spark_version``. ``DataFlowRuntime`` is populated with ``script_uri`` and
-``script_bucket``. The ``script_uri`` specifies the path to the script. It can be
-local or remote (an Object Storage path). If the path is local, then
-``script_bucket`` must be specified additionally because Data Flow
+, ``spark_version``, ``defined_tags`` and ``freeform_tags``. ``DataFlowRuntime`` is 
+populated with ``script_uri`` and ``script_bucket``. The ``script_uri`` specifies the 
+path to the script. It can be local or remote (an Object Storage path). If the path
+is local, then ``script_bucket`` must be specified additionally because Data Flow
 requires a script to be available in Object Storage. ADS
 performs the upload step for you, as long as you give the bucket name
 or the Object Storage path prefix to upload the script. Either can be
@@ -272,6 +274,10 @@ accepted. In the next example, the prefix is given for ``script_bucket``.
 		    .with_executor_shape("VM.Standard.E4.Flex")
 		    .with_executor_shape_config(ocpus=4, memory_in_gbs=64)
             .with_spark_version("3.0.2")
+            .with_defined_tag(
+                **{"Oracle-Tags": {"CreatedBy": "test_name@oracle.com"}}
+            )
+            .with_freeform_tag(test_freeform_key="test_freeform_value")
         )
         runtime_config = (
             DataFlowRuntime()
@@ -393,6 +399,10 @@ In the next example, ``archive_uri`` is given as an Object Storage location.
                 "spark.driverEnv.myEnvVariable": "value1",
                 "spark.executorEnv.myEnvVariable": "value2",
             })
+            .with_defined_tag(
+                **{"Oracle-Tags": {"CreatedBy": "test_name@oracle.com"}}
+            )
+            .with_freeform_tag(test_freeform_key="test_freeform_value")
         )
         runtime_config = (
             DataFlowRuntime()
@@ -566,6 +576,11 @@ into the ``Job.from_yaml()`` function to build a Data Flow job:
         numExecutors: 1
         sparkVersion: 3.2.1
         privateEndpointId: <private_endpoint_ocid>
+        definedTags:
+          Oracle-Tags:
+            CreatedBy: test_name@oracle.com
+        freeformTags:
+          test_freeform_key: test_freeform_value
       type: dataFlow
     name: dataflow_app_name
     runtime:
@@ -647,6 +662,12 @@ into the ``Job.from_yaml()`` function to build a Data Flow job:
             configuration:
                 required: false
                 type: dict
+            definedTags:
+                required: false
+                type: dict
+            freeformTags:
+                required: false
+                type: dict
     type:
         allowed:
             - dataFlow
@@ -694,7 +715,10 @@ into the ``Job.from_yaml()`` function to build a Data Flow job:
             configuration:
                 required: false
                 type: dict
-            freeform_tag:
+            definedTags:
+                required: false
+                type: dict
+            freeformTags:
                 required: false
                 type: dict
             scriptBucket:

docs/source/user_guide/apachespark/quickstart.rst

@@ -45,6 +45,10 @@ followed by the spark version, 3.2.1.
             .with_executor_shape("VM.Standard.E4.Flex")
             .with_executor_shape_config(ocpus=4, memory_in_gbs=64)
             .with_spark_version("3.2.1")
+            .with_defined_tag(
+                **{"Oracle-Tags": {"CreatedBy": "test_name@oracle.com"}}
+            )
+            .with_freeform_tag(test_freeform_key="test_freeform_value")
         )
         runtime_config = (
             DataFlowRuntime()
@@ -95,6 +99,11 @@ Assuming you have the following two files written in your current directory as `
                   memory_in_gbs: 64
                 sparkVersion: 3.2.1
                 numExecutors: 1
+                definedTags:
+                  Oracle-Tags:
+                    CreatedBy: test_name@oracle.com
+                freeformTags:
+                  test_freeform_key: test_freeform_value
             type: dataFlow
         runtime:
             kind: runtime
@@ -185,6 +194,10 @@ From a Python Environment
             .with_executor_shape("VM.Standard.E4.Flex")
             .with_executor_shape_config(ocpus=4, memory_in_gbs=64)
             .with_spark_version("3.2.1")
+            .with_defined_tag(
+                **{"Oracle-Tags": {"CreatedBy": "test_name@oracle.com"}}
+            )
+            .with_freeform_tag(test_freeform_key="test_freeform_value")
         )
         runtime_config = (
             DataFlowRuntime()
@@ -275,6 +288,11 @@ Again, assume you have the following two files written in your current directory
                     memory_in_gbs: 64
                 sparkVersion: 3.2.1
                 numExecutors: 1
+                definedTags:
+                    Oracle-Tags:
+                        CreatedBy: test_name@oracle.com
+                freeformTags:
+                    test_freeform_key: test_freeform_value
             type: dataFlow
         runtime:
             kind: runtime

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")