Added logging_config,snapstart,ephemeral_storage parameters to aws la…

…mbdacreatefunctionoperator (apache#39300)
RodrigoGanancia · May 10, 2024 · 331a8f7 · 331a8f7
1 parent 7cde0ae
commit 331a8f7
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 0 deletions.
diff --git 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
@@ -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)