Repository: TobikoData/sqlmesh
Pull request #3488 (2024-12-08 23:49:30)

Fix(athena): Drop partitions in batches of 25
Prior to this change, if the Athena adapter needed to replace more than 25 partitions in a partitioned model, the following boto3 error would occur:
```
'partitionsToDelete' failed to satisfy constraint: Member must have length less than or
equal to 25
```

This PR changes the adapter to call `batch_delete_partition` in batches of 25, which is the [maximum number](https://docs.aws.amazon.com/glue/latest/webapi/API_BatchDeletePartition.html#API_BatchDeletePartition_RequestParameters) allowed at once.

Note that it's just a dumb loop, technically it could be called in parallel from multiple worker threads to reduce latency if there are a lot of partitions to drop but we could add that later if it turns out to be a problem


Patch:
diff --git a/sqlmesh/core/engine_adapter/athena.py b/sqlmesh/core/engine_adapter/athena.py
index 2cb500e5..f257b6a5 100644
--- a/sqlmesh/core/engine_adapter/athena.py
+++ b/sqlmesh/core/engine_adapter/athena.py
@@ -527,13 +527,21 @@ class AthenaEngineAdapter(PandasNativeFetchDFSupportMixin, RowDiffMixin):
         raise SQLMeshError(f"Table {table} has no location set in the metastore!")
 
     def _drop_partitions_from_metastore(
-        self, table: exp.Table, partition_values: t.List[t.List[t.Any]]
+        self, table: exp.Table, partition_values: t.List[t.List[str]]
     ) -> None:
-        self._glue_client.batch_delete_partition(
-            DatabaseName=table.db,
-            TableName=table.name,
-            PartitionsToDelete=[{"Values": v} for v in partition_values],
-        )
+        # todo: switch to itertools.batched when our minimum supported Python is 3.12
+        # 25 = maximum number of partitions that batch_delete_partition can process at once
+        # ref: https://docs.aws.amazon.com/glue/latest/webapi/API_BatchDeletePartition.html#API_BatchDeletePartition_RequestParameters
+        def _chunks() -> t.Iterable[t.List[t.List[str]]]:
+            for i in range(0, len(partition_values), 25):
+                yield partition_values[i : i + 25]
+
+        for batch in _chunks():
+            self._glue_client.batch_delete_partition(
+                DatabaseName=table.db,
+                TableName=table.name,
+                PartitionsToDelete=[{"Values": v} for v in batch],
+            )
 
     def delete_from(self, table_name: TableName, where: t.Union[str, exp.Expression]) -> None:
         table = exp.to_table(table_name)


Test patch:
diff --git a/tests/core/engine_adapter/test_athena.py b/tests/core/engine_adapter/test_athena.py
index fc1b148b..29697469 100644
--- a/tests/core/engine_adapter/test_athena.py
+++ b/tests/core/engine_adapter/test_athena.py
@@ -365,3 +365,42 @@ def test_create_state_table(adapter: AthenaEngineAdapter):
     assert to_sql_calls(adapter) == [
         "CREATE TABLE IF NOT EXISTS `_snapshots` (`name` STRING) LOCATION 's3://base/_snapshots/' TBLPROPERTIES ('table_type'='iceberg')"
     ]
+
+
+def test_drop_partitions_from_metastore_uses_batches(
+    adapter: AthenaEngineAdapter, mocker: MockerFixture
+):
+    glue_client_mock = mocker.patch.object(AthenaEngineAdapter, "_glue_client", autospec=True)
+
+    glue_client_mock.batch_delete_partition.assert_not_called()
+
+    partition_values = []
+
+    for i in range(63):
+        partition_values.append([str(i)])
+
+    adapter._drop_partitions_from_metastore(
+        table=exp.table_("foo"), partition_values=partition_values
+    )
+
+    glue_client_mock.batch_delete_partition.assert_called()
+
+    # should have been called in batches of 25
+    calls = glue_client_mock.batch_delete_partition.call_args_list
+    assert len(calls) == 3
+
+    assert len(calls[0][1]["PartitionsToDelete"]) == 25
+    assert len(calls[1][1]["PartitionsToDelete"]) == 25
+    assert len(calls[2][1]["PartitionsToDelete"]) == 13
+
+    # first call 0-24
+    assert calls[0][1]["PartitionsToDelete"][0]["Values"][0] == "0"
+    assert calls[0][1]["PartitionsToDelete"][-1]["Values"][0] == "24"
+
+    # second call 25-49
+    assert calls[1][1]["PartitionsToDelete"][0]["Values"][0] == "25"
+    assert calls[1][1]["PartitionsToDelete"][-1]["Values"][0] == "49"
+
+    # third call 50-62
+    assert calls[2][1]["PartitionsToDelete"][0]["Values"][0] == "50"
+    assert calls[2][1]["PartitionsToDelete"][-1]["Values"][0] == "62"