From effb075b62cdb8862b9fa58ff272863c67e84635 Mon Sep 17 00:00:00 2001
From: Don Kuntz <dkuntz@journera.com>
Date: Thu, 10 Dec 2020 14:03:37 -0600
Subject: [PATCH] Add glue.batch_update_partition (#3534)

* Add glue.batch_update_partition

* Fix error output for glue.batch_update_partition and add test case for non-existent partition in the batch update
---
 moto/glue/responses.py              |  29 +++++++
 tests/test_glue/test_datacatalog.py | 113 ++++++++++++++++++++++++++++
 2 files changed, 142 insertions(+)

diff --git a/moto/glue/responses.py b/moto/glue/responses.py
index 66185e099..ba9cee8fc 100644
--- a/moto/glue/responses.py
+++ b/moto/glue/responses.py
@@ -208,6 +208,35 @@ class GlueResponse(BaseResponse):
 
         return ""
 
+    def batch_update_partition(self):
+        database_name = self.parameters.get("DatabaseName")
+        table_name = self.parameters.get("TableName")
+        table = self.glue_backend.get_table(database_name, table_name)
+
+        errors_output = []
+        for entry in self.parameters.get("Entries"):
+            part_to_update = entry["PartitionValueList"]
+            part_input = entry["PartitionInput"]
+
+            try:
+                table.update_partition(part_to_update, part_input)
+            except PartitionNotFoundException:
+                errors_output.append(
+                    {
+                        "PartitionValueList": part_to_update,
+                        "ErrorDetail": {
+                            "ErrorCode": "EntityNotFoundException",
+                            "ErrorMessage": "Partition not found.",
+                        },
+                    }
+                )
+
+        out = {}
+        if errors_output:
+            out["Errors"] = errors_output
+
+        return json.dumps(out)
+
     def delete_partition(self):
         database_name = self.parameters.get("DatabaseName")
         table_name = self.parameters.get("TableName")
diff --git a/tests/test_glue/test_datacatalog.py b/tests/test_glue/test_datacatalog.py
index 38a3831d5..46ef910b7 100644
--- a/tests/test_glue/test_datacatalog.py
+++ b/tests/test_glue/test_datacatalog.py
@@ -663,6 +663,119 @@ def test_update_partition_move():
     )
 
 
+@mock_glue
+def test_batch_update_partition():
+    client = boto3.client("glue", region_name="us-east-1")
+    database_name = "myspecialdatabase"
+    table_name = "myfirsttable"
+
+    values = [
+        ["2020-12-04"],
+        ["2020-12-05"],
+        ["2020-12-06"],
+    ]
+
+    new_values = [
+        ["2020-11-04"],
+        ["2020-11-05"],
+        ["2020-11-06"],
+    ]
+
+    helpers.create_database(client, database_name)
+    helpers.create_table(client, database_name, table_name)
+
+    batch_update_values = []
+    for idx, value in enumerate(values):
+        helpers.create_partition(client, database_name, table_name, values=value)
+        batch_update_values.append(
+            {
+                "PartitionValueList": value,
+                "PartitionInput": helpers.create_partition_input(
+                    database_name,
+                    table_name,
+                    values=new_values[idx],
+                    columns=[{"Name": "country", "Type": "string"}],
+                ),
+            }
+        )
+
+    response = client.batch_update_partition(
+        DatabaseName=database_name, TableName=table_name, Entries=batch_update_values,
+    )
+
+    for value in values:
+        with pytest.raises(ClientError) as exc:
+            helpers.get_partition(client, database_name, table_name, value)
+        exc.value.response["Error"]["Code"].should.equal("EntityNotFoundException")
+
+    for value in new_values:
+        response = client.get_partition(
+            DatabaseName=database_name, TableName=table_name, PartitionValues=value
+        )
+        partition = response["Partition"]
+
+        partition["TableName"].should.equal(table_name)
+        partition["StorageDescriptor"]["Columns"].should.equal(
+            [{"Name": "country", "Type": "string"}]
+        )
+
+
+@mock_glue
+def test_batch_update_partition_missing_partition():
+    client = boto3.client("glue", region_name="us-east-1")
+    database_name = "myspecialdatabase"
+    table_name = "myfirsttable"
+
+    values = [
+        ["2020-12-05"],
+        ["2020-12-06"],
+    ]
+
+    new_values = [
+        ["2020-11-05"],
+        ["2020-11-06"],
+    ]
+
+    helpers.create_database(client, database_name)
+    helpers.create_table(client, database_name, table_name)
+
+    batch_update_values = []
+    for idx, value in enumerate(values):
+        helpers.create_partition(client, database_name, table_name, values=value)
+        batch_update_values.append(
+            {
+                "PartitionValueList": value,
+                "PartitionInput": helpers.create_partition_input(
+                    database_name,
+                    table_name,
+                    values=new_values[idx],
+                    columns=[{"Name": "country", "Type": "string"}],
+                ),
+            }
+        )
+
+    # add a non-existent partition to the batch update values
+    batch_update_values.append(
+        {
+            "PartitionValueList": ["2020-10-10"],
+            "PartitionInput": helpers.create_partition_input(
+                database_name,
+                table_name,
+                values=["2019-09-09"],
+                columns=[{"Name": "country", "Type": "string"}],
+            ),
+        }
+    )
+
+    response = client.batch_update_partition(
+        DatabaseName=database_name, TableName=table_name, Entries=batch_update_values,
+    )
+
+    response.should.have.key("Errors")
+    response["Errors"].should.have.length_of(1)
+    response["Errors"][0]["PartitionValueList"].should.equal(["2020-10-10"])
+
+
 @mock_glue
 def test_delete_partition():
     client = boto3.client("glue", region_name="us-east-1")