From 4d0099499faf54b940f807781c7e821ea69b8f70 Mon Sep 17 00:00:00 2001 From: Jack Danger Date: Tue, 22 Oct 2019 14:37:29 -0700 Subject: [PATCH] Athena - implementing first two endpoints (#2506) This implements create_work_group() and list_work_groups() --- moto/__init__.py | 1 + moto/athena/__init__.py | 7 +++ moto/athena/exceptions.py | 18 ++++++++ moto/athena/models.py | 79 ++++++++++++++++++++++++++++++++ moto/athena/responses.py | 35 ++++++++++++++ moto/athena/urls.py | 10 ++++ moto/athena/utils.py | 1 + moto/backends.py | 4 +- tests/test_athena/test_athena.py | 59 ++++++++++++++++++++++++ 9 files changed, 213 insertions(+), 1 deletion(-) create mode 100644 moto/athena/__init__.py create mode 100644 moto/athena/exceptions.py create mode 100644 moto/athena/models.py create mode 100644 moto/athena/responses.py create mode 100644 moto/athena/urls.py create mode 100644 moto/athena/utils.py create mode 100644 tests/test_athena/test_athena.py diff --git a/moto/__init__.py b/moto/__init__.py index f82a411cf..8e915933a 100644 --- a/moto/__init__.py +++ b/moto/__init__.py @@ -7,6 +7,7 @@ __version__ = '1.3.14.dev' from .acm import mock_acm # flake8: noqa from .apigateway import mock_apigateway, mock_apigateway_deprecated # flake8: noqa +from .athena import mock_athena # flake8: noqa from .autoscaling import mock_autoscaling, mock_autoscaling_deprecated # flake8: noqa from .awslambda import mock_lambda, mock_lambda_deprecated # flake8: noqa from .cloudformation import mock_cloudformation, mock_cloudformation_deprecated # flake8: noqa diff --git a/moto/athena/__init__.py b/moto/athena/__init__.py new file mode 100644 index 000000000..c7bfa2b1f --- /dev/null +++ b/moto/athena/__init__.py @@ -0,0 +1,7 @@ +from __future__ import unicode_literals +from .models import athena_backends +from ..core.models import base_decorator, deprecated_base_decorator + +athena_backend = athena_backends['us-east-1'] +mock_athena = base_decorator(athena_backends) +mock_athena_deprecated = deprecated_base_decorator(athena_backends) diff --git a/moto/athena/exceptions.py b/moto/athena/exceptions.py new file mode 100644 index 000000000..1faa54731 --- /dev/null +++ b/moto/athena/exceptions.py @@ -0,0 +1,18 @@ +from __future__ import unicode_literals + +import json +from werkzeug.exceptions import BadRequest + + +class AthenaClientError(BadRequest): + + def __init__(self, code, message): + super(AthenaClientError, self).__init__() + self.description = json.dumps({ + "Error": { + "Code": code, + "Message": message, + 'Type': "InvalidRequestException", + }, + 'RequestId': '6876f774-7273-11e4-85dc-39e55ca848d1', + }) diff --git a/moto/athena/models.py b/moto/athena/models.py new file mode 100644 index 000000000..d6d88cde8 --- /dev/null +++ b/moto/athena/models.py @@ -0,0 +1,79 @@ +from __future__ import unicode_literals +import time + +import boto3 +from moto.core import BaseBackend, BaseModel + +ACCOUNT_ID = 123456789012 + + +class TaggableResourceMixin(object): + # This mixing was copied from Redshift when initially implementing + # Athena. TBD if it's worth the overhead. + + def __init__(self, region_name, resource_name, tags): + self.region = region_name + self.resource_name = resource_name + self.tags = tags or [] + + @property + def arn(self): + return "arn:aws:athena:{region}:{account_id}:{resource_name}".format( + region=self.region, + account_id=ACCOUNT_ID, + resource_name=self.resource_name) + + def create_tags(self, tags): + new_keys = [tag_set['Key'] for tag_set in tags] + self.tags = [tag_set for tag_set in self.tags + if tag_set['Key'] not in new_keys] + self.tags.extend(tags) + return self.tags + + def delete_tags(self, tag_keys): + self.tags = [tag_set for tag_set in self.tags + if tag_set['Key'] not in tag_keys] + return self.tags + + +class WorkGroup(TaggableResourceMixin, BaseModel): + + resource_type = 'workgroup' + state = 'ENABLED' + + def __init__(self, athena_backend, name, configuration, description, tags): + self.region_name = athena_backend.region_name + super(WorkGroup, self).__init__(self.region_name, "workgroup/{}".format(name), tags) + self.athena_backend = athena_backend + self.name = name + self.description = description + self.configuration = configuration + + +class AthenaBackend(BaseBackend): + region_name = None + + def __init__(self, region_name=None): + if region_name is not None: + self.region_name = region_name + self.work_groups = {} + + def create_work_group(self, name, configuration, description, tags): + if name in self.work_groups: + return None + work_group = WorkGroup(self, name, configuration, description, tags) + self.work_groups[name] = work_group + return work_group + + def list_work_groups(self): + return [{ + 'Name': wg.name, + 'State': wg.state, + 'Description': wg.description, + 'CreationTime': time.time(), + } for wg in self.work_groups.values()] + + +athena_backends = {} +for region in boto3.Session().get_available_regions('athena'): + athena_backends[region] = AthenaBackend(region) diff --git a/moto/athena/responses.py b/moto/athena/responses.py new file mode 100644 index 000000000..13d33c129 --- /dev/null +++ b/moto/athena/responses.py @@ -0,0 +1,35 @@ +import json + +from moto.core.responses import BaseResponse +from .models import athena_backends + + +class AthenaResponse(BaseResponse): + + @property + def athena_backend(self): + return athena_backends[self.region] + + def create_work_group(self): + name = self._get_param('Name') + description = self._get_param('Description') + configuration = self._get_param('Configuration') + tags = self._get_param('Tags') + work_group = self.athena_backend.create_work_group(name, configuration, description, tags) + if not work_group: + return json.dumps({ + '__type': 'InvalidRequestException', + 'Message': 'WorkGroup already exists', + }), dict(status=400) + return json.dumps({ + "CreateWorkGroupResponse": { + "ResponseMetadata": { + "RequestId": "384ac68d-3775-11df-8963-01868b7c937a", + } + } + }) + + def list_work_groups(self): + return json.dumps({ + "WorkGroups": self.athena_backend.list_work_groups() + }) diff --git a/moto/athena/urls.py b/moto/athena/urls.py new file mode 100644 index 000000000..bdd4ebc1e --- /dev/null +++ b/moto/athena/urls.py @@ -0,0 +1,10 @@ +from __future__ import unicode_literals +from .responses import AthenaResponse + +url_bases = [ + "https?://athena.(.+).amazonaws.com", +] + +url_paths = { + '{0}/$': AthenaResponse.dispatch, +} diff --git a/moto/athena/utils.py b/moto/athena/utils.py new file mode 100644 index 000000000..baffc4882 --- /dev/null +++ b/moto/athena/utils.py @@ -0,0 +1 @@ +from __future__ import unicode_literals diff --git a/moto/backends.py b/moto/backends.py index 8a20697c2..0a387ac7e 100644 --- a/moto/backends.py +++ b/moto/backends.py @@ -2,6 +2,7 @@ from __future__ import unicode_literals from moto.acm import acm_backends from moto.apigateway import apigateway_backends +from moto.athena import athena_backends from moto.autoscaling import autoscaling_backends from moto.awslambda import lambda_backends from moto.cloudformation import cloudformation_backends @@ -35,8 +36,8 @@ from moto.redshift import redshift_backends from moto.resourcegroups import resourcegroups_backends from moto.route53 import route53_backends from moto.s3 import s3_backends -from moto.ses import ses_backends from moto.secretsmanager import secretsmanager_backends +from moto.ses import ses_backends from moto.sns import sns_backends from moto.sqs import sqs_backends from moto.ssm import ssm_backends @@ -53,6 +54,7 @@ from moto.config import config_backends BACKENDS = { 'acm': acm_backends, 'apigateway': apigateway_backends, + 'athena': athena_backends, 'autoscaling': autoscaling_backends, 'batch': batch_backends, 'cloudformation': cloudformation_backends, diff --git a/tests/test_athena/test_athena.py b/tests/test_athena/test_athena.py new file mode 100644 index 000000000..ad6ef908f --- /dev/null +++ b/tests/test_athena/test_athena.py @@ -0,0 +1,59 @@ +from __future__ import unicode_literals + +import datetime + +from botocore.exceptions import ClientError +import boto3 +import sure # noqa + +from moto import mock_athena + + +@mock_athena +def test_create_work_group(): + client = boto3.client('athena', region_name='us-east-1') + + response = client.create_work_group( + Name='athena_workgroup', + Description='Test work group', + Configuration={ + 'ResultConfiguration': { + 'OutputLocation': 's3://bucket-name/prefix/', + 'EncryptionConfiguration': { + 'EncryptionOption': 'SSE_KMS', + 'KmsKey': 'aws:arn:kms:1233456789:us-east-1:key/number-1', + }, + }, + }, + Tags=[], + ) + + try: + # The second time should throw an error + response = client.create_work_group( + Name='athena_workgroup', + Description='duplicate', + Configuration={ + 'ResultConfiguration': { + 'OutputLocation': 's3://bucket-name/prefix/', + 'EncryptionConfiguration': { + 'EncryptionOption': 'SSE_KMS', + 'KmsKey': 'aws:arn:kms:1233456789:us-east-1:key/number-1', + }, + }, + }, + ) + except ClientError as err: + err.response['Error']['Code'].should.equal('InvalidRequestException') + err.response['Error']['Message'].should.equal('WorkGroup already exists') + else: + raise RuntimeError('Should have raised ResourceNotFoundException') + + # Then test the work group appears in the work group list + response = client.list_work_groups() + + response['WorkGroups'].should.have.length_of(1) + work_group = response['WorkGroups'][0] + work_group['Name'].should.equal('athena_workgroup') + work_group['Description'].should.equal('Test work group') + work_group['State'].should.equal('ENABLED')