Merge pull request #1848 from ashb/expand-glue-catalog-mocking

Mock more of the Glue Data Catalog APIs
This commit is contained in:
Steve Pulec 2018-10-15 01:14:03 -04:00 committed by GitHub
commit dfc793916d
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 673 additions and 61 deletions

View file

@ -6,19 +6,56 @@ class GlueClientError(JsonRESTError):
code = 400
class DatabaseAlreadyExistsException(GlueClientError):
def __init__(self):
self.code = 400
super(DatabaseAlreadyExistsException, self).__init__(
'DatabaseAlreadyExistsException',
'Database already exists.'
class AlreadyExistsException(GlueClientError):
def __init__(self, typ):
super(GlueClientError, self).__init__(
'AlreadyExistsException',
'%s already exists.' % (typ),
)
class TableAlreadyExistsException(GlueClientError):
class DatabaseAlreadyExistsException(AlreadyExistsException):
def __init__(self):
self.code = 400
super(TableAlreadyExistsException, self).__init__(
'TableAlreadyExistsException',
'Table already exists.'
super(DatabaseAlreadyExistsException, self).__init__('Database')
class TableAlreadyExistsException(AlreadyExistsException):
def __init__(self):
super(TableAlreadyExistsException, self).__init__('Table')
class PartitionAlreadyExistsException(AlreadyExistsException):
def __init__(self):
super(PartitionAlreadyExistsException, self).__init__('Partition')
class EntityNotFoundException(GlueClientError):
def __init__(self, msg):
super(GlueClientError, self).__init__(
'EntityNotFoundException',
msg,
)
class DatabaseNotFoundException(EntityNotFoundException):
def __init__(self, db):
super(DatabaseNotFoundException, self).__init__(
'Database %s not found.' % db,
)
class TableNotFoundException(EntityNotFoundException):
def __init__(self, tbl):
super(TableNotFoundException, self).__init__(
'Table %s not found.' % tbl,
)
class PartitionNotFoundException(EntityNotFoundException):
def __init__(self):
super(PartitionNotFoundException, self).__init__("Cannot find partition.")
class VersionNotFoundException(EntityNotFoundException):
def __init__(self):
super(VersionNotFoundException, self).__init__("Version not found.")

View file

@ -1,8 +1,19 @@
from __future__ import unicode_literals
import time
from moto.core import BaseBackend, BaseModel
from moto.compat import OrderedDict
from.exceptions import DatabaseAlreadyExistsException, TableAlreadyExistsException
from.exceptions import (
JsonRESTError,
DatabaseAlreadyExistsException,
DatabaseNotFoundException,
TableAlreadyExistsException,
TableNotFoundException,
PartitionAlreadyExistsException,
PartitionNotFoundException,
VersionNotFoundException,
)
class GlueBackend(BaseBackend):
@ -19,7 +30,10 @@ class GlueBackend(BaseBackend):
return database
def get_database(self, database_name):
return self.databases[database_name]
try:
return self.databases[database_name]
except KeyError:
raise DatabaseNotFoundException(database_name)
def create_table(self, database_name, table_name, table_input):
database = self.get_database(database_name)
@ -33,7 +47,10 @@ class GlueBackend(BaseBackend):
def get_table(self, database_name, table_name):
database = self.get_database(database_name)
return database.tables[table_name]
try:
return database.tables[table_name]
except KeyError:
raise TableNotFoundException(table_name)
def get_tables(self, database_name):
database = self.get_database(database_name)
@ -52,9 +69,84 @@ class FakeTable(BaseModel):
def __init__(self, database_name, table_name, table_input):
self.database_name = database_name
self.name = table_name
self.table_input = table_input
self.storage_descriptor = self.table_input.get('StorageDescriptor', {})
self.partition_keys = self.table_input.get('PartitionKeys', [])
self.partitions = OrderedDict()
self.versions = []
self.update(table_input)
def update(self, table_input):
self.versions.append(table_input)
def get_version(self, ver):
try:
if not isinstance(ver, int):
# "1" goes to [0]
ver = int(ver) - 1
except ValueError as e:
raise JsonRESTError("InvalidInputException", str(e))
try:
return self.versions[ver]
except IndexError:
raise VersionNotFoundException()
def as_dict(self, version=-1):
obj = {
'DatabaseName': self.database_name,
'Name': self.name,
}
obj.update(self.get_version(version))
return obj
def create_partition(self, partiton_input):
partition = FakePartition(self.database_name, self.name, partiton_input)
key = str(partition.values)
if key in self.partitions:
raise PartitionAlreadyExistsException()
self.partitions[str(partition.values)] = partition
def get_partitions(self):
return [p for str_part_values, p in self.partitions.items()]
def get_partition(self, values):
try:
return self.partitions[str(values)]
except KeyError:
raise PartitionNotFoundException()
def update_partition(self, old_values, partiton_input):
partition = FakePartition(self.database_name, self.name, partiton_input)
key = str(partition.values)
if old_values == partiton_input['Values']:
# Altering a partition in place. Don't remove it so the order of
# returned partitions doesn't change
if key not in self.partitions:
raise PartitionNotFoundException()
else:
removed = self.partitions.pop(str(old_values), None)
if removed is None:
raise PartitionNotFoundException()
if key in self.partitions:
# Trying to update to overwrite a partition that exists
raise PartitionAlreadyExistsException()
self.partitions[key] = partition
class FakePartition(BaseModel):
def __init__(self, database_name, table_name, partiton_input):
self.creation_time = time.time()
self.database_name = database_name
self.table_name = table_name
self.partition_input = partiton_input
self.values = self.partition_input.get('Values', [])
def as_dict(self):
obj = {
'DatabaseName': self.database_name,
'TableName': self.table_name,
'CreationTime': self.creation_time,
}
obj.update(self.partition_input)
return obj
glue_backend = GlueBackend()

View file

@ -37,27 +37,94 @@ class GlueResponse(BaseResponse):
database_name = self.parameters.get('DatabaseName')
table_name = self.parameters.get('Name')
table = self.glue_backend.get_table(database_name, table_name)
return json.dumps({'Table': table.as_dict()})
def update_table(self):
database_name = self.parameters.get('DatabaseName')
table_input = self.parameters.get('TableInput')
table_name = table_input.get('Name')
table = self.glue_backend.get_table(database_name, table_name)
table.update(table_input)
return ""
def get_table_versions(self):
database_name = self.parameters.get('DatabaseName')
table_name = self.parameters.get('TableName')
table = self.glue_backend.get_table(database_name, table_name)
return json.dumps({
'Table': {
'DatabaseName': table.database_name,
'Name': table.name,
'PartitionKeys': table.partition_keys,
'StorageDescriptor': table.storage_descriptor
}
"TableVersions": [
{
"Table": table.as_dict(version=n),
"VersionId": str(n + 1),
} for n in range(len(table.versions))
],
})
def get_table_version(self):
database_name = self.parameters.get('DatabaseName')
table_name = self.parameters.get('TableName')
table = self.glue_backend.get_table(database_name, table_name)
ver_id = self.parameters.get('VersionId')
return json.dumps({
"TableVersion": {
"Table": table.as_dict(version=ver_id),
"VersionId": ver_id,
},
})
def get_tables(self):
database_name = self.parameters.get('DatabaseName')
tables = self.glue_backend.get_tables(database_name)
return json.dumps(
{
'TableList': [
{
'DatabaseName': table.database_name,
'Name': table.name,
'PartitionKeys': table.partition_keys,
'StorageDescriptor': table.storage_descriptor
} for table in tables
]
}
)
return json.dumps({
'TableList': [
table.as_dict() for table in tables
]
})
def get_partitions(self):
database_name = self.parameters.get('DatabaseName')
table_name = self.parameters.get('TableName')
if 'Expression' in self.parameters:
raise NotImplementedError("Expression filtering in get_partitions is not implemented in moto")
table = self.glue_backend.get_table(database_name, table_name)
return json.dumps({
'Partitions': [
p.as_dict() for p in table.get_partitions()
]
})
def get_partition(self):
database_name = self.parameters.get('DatabaseName')
table_name = self.parameters.get('TableName')
values = self.parameters.get('PartitionValues')
table = self.glue_backend.get_table(database_name, table_name)
p = table.get_partition(values)
return json.dumps({'Partition': p.as_dict()})
def create_partition(self):
database_name = self.parameters.get('DatabaseName')
table_name = self.parameters.get('TableName')
part_input = self.parameters.get('PartitionInput')
table = self.glue_backend.get_table(database_name, table_name)
table.create_partition(part_input)
return ""
def update_partition(self):
database_name = self.parameters.get('DatabaseName')
table_name = self.parameters.get('TableName')
part_input = self.parameters.get('PartitionInput')
part_to_update = self.parameters.get('PartitionValueList')
table = self.glue_backend.get_table(database_name, table_name)
table.update_partition(part_to_update, part_input)
return ""