Merge branch 'main' into zarir.hamza/disable-rc-in-lambda

Author: zarirhamza

Dockerfile

@@ -33,7 +33,13 @@ RUN pip install --no-cache-dir . -t ./python/lib/$runtime/site-packages
 RUN rm -rf ./python/lib/$runtime/site-packages/botocore*
 RUN rm -rf ./python/lib/$runtime/site-packages/setuptools
 RUN rm -rf ./python/lib/$runtime/site-packages/jsonschema/tests
-RUN rm -rf ./python/lib/$runtime/site-packages/ddtrace/appsec/_iast
+
+# Remove unsupported appsec modules
+RUN rm -rf \
+    ./python/lib/$runtime/site-packages/ddtrace/appsec/_iast \
+    ./python/lib/$runtime/site-packages/ddtrace/appsec/sca \
+    ./python/lib/$runtime/site-packages/ddtrace/appsec/_shared
+
 # CI Visibility paths/integrations
 RUN rm -rf \
     ./python/lib/$runtime/site-packages/ddtrace/contrib/coverage/ \

ci/get_secrets.sh

@@ -21,21 +21,11 @@ fi
 
 printf "Getting AWS External ID...\n"
 
-EXTERNAL_ID=$(aws ssm get-parameter \
-    --region us-east-1 \
-    --name "ci.datadog-lambda-python.$EXTERNAL_ID_NAME" \
-    --with-decryption \
-    --query "Parameter.Value" \
-    --out text)
+EXTERNAL_ID=$(vault kv get -field="$EXTERNAL_ID_NAME" kv/k8s/gitlab-runner/datadog-lambda-python/secrets)
 
 printf "Getting DD API KEY...\n"
 
-export DD_API_KEY=$(aws ssm get-parameter \
-    --region us-east-1 \
-    --name ci.datadog-lambda-python.dd-api-key \
-    --with-decryption \
-    --query "Parameter.Value" \
-    --out text)
+export DD_API_KEY=$(vault kv get -field=dd-api-key kv/k8s/gitlab-runner/datadog-lambda-python/secrets)
 
 printf "Assuming role...\n"
 

ci/input_files/build.yaml.tpl

@@ -285,30 +285,5 @@ e2e-test-status:
     - if: '$SKIP_E2E_TESTS == "true"'
       when: never
     - when: on_success
-  script: |
-      GITLAB_API_TOKEN=$(aws ssm get-parameter --region us-east-1 --name "ci.${CI_PROJECT_NAME}.serverless-e2e-gitlab-token" --with-decryption --query "Parameter.Value" --out text)
-      URL="${CI_API_V4_URL}/projects/${CI_PROJECT_ID}/pipelines/${CI_PIPELINE_ID}/bridges"
-      echo "Fetching E2E job status from: $URL"
-      while true; do
-        RESPONSE=$(curl -s --header "PRIVATE-TOKEN: ${GITLAB_API_TOKEN}" "$URL")
-        E2E_JOB_STATUS=$(echo "$RESPONSE" | jq -r '.[] | select(.name=="e2e-test") | .downstream_pipeline.status')
-        echo -n "E2E job status: $E2E_JOB_STATUS, "
-        if [ "$E2E_JOB_STATUS" == "success" ]; then
-          echo "✅ E2E tests completed successfully"
-          exit 0
-        elif [ "$E2E_JOB_STATUS" == "failed" ]; then
-          echo "❌ E2E tests failed"
-          exit 1
-        elif [ "$E2E_JOB_STATUS" == "running" ]; then
-          echo "⏳ E2E tests are still running, retrying in 1 minute..."
-        elif [ "$E2E_JOB_STATUS" == "canceled" ]; then
-          echo "🚫 E2E tests were canceled"
-          exit 1
-        elif [ "$E2E_JOB_STATUS" == "skipped" ]; then
-          echo "⏭️ E2E tests were skipped"
-          exit 0
-        else
-          echo "❓ Unknown E2E test status: $E2E_JOB_STATUS, retrying in 1 minute..."
-        fi
-        sleep 60
-      done
+  script: 
+      - ci/poll_e2e.sh

ci/poll_e2e.sh

@@ -0,0 +1,38 @@
+curl -OL "binaries.ddbuild.io/dd-source/authanywhere/LATEST/authanywhere-linux-amd64" && mv "authanywhere-linux-amd64" /bin/authanywhere && chmod +x /bin/authanywhere
+
+BTI_CI_API_TOKEN=$(authanywhere --audience rapid-devex-ci)
+
+BTI_RESPONSE=$(curl --silent --request GET \
+    --header "$BTI_CI_API_TOKEN" \
+    --header "Content-Type: application/vnd.api+json" \
+    "https://bti-ci-api.us1.ddbuild.io/internal/ci/gitlab/token?owner=DataDog&repository=datadog-lambda-python")
+
+GITLAB_TOKEN=$(echo "$BTI_RESPONSE" | jq -r '.token // empty') 
+
+URL="${CI_API_V4_URL}/projects/${CI_PROJECT_ID}/pipelines/${CI_PIPELINE_ID}/bridges"
+
+echo "Fetching E2E job status from: $URL"
+
+while true; do
+    RESPONSE=$(curl -s --header "PRIVATE-TOKEN: ${GITLAB_TOKEN}" "$URL")
+    E2E_JOB_STATUS=$(echo "$RESPONSE" | jq -r '.[] | select(.name=="e2e-test") | .downstream_pipeline.status')
+    echo -n "E2E job status: $E2E_JOB_STATUS, "
+    if [ "$E2E_JOB_STATUS" == "success" ]; then
+        echo "✅ E2E tests completed successfully"
+        exit 0
+    elif [ "$E2E_JOB_STATUS" == "failed" ]; then
+        echo "❌ E2E tests failed"
+        exit 1
+    elif [ "$E2E_JOB_STATUS" == "running" ]; then
+        echo "⏳ E2E tests are still running, retrying in 2 minutes..."
+    elif [ "$E2E_JOB_STATUS" == "canceled" ]; then
+        echo "🚫 E2E tests were canceled"
+        exit 1
+    elif [ "$E2E_JOB_STATUS" == "skipped" ]; then
+        echo "⏭️ E2E tests were skipped"
+        exit 0
+    else
+        echo "❓ Unknown E2E test status: $E2E_JOB_STATUS, retrying in 2 minutes..."
+    fi
+    sleep 120
+done

datadog_lambda/durable.py

@@ -41,11 +41,17 @@ def extract_durable_function_tags(event):
     if not parsed:
         logger.error("Failed to parse DurableExecutionArn: %s", durable_execution_arn)
         return {}
-
     execution_name, execution_id = parsed
+    # Use the number of operations to determine if it's the first invocation. This is
+    # what the durable execution SDK does to determine the replay status.
+    operations = event.get("InitialExecutionState", {}).get("Operations", [])
+    is_first_invocation = len(operations) == 1
     return {
         "aws_lambda.durable_function.execution_name": execution_name,
         "aws_lambda.durable_function.execution_id": execution_id,
+        "aws_lambda.durable_function.first_invocation": str(
+            is_first_invocation
+        ).lower(),
     }
 
 

datadog_lambda/tracing.py

@@ -977,16 +977,12 @@ def process_injected_data(event, request_time_epoch_ms, args, tags):
             start_time_ns = int(
                 injected_authorizer_data.get(Headers.Parent_Span_Finish_Time)
             )
-            finish_time_ns = (
-                request_time_epoch_ms
-                + (
-                    int(
-                        event["requestContext"]["authorizer"].get(
-                            "integrationLatency", 0
-                        )
-                    )
-                )
-            ) * 1e6
+            integration_latency = int(
+                event["requestContext"]["authorizer"].get("integrationLatency", 0)
+            )
+            finish_time_ns = max(
+                start_time_ns, (request_time_epoch_ms + integration_latency) * 1e6
+            )
             upstream_authorizer_span = insert_upstream_authorizer_span(
                 args, tags, start_time_ns, finish_time_ns
             )

datadog_lambda/wrapper.py

@@ -191,7 +191,6 @@ def __call__(self, event, context, **kwargs):
             if self.blocking_response:
                 return self.blocking_response
             self.response = self.func(event, context, **kwargs)
-            return self.response
         except BlockingException:
             self.blocking_response = get_asm_blocked_response(self.event_source)
         except Exception:
@@ -204,8 +203,9 @@ def __call__(self, event, context, **kwargs):
             raise
         finally:
             self._after(event, context)
-            if self.blocking_response:
-                return self.blocking_response
+        if self.blocking_response:
+            return self.blocking_response
+        return self.response
 
     def _inject_authorizer_span_headers(self, request_id):
         reference_span = self.inferred_span if self.inferred_span else self.span

pyproject.toml

@@ -33,7 +33,10 @@ ddtrace = [
     {version = ">=3.19.1,<4", python = ">=3.8,<3.10"},
     {version = ">=4.1.1,<5,!=4.6.*", python = ">=3.10"}
 ]
-ujson = ">=5.9.0"
+ujson = [
+    {version = ">=5.10.0,<5.12.0", python = ">=3.8,<3.10"},
+    {version = ">=5.12.0", python = ">=3.10"}
+]
 botocore = { version = "^1.34.0", optional = true }
 requests = { version ="^2.22.0", optional = true }
 pytest = { version= "^8.0.0", optional = true }

scripts/build_layers.sh

@@ -91,29 +91,49 @@ trap cleanup EXIT
 # Helper: replace the multi-line ddtrace dependency in pyproject.toml.
 # Uses perl instead of sed -z for macOS/Linux portability.
 replace_ddtrace_dep() {
+    echo "Replacing dep with $1"
     perl -i -0777 -pe "s|ddtrace = \[[^\]]*\]|$1|gs" pyproject.toml
 }
 
 function make_path_absolute {
     echo "$(cd "$(dirname "$1")"; pwd)/$(basename "$1")"
 }
 
-function docker_build_zip {
-    # Args: [python version] [zip destination]
+function search_wheel {
+    # Args: [wheel base name] [index]
 
-    destination=$(make_path_absolute $2)
-    arch=$3
+    WHEEL_BASENAME=$1
+    INDEX=$2
+
+    SEARCH_PATTERN="${WHEEL_BASENAME}-[^\"]*${PY_TAG}[^\"]*${PLATFORM}[^\"]*\.whl"
+    INDEX_URL="${S3_BASE}/index-${INDEX}.html" 
+    echo "Searching for wheel ${SEARCH_PATTERN}"
+    export WHEEL_FILE=$(curl -sSfL ${INDEX_URL} | grep -o "$SEARCH_PATTERN" | head -n 1)
+    if [ ! -z "${WHEEL_FILE}" ]; then
+        curl -sSfL "${S3_BASE}/${WHEEL_FILE}" -o "${WHEEL_FILE}"
+        echo "Using S3 wheel: ${WHEEL_FILE}"
+        replace_ddtrace_dep "${WHEEL_BASENAME} = { file = \"${WHEEL_FILE}\" }"
+    fi
+}
+
+function find_and_spec_wheel {
+    # Args: [python version] [wheel base name] [index]
+
+    arch=$2
+    wheel_basename=$3
+    index=$4
 
     # Restore pyproject.toml to a clean state for each build iteration
     cp pyproject.toml.bak pyproject.toml
 
     # Replace ddtrace source if necessary
     if [ -n "$DD_TRACE_COMMIT" ]; then
-        replace_ddtrace_dep "ddtrace = { git = \"https://github.com/DataDog/dd-trace-py.git\", rev = \"$DD_TRACE_COMMIT\" }"
+        replace_ddtrace_dep "${wheel_basename} = { git = \"https://github.com/DataDog/dd-trace-py.git\", rev = \"$DD_TRACE_COMMIT\" }"
     elif [ -n "$DD_TRACE_COMMIT_BRANCH" ]; then
-        replace_ddtrace_dep "ddtrace = { git = \"https://github.com/DataDog/dd-trace-py.git\", branch = \"$DD_TRACE_COMMIT_BRANCH\" }"
+        replace_ddtrace_dep "${wheel_basename} = { git = \"https://github.com/DataDog/dd-trace-py.git\", branch = \"$DD_TRACE_COMMIT_BRANCH\" }"
     elif [ -n "$DD_TRACE_WHEEL" ]; then
-        replace_ddtrace_dep "ddtrace = { file = \"$DD_TRACE_WHEEL\" }"
+        wheel_basename=$(sed 's/^.*\///' <<< ${DD_TRACE_WHEEL%%-*})
+        replace_ddtrace_dep "${wheel_basename} = { file = \"$DD_TRACE_WHEEL\" }"
     elif [ -n "$UPSTREAM_PIPELINE_ID" ]; then
         S3_BASE="https://dd-trace-py-builds.s3.amazonaws.com/${UPSTREAM_PIPELINE_ID}"
         if [ "${arch}" = "amd64" ]; then
@@ -122,18 +142,19 @@ function docker_build_zip {
             PLATFORM="manylinux2014_aarch64"
         fi
         PY_TAG="cp$(echo "$1" | tr -d '.')"
-        WHEEL_FILE=$(curl -sSfL "${S3_BASE}/index-manylinux2014.html" \
-            | grep -o "ddtrace-[^\"]*${PY_TAG}[^\"]*${PLATFORM}[^\"]*\.whl" \
-            | head -n 1)
+        search_wheel ${wheel_basename} ${index}
         if [ -z "${WHEEL_FILE}" ]; then
             echo "No S3 wheel found for ${PY_TAG} ${PLATFORM}, using default pyproject.toml version"
-        else
-            curl -sSfL "${S3_BASE}/${WHEEL_FILE}" -o "${WHEEL_FILE}"
-            echo "Using S3 wheel: ${WHEEL_FILE}"
-            replace_ddtrace_dep "ddtrace = { file = \"${WHEEL_FILE}\" }"
+            return 1
         fi
     fi
+}
+
+function docker_build_zip {
+    # Args: [python version] [zip destination]
 
+    destination=$(make_path_absolute $2)
+    arch=$3
     # Install datadogpy in a docker container to avoid the mess from switching
     # between different python runtimes.
     temp_dir=$(mktemp -d)
@@ -159,6 +180,14 @@ do
     for architecture in "${ARCHS[@]}"
     do
         echo "Building layer for Python ${python_version} arch=${architecture}"
+        set +e
+        find_and_spec_wheel ${python_version} ${architecture} "ddtrace_serverless" "serverless"
+        FAILURE=$?
+        if [ $FAILURE != 0 ]; then
+            echo "Attempting layer build again with package ddtrace"
+            find_and_spec_wheel ${python_version} ${architecture} "ddtrace" "manylinux2014"
+        fi
+        set -e
         docker_build_zip ${python_version} $LAYER_DIR/${LAYER_FILES_PREFIX}-${architecture}-${python_version}.zip ${architecture}
     done
 done

tests/event_samples/authorizer-request-api-gateway-v1.json

@@ -55,7 +55,7 @@
   "requestContext": {
     "resourceId": "0et54l",
     "authorizer": {
-      "_datadog": "eyJ4LWRhdGFkb2ctdHJhY2UtaWQiOiAiMTM0Nzg3MDU5OTU3OTcyMjEyMDkiLCAieC1kYXRhZG9nLXBhcmVudC1pZCI6ICI4NDcxMjg4MjYzMzg0MjE2ODk2IiwgIngtZGF0YWRvZy1zYW1wbGluZy1wcmlvcml0eSI6ICIxIiwgIngtZGF0YWRvZy1wYXJlbnQtc3Bhbi1maW5pc2gtdGltZSI6IDE2NjMyOTUwMjE4MjcuNTIxLCAieC1kYXRhZG9nLWF1dGhvcml6aW5nLXJlcXVlc3RpZCI6ICJhYmMxMjMifQ==",
+      "_datadog": "eyJ4LWRhdGFkb2ctdHJhY2UtaWQiOiAiMTM0Nzg3MDU5OTU3OTcyMjEyMDkiLCAieC1kYXRhZG9nLXBhcmVudC1pZCI6ICI4NDcxMjg4MjYzMzg0MjE2ODk2IiwgIngtZGF0YWRvZy1zYW1wbGluZy1wcmlvcml0eSI6ICIxIiwgIngtZGF0YWRvZy1wYXJlbnQtc3Bhbi1maW5pc2gtdGltZSI6IDE2NjMyOTUwMjE4Mjc1MjEwMDAsICJ4LWRhdGFkb2ctYXV0aG9yaXppbmctcmVxdWVzdGlkIjogImFiYzEyMyJ9",
       "scope": "this is just a string",
       "principalId": "foo",
       "integrationLatency": 1897