O hai, Orchestra!

Author: marcua

.gitignore

@@ -0,0 +1,24 @@
+# Byte-compiled / optimized / DLL files
+*.py[cod]
+*~
+__pycache__/
+
+# Sqlite
+*.sqlite3
+
+# Mac related
+.DS_Store
+
+# Emacs temp files
+[#]*[#]
+.\#*
+
+# Coverage data
+coverage_annotations/
+.coverage
+
+# Database backups data
+backup.sql
+
+# Compressed staticfiles
+/orchestra/staticfiles

LICENSE

@@ -186,7 +186,7 @@
       same "printed page" as the copyright notice for easier
       identification within third-party archives.
 
-   Copyright {yyyy} {name of copyright owner}
+   Copyright 2015 Unlimited Labs, Inc.
 
    Licensed under the Apache License, Version 2.0 (the "License");
    you may not use this file except in compliance with the License.

NOTICE

@@ -0,0 +1,6 @@
+Orchestra
+Copyright 2015 Unlimited Labs, Inc.
+
+We based the Orchestra html/css/js theme off of Dashgum (http://www.blacktie.co/demo/dashgumfree/)
+Creative commons license with attribution (http://creativecommons.org/licenses/by/3.0/)
+Source: http://www.blacktie.co/about-themes/

beanstalk_dispatch/README.md

@@ -0,0 +1,53 @@
+The Beanstalk dispatcher is a Django app that runs functions that have
+been scheduled to run an AWS SQS queue and executes them on Elastic
+Beanstalk Worker machines that are listening to that queue.
+
+## Configuration
+
+To install:
+
+  * create an Elastic Beanstalk environment for an application
+that has the following two parameters in `settings.py`:
+
+     ```python
+     BEANSTALK_DISPATCH_SQS_KEY = 'your AWS key for accessing SQS'
+     BEANSTALK_DISPATCH_SQS_SECRET = 'your AWS secret for accessing SQS'
+     ```
+
+  * add `beanstalk_dispatch` to settings.py's `INSTALLED_APPS`
+
+  * Add `url(r'^beanstalk_dispatch/',
+    include('beanstalk_dispatch.urls')),` to your main `urls.py`
+
+  * Add `/beanstalk_dispatch/dispatcher` as the HTTP endpoint or your
+    beanstalk worker configuration in the AWS console.
+
+  * Add a dispatch table.  The dispatcher works by creating an HTTP
+endpoint that a local SQS/Beanstalk daemon POSTs requests to.  That
+endpoint consults a `BEANSTALK_DISPATCH_TABLE`, which maps function
+names onto functions to run.  Here's an example:
+
+      ```python
+      if os.environ.get('BEANSTALK_WORKER') == 'True':
+        BEANSTALK_DISPATCH_TABLE = {
+            'a_function_to_dispatch': ('some_package.beanstalk_tasks', 'the_name_of_the_function_in_the_module')}
+      ```
+
+   The first line is a check we have that ensures this type of machine
+should be a beanstalk worker.  We set a BEANSTALK_WORKER environment
+variable to True in the environment's configuration only on our worker
+machines.  This avoids other environments (e.g., our web servers) from
+serving as open proxies for running arbitrary code.
+
+   The second line is the dispatch table.  It maps function name (e.g.,
+`a_function_to_dispatch`) to a module (e.g.,
+`some_package.beanstalk_tasks` and function in that module (e.g.,
+`the_name_of_the_function_in_the_module`).
+
+
+## Scheduling a function to run
+
+The `beanstalk_dispatch.client.schedule_function` schedules a function
+to run on a given SQS queue.  The function name you pass it must be a
+key in the `BEANSTALK_DISPATCH_TABLE`, and the queue name you pass it
+must be a queue for which a beanstalk worker is configured.

beanstalk_dispatch/__init__.py

@@ -0,0 +1,6 @@
+class BeanstalkDispatchError(Exception):
+    pass
+
+ARGS = 'args'
+FUNCTION = 'function'
+KWARGS = 'kwargs'

beanstalk_dispatch/admin.py

@@ -0,0 +1 @@
+# Register your models here.

beanstalk_dispatch/client.py

@@ -0,0 +1,22 @@
+import boto.sqs
+
+from beanstalk_dispatch.common import create_request_body
+from django.conf import settings
+
+
+def schedule_function(queue_name, function_name, *args, **kwargs):
+    """
+    Schedule a function named `function_name` to be run by workers on
+    the queue `queue_name` with *args and **kwargs as specified by that
+    function.
+    """
+    body = create_request_body(function_name, *args, **kwargs)
+    connection = boto.connect_sqs(
+        settings.BEANSTALK_DISPATCH_SQS_KEY,
+        settings.BEANSTALK_DISPATCH_SQS_SECRET)
+    queue = connection.get_queue(queue_name)
+    if not queue:
+        queue = connection.create_queue(queue_name)
+    message = boto.sqs.message.Message()
+    message.set_body(body)
+    queue.write(message)

beanstalk_dispatch/common.py

@@ -0,0 +1,9 @@
+import json
+
+from beanstalk_dispatch import ARGS
+from beanstalk_dispatch import FUNCTION
+from beanstalk_dispatch import KWARGS
+
+
+def create_request_body(function_name, *args, **kwargs):
+    return json.dumps({FUNCTION: function_name, ARGS: args, KWARGS: kwargs})

beanstalk_dispatch/execution.py

@@ -0,0 +1,35 @@
+from beanstalk_dispatch import ARGS
+from beanstalk_dispatch import FUNCTION
+from beanstalk_dispatch import KWARGS
+from beanstalk_dispatch import BeanstalkDispatchError
+from django.conf import settings
+from importlib import import_module
+
+
+def execute_function(function_request):
+    """
+    Given a request created by
+    `beanstalk_dispatch.common.create_request_body`, executes the
+    request.  This function is to be run on a beanstalk worker.
+    """
+    dispatch_table = getattr(settings, 'BEANSTALK_DISPATCH_TABLE', None)
+
+    if dispatch_table is None:
+        raise BeanstalkDispatchError('No beanstalk dispatch table configured')
+    for key in (FUNCTION, ARGS, KWARGS):
+        if key not in function_request.keys():
+            raise BeanstalkDispatchError(
+                'Please provide a {} argument'.format(key))
+
+    module_name, function_name = (
+        dispatch_table.get(function_request[FUNCTION], (None, None)))
+    if module_name and function_name:
+        # TODO(marcua): Catch import errors and rethrow them as
+        # BeanstalkDispatchErrors.
+        module = import_module(module_name)
+        function = getattr(module, function_name)
+        function(*function_request[ARGS], **function_request[KWARGS])
+    else:
+        raise BeanstalkDispatchError(
+            'Requested function not found: {}'.format(
+                function_request[FUNCTION]))

beanstalk_dispatch/migrations/__init__.py

@@ -0,0 +1 @@
+# Create your models here.

beanstalk_dispatch/models.py

@@ -0,0 +1,41 @@
+import boto
+import json
+
+from beanstalk_dispatch import ARGS
+from beanstalk_dispatch import FUNCTION
+from beanstalk_dispatch import KWARGS
+from beanstalk_dispatch.client import schedule_function
+from django.test import TestCase
+from django.test import override_settings
+from moto import mock_sqs
+
+
+@mock_sqs
+@override_settings(
+    BEANSTALK_DISPATCH_SQS_KEY='', BEANSTALK_DISPATCH_SQS_SECRET='')
+class ClientTestCase(TestCase):
+
+    def setUp(self):
+        self.queue_name = 'testing-queue'
+
+    def test_function_scheduling(self):
+        # Check the message on the queue.
+        sqs_connection = boto.connect_sqs('', '')
+        sqs_connection.create_queue(self.queue_name)
+        queue = sqs_connection.get_queue(self.queue_name)
+        messages = queue.get_messages()
+        self.assertEquals(len(messages), 0)
+
+        # Schedule a function.
+        schedule_function(
+            self.queue_name, 'a-function', '1', '2', kwarg1=1, kwarg2=2)
+
+        messages = queue.get_messages()
+        self.assertEquals(len(messages), 1)
+
+        # For some reason, boto base64-encodes the messages, but moto does
+        # not.  Life.
+        self.assertEquals(
+            json.loads(messages[0].get_body()),
+            {FUNCTION: 'a-function', ARGS: ['1', '2'], KWARGS: {
+                'kwarg1': 1, 'kwarg2': 2}})

beanstalk_dispatch/tests/__init__.py

@@ -0,0 +1,114 @@
+import json
+
+from base64 import b64encode
+from beanstalk_dispatch import ARGS
+from beanstalk_dispatch import FUNCTION
+from beanstalk_dispatch import KWARGS
+from beanstalk_dispatch.common import create_request_body
+from django.core.urlresolvers import reverse
+from django.test import Client
+from django.test import TestCase
+from django.test import override_settings
+
+# Don't log logger errors.
+import logging
+logging.disable(logging.CRITICAL)
+
+
+CALL_COUNTER = 0
+
+
+def counter_incrementer(first_arg, second_arg=None):
+    global CALL_COUNTER
+    CALL_COUNTER += first_arg
+    if second_arg:
+        CALL_COUNTER += second_arg
+
+
+DISPATCH_SETTINGS = {
+    'BEANSTALK_DISPATCH_TABLE': {
+        'the_counter': (
+            'beanstalk_dispatch.tests.test_dispatcher',
+            'counter_incrementer')}}
+
+
+class DispatcherTestCase(TestCase):
+
+    """ Test the server-side function dispatcher.
+
+    In these tests, we base64-encode every message we send to the server
+    because this is what boto does.
+    """
+
+    def setUp(self):
+        global CALL_COUNTER
+        CALL_COUNTER = 0
+        self.client = Client()
+        self.url = reverse('beanstalk_dispatcher')
+
+    @override_settings(**DISPATCH_SETTINGS)
+    def test_no_get(self):
+        response = self.client.get(self.url)
+        self.assertEquals(response.status_code, 405)
+
+    @override_settings(BEANSTALK_DISPATCH_TABLE=None)
+    def test_no_dispatch(self):
+        response = self.client.post(
+            self.url, b64encode(
+                create_request_body('some_func').encode('ascii')),
+            content_type='application/json')
+        self.assertEquals(response.status_code, 400)
+        self.assertEquals(json.loads(response.content.decode()),
+                          {'message': 'No beanstalk dispatch table configured',
+                           'error': 400})
+
+    @override_settings(**DISPATCH_SETTINGS)
+    def test_missing_function(self):
+        response = self.client.post(
+            self.url,
+            b64encode(create_request_body('nonexistent_func').encode('ascii')),
+            content_type='application/json')
+        self.assertEquals(response.status_code, 400)
+        self.assertEquals(
+            json.loads(response.content.decode()),
+            {'message': 'Requested function not found: nonexistent_func',
+                'error': 400})
+
+    @override_settings(**DISPATCH_SETTINGS)
+    def test_malformed_request(self):
+        keys = {FUNCTION, ARGS, KWARGS}
+        for missing_key in keys:
+            request_body = {key: 'test' for key in
+                            keys - {missing_key}}
+            response = self.client.post(
+                self.url,
+                b64encode(json.dumps(request_body).encode('ascii')),
+                content_type='application/json')
+            self.assertEquals(response.status_code, 400)
+            self.assertEquals(json.loads(response.content.decode()), {
+                'message': 'Please provide a {} argument'.format(missing_key),
+                'error': 400})
+
+    @override_settings(**DISPATCH_SETTINGS)
+    def test_both_args_kwargs(self):
+        body = b64encode(
+            create_request_body('the_counter', 1, second_arg=5)
+            .encode('ascii'))
+        response = self.client.post(self.url,
+                                    body,
+                                    content_type='application/json')
+        self.assertEquals(response.status_code, 200)
+        self.assertEquals(json.loads(response.content.decode()),
+                          {})
+        self.assertEquals(CALL_COUNTER, 6)
+
+    @override_settings(**DISPATCH_SETTINGS)
+    def test_just_args(self):
+        body = b64encode(create_request_body('the_counter', 2).encode('ascii'))
+        response = self.client.post(self.url,
+                                    body,
+                                    content_type='application/json')
+        self.assertEquals(response.status_code, 200)
+        self.assertEquals(json.loads(response.content.decode()),
+                          {})
+        self.assertEquals(CALL_COUNTER, 2)

beanstalk_dispatch/tests/test_client.py

@@ -0,0 +1,7 @@
+from django.conf.urls import patterns
+from django.conf.urls import url
+from beanstalk_dispatch.views import dispatcher
+
+urlpatterns = patterns(
+    '',
+    url(r'^', dispatcher, name='beanstalk_dispatcher'))

beanstalk_dispatch/tests/test_dispatcher.py

@@ -0,0 +1,24 @@
+import json
+
+from base64 import b64decode
+from beanstalk_dispatch.execution import execute_function
+from jsonview.decorators import json_view
+from jsonview.exceptions import BadRequest
+from django.views.decorators.csrf import csrf_exempt
+from django.views.decorators.http import require_http_methods
+
+import logging
+logger = logging.getLogger(__name__)
+
+
+@require_http_methods(['POST'])
+@json_view
+@csrf_exempt
+def dispatcher(request):
+    function_request = json.loads(b64decode(request.body.decode()).decode())
+    try:
+        execute_function(function_request)
+        return {}
+    except Exception as e:
+        logger.error('Failure running function', exc_info=True)
+        raise BadRequest(e)