[AIRFLOW-1571] Add AWS Lambda Hook

airflow/contrib/hooks/aws_lambda_hook.py

@@ -0,0 +1,62 @@
+# -*- coding: utf-8 -*-
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from airflow.contrib.hooks.aws_hook import AwsHook
+
+
+class AwsLambdaHook(AwsHook):
+    """
+    Interact with AWS Lambda
+
+    :param function_name: AWS Lambda Function Name
+    :type function_name: str
+    :param region_name: AWS Region Name (example: us-west-2)
+    :type region_name: str
+    :param log_type: Tail Invocation Request
+    :type log_type: str
+    :param qualifier: AWS Lambda Function Version or Alias Name
+    :type qualifier: str
+    :param invocation_type: AWS Lambda Invocation Type (RequestResponse, Event etc)
+    :type invocation_type: str
+    """
+
+    def __init__(self, function_name, region_name=None, log_type='None', qualifier='$LATEST',
+                 invocation_type='RequestResponse', *args, **kwargs):
+        self.function_name = function_name
+        self.region_name = region_name
+        self.log_type = log_type
+        self.invocation_type = invocation_type
+        self.qualifier = qualifier
+        super(AwsLambdaHook, self).__init__(*args, **kwargs)
+
+    def get_conn(self):
+        self.conn = self.get_client_type('lambda', self.region_name)
+        return self.conn
+
+    def invoke_lambda(self, payload):
+        """
+        Invoke Lambda Function
+        """
+
+        awslambda_conn = self.get_conn()
+
+        response = awslambda_conn.invoke(
+            FunctionName=self.function_name,
+            InvocationType=self.invocation_type,
+            LogType=self.log_type,
+            Payload=payload,
+            Qualifier=self.qualifier
+        )
+
+        return response

tests/contrib/hooks/test_aws_lambda_hook.py

@@ -0,0 +1,82 @@
+# -*- coding: utf-8 -*-
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import unittest
+import io
+import json
+import textwrap
+import zipfile
+import base64
+
+from airflow.contrib.hooks.aws_lambda_hook import AwsLambdaHook
+
+try:
+    from moto import mock_lambda
+except ImportError:
+    mock_lambda = None
+
+
+class TestAwsLambdaHook(unittest.TestCase):
+
+    @unittest.skipIf(mock_lambda is None, 'mock_lambda package not present')
+    @mock_lambda
+    def test_get_conn_returns_a_boto3_connection(self):
+        hook = AwsLambdaHook(aws_conn_id='aws_default',
+                             function_name="test_function", region_name="us-east-1")
+        self.assertIsNotNone(hook.get_conn())
+
+    def lambda_function(self):
+        code = textwrap.dedent("""
+def lambda_handler(event, context):
+    return event
+        """)
+        zip_output = io.BytesIO()
+        zip_file = zipfile.ZipFile(zip_output, 'w', zipfile.ZIP_DEFLATED)
+        zip_file.writestr('lambda_function.zip', code)
+        zip_file.close()
+        zip_output.seek(0)
+        return zip_output.read()
+
+    @unittest.skipIf(mock_lambda is None, 'mock_lambda package not present')
+    @mock_lambda
+    def test_invoke_lambda_function(self):
+
+        hook = AwsLambdaHook(aws_conn_id='aws_default',
+                             function_name="test_function", region_name="us-east-1")
+
+        hook.get_conn().create_function(
+            FunctionName='test_function',
+            Runtime='python2.7',
+            Role='test-iam-role',
+            Handler='lambda_function.lambda_handler',
+            Code={
+                'ZipFile': self.lambda_function(),
+            },
+            Description='test lambda function',
+            Timeout=3,
+            MemorySize=128,
+            Publish=True,
+        )
+
+        input = {'hello': 'airflow'}
+        response = hook.invoke_lambda(payload=json.dumps(input))
+
+        self.assertEquals(response["StatusCode"], 202)
+        self.assertEquals(base64.b64decode(response["LogResult"]).decode(
+            'utf-8'), json.dumps(input))
+
+
+if __name__ == '__main__':
+    unittest.main()