From c8d8aa4dd0acf3ba6c62bbf57fa81943d6700f0c Mon Sep 17 00:00:00 2001 From: Don Kuntz Date: Wed, 17 Jul 2019 14:07:19 -0500 Subject: [PATCH] Add glue.batch_get_partition endpoint --- moto/glue/responses.py | 17 +++++++++ tests/test_glue/test_datacatalog.py | 57 +++++++++++++++++++++++++++++ 2 files changed, 74 insertions(+) diff --git a/moto/glue/responses.py b/moto/glue/responses.py index cb1ecf519..875513e7f 100644 --- a/moto/glue/responses.py +++ b/moto/glue/responses.py @@ -141,6 +141,23 @@ class GlueResponse(BaseResponse): return json.dumps({'Partition': p.as_dict()}) + def batch_get_partition(self): + database_name = self.parameters.get('DatabaseName') + table_name = self.parameters.get('TableName') + partitions_to_get = self.parameters.get('PartitionsToGet') + + table = self.glue_backend.get_table(database_name, table_name) + + partitions = [] + for values in partitions_to_get: + try: + p = table.get_partition(values=values["Values"]) + partitions.append(p.as_dict()) + except PartitionNotFoundException: + continue + + return json.dumps({'Partitions': partitions}) + def create_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 232ab3019..9034feb55 100644 --- a/tests/test_glue/test_datacatalog.py +++ b/tests/test_glue/test_datacatalog.py @@ -419,6 +419,63 @@ def test_get_partition(): partition['Values'].should.equal(values[1]) +@mock_glue +def test_batch_get_partition(): + client = boto3.client('glue', region_name='us-east-1') + database_name = 'myspecialdatabase' + table_name = 'myfirsttable' + helpers.create_database(client, database_name) + + helpers.create_table(client, database_name, table_name) + + values = [['2018-10-01'], ['2018-09-01']] + + helpers.create_partition(client, database_name, table_name, values=values[0]) + helpers.create_partition(client, database_name, table_name, values=values[1]) + + partitions_to_get = [ + {'Values': values[0]}, + {'Values': values[1]}, + ] + response = client.batch_get_partition(DatabaseName=database_name, TableName=table_name, PartitionsToGet=partitions_to_get) + + partitions = response['Partitions'] + partitions.should.have.length_of(2) + + partition = partitions[1] + partition['TableName'].should.equal(table_name) + partition['Values'].should.equal(values[1]) + + +@mock_glue +def test_batch_get_partition_missing_partition(): + client = boto3.client('glue', region_name='us-east-1') + database_name = 'myspecialdatabase' + table_name = 'myfirsttable' + helpers.create_database(client, database_name) + + helpers.create_table(client, database_name, table_name) + + values = [['2018-10-01'], ['2018-09-01'], ['2018-08-01']] + + helpers.create_partition(client, database_name, table_name, values=values[0]) + helpers.create_partition(client, database_name, table_name, values=values[2]) + + partitions_to_get = [ + {'Values': values[0]}, + {'Values': values[1]}, + {'Values': values[2]}, + ] + response = client.batch_get_partition(DatabaseName=database_name, TableName=table_name, PartitionsToGet=partitions_to_get) + + partitions = response['Partitions'] + partitions.should.have.length_of(2) + + partitions[0]['Values'].should.equal(values[0]) + partitions[1]['Values'].should.equal(values[2]) + + + @mock_glue def test_update_partition_not_found_moving(): client = boto3.client('glue', region_name='us-east-1')