Athena - implementing first two endpoints (#2506)

This implements create_work_group() and list_work_groups()

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)

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',
+        })

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)

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()
+        })

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,
+}

moto/athena/utils.py

@@ -0,0 +1 @@
+from __future__ import unicode_literals