catalin/moto
1
0
Fork 0
Code Issues Pull requests Projects Releases Wiki Activity

Run black on moto & test directories.

Browse source
This commit is contained in:
Asher Foa 2019-10-31 08:44:26 -07:00
commit 96e5b1993d
507 changed files with 52541 additions and 47814 deletions
Download patch file Download diff file Expand all files Collapse all files

208
moto/kinesis/models.py
View file

@ -13,9 +13,18 @@ from hashlib import md5
from moto.compat import OrderedDict
from moto.core import BaseBackend, BaseModel
from moto.core.utils import unix_time
from .exceptions import StreamNotFoundError, ShardNotFoundError, ResourceInUseError, \
ResourceNotFoundError, InvalidArgumentError
from .utils import compose_shard_iterator, compose_new_shard_iterator, decompose_shard_iterator
from .exceptions import (
StreamNotFoundError,
ShardNotFoundError,
ResourceInUseError,
ResourceNotFoundError,
InvalidArgumentError,
)
from .utils import (
compose_shard_iterator,
compose_new_shard_iterator,
decompose_shard_iterator,
)
class Record(BaseModel):
@ -32,12 +41,11 @@ class Record(BaseModel):
"Data": self.data,
"PartitionKey": self.partition_key,
"SequenceNumber": str(self.sequence_number),
"ApproximateArrivalTimestamp": self.created_at_datetime.isoformat()
"ApproximateArrivalTimestamp": self.created_at_datetime.isoformat(),
}
class Shard(BaseModel):
def __init__(self, shard_id, starting_hash, ending_hash):
self._shard_id = shard_id
self.starting_hash = starting_hash
@ -75,7 +83,8 @@ class Shard(BaseModel):
last_sequence_number = 0
sequence_number = last_sequence_number + 1
self.records[sequence_number] = Record(
partition_key, data, sequence_number, explicit_hash_key)
partition_key, data, sequence_number, explicit_hash_key
)
return sequence_number
def get_min_sequence_number(self):
@ -94,25 +103,31 @@ class Shard(BaseModel):
else:
# find the last item in the list that was created before
# at_timestamp
r = next((r for r in reversed(self.records.values()) if r.created_at < at_timestamp), None)
r = next(
(
r
for r in reversed(self.records.values())
if r.created_at < at_timestamp
),
None,
)
return r.sequence_number
def to_json(self):
return {
"HashKeyRange": {
"EndingHashKey": str(self.ending_hash),
"StartingHashKey": str(self.starting_hash)
"StartingHashKey": str(self.starting_hash),
},
"SequenceNumberRange": {
"EndingSequenceNumber": self.get_max_sequence_number(),
"StartingSequenceNumber": self.get_min_sequence_number(),
},
"ShardId": self.shard_id
"ShardId": self.shard_id,
}
class Stream(BaseModel):
def __init__(self, stream_name, shard_count, region):
self.stream_name = stream_name
self.shard_count = shard_count
@ -123,10 +138,11 @@ class Stream(BaseModel):
self.tags = {}
self.status = "ACTIVE"
step = 2**128 // shard_count
hash_ranges = itertools.chain(map(lambda i: (i, i * step, (i + 1) * step),
range(shard_count - 1)),
[(shard_count - 1, (shard_count - 1) * step, 2**128)])
step = 2 ** 128 // shard_count
hash_ranges = itertools.chain(
map(lambda i: (i, i * step, (i + 1) * step), range(shard_count - 1)),
[(shard_count - 1, (shard_count - 1) * step, 2 ** 128)],
)
for index, start, end in hash_ranges:
shard = Shard(index, start, end)
@ -137,7 +153,7 @@ class Stream(BaseModel):
return "arn:aws:kinesis:{region}:{account_number}:{stream_name}".format(
region=self.region,
account_number=self.account_number,
stream_name=self.stream_name
stream_name=self.stream_name,
)
def get_shard(self, shard_id):
@ -158,21 +174,22 @@ class Stream(BaseModel):
key = int(explicit_hash_key)
if key >= 2**128:
if key >= 2 ** 128:
raise InvalidArgumentError("explicit_hash_key")
else:
key = int(md5(partition_key.encode('utf-8')).hexdigest(), 16)
key = int(md5(partition_key.encode("utf-8")).hexdigest(), 16)
for shard in self.shards.values():
if shard.starting_hash <= key < shard.ending_hash:
return shard
def put_record(self, partition_key, explicit_hash_key, sequence_number_for_ordering, data):
def put_record(
self, partition_key, explicit_hash_key, sequence_number_for_ordering, data
):
shard = self.get_shard_for_key(partition_key, explicit_hash_key)
sequence_number = shard.put_record(
partition_key, data, explicit_hash_key)
sequence_number = shard.put_record(partition_key, data, explicit_hash_key)
return sequence_number, shard.shard_id
def to_json(self):
@ -198,64 +215,69 @@ class Stream(BaseModel):
}
@classmethod
def create_from_cloudformation_json(cls, resource_name, cloudformation_json, region_name):
properties = cloudformation_json['Properties']
region = properties.get('Region', 'us-east-1')
shard_count = properties.get('ShardCount', 1)
return Stream(properties['Name'], shard_count, region)
def create_from_cloudformation_json(
cls, resource_name, cloudformation_json, region_name
):
properties = cloudformation_json["Properties"]
region = properties.get("Region", "us-east-1")
shard_count = properties.get("ShardCount", 1)
return Stream(properties["Name"], shard_count, region)
class FirehoseRecord(BaseModel):
def __init__(self, record_data):
self.record_id = 12345678
self.record_data = record_data
class DeliveryStream(BaseModel):
def __init__(self, stream_name, **stream_kwargs):
self.name = stream_name
self.redshift_username = stream_kwargs.get('redshift_username')
self.redshift_password = stream_kwargs.get('redshift_password')
self.redshift_jdbc_url = stream_kwargs.get('redshift_jdbc_url')
self.redshift_role_arn = stream_kwargs.get('redshift_role_arn')
self.redshift_copy_command = stream_kwargs.get('redshift_copy_command')
self.redshift_username = stream_kwargs.get("redshift_username")
self.redshift_password = stream_kwargs.get("redshift_password")
self.redshift_jdbc_url = stream_kwargs.get("redshift_jdbc_url")
self.redshift_role_arn = stream_kwargs.get("redshift_role_arn")
self.redshift_copy_command = stream_kwargs.get("redshift_copy_command")
self.s3_config = stream_kwargs.get('s3_config')
self.extended_s3_config = stream_kwargs.get('extended_s3_config')
self.s3_config = stream_kwargs.get("s3_config")
self.extended_s3_config = stream_kwargs.get("extended_s3_config")
self.redshift_s3_role_arn = stream_kwargs.get('redshift_s3_role_arn')
self.redshift_s3_bucket_arn = stream_kwargs.get(
'redshift_s3_bucket_arn')
self.redshift_s3_prefix = stream_kwargs.get('redshift_s3_prefix')
self.redshift_s3_role_arn = stream_kwargs.get("redshift_s3_role_arn")
self.redshift_s3_bucket_arn = stream_kwargs.get("redshift_s3_bucket_arn")
self.redshift_s3_prefix = stream_kwargs.get("redshift_s3_prefix")
self.redshift_s3_compression_format = stream_kwargs.get(
'redshift_s3_compression_format', 'UNCOMPRESSED')
"redshift_s3_compression_format", "UNCOMPRESSED"
)
self.redshift_s3_buffering_hints = stream_kwargs.get(
'redshift_s3_buffering_hints')
"redshift_s3_buffering_hints"
)
self.records = []
self.status = 'ACTIVE'
self.status = "ACTIVE"
self.created_at = datetime.datetime.utcnow()
self.last_updated = datetime.datetime.utcnow()
@property
def arn(self):
return 'arn:aws:firehose:us-east-1:123456789012:deliverystream/{0}'.format(self.name)
return "arn:aws:firehose:us-east-1:123456789012:deliverystream/{0}".format(
self.name
)
def destinations_to_dict(self):
if self.s3_config:
return [{
'DestinationId': 'string',
'S3DestinationDescription': self.s3_config,
}]
return [
{"DestinationId": "string", "S3DestinationDescription": self.s3_config}
]
elif self.extended_s3_config:
return [{
'DestinationId': 'string',
'ExtendedS3DestinationDescription': self.extended_s3_config,
}]
return [
{
"DestinationId": "string",
"ExtendedS3DestinationDescription": self.extended_s3_config,
}
]
else:
return [{
return [
{
"DestinationId": "string",
"RedshiftDestinationDescription": {
"ClusterJDBCURL": self.redshift_jdbc_url,
@ -266,12 +288,12 @@ class DeliveryStream(BaseModel):
"BufferingHints": self.redshift_s3_buffering_hints,
"CompressionFormat": self.redshift_s3_compression_format,
"Prefix": self.redshift_s3_prefix,
"RoleARN": self.redshift_s3_role_arn
"RoleARN": self.redshift_s3_role_arn,
},
"Username": self.redshift_username,
},
}
]
}
]
def to_dict(self):
return {
@ -294,7 +316,6 @@ class DeliveryStream(BaseModel):
class KinesisBackend(BaseBackend):
def __init__(self):
self.streams = OrderedDict()
self.delivery_streams = {}
@ -323,14 +344,24 @@ class KinesisBackend(BaseBackend):
return self.streams.pop(stream_name)
raise StreamNotFoundError(stream_name)
def get_shard_iterator(self, stream_name, shard_id, shard_iterator_type, starting_sequence_number,
at_timestamp):
def get_shard_iterator(
self,
stream_name,
shard_id,
shard_iterator_type,
starting_sequence_number,
at_timestamp,
):
# Validate params
stream = self.describe_stream(stream_name)
shard = stream.get_shard(shard_id)
shard_iterator = compose_new_shard_iterator(
stream_name, shard, shard_iterator_type, starting_sequence_number, at_timestamp
stream_name,
shard,
shard_iterator_type,
starting_sequence_number,
at_timestamp,
)
return shard_iterator
@ -341,14 +372,24 @@ class KinesisBackend(BaseBackend):
stream = self.describe_stream(stream_name)
shard = stream.get_shard(shard_id)
records, last_sequence_id, millis_behind_latest = shard.get_records(last_sequence_id, limit)
records, last_sequence_id, millis_behind_latest = shard.get_records(
last_sequence_id, limit
)
next_shard_iterator = compose_shard_iterator(
stream_name, shard, last_sequence_id)
stream_name, shard, last_sequence_id
)
return next_shard_iterator, records, millis_behind_latest
def put_record(self, stream_name, partition_key, explicit_hash_key, sequence_number_for_ordering, data):
def put_record(
self,
stream_name,
partition_key,
explicit_hash_key,
sequence_number_for_ordering,
data,
):
stream = self.describe_stream(stream_name)
sequence_number, shard_id = stream.put_record(
@ -360,10 +401,7 @@ class KinesisBackend(BaseBackend):
def put_records(self, stream_name, records):
stream = self.describe_stream(stream_name)
response = {
"FailedRecordCount": 0,
"Records": []
}
response = {"FailedRecordCount": 0, "Records": []}
for record in records:
partition_key = record.get("PartitionKey")
@ -373,10 +411,9 @@ class KinesisBackend(BaseBackend):
sequence_number, shard_id = stream.put_record(
partition_key, explicit_hash_key, None, data
)
response['Records'].append({
"SequenceNumber": sequence_number,
"ShardId": shard_id
})
response["Records"].append(
{"SequenceNumber": sequence_number, "ShardId": shard_id}
)
return response
@ -386,18 +423,18 @@ class KinesisBackend(BaseBackend):
if shard_to_split not in stream.shards:
raise ResourceNotFoundError(shard_to_split)
if not re.match(r'0|([1-9]\d{0,38})', new_starting_hash_key):
if not re.match(r"0|([1-9]\d{0,38})", new_starting_hash_key):
raise InvalidArgumentError(new_starting_hash_key)
new_starting_hash_key = int(new_starting_hash_key)
shard = stream.shards[shard_to_split]
last_id = sorted(stream.shards.values(),
key=attrgetter('_shard_id'))[-1]._shard_id
last_id = sorted(stream.shards.values(), key=attrgetter("_shard_id"))[
-1
]._shard_id
if shard.starting_hash < new_starting_hash_key < shard.ending_hash:
new_shard = Shard(
last_id + 1, new_starting_hash_key, shard.ending_hash)
new_shard = Shard(last_id + 1, new_starting_hash_key, shard.ending_hash)
shard.ending_hash = new_starting_hash_key
stream.shards[new_shard.shard_id] = new_shard
else:
@ -434,10 +471,11 @@ class KinesisBackend(BaseBackend):
del stream.shards[shard2.shard_id]
for index in shard2.records:
record = shard2.records[index]
shard1.put_record(record.partition_key,
record.data, record.explicit_hash_key)
shard1.put_record(
record.partition_key, record.data, record.explicit_hash_key
)
''' Firehose '''
""" Firehose """
def create_delivery_stream(self, stream_name, **stream_kwargs):
stream = DeliveryStream(stream_name, **stream_kwargs)
@ -461,25 +499,21 @@ class KinesisBackend(BaseBackend):
record = stream.put_record(record_data)
return record
def list_tags_for_stream(self, stream_name, exclusive_start_tag_key=None, limit=None):
def list_tags_for_stream(
self, stream_name, exclusive_start_tag_key=None, limit=None
):
stream = self.describe_stream(stream_name)
tags = []
result = {
'HasMoreTags': False,
'Tags': tags
}
result = {"HasMoreTags": False, "Tags": tags}
for key, val in sorted(stream.tags.items(), key=lambda x: x[0]):
if limit and len(tags) >= limit:
result['HasMoreTags'] = True
result["HasMoreTags"] = True
break
if exclusive_start_tag_key and key < exclusive_start_tag_key:
continue
tags.append({
'Key': key,
'Value': val
})
tags.append({"Key": key, "Value": val})
return result