apache-airflow-providers-amazon 9.8.0__py3-none-any.whl → 9.9.0rc1__py3-none-any.whl

airflow/providers/amazon/__init__.py

@@ -29,7 +29,7 @@ from airflow import __version__ as airflow_version
 
 __all__ = ["__version__"]
 
-__version__ = "9.8.0"
+__version__ = "9.9.0"
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
     "2.10.0"

airflow/providers/amazon/aws/executors/aws_lambda/__init__.py

@@ -0,0 +1,21 @@
+# 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
+
+__all__ = ["AwsLambdaExecutor"]
+
+from airflow.providers.amazon.aws.executors.aws_lambda.lambda_executor import AwsLambdaExecutor

airflow/providers/amazon/aws/executors/aws_lambda/docker/Dockerfile

@@ -0,0 +1,107 @@
+# 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.
+
+# Use the official AWS Lambda Python base image.
+# If you wish to use a different Python version, please update the line below
+FROM public.ecr.aws/lambda/python:3.12
+
+# hadolint ignore=DL3041
+RUN dnf -y install unzip \
+    && dnf clean all \
+    && rm -rf /var/cache/dnf
+
+# Install the AWS CLI.
+RUN curl "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o "awscliv2.zip" && \
+    unzip awscliv2.zip \
+    && ./aws/install \
+    && rm -rf ./aws awscliv2.zip
+
+## Install Airflow and dependencies.
+# The most current version of Airflow is installed by default, along with the amazon and postgres
+# provider packages.
+# If you would like to install a specific version, you can use the requirements.txt to change the
+# version along with installing your dependencies or update this Dockerfile to install a specific
+# version of Airflow.
+
+# NOTE: If you change the below line, specifically removing the amazon extra, please ensure boto3 is
+# installed via another method. Boto3 is required for the Lambda executor to function properly.
+# hadolint ignore=SC2102
+RUN pip install --no-cache-dir --upgrade pip && pip install --no-cache-dir apache-airflow[amazon,postgres]
+# /tmp is the only writable directory in Lambda, so we need to set the AIRFLOW_HOME there.
+ENV AIRFLOW_HOME=/tmp/airflow
+# Dags are read-only, so they can be stored in opt (or another path provided).
+ARG container_dag_path=/opt/airflow/dags
+ENV AIRFLOW__CORE__DAGS_FOLDER=$container_dag_path
+RUN mkdir -p /opt/airflow/dags
+
+# Python dependencies can be installed by providing a requirements.txt.
+# If the file is in a different location, use the requirements_path build argument to specify
+# the file path.
+ARG requirements_path=./requirements.txt
+ENV REQUIREMENTS_PATH=$requirements_path
+# Uncomment the two lines below to copy the requirements.txt file to the container, and
+# install the dependencies.
+# COPY --chown=airflow:root $REQUIREMENTS_PATH /opt/airflow/requirements.txt
+# RUN pip install --no-cache-dir -r /opt/airflow/requirements.txt
+
+## AWS Authentication
+# The image requires access to AWS services. This Dockerfile supports 2 ways to authenticate with AWS.
+# The first is using build arguments where you can provide the AWS credentials as arguments
+# passed when building the image. The other option is to leverage the Lambda execution role. Airflow
+# will default to using Boto credential strategy which will look for roles from Lambda, this is the
+# preferred approach. See the Lambda Executor Airflow documentation for more details.
+
+# If you would like to use an alternative method of authentication, feel free to make the
+# necessary changes to this file.
+
+# Uncomment to use these arguments to provide AWS authentication information if not using the Lambda
+# execution role.
+#ARG aws_access_key_id
+#ARG aws_secret_access_key
+#ARG aws_default_region
+#ARG aws_session_token
+
+#ENV AWS_ACCESS_KEY_ID=$aws_access_key_id
+#ENV AWS_SECRET_ACCESS_KEY=$aws_secret_access_key
+#ENV AWS_DEFAULT_REGION=$aws_default_region
+#ENV AWS_SESSION_TOKEN=$aws_session_token
+
+## Loading DAGs
+# This Dockerfile supports 2 ways to load DAGs onto the container.
+# One is to download them from S3 at runtime during the Lambda app invocation. The other
+# is to copy the dags into the image at build time, this will make task execution
+# much faster, since the images will already be present but the image will need to be rebuilt
+# every time the DAGs are updated.
+# If you would like to use an alternative method of loading DAGs, feel free to make the
+# necessary changes to this file.
+
+ARG host_dag_path=./dags
+ENV HOST_DAG_PATH=$host_dag_path
+# Uncomment the line below to copy the DAGs from the host to the container.
+# COPY $HOST_DAG_PATH $AIRFLOW__CORE__DAGS_FOLDER
+
+# Use these arguments to load DAGs at runtime. If you are using the provided Lambda function (app.py),
+# it will check for this environment variable and download the DAGs from S3. See the example app in Lambda
+# Executor documentation for details.
+ARG s3_uri
+ENV S3_URI=$s3_uri
+
+# Copy your Lambda function code into the Docker build directory/context.
+COPY app.py ${LAMBDA_TASK_ROOT}/
+
+# Specify the Lambda function handler (update if you provide a different handler with a different name).
+CMD ["app.lambda_handler"]

airflow/providers/amazon/aws/executors/aws_lambda/docker/__init__.py

@@ -0,0 +1,16 @@
+# 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.

airflow/providers/amazon/aws/executors/aws_lambda/docker/app.py

@@ -0,0 +1,129 @@
+# 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
+
+import json
+import logging
+import os
+import subprocess
+from tempfile import mkdtemp
+
+import boto3
+
+"""
+Example Lambda function to execute an Airflow command or workload. Use or modify this code as needed.
+"""
+
+log = logging.getLogger()
+log.setLevel(logging.INFO)
+
+
+# Get the S3 URI from the environment variable. Set either on the Lambda function or in the
+# docker image used for the lambda invocations.
+S3_URI = os.environ.get("S3_URI", None)
+# Input and output keys
+TASK_KEY_KEY = "task_key"
+COMMAND_KEY = "command"
+RETURN_CODE_KEY = "return_code"
+
+
+def lambda_handler(event, context):
+    log.info("Received event: %s", event)
+    log.info("Received context: %s", context)
+
+    command = event.get(COMMAND_KEY)
+    task_key = event.get(TASK_KEY_KEY)
+
+    # Any pre-processing or validation of the command or use of the context can be done here or above.
+
+    # Sync dags from s3 to the local dags directory
+    if S3_URI:
+        fetch_dags_from_s3(S3_URI)
+    # This function must be called, it executes the Airflow command and reports to SQS.
+    run_and_report(command, task_key)
+
+    # Any post-processing or cleanup can be done here.
+
+
+def run_and_report(command, task_key):
+    """Execute the provided Airflow command or workload and report the result via SQS."""
+    try:
+        log.info("Starting execution for task: %s", task_key)
+        result = subprocess.run(
+            command, shell=isinstance(command, str), stdout=subprocess.PIPE, stderr=subprocess.STDOUT
+        )
+        return_code = result.returncode
+        log.info("Execution completed for task %s with return code %s", task_key, return_code)
+        log.info("Output:")
+        log.info("%s", result.stdout.decode())
+    except Exception:
+        log.exception("Error executing task %s: ", task_key)
+        return_code = 1  # Non-zero indicates failure to run the task
+
+    queue_url = get_queue_url()
+    message = json.dumps({TASK_KEY_KEY: task_key, RETURN_CODE_KEY: return_code})
+    try:
+        sqs_client = get_sqs_client()
+        sqs_client.send_message(QueueUrl=queue_url, MessageBody=message)
+        log.info("Sent result to SQS %s", message)
+    except Exception:
+        log.exception("Failed to send message to SQS for task %s", task_key)
+
+
+def get_sqs_client():
+    """Create an SQS client. Credentials and region are automatically picked up from the environment."""
+    return boto3.client("sqs")
+
+
+def get_queue_url():
+    """
+    Get the SQS queue URL from the environment variable.
+
+    Set either on the Lambda function or in the image used for the lambda invocations.
+    """
+    queue_url = os.environ.get("AIRFLOW__AWS_LAMBDA_EXECUTOR__QUEUE_URL", os.environ.get("QUEUE_URL", None))
+    if not queue_url:
+        raise RuntimeError(
+            "No Queue URL detected (either AIRFLOW__AWS_LAMBDA_EXECUTOR__QUEUE_URL or "
+            "QUEUE_URL); Will be unable to send task results. Exiting!"
+        )
+    return queue_url
+
+
+def fetch_dags_from_s3(s3_uri):
+    """Fetch DAGs from S3 and sync them to the local dags directory."""
+    log.info("Fetching DAGs from S3 URI: %s", s3_uri)
+    # Use a named temporary directory for the local dags folder, only tmp is writeable in Lambda
+    local_dags_dir = mkdtemp(prefix="airflow_dags_")
+    log.info("Setting AIRFLOW__CORE__DAGS_FOLDER to: %s", local_dags_dir)
+    os.environ["AIRFLOW__CORE__DAGS_FOLDER"] = local_dags_dir
+
+    # S3 URI format s3://bucket-name/path/to/dags/
+    bucket_name = s3_uri.split("/")[2]
+    prefix = "/".join(s3_uri.split("/")[3:])
+
+    s3_resource = boto3.resource("s3")
+    bucket = s3_resource.Bucket(bucket_name)
+
+    for obj in bucket.objects.filter(Prefix=prefix):
+        if obj.key.endswith("/"):
+            # Skip directories
+            continue
+        key = obj.key
+        local_path = os.path.join(local_dags_dir, os.path.basename(key))
+        log.info("Downloading %s to %s", key, local_path)
+        bucket.download_file(key, local_path)