Merge branch 'main' into avara1986-patch-1

Author: avara1986

Dockerfile

@@ -37,6 +37,22 @@ RUN rm -rf ./python/lib/$runtime/site-packages/ddtrace/appsec/_iast
 RUN rm -rf ./python/lib/$runtime/site-packages/ddtrace/appsec/sca
 RUN rm -rf ./python/lib/$runtime/site-packages/ddtrace/appsec/_shared
 RUN rm -rf ./python/lib/$runtime/site-packages/ddtrace/internal/test_visibility
+
+# CI Visibility paths/integrations
+RUN rm -rf \
+    ./python/lib/$runtime/site-packages/ddtrace/contrib/coverage/ \
+    ./python/lib/$runtime/site-packages/ddtrace/contrib/pytest/ \
+    ./python/lib/$runtime/site-packages/ddtrace/contrib/pytest_bdd/ \
+    ./python/lib/$runtime/site-packages/ddtrace/contrib/pytest_benchmark/ \
+    ./python/lib/$runtime/site-packages/ddtrace/contrib/selenium/ \
+    ./python/lib/$runtime/site-packages/ddtrace/contrib/unittest/ \
+    ./python/lib/$runtime/site-packages/ddtrace/ext/ci_visibility \
+    ./python/lib/$runtime/site-packages/ddtrace/ext/test_visibility \
+    ./python/lib/$runtime/site-packages/ddtrace/internal/ci_visibility \
+    ./python/lib/$runtime/site-packages/ddtrace/internal/coverage \
+    ./python/lib/$runtime/site-packages/ddtrace/internal/test_visibility \
+    ./python/lib/$runtime/site-packages/ddtrace/testing/
+
 # Dogshell
 RUN rm -rf ./python/lib/$runtime/site-packages/datadog/dogshell
 RUN rm -rf ./python/lib/$runtime/site-packages/bin/dog*
@@ -64,10 +80,18 @@ RUN rm -rf \
 # https://docs.python.org/3.11/using/cmdline.html#cmdoption-O
 # https://docs.python.org/3/using/cmdline.html#envvar-PYTHONNODEBUGRANGES
 RUN PYTHONNODEBUGRANGES=1 python -OO -m compileall -b ./python/lib/$runtime/site-packages
-# remove all .py files except ddtrace/contrib/*/patch.py which are necessary
-# for ddtrace.patch to discover instrumationation packages.
-RUN find ./python/lib/$runtime/site-packages -name \*.py | grep -v ddtrace/contrib | xargs rm -rf
-RUN find ./python/lib/$runtime/site-packages/ddtrace/contrib -name \*.py | grep -v patch.py | xargs rm -rf
+# remove all .py files
+# DEV: ddtrace>=4.7.0rc3 checks for .pyc files in addition to .py files for instrumentation
+# discovery (DataDog/dd-trace-py#17196), so we can safely remove all .py files.
+# For older versions, we need to keep patch.py files for instrumentation discovery.
+RUN pip install --quiet packaging && \
+    DDTRACE_VERSION=$(grep "^Version:" ./python/lib/$runtime/site-packages/ddtrace-*.dist-info/METADATA | awk '{print $2}') && \
+    if python -c "from packaging.version import Version; exit(0 if Version('$DDTRACE_VERSION') >= Version('4.7.0rc3') else 1)"; then \
+        find ./python/lib/$runtime/site-packages -name \*.py | xargs rm -rf; \
+    else \
+        find ./python/lib/$runtime/site-packages -name \*.py | grep -v ddtrace/contrib | xargs rm -rf && \
+        find ./python/lib/$runtime/site-packages/ddtrace/contrib -name \*.py | grep -v patch.py | xargs rm -rf; \
+    fi
 RUN find ./python/lib/$runtime/site-packages -name __pycache__ -type d -exec rm -r {} \+
 
 # When building ddtrace from branch, remove extra source files.  These are

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"
 

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 {
-        "durable_function_execution_name": execution_name,
-        "durable_function_execution_id": execution_id,
+        "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(),
     }
 
 

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

scripts/check_layer_size.sh

@@ -8,7 +8,7 @@
 # Compares layer size to threshold, and fails if below that threshold
 
 set -e
-MAX_LAYER_COMPRESSED_SIZE_KB=$(expr 9 \* 1024)      # 9216 KB
+MAX_LAYER_COMPRESSED_SIZE_KB=$(expr 9 \* 1024 + 15)      # 9231 KB
 MAX_LAYER_UNCOMPRESSED_SIZE_KB=$(expr 25 \* 1024)   # 25600 KB
 
 

tests/test_durable.py

@@ -46,18 +46,42 @@ def test_works_with_numeric_version_qualifier(self):
 
 
 class TestExtractDurableFunctionTags(unittest.TestCase):
-    def test_extracts_tags_from_event_with_durable_execution_arn(self):
+    def test_sets_first_invocation_true_when_only_execution_operation(self):
+        # One operation (the current EXECUTION operation itself) → not replaying → first invocation
         event = {
             "DurableExecutionArn": "arn:aws:lambda:us-east-1:123456789012:function:my-func:1/durable-execution/my-execution/550e8400-e29b-41d4-a716-446655440004",
             "CheckpointToken": "some-token",
-            "InitialExecutionState": {"Operations": []},
+            "InitialExecutionState": {"Operations": [{"OperationType": "EXECUTION"}]},
         }
         result = extract_durable_function_tags(event)
         self.assertEqual(
             result,
             {
-                "durable_function_execution_name": "my-execution",
-                "durable_function_execution_id": "550e8400-e29b-41d4-a716-446655440004",
+                "aws_lambda.durable_function.execution_name": "my-execution",
+                "aws_lambda.durable_function.execution_id": "550e8400-e29b-41d4-a716-446655440004",
+                "aws_lambda.durable_function.first_invocation": "true",
+            },
+        )
+
+    def test_sets_first_invocation_false_when_multiple_operations(self):
+        # More than one operation → replaying → not first invocation
+        event = {
+            "DurableExecutionArn": "arn:aws:lambda:us-east-1:123456789012:function:my-func:1/durable-execution/my-execution/550e8400-e29b-41d4-a716-446655440004",
+            "CheckpointToken": "some-token",
+            "InitialExecutionState": {
+                "Operations": [
+                    {"OperationType": "EXECUTION"},
+                    {"OperationType": "STEP"},
+                ]
+            },
+        }
+        result = extract_durable_function_tags(event)
+        self.assertEqual(
+            result,
+            {
+                "aws_lambda.durable_function.execution_name": "my-execution",
+                "aws_lambda.durable_function.execution_id": "550e8400-e29b-41d4-a716-446655440004",
+                "aws_lambda.durable_function.first_invocation": "false",
             },
         )