Merge pull request #12 from nrccua/ARCH-526-add-ability-to-pass-more-…

…regions-to-aws-services

Add option for adding more regions.

HISTORY.rst

@@ -3,6 +3,12 @@ History
 =======
 
 
+v0.9.8 (2021-02-01)
+-----------------------
+
+* Add ability to add more regions besides us-east-1 & us-east-2.
+
+
 v0.9.7 (2021-01-06)
 -----------------------
 

aioradio/aws/dynamodb.py

@@ -6,10 +6,20 @@
 
 from aioradio.aws.utils import AwsServiceManager
 
-AWS_SERVICE = AwsServiceManager(service='dynamodb', regions=['us-east-1', 'us-east-2'], module='aioboto3')
+AWS_SERVICE = AwsServiceManager(service='dynamodb', regions=['us-east-1'], module='aioboto3')
 DYNAMO = AWS_SERVICE.service_dict
 
 
+async def add_regions(regions: List[str]):
+    """Add regions to DynamoDB AWS service.
+
+    Args:
+        regions (List[str]): List of AWS regions
+    """
+
+    AWS_SERVICE.add_regions(regions)
+
+
 @AWS_SERVICE.active
 async def create_dynamo_table(
         table_name: str,

aioradio/aws/secrets.py

@@ -1,13 +1,24 @@
 """Generic async AWS functions for Secrets Manager."""
 
 import base64
+from typing import List
 
 from aioradio.aws.utils import AwsServiceManager
 
-AWS_SERVICE = AwsServiceManager(service='secretsmanager', regions=['us-east-1', 'us-east-2'])
+AWS_SERVICE = AwsServiceManager(service='secretsmanager', regions=['us-east-1'])
 SECRETS = AWS_SERVICE.service_dict
 
 
+async def add_regions(regions: List[str]):
+    """Add regions to Secret Manager AWS service.
+
+    Args:
+        regions (List[str]): List of AWS regions
+    """
+
+    AWS_SERVICE.add_regions(regions)
+
+
 @AWS_SERVICE.active
 async def get_secret(secret_name: str, region: str) -> str:
     """Get secret from AWS Secrets Manager.

aioradio/aws/sqs.py

@@ -11,10 +11,20 @@
 from aioradio.aws.utils import AwsServiceManager
 
 LOG = logging.getLogger(__name__)
-AWS_SERVICE = AwsServiceManager(service='sqs', regions=['us-east-1', 'us-east-2'])
+AWS_SERVICE = AwsServiceManager(service='sqs', regions=['us-east-1'])
 SQS = AWS_SERVICE.service_dict
 
 
+async def add_regions(regions: List[str]):
+    """Add regions to SQS AWS service.
+
+    Args:
+        regions (List[str]): List of AWS regions
+    """
+
+    AWS_SERVICE.add_regions(regions)
+
+
 @AWS_SERVICE.active
 async def create_queue(queue: str, region: str, attributes: Dict[str, str]) -> Dict[str, Any]:
     """Create SQS queue in region defined.

aioradio/aws/utils.py

@@ -54,6 +54,23 @@ def __post_init__(self):
         except RuntimeError:
             loop = asyncio.new_event_loop().run_until_complete(self.create_scheduler())
 
+    def add_regions(self, regions: List[str]):
+        """Add regions to AWS service.
+
+        Args:
+            regions (List[str]): [description]
+        """
+
+        data = {'client': {'obj': None, 'session': None, 'busy': True}, 'active': 0}
+        if self.service == 'dynamodb':
+            data['resource'] = {'obj': None, 'session': None, 'busy': True}
+
+        for region in regions:
+            if region not in self.regions:
+                self.regions.append(region)
+                self.service_dict[region] = deepcopy(data)
+                asyncio.get_event_loop().create_task(self.spawn_aws_service_process(region))
+
     def __del__(self):
         self.scheduler.close()
 
@@ -63,9 +80,21 @@ async def create_scheduler(self):
         self.scheduler = await aiojobs.create_scheduler()
         if self.regions:
             for region in self.regions:
-                await self.scheduler.spawn(self.aio_server(item='client', region=region))
-                if self.module == 'aioboto3':
-                    await self.scheduler.spawn(self.aio_server(item='resource', region=region))
+                await self.spawn_aws_service_process(region)
+        else:
+            await self.spawn_aws_service_process()
+
+    async def spawn_aws_service_process(self, region: str=''):
+        """Spawn scheduler process on a per region basis.
+
+        Args:
+            region (str, optional): AWS region. Defaults to ''.
+        """
+
+        if region:
+            await self.scheduler.spawn(self.aio_server(item='client', region=region))
+            if self.module == 'aioboto3':
+                await self.scheduler.spawn(self.aio_server(item='resource', region=region))
         else:
             await self.scheduler.spawn(self.aio_server(item='client'))
             if self.module == 'aioboto3':

aioradio/tests/aws_secrets_test.py

@@ -4,11 +4,17 @@
 
 import pytest
 
-from aioradio.aws.secrets import get_secret
+from aioradio.aws.secrets import add_regions, get_secret
 
 pytestmark = pytest.mark.asyncio
 
 
+async def test_add_regions():
+    """Add us-east-2 region."""
+
+    await add_regions(['us-east-2'])
+
+
 async def test_secrets_get_secret(create_secret):
     """Test getting secret from Secrets Manager."""
 

aioradio/tests/dynamodb_test.py

@@ -7,7 +7,7 @@
 import pytest
 from boto3.dynamodb.conditions import Attr, Key
 
-from aioradio.aws.dynamodb import (batch_get_items_from_dynamo,
+from aioradio.aws.dynamodb import (add_regions, batch_get_items_from_dynamo,
                                    batch_write_to_dynamo,
                                    get_list_of_dynamo_tables,
                                    put_item_in_dynamo, query_dynamo,
@@ -22,6 +22,12 @@
 pytestmark = pytest.mark.asyncio
 
 
+async def test_add_regions():
+    """Add us-east-2 region."""
+
+    await add_regions(['us-east-2'])
+
+
 async def test_dynamodb_create_table(create_table):
     """Test creating a DynamoDB table."""
 

aioradio/tests/sqs_test.py

@@ -8,8 +8,8 @@
 import orjson
 import pytest
 
-from aioradio.aws.sqs import (delete_messages, get_messages, purge_messages,
-                              send_messages)
+from aioradio.aws.sqs import (add_regions, delete_messages, get_messages,
+                              purge_messages, send_messages)
 
 QUEUE = 'pytest'
 REGION = 'us-east-2'
@@ -18,6 +18,12 @@
 pytestmark = pytest.mark.asyncio
 
 
+async def test_add_regions():
+    """Add us-east-2 region."""
+
+    await add_regions(['us-east-2'])
+
+
 async def test_sqs_creating_queue(sqs_queue_url):
     """Create mock SQS queue."""
 

setup.py

@@ -7,7 +7,7 @@
     long_description = fileobj.read()
 
 setup(name='aioradio',
-    version='0.9.7',
+    version='0.9.8',
     description='Generic asynchronous i/o python utilities for AWS services (SQS, S3, DynamoDB, Secrets Manager), Redis, MSSQL (pyodbc), JIRA and more',
     long_description=long_description,
     long_description_content_type="text/markdown",