diff --git a/airflow/providers/amazon/aws/hooks/lambda_function.py b/airflow/providers/amazon/aws/hooks/lambda_function.py index 771a01aeb7fd1..dea2a4272b2f1 100644 --- a/airflow/providers/amazon/aws/hooks/lambda_function.py +++ b/airflow/providers/amazon/aws/hooks/lambda_function.py @@ -106,6 +106,9 @@ def create_lambda( image_config: Any | None = None, code_signing_config_arn: str | None = None, architectures: list[str] | None = None, + ephemeral_storage: Any | None = None, + snap_start: Any | None = None, + logging_config: Any | None = None, ) -> dict: """ Create a Lambda function. @@ -151,6 +154,10 @@ def create_lambda( A code-signing configuration includes a set of signing profiles, which define the trusted publishers for this function. :param architectures: The instruction set architecture that the function supports. + :param ephemeral_storage: The size of the function's /tmp directory in MB. + The default value is 512, but can be any whole number between 512 and 10,240 MB + :param snap_start: The function's SnapStart setting + :param logging_config: The function's Amazon CloudWatch Logs configuration settings """ if package_type == "Zip": if handler is None: @@ -181,6 +188,9 @@ def create_lambda( "ImageConfig": image_config, "CodeSigningConfigArn": code_signing_config_arn, "Architectures": architectures, + "EphemeralStorage": ephemeral_storage, + "SnapStart": snap_start, + "LoggingConfig": logging_config, } return self.conn.create_function(**trim_none_values(create_function_args)) diff --git a/tests/providers/amazon/aws/operators/test_lambda_function.py b/tests/providers/amazon/aws/operators/test_lambda_function.py index 723f7584d9937..9c977ee883ed7 100644 --- a/tests/providers/amazon/aws/operators/test_lambda_function.py +++ b/tests/providers/amazon/aws/operators/test_lambda_function.py @@ -125,6 +125,41 @@ def test_create_lambda_deferrable(self, _): with pytest.raises(TaskDeferred): operator.execute(None) + @mock.patch.object(LambdaHook, "create_lambda") + @mock.patch.object(LambdaHook, "conn") + @pytest.mark.parametrize( + "config", + [ + pytest.param( + { + "architectures": ["arm64"], + "logging_config": {"LogFormat": "Text", "LogGroup": "/custom/log-group/"}, + "snap_start": {"ApplyOn": "PublishedVersions"}, + "ephemeral_storage": {"Size": 1024}, + }, + id="with-config-argument", + ), + ], + ) + def test_create_lambda_using_config_argument(self, mock_hook_conn, mock_hook_create_lambda, config): + operator = LambdaCreateFunctionOperator( + task_id="task_test", + function_name=FUNCTION_NAME, + role=ROLE_ARN, + code={ + "ImageUri": IMAGE_URI, + }, + config=config, + ) + operator.execute(None) + + mock_hook_create_lambda.assert_called_once() + mock_hook_conn.get_waiter.assert_not_called() + assert operator.config.get("logging_config") == config.get("logging_config") + assert operator.config.get("architectures") == config.get("architectures") + assert operator.config.get("snap_start") == config.get("snap_start") + assert operator.config.get("ephemeral_storage") == config.get("ephemeral_storage") + class TestLambdaInvokeFunctionOperator: @pytest.mark.parametrize("payload", PAYLOADS)