Support AWS DynamoDB Stream ARN lookup (#201)

crunchyroll · Sep 15, 2020 · facf1ca · facf1ca
1 parent fad3596
commit facf1ca
Show file tree

Hide file tree

Showing 3 changed files with 53 additions and 9 deletions.
diff --git a/efopen/ef_aws_resolver.py b/efopen/ef_aws_resolver.py
@@ -878,6 +878,23 @@ def ecr_repository_uri(self, lookup, default=None):
     except ClientError:
       return default
 
+  def dynamodb_stream_arn(self, lookup, default=None):
+    """
+    Args:
+      lookup: the name of the DynamoDB table to look up
+      default: the optional value to return if lookup failed; returns None if not set
+    Returns:
+      The DynamoDB Stream ARN with a label matching 'lookup' or default/None if no match found
+    """
+    try:
+      dynamodb_table = EFAwsResolver.__CLIENTS["dynamodb"].describe_table(TableName=lookup)
+      if dynamodb_table:
+        return dynamodb_table["Table"]["LatestStreamArn"]
+      else:
+        return default
+    except ClientError:
+      return default
+
   def lookup(self, token):
     try:
       kv = token.split(",")
@@ -899,6 +916,8 @@ def lookup(self, token):
       return self.cognito_idp_user_pool_arn(*kv[1:])
     elif kv[0] == "cognito-idp:user-pool-id":
       return self.cognito_idp_user_pool_id(*kv[1:])
+    elif kv[0] == "dynamodb:stream-arn":
+      return self.dynamodb_stream_arn(*kv[1:])
     elif kv[0] == "ec2:elasticip/elasticip-id":
       return self.ec2_elasticip_elasticip_id(*kv[1:])
     elif kv[0] == "ec2:elasticip/elasticip-ipaddress":

diff --git a/efopen/ef_template_resolver.py b/efopen/ef_template_resolver.py
@@ -234,6 +234,7 @@ def __init__(self,
       "cloudfront",
       "cognito-identity",
       "cognito-idp",
+      "dynamodb",
       "ec2",
       "ecr",
       "elbv2",

diff --git a/tests/unit_tests/test_ef_aws_resolver.py b/tests/unit_tests/test_ef_aws_resolver.py
@@ -45,28 +45,30 @@ def setUp(self):
     mock_cloud_front_client = Mock(name="Mock CloudFront Client")
     mock_cognito_identity_client = Mock(name="Mock Cognito Identity Client")
     mock_cognito_idp_client = Mock(name="Mock Cognito IDP Client")
+    mock_dynamodb_client = Mock(name="Mock DynamoDB Client")
     mock_ec2_client = Mock(name="Mock EC2 Client")
     mock_ecr_client = Mock(name="Mock ECR Client")
-    mock_route_53_client = Mock(name="Mock Route 53 Client")
-    mock_waf_client = Mock(name="Mock WAF Client")
-    mock_session = Mock(name="Mock Client")
-    mock_kms_client = Mock(name="Mock KMS Client")
     mock_elbv2_client = Mock(name="Mock ELBV2 Client")
+    mock_kms_client = Mock(name="Mock KMS Client")
     mock_ram_client = Mock(name="Mock RAM Client")
+    mock_route_53_client = Mock(name="Mock Route 53 Client")
+    mock_session = Mock(name="Mock Client")
+    mock_waf_client = Mock(name="Mock WAF Client")
 
     self._clients = {
+      "SESSION": mock_session,
       "cloudformation": mock_cloud_formation_client,
       "cloudfront": mock_cloud_front_client,
       "cognito-identity": mock_cognito_identity_client,
       "cognito-idp": mock_cognito_idp_client,
+      "dynamodb": mock_dynamodb_client,
       "ec2": mock_ec2_client,
       "ecr": mock_ecr_client,
-      "route53": mock_route_53_client,
-      "waf": mock_waf_client,
-      "SESSION": mock_session,
-      "kms": mock_kms_client,
       "elbv2": mock_elbv2_client,
-      "ram": mock_ram_client
+      "kms": mock_kms_client,
+      "ram": mock_ram_client,
+      "route53": mock_route_53_client,
+      "waf": mock_waf_client
     }
 
   def tearDown(self):
@@ -2926,3 +2928,25 @@ def test_ram_resource_arn_no_match(self):
     ef_aws_resolver = EFAwsResolver(self._clients)
     result = ef_aws_resolver.lookup("ram:resource-share/resource-arn,cant_possibly_match")
     self.assertIsNone(result)
+
+  def test_dynamodb_stream_arn(self):
+    # Mock the return values involved with this lookup
+    resource_response = {
+      "Table": {
+        "LatestStreamArn": "arn:aws:dynamodb:us-west-2:490645551402:table/alpha0-dynamodb-dummy-table/stream/2020-09-14T23:42:19.796"
+      }
+    }
+    self._clients["dynamodb"].describe_table.return_value = resource_response
+
+    ef_aws_resolver = EFAwsResolver(self._clients)
+    result = ef_aws_resolver.lookup("dynamodb:stream-arn,alpha0-dynamodb-dummy-table")
+    self.assertEqual("arn:aws:dynamodb:us-west-2:490645551402:table/alpha0-dynamodb-dummy-table/stream/2020-09-14T23:42:19.796", result)
+
+  def test_dynamodb_stream_arn_no_match(self):
+    # Mock the return values involved with this lookup
+    resource_response = None
+    self._clients["dynamodb"].describe_table.return_value = resource_response
+
+    ef_aws_resolver = EFAwsResolver(self._clients)
+    result = ef_aws_resolver.lookup("dynamodb:stream-arn,alpha0-dynamodb-dummy-table")
+    self.assertIsNone(result)