Only allow a single create-annotation per script

examples/customer_table.py

@@ -3,7 +3,6 @@
 import boto3
 from migrator.dynamodb_migrator import Migrator
 from time import sleep
-from uuid import uuid4
 
 
 dynamodb = boto3.client('dynamodb')

src/migrator/dynamodb_migrator.py

@@ -3,6 +3,7 @@
 import os
 from functools import wraps
 from time import sleep
+from migrator.exceptions.MigratorScriptException import MigratorScriptException
 
 
 class Migrator():
@@ -19,6 +20,7 @@ def __init__(self, identifier = None):
         self._function_list = []
         self._current_identifier = identifier if identifier else os.path.basename(__file__)
         self._get_or_create_metadata_table()
+        self._table_created = False
 
     def _get_or_create_metadata_table(self):
         try:
@@ -53,10 +55,17 @@ def wrapper(*args, **kwargs):
         return inner_function
 
     def create(self, **kwargs):
+        if self._table_created:
+            self._logger.error("Unable to execute script")
+            self._logger.error("Ensure that you have only one create-annotation per script")
+            self._logger.error("Each table should have it's own script")
+            raise MigratorScriptException("Unable to create multiple tables per script")
+
         def inner_function(function):
             self._function_list.append({'identifier': self._current_identifier,
                                         'table_properties': kwargs,
                                         'func': function})
+            self._table_created = True
         return inner_function
 
     def migrate(self):

src/migrator/exceptions/MigratorException.py

@@ -0,0 +1,2 @@
+class MigratorException(Exception):
+    pass

src/migrator/exceptions/MigratorScriptException.py

@@ -0,0 +1,5 @@
+from migrator.exceptions.MigratorException import MigratorException
+
+
+class MigratorScriptException(MigratorException):
+    pass

tests/migration_scripts/multiple_create_table.py

@@ -0,0 +1,47 @@
+#!/usr/bin/python
+
+from migrator.dynamodb_migrator import Migrator
+from uuid import uuid4
+
+
+migrator = Migrator()
+
+
+@migrator.version(1)
+@migrator.create(
+    AttributeDefinitions=[{
+        'AttributeName': 'hash_key',
+        'AttributeType': 'N'
+    }],
+    TableName=str(uuid4()),
+    KeySchema=[{
+        'AttributeName': 'hash_key',
+        'KeyType': 'HASH'
+    }],
+    BillingMode='PAY_PER_REQUEST')
+def create_table(created_table):
+    # There are two create-statements in this file
+    # There should always only be one
+    # The table should never be created
+    # And we should never get here
+    assert False
+
+
+@migrator.version(1)
+@migrator.create(
+    AttributeDefinitions=[{
+        'AttributeName': 'hash_key',
+        'AttributeType': 'N'
+    }],
+    TableName=str(uuid4()),
+    KeySchema=[{
+        'AttributeName': 'hash_key',
+        'KeyType': 'HASH'
+    }],
+    BillingMode='PAY_PER_REQUEST')
+def create_another_table(created_table):
+    # There are two create-statements in this file
+    # There should always only be one
+    # The table should never be created
+    # And we should never get here
+    assert False

tests/test_table_creation.py

@@ -1,6 +1,7 @@
 from importlib import reload
 from time import sleep
 from .mock_wrapper import mock_aws
+from migrator.exceptions.MigratorScriptException import MigratorScriptException
 
 
 @mock_aws
@@ -64,6 +65,16 @@ def test_create_table_script__assert_metadata_table_is_created(dynamodb):
     delete_metadata_table(dynamodb)
 
 
+@mock_aws
+def test_create_table_script__assert_error_when_there_are_multiple_create_statements(dynamodb):
+    try:
+        from .migration_scripts import multiple_create_table
+        reload(multiple_create_table)  # Ensure that this script isnt already loaded
+        assert True, "Script execution should fail, as only a single create-annotation per script is allowed"
+    except MigratorScriptException:
+        delete_metadata_table(dynamodb)
+
+
 def delete_metadata_table(dynamodb):
     try:
         dynamodb.delete_table(TableName='dynamodb_migrator_metadata')