Repository: bimpression/sosw

Issue:
Scheduler: allow lambda_name in job to be ARN
If the passed in the job `lambda_name` (aka Labourer.id) is a full ARN trim it using regex to leave just the name.

Discussion:

Patch:
diff --git a/sosw/components/helpers.py b/sosw/components/helpers.py
index 8d90320..76aef77 100644
--- a/sosw/components/helpers.py
+++ b/sosw/components/helpers.py
@@ -25,6 +25,7 @@ __all__ = ['validate_account_to_dashed',
            'validate_list_of_words_from_csv_or_list',
            'first_or_none',
            'recursive_update',
+           'trim_arn_to_name',
            ]
 
 import re
@@ -668,3 +669,22 @@ def recursive_update(d: Dict, u: Mapping) -> Dict:
             new[k] = v
 
     return new
+
+
+def trim_arn_to_name(arn: str) -> str:
+    """
+    Extract just the name of function from full ARN. Supports versions, aliases or raw name (without ARN).
+
+    More information about ARN Format:
+    https://docs.aws.amazon.com/general/latest/gr/aws-arns-and-namespaces.html#genref-arns
+    """
+
+    # Special handling for super global services (e.g. S3 buckets)
+    if arn.count(':') < 6 and '/' not in arn:
+        return arn.split(':')[-1]
+
+    # Seems a little messy, but passes more/less any test of different ARNs we tried.
+    pattern = "(arn:aws:[0-9a-zA-Z-]{2,20}:[0-9a-zA-Z-]{0,12}:[0-9]{12}:[0-9a-zA-Z-]{2,20}[:/])?" \
+              "(?P<name>[0-9a-zA-Z_=,.@-]*)(:)?([0-9a-zA-Z$]*)?"
+
+    return re.search(pattern, arn).group('name')
diff --git a/sosw/scheduler.py b/sosw/scheduler.py
index 98074a2..2756d89 100644
--- a/sosw/scheduler.py
+++ b/sosw/scheduler.py
@@ -22,7 +22,7 @@ from copy import deepcopy
 from typing import List, Set, Tuple, Union, Optional, Dict
 
 from sosw.app import Processor
-from sosw.components.helpers import get_list_of_multiple_or_one_or_empty_from_dict
+from sosw.components.helpers import get_list_of_multiple_or_one_or_empty_from_dict, trim_arn_to_name
 from sosw.labourer import Labourer
 from sosw.managers.task import TaskManager
 
@@ -331,20 +331,11 @@ class Scheduler(Processor):
         job = load(jh['job']) if 'job' in jh else jh
 
         assert 'lambda_name' in job, f"Job is missing required parameter 'lambda_name': {job}"
-        job['lambda_name'] = job['lambda_name']
+        job['lambda_name'] = trim_arn_to_name(job['lambda_name'])
 
         return job
 
 
-    def get_name_from_arn(self, arn):
-        """ Extract just the name of function from full ARN. Supports versions, aliases or raw name (without ARN). """
-
-        pattern = "(arn:aws:lambda:[0-9a-zA-Z-]{6,12}:[0-9]{12}:function:)?" \
-                  "(?P<name>[0-9a-zA-Z_=,.@-]*)(:)?([0-9a-zA-Z$]*)?"
-
-        return re.search(pattern, arn).group('name')
-
-
     def process_file(self):
 
         file_name = self.get_and_lock_queue_file()


Test patch:
diff --git a/sosw/components/test/unit/test_helpers.py b/sosw/components/test/unit/test_helpers.py
index 2949462..4563593 100644
--- a/sosw/components/test/unit/test_helpers.py
+++ b/sosw/components/test/unit/test_helpers.py
@@ -545,5 +545,26 @@ class helpers_UnitTestCase(unittest.TestCase):
         self.assertIsNone(recursive_update(a, b)['b'])
 
 
+    def test_trim_arn_to_name(self):
+
+        TESTS = [
+            ('bar_with_no_arn', 'bar_with_no_arn'),
+            ('arn:aws:lambda:us-west-2:000000000000:function:bar', 'bar'),
+            ('arn:aws:lambda:us-west-2:000000000000:function:bar:', 'bar'),
+            ('arn:aws:lambda:us-west-2:000000000000:function:bar:$LATEST', 'bar'),
+            ('arn:aws:lambda:us-west-2:000000000000:function:bar:12', 'bar'),
+            ('arn:aws:lambda:us-west-2:000000000000:function:bar:12', 'bar'),
+            ('arn:aws:s3:::autotest-sosw', 'autotest-sosw'),
+            ('arn:aws:iam::000000000000:role/aws-code-deploy-role', 'aws-code-deploy-role'),
+            ('arn:aws:rds:us-west-2:000000000000:cluster:aws-cluster-01', 'aws-cluster-01'),
+            ('arn:aws:rds:us-west-2:000000000000:db:aws-01-00', 'aws-01-00'),
+            ('arn:aws:events:us-west-2:123456000000:rule/aws-sr-01', 'aws-sr-01'),
+            ('arn:aws:dynamodb:us-west-2:123456000321:table/sosw_tasks', 'sosw_tasks'),
+        ]
+
+        for test, expected in TESTS:
+            self.assertEqual(trim_arn_to_name(test), expected)
+
+
 if __name__ == '__main__':
     unittest.main()
diff --git a/sosw/test/unit/test_scheduler.py b/sosw/test/unit/test_scheduler.py
index f590877..11c1479 100644
--- a/sosw/test/unit/test_scheduler.py
+++ b/sosw/test/unit/test_scheduler.py
@@ -211,6 +211,8 @@ class Scheduler_UnitTestCase(unittest.TestCase):
         TESTS = [
             ({'job': {'lambda_name': 'foo', 'payload_attr': 'val'}}, {'lambda_name': 'foo', 'payload_attr': 'val'}),
             ({'lambda_name': 'foo', 'payload_attr': 'val'}, {'lambda_name': 'foo', 'payload_attr': 'val'}),
+            ({'lambda_name': 'arn:aws:lambda:us-west-2:000000000000:function:foo', 'payload_attr': 'val'},
+             {'lambda_name': 'foo', 'payload_attr': 'val'}),
             ({'job': {'lambda_name': 'foo', 'payload_attr': 'val'}}, {'lambda_name': 'foo', 'payload_attr': 'val'}),
 
             # JSONs
@@ -239,20 +241,6 @@ class Scheduler_UnitTestCase(unittest.TestCase):
             self.assertRaises(Exception, self.scheduler.extract_job_from_payload, test)
 
 
-    def test_get_name_from_arn(self):
-
-        TESTS = [
-            ('bar_with_no_arn', 'bar_with_no_arn'),
-            ('arn:aws:lambda:us-west-2:000000000000:function:bar', 'bar'),
-            ('arn:aws:lambda:us-west-2:000000000000:function:bar:', 'bar'),
-            ('arn:aws:lambda:us-west-2:000000000000:function:bar:$LATEST', 'bar'),
-            ('arn:aws:lambda:us-west-2:000000000000:function:bar:12', 'bar'),
-        ]
-
-        for test, expected in TESTS:
-            self.assertEqual(self.scheduler.get_name_from_arn(test), expected)
-
-
     def test_needs_chunking__isolate_root(self):
 
         pl = deepcopy(self.PAYLOAD)