Glue: Add MLTransform support and tagging for DevEndpoints/Workflows

IMPLEMENTATION_COVERAGE.md

@@ -5129,7 +5129,7 @@
 - [ ] create_integration_resource_property
 - [ ] create_integration_table_properties
 - [X] create_job
-- [ ] create_ml_transform
+- [X] create_ml_transform
 - [X] create_partition
 - [ ] create_partition_index
 - [X] create_registry
@@ -5161,7 +5161,7 @@
 - [ ] delete_integration_resource_property
 - [ ] delete_integration_table_properties
 - [X] delete_job
-- [ ] delete_ml_transform
+- [X] delete_ml_transform
 - [X] delete_partition
 - [ ] delete_partition_index
 - [X] delete_registry
@@ -5226,8 +5226,8 @@
 - [ ] get_materialized_view_refresh_task_run
 - [ ] get_ml_task_run
 - [ ] get_ml_task_runs
-- [ ] get_ml_transform
-- [ ] get_ml_transforms
+- [X] get_ml_transform
+- [X] get_ml_transforms
 - [X] get_partition
 - [ ] get_partition_indexes
 - [X] get_partitions

docs/docs/services/glue.rst

@@ -52,7 +52,7 @@ glue
 - [ ] create_integration_resource_property
 - [ ] create_integration_table_properties
 - [X] create_job
-- [ ] create_ml_transform
+- [X] create_ml_transform
 - [X] create_partition
 - [ ] create_partition_index
 - [X] create_registry
@@ -88,7 +88,7 @@ The following parameters/features are not yet implemented: Glue Schema Registry:
 - [ ] delete_integration_resource_property
 - [ ] delete_integration_table_properties
 - [X] delete_job
-- [ ] delete_ml_transform
+- [X] delete_ml_transform
 - [X] delete_partition
 - [ ] delete_partition_index
 - [X] delete_registry
@@ -153,8 +153,8 @@ The following parameters/features are not yet implemented: Glue Schema Registry:
 - [ ] get_materialized_view_refresh_task_run
 - [ ] get_ml_task_run
 - [ ] get_ml_task_runs
-- [ ] get_ml_transform
-- [ ] get_ml_transforms
+- [X] get_ml_transform
+- [X] get_ml_transforms
 - [X] get_partition
 - [ ] get_partition_indexes
 - [X] get_partitions

moto/glue/models.py

@@ -153,6 +153,70 @@ def as_dict(self) -> dict[str, Any]:
         return response
 
 
+class FakeMLTransform(BaseModel):
+    def __init__(
+        self,
+        name: str,
+        input_record_tables: list[dict[str, Any]],
+        parameters: dict[str, Any],
+        role: str,
+        backend: "GlueBackend",
+        description: str | None = None,
+        glue_version: str | None = None,
+        max_capacity: float | None = None,
+        worker_type: str | None = None,
+        number_of_workers: int | None = None,
+        timeout: int | None = None,
+        max_retries: int | None = None,
+        transform_encryption: dict[str, Any] | None = None,
+    ):
+        self.transform_id = f"tfm-{mock_random.get_random_hex(32)}"
+        self.name = name
+        self.input_record_tables = input_record_tables or []
+        self.parameters = parameters or {}
+        self.role = role
+        self.description = description
+        self.glue_version = glue_version or "1.0"
+        self.max_capacity = max_capacity
+        self.worker_type = worker_type
+        self.number_of_workers = number_of_workers
+        self.timeout = timeout or 2880
+        self.max_retries = max_retries if max_retries is not None else 0
+        self.transform_encryption = transform_encryption
+        self.status = "READY"
+        self.created_on = utcnow()
+        self.last_modified_on = self.created_on
+        self.label_count = 0
+        self.backend = backend
+        self.arn = f"arn:{get_partition(backend.region_name)}:glue:{backend.region_name}:{backend.account_id}:mlTransform/{self.transform_id}"
+
+    def as_dict(self) -> dict[str, Any]:
+        response: dict[str, Any] = {
+            "TransformId": self.transform_id,
+            "Name": self.name,
+            "Description": self.description,
+            "Status": self.status,
+            "CreatedOn": self.created_on,
+            "LastModifiedOn": self.last_modified_on,
+            "InputRecordTables": self.input_record_tables,
+            "Parameters": self.parameters,
+            "Role": self.role,
+            "GlueVersion": self.glue_version,
+            "Timeout": self.timeout,
+            "MaxRetries": self.max_retries,
+            "LabelCount": self.label_count,
+        }
+        if self.max_capacity is not None:
+            response["MaxCapacity"] = self.max_capacity
+        if self.worker_type is not None:
+            response["WorkerType"] = self.worker_type
+        if self.number_of_workers is not None:
+            response["NumberOfWorkers"] = self.number_of_workers
+        if self.transform_encryption is not None:
+            response["TransformEncryption"] = self.transform_encryption
+        return response
+
+
 class GlueBackend(BaseBackend, TaggableResourcesMixin):
     SERVICE_NAMESPACE = "glue"
 
@@ -229,6 +293,12 @@ class GlueBackend(BaseBackend, TaggableResourcesMixin):
             "limit_default": 100,
             "unique_attribute": "workflow_run_id",
         },
+        "get_ml_transforms": {
+            "input_token": "next_token",
+            "limit_key": "max_results",
+            "limit_default": 100,
+            "unique_attribute": "transform_id",
+        },
     }
 
     def __init__(self, region_name: str, account_id: str):
@@ -251,6 +321,7 @@ def __init__(self, region_name: str, account_id: str):
         self.num_schemas = 0
         self.num_schema_versions = 0
         self.dev_endpoints: dict[str, FakeDevEndpoint] = OrderedDict()
+        self.ml_transforms: dict[str, FakeMLTransform] = OrderedDict()
         self.data_catalog_encryption_settings: dict[str, dict[str, Any]] = {}
         self.resource_policies: dict[str, dict[str, Any]] = {}
         self.default_catalog_arn = f"arn:{get_partition(self.region_name)}:glue:{self.region_name}:{self.account_id}:catalog"
@@ -1319,6 +1390,8 @@ def create_dev_endpoint(
             dev_endpoint.public_keys = public_keys
 
         self.dev_endpoints[endpoint_name] = dev_endpoint
+        resource_arn = f"arn:{get_partition(self.region_name)}:glue:{self.region_name}:{self.account_id}:devEndpoint/{endpoint_name}"
+        self.tag_resource(resource_arn, tags)
         return dev_endpoint
 
     def get_dev_endpoint(self, endpoint_name: str) -> FakeDevEndpoint:
@@ -1332,6 +1405,60 @@ def delete_dev_endpoint(self, endpoint_name: str) -> None:
             del self.dev_endpoints[endpoint_name]
         except KeyError:
             raise EntityNotFoundException(f"DevEndpoint {endpoint_name} not found")
+        resource_arn = f"arn:{get_partition(self.region_name)}:glue:{self.region_name}:{self.account_id}:devEndpoint/{endpoint_name}"
+        self.tagger.delete_all_tags_for_resource(resource_arn)
+
+    def create_ml_transform(
+        self,
+        name: str,
+        input_record_tables: list[dict[str, Any]],
+        parameters: dict[str, Any],
+        role: str,
+        description: str | None = None,
+        glue_version: str | None = None,
+        max_capacity: float | None = None,
+        worker_type: str | None = None,
+        number_of_workers: int | None = None,
+        timeout: int | None = None,
+        max_retries: int | None = None,
+        tags: dict[str, str] | None = None,
+        transform_encryption: dict[str, Any] | None = None,
+    ) -> str:
+        ml_transform = FakeMLTransform(
+            name=name,
+            input_record_tables=input_record_tables,
+            parameters=parameters,
+            role=role,
+            description=description,
+            glue_version=glue_version,
+            max_capacity=max_capacity,
+            worker_type=worker_type,
+            number_of_workers=number_of_workers,
+            timeout=timeout,
+            max_retries=max_retries,
+            transform_encryption=transform_encryption,
+            backend=self,
+        )
+        self.ml_transforms[ml_transform.transform_id] = ml_transform
+        self.tag_resource(ml_transform.arn, tags)
+        return ml_transform.transform_id
+
+    def get_ml_transform(self, transform_id: str) -> "FakeMLTransform":
+        try:
+            return self.ml_transforms[transform_id]
+        except KeyError:
+            raise EntityNotFoundException(f"MLTransform {transform_id} not found")
+
+    @paginate(pagination_model=PAGINATION_MODEL)
+    def get_ml_transforms(self) -> list["FakeMLTransform"]:
+        return list(self.ml_transforms.values())
+
+    def delete_ml_transform(self, transform_id: str) -> str:
+        if transform_id not in self.ml_transforms:
+            raise EntityNotFoundException(f"MLTransform {transform_id} not found")
+        ml_transform = self.ml_transforms.pop(transform_id)
+        self.tagger.delete_all_tags_for_resource(ml_transform.arn)
+        return transform_id
 
     def create_connection(
         self,
@@ -1516,6 +1643,8 @@ def create_workflow(
         self.workflows[name] = FakeWorkflow(
             name, default_run_properties, description, max_concurrent_runs, tags
         )
+        resource_arn = f"arn:{get_partition(self.region_name)}:glue:{self.region_name}:{self.account_id}:workflow/{name}"
+        self.tag_resource(resource_arn, tags)
         return name
 
     def get_workflow(
@@ -1558,6 +1687,8 @@ def delete_workflow(
     ) -> str:
         if self.workflows.get(name):
             del self.workflows[name]
+            resource_arn = f"arn:{get_partition(self.region_name)}:glue:{self.region_name}:{self.account_id}:workflow/{name}"
+            self.tagger.delete_all_tags_for_resource(resource_arn)
         return name
 
     def get_workflow_run(

moto/glue/responses.py

@@ -870,6 +870,61 @@ def delete_dev_endpoint(self) -> ActionResult:
         self.glue_backend.delete_dev_endpoint(endpoint_name)
         return EmptyResult()
 
+    def create_ml_transform(self) -> ActionResult:
+        name = self._get_param("Name")
+        input_record_tables = self._get_param("InputRecordTables")
+        parameters = self._get_param("Parameters")
+        role = self._get_param("Role")
+        description = self._get_param("Description")
+        glue_version = self._get_param("GlueVersion")
+        max_capacity = self._get_param("MaxCapacity")
+        worker_type = self._get_param("WorkerType")
+        number_of_workers = self._get_int_param("NumberOfWorkers")
+        timeout = self._get_int_param("Timeout")
+        max_retries = self._get_int_param("MaxRetries")
+        tags = self._get_param("Tags")
+        transform_encryption = self._get_param("TransformEncryption")
+
+        transform_id = self.glue_backend.create_ml_transform(
+            name=name,
+            input_record_tables=input_record_tables,
+            parameters=parameters,
+            role=role,
+            description=description,
+            glue_version=glue_version,
+            max_capacity=max_capacity,
+            worker_type=worker_type,
+            number_of_workers=number_of_workers,
+            timeout=timeout,
+            max_retries=max_retries,
+            tags=tags,
+            transform_encryption=transform_encryption,
+        )
+        return ActionResult({"TransformId": transform_id})
+
+    def get_ml_transform(self) -> ActionResult:
+        transform_id = self._get_param("TransformId")
+        ml_transform = self.glue_backend.get_ml_transform(transform_id)
+        return ActionResult(ml_transform.as_dict())
+
+    def get_ml_transforms(self) -> ActionResult:
+        next_token = self._get_param("NextToken")
+        max_results = self._get_int_param("MaxResults")
+        ml_transforms, next_token = self.glue_backend.get_ml_transforms(
+            next_token=next_token, max_results=max_results
+        )
+        return ActionResult(
+            {
+                "Transforms": [transform.as_dict() for transform in ml_transforms],
+                "NextToken": next_token,
+            }
+        )
+
+    def delete_ml_transform(self) -> ActionResult:
+        transform_id = self._get_param("TransformId")
+        transform_id = self.glue_backend.delete_ml_transform(transform_id)
+        return ActionResult({"TransformId": transform_id})
+
     def create_connection(self) -> ActionResult:
         catalog_id = self._get_param("CatalogId")
         connection_input = self._get_param("ConnectionInput")