Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Added logging_config,snapstart,ephemeral_storage parameters to aws lambdacreatefunctionoperator #39300

Merged
10 changes: 10 additions & 0 deletions airflow/providers/amazon/aws/hooks/lambda_function.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
gopidesupavan marked this conversation as resolved.
Show resolved Hide resolved
snap_start: Any | None = None,
logging_config: Any | None = None,
) -> dict:
"""
Create a Lambda function.
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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))

Expand Down
32 changes: 32 additions & 0 deletions tests/providers/amazon/aws/operators/test_lambda_function.py
Original file line number Diff line number Diff line change
Expand Up @@ -125,6 +125,38 @@ 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/"
}
}, 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")

gopidesupavan marked this conversation as resolved.
Show resolved Hide resolved

class TestLambdaInvokeFunctionOperator:
@pytest.mark.parametrize("payload", PAYLOADS)
Expand Down