Create Lambda create operator and sensor

airflow/providers/amazon/aws/hooks/lambda_function.py

@@ -21,6 +21,7 @@
 from typing import Any
 
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+from airflow.providers.amazon.aws.utils import trim_none_values
 
 
 class LambdaHook(AwsBaseHook):
@@ -70,7 +71,7 @@ def invoke_lambda(
             "Payload": payload,
             "Qualifier": qualifier,
         }
-        return self.conn.invoke(**{k: v for k, v in invoke_args.items() if v is not None})
+        return self.conn.invoke(**trim_none_values(invoke_args))
 
     def create_lambda(
         self,
@@ -127,6 +128,4 @@ def create_lambda(
             "CodeSigningConfigArn": code_signing_config_arn,
             "Architectures": architectures,
         }
-        return self.conn.create_function(
-            **{k: v for k, v in create_function_args.items() if v is not None},
-        )
+        return self.conn.create_function(**trim_none_values(create_function_args))

airflow/providers/amazon/aws/operators/lambda_function.py

@@ -20,35 +20,122 @@
 import json
 from typing import TYPE_CHECKING, Sequence
 
+from airflow.compat.functools import cached_property
 from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.lambda_function import LambdaHook
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
 
 
-class AwsLambdaInvokeFunctionOperator(BaseOperator):
+class LambdaCreateFunctionOperator(BaseOperator):
     """
-    Invokes an AWS Lambda function.
-    You can invoke a function synchronously (and wait for the response),
-    or asynchronously.
-    To invoke a function asynchronously,
-    set `invocation_type` to `Event`. For more details,
-    review the boto3 Lambda invoke docs.
+    Creates an AWS Lambda function.
+
+    More information regarding parameters of this operator can be found here
+    https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/lambda.html#Lambda.Client.create_function
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:LambdaCreateFunctionOperator`
 
     :param function_name: The name of the AWS Lambda function, version, or alias.
-    :param log_type: Set to Tail to include the execution log in the response. Otherwise, set to "None".
-    :param qualifier: Specify a version or alias to invoke a published version of the function.
-    :param invocation_type: One of RequestResponse / Event / DryRun
-    :param client_context: Up to 3,583 bytes of base64-encoded data about the invoking client
-        to pass to the function in the context object.
-    :param payload: The JSON string that you want to provide to your Lambda function as input.
+    :param runtime: The identifier of the function's runtime. Runtime is required if the deployment package
+        is a .zip file archive.
+    :param role: The Amazon Resource Name (ARN) of the function's execution role.
+    :param handler: The name of the method within your code that Lambda calls to run your function.
+        Handler is required if the deployment package is a .zip file archive.
+    :param code: The code for the function.
+    :param description: A description of the function.
+    :param timeout: The amount of time (in seconds) that Lambda allows a function to run before stopping it.
+    :param config: Optional dictionary for arbitrary parameters to the boto API create_lambda call.
+    :param wait_for_completion: If True, the operator will wait until the function is active.
     :param aws_conn_id: The AWS connection ID to use
+    """
+
+    template_fields: Sequence[str] = (
+        "function_name",
+        "runtime",
+        "role",
+        "handler",
+        "code",
+        "config",
+    )
+    ui_color = "#ff7300"
+
+    def __init__(
+        self,
+        *,
+        function_name: str,
+        runtime: str | None = None,
+        role: str,
+        handler: str | None = None,
+        code: dict,
+        description: str | None = None,
+        timeout: int | None = None,
+        config: dict = {},
+        wait_for_completion: bool = False,
+        aws_conn_id: str = "aws_default",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.function_name = function_name
+        self.runtime = runtime
+        self.role = role
+        self.handler = handler
+        self.code = code
+        self.description = description
+        self.timeout = timeout
+        self.config = config
+        self.wait_for_completion = wait_for_completion
+        self.aws_conn_id = aws_conn_id
+
+    @cached_property
+    def hook(self) -> LambdaHook:
+        return LambdaHook(aws_conn_id=self.aws_conn_id)
+
+    def execute(self, context: Context):
+        self.log.info("Creating AWS Lambda function: %s", self.function_name)
+        response = self.hook.create_lambda(
+            function_name=self.function_name,
+            runtime=self.runtime,
+            role=self.role,
+            handler=self.handler,
+            code=self.code,
+            description=self.description,
+            timeout=self.timeout,
+            **self.config,
+        )
+        self.log.info("Lambda response: %r", response)
+
+        if self.wait_for_completion:
+            self.log.info("Wait for Lambda function to be active")
+            waiter = self.hook.conn.get_waiter("function_active_v2")
+            waiter.wait(
+                FunctionName=self.function_name,
+            )
+
+        return response.get("FunctionArn")
+
+
+class AwsLambdaInvokeFunctionOperator(BaseOperator):
+    """
+    Invokes an AWS Lambda function. You can invoke a function synchronously (and wait for the response),
+    or asynchronously.
+    To invoke a function asynchronously, set `invocation_type` to `Event`. For more details,
+    review the boto3 Lambda invoke docs.
 
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
         :ref:`howto/operator:AwsLambdaInvokeFunctionOperator`
 
+    :param function_name: The name of the AWS Lambda function, version, or alias.
+    :param log_type: Set to Tail to include the execution log in the response. Otherwise, set to "None".
+    :param qualifier: Specify a version or alias to invoke a published version of the function.
+    :param invocation_type: AWS Lambda invocation type (RequestResponse, Event, DryRun)
+    :param client_context: Data about the invoking client to pass to the function in the context object
+    :param payload: JSON provided as input to the Lambda function
+    :param aws_conn_id: The AWS connection ID to use
     """
 
     template_fields: Sequence[str] = ("function_name", "payload", "qualifier", "invocation_type")
@@ -75,16 +162,19 @@ def __init__(
         self.client_context = client_context
         self.aws_conn_id = aws_conn_id
 
+    @cached_property
+    def hook(self) -> LambdaHook:
+        return LambdaHook(aws_conn_id=self.aws_conn_id)
+
     def execute(self, context: Context):
         """
         Invokes the target AWS Lambda function from Airflow.
 
         :return: The response payload from the function, or an error object.
         """
-        hook = LambdaHook(aws_conn_id=self.aws_conn_id)
         success_status_codes = [200, 202, 204]
         self.log.info("Invoking AWS Lambda function: %s with payload: %s", self.function_name, self.payload)
-        response = hook.invoke_lambda(
+        response = self.hook.invoke_lambda(
             function_name=self.function_name,
             invocation_type=self.invocation_type,
             log_type=self.log_type,

airflow/providers/amazon/aws/sensors/lambda_function.py

@@ -0,0 +1,86 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you 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 __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Sequence
+
+from airflow.providers.amazon.aws.hooks.lambda_function import LambdaHook
+from airflow.providers.amazon.aws.utils import trim_none_values
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+from airflow.compat.functools import cached_property
+from airflow.exceptions import AirflowException
+from airflow.sensors.base import BaseSensorOperator
+
+
+class LambdaFunctionStateSensor(BaseSensorOperator):
+    """
+    Asks for the state of the Lambda until it reaches a target state.
+    If the query fails, the task will fail.
+
+    .. seealso::
+        For more information on how to use this sensor, take a look at the guide:
+        :ref:`howto/sensor:LambdaFunctionStateSensor`
+
+    :param function_name: The name of the AWS Lambda function, version, or alias.
+    :param qualifier: Specify a version or alias to get details about a published version of the function.
+    :param target_states: The Lambda states desired.
+    :param aws_conn_id: aws connection to use, defaults to 'aws_default'
+    """
+
+    FAILURE_STATES = ("Failed",)
+
+    template_fields: Sequence[str] = (
+        "function_name",
+        "qualifier",
+    )
+
+    def __init__(
+        self,
+        *,
+        function_name: str,
+        qualifier: str | None = None,
+        target_states: list = ["Active"],
+        aws_conn_id: str = "aws_default",
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.aws_conn_id = aws_conn_id
+        self.function_name = function_name
+        self.qualifier = qualifier
+        self.target_states = target_states
+
+    def poke(self, context: Context) -> bool:
+        get_function_args = {
+            "FunctionName": self.function_name,
+            "Qualifier": self.qualifier,
+        }
+        state = self.hook.conn.get_function(**trim_none_values(get_function_args))["Configuration"]["State"]
+
+        if state in self.FAILURE_STATES:
+            raise AirflowException(
+                "Lambda function state sensor failed because the Lambda is in a failed state"
+            )
+
+        return state in self.target_states
+
+    @cached_property
+    def hook(self) -> LambdaHook:
+        return LambdaHook(aws_conn_id=self.aws_conn_id)

airflow/providers/amazon/provider.yaml

@@ -354,6 +354,9 @@ sensors:
       - airflow.providers.amazon.aws.sensors.glue
       - airflow.providers.amazon.aws.sensors.glue_crawler
       - airflow.providers.amazon.aws.sensors.glue_catalog_partition
+  - integration-name: AWS Lambda
+    python-modules:
+      - airflow.providers.amazon.aws.sensors.lambda_function
   - integration-name: Amazon RDS
     python-modules:
       - airflow.providers.amazon.aws.sensors.rds

docs/apache-airflow-providers-amazon/operators/lambda.rst

@@ -33,6 +33,20 @@ Prerequisite Tasks
 Operators
 ---------
 
+.. _howto/operator:LambdaCreateFunctionOperator:
+
+Create an AWS Lambda function
+=============================
+
+To create an AWS lambda function you can use
+:class:`~airflow.providers.amazon.aws.operators.lambda_function.LambdaCreateFunctionOperator`.
+
+.. exampleinclude:: /../../tests/system/providers/amazon/aws/example_lambda.py
+    :language: python
+    :dedent: 4
+    :start-after: [START howto_operator_create_lambda_function]
+    :end-before: [END howto_operator_create_lambda_function]
+
 .. _howto/operator:AwsLambdaInvokeFunctionOperator:
 
 Invoke an AWS Lambda function
@@ -41,12 +55,29 @@ Invoke an AWS Lambda function
 To invoke an AWS lambda function you can use
 :class:`~airflow.providers.amazon.aws.operators.lambda_function.AwsLambdaInvokeFunctionOperator`.
 
+.. exampleinclude:: /../../tests/system/providers/amazon/aws/example_lambda.py
+    :language: python
+    :dedent: 4
+    :start-after: [START howto_operator_invoke_lambda_function]
+    :end-before: [END howto_operator_invoke_lambda_function]
+
+Sensors
+---------
+
+.. _howto/sensor:LambdaFunctionStateSensor:
+
+Wait on an Amazon Lambda function state
+=======================================
+
+To check the state of an Amazon Lambda function until it reaches the target state or another terminal
+state you can use :class:`~airflow.providers.amazon.aws.sensors.lambda_function.LambdaFunctionStateSensor`.
 
 .. exampleinclude:: /../../tests/system/providers/amazon/aws/example_lambda.py
     :language: python
     :dedent: 4
-    :start-after: [START howto_operator_lambda]
-    :end-before: [END howto_operator_lambda]
+    :start-after: [START howto_sensor_lambda_function_state]
+    :end-before: [END howto_sensor_lambda_function_state]
+
 
 Reference
 ---------