fix: psycopg2 register_default_jsonb TypeError with InstrumentedConnection + stack tests (#52)

Author: jy-tan

.github/workflows/e2e.yml

@@ -15,35 +15,43 @@ on:
 
 jobs:
   discover:
-    name: Discover E2E Tests
+    name: Discover Tests
     runs-on: ubuntu-latest
     outputs:
-      matrix: ${{ steps.set-matrix.outputs.matrix }}
+      e2e_matrix: ${{ steps.set-matrix.outputs.e2e_matrix }}
+      stack_matrix: ${{ steps.set-matrix.outputs.stack_matrix }}
     steps:
       - name: Checkout
         uses: actions/checkout@v4
 
-      - name: Find all e2e-tests directories
+      - name: Find all test directories
         id: set-matrix
         run: |
-          # Find all directories with e2e-tests and convert to JSON array
-          LIBRARIES=$(find drift/instrumentation -type d -name "e2e-tests" \
+          # Find all e2e-tests directories (single instrumentation)
+          E2E_TESTS=$(find drift/instrumentation -type d -name "e2e-tests" \
             | sed 's|drift/instrumentation/||' | sed 's|/e2e-tests||' | sort \
             | jq -R -s -c 'split("\n") | map(select(length > 0))')
           
-          echo "Found libraries with e2e-tests: $LIBRARIES"
-          echo "matrix=$LIBRARIES" >> $GITHUB_OUTPUT
+          # Find all stack-tests directories (multi-instrumentation)
+          STACK_TESTS=$(find drift/stack-tests -mindepth 1 -maxdepth 1 -type d 2>/dev/null \
+            | xargs -I {} basename {} | sort \
+            | jq -R -s -c 'split("\n") | map(select(length > 0))') || echo "[]"
+          
+          echo "Found e2e-tests: $E2E_TESTS"
+          echo "Found stack-tests: $STACK_TESTS"
+          echo "e2e_matrix=$E2E_TESTS" >> $GITHUB_OUTPUT
+          echo "stack_matrix=$STACK_TESTS" >> $GITHUB_OUTPUT
 
   e2e:
-    name: E2E Tests - ${{ matrix.library }}
+    name: E2E - ${{ matrix.library }}
     needs: discover
     runs-on: ubuntu-latest
     timeout-minutes: 30
     strategy:
       fail-fast: false
       max-parallel: 6
       matrix:
-        library: ${{ fromJSON(needs.discover.outputs.matrix) }}
+        library: ${{ fromJSON(needs.discover.outputs.e2e_matrix) }}
     steps:
       - name: Checkout
         uses: actions/checkout@v4
@@ -109,3 +117,80 @@ jobs:
           docker volume prune -f || true
           # Clean up networks
           docker network prune -f || true
+
+  stack:
+    name: Stack - ${{ matrix.test }}
+    needs: discover
+    if: ${{ needs.discover.outputs.stack_matrix != '[]' && needs.discover.outputs.stack_matrix != '' }}
+    runs-on: ubuntu-latest
+    timeout-minutes: 30
+    strategy:
+      fail-fast: false
+      max-parallel: 3
+      matrix:
+        test: ${{ fromJSON(needs.discover.outputs.stack_matrix) }}
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v4
+        with:
+          version: "latest"
+
+      - name: Setup Python
+        run: uv python install 3.9
+
+      - name: Setup Docker Buildx
+        uses: docker/setup-buildx-action@v3
+        with:
+          driver: docker
+
+      - name: Install SDK dependencies
+        run: uv sync --all-extras
+
+      - name: Build SDK
+        run: uv build
+
+      - name: Verify SDK build
+        run: |
+          ls -la dist/ || (echo "dist folder not found!" && exit 1)
+          test -f dist/*.whl || (echo "SDK build incomplete!" && exit 1)
+
+      - name: Get latest Tusk CLI version
+        id: tusk-version
+        run: |
+          VERSION=$(curl -s -H "Authorization: Bearer ${{ secrets.GITHUB_TOKEN }}" \
+            "https://api.github.com/repos/Use-Tusk/tusk-drift-cli/releases/latest" \
+            | grep '"tag_name":' | sed -E 's/.*"([^"]+)".*/\1/')
+          echo "version=$VERSION" >> $GITHUB_OUTPUT
+          echo "Latest Tusk CLI version: $VERSION"
+
+      - name: Build base image
+        env:
+          DOCKER_DEFAULT_PLATFORM: linux/amd64
+        run: |
+          docker build \
+            --build-arg TUSK_CLI_VERSION=${{ steps.tusk-version.outputs.version }} \
+            -t python-e2e-base:latest \
+            -f drift/instrumentation/e2e_common/Dockerfile.base \
+            .
+
+      - name: Run stack tests for ${{ matrix.test }}
+        env:
+          DOCKER_DEFAULT_PLATFORM: linux/amd64
+          TUSK_CLI_VERSION: ${{ steps.tusk-version.outputs.version }}
+        run: |
+          chmod +x ./drift/stack-tests/${{ matrix.test }}/run.sh
+          cd ./drift/stack-tests/${{ matrix.test }} && ./run.sh 8000
+
+      - name: Cleanup Docker resources
+        if: always()
+        run: |
+          # Stop all running containers
+          docker ps -aq | xargs -r docker stop || true
+          docker ps -aq | xargs -r docker rm || true
+          # Clean up volumes
+          docker volume prune -f || true
+          # Clean up networks
+          docker network prune -f || true

CONTRIBUTING.md

@@ -68,12 +68,14 @@ E2E tests validate full instrumentation workflows using Docker containers. They
 
 #### Running E2E Tests
 
-Run all e2e tests:
+Run all tests (e2e + stack):
 
 ```bash
-./run-all-e2e-tests.sh        # Sequential (default)
-./run-all-e2e-tests.sh 2      # 2 tests in parallel
-./run-all-e2e-tests.sh 0      # All tests in parallel
+./run-all-e2e-tests.sh                         # Run all tests sequentially
+./run-all-e2e-tests.sh -c 2                    # Run 2 tests concurrently
+./run-all-e2e-tests.sh -c 0                    # Run all tests in parallel
+./run-all-e2e-tests.sh --instrumentation-only  # Run only e2e tests
+./run-all-e2e-tests.sh --stack-only            # Run only stack tests
 ```
 
 Run a single instrumentation's e2e test:
@@ -143,12 +145,28 @@ python src/test_requests.py
 
 For more details, see `drift/instrumentation/README-e2e-tests.md`.
 
+### Stack Tests
+
+Stack tests validate multiple instrumentations working together in realistic application architectures (e.g., Django + PostgreSQL, FastAPI + Redis). They catch bugs at integration points that don't surface in isolated e2e testing.
+
+```bash
+# Run a specific stack test
+cd drift/stack-tests/django-postgres
+./run.sh
+
+# Or run all tests (including stack tests) from the root
+./run-all-e2e-tests.sh
+```
+
+For available tests and details, see `drift/stack-tests/README.md`.
+
 ## Documentation
 
 | Document | Description |
 |----------|-------------|
 | `docs/context-propagation.md` | Context propagation behavior, edge cases, and patterns |
 | `drift/instrumentation/README-e2e-tests.md` | E2E test architecture and debugging |
+| `drift/stack-tests/README.md` | Stack tests for multi-instrumentation scenarios |
 
 ## For Maintainers
 

E2E_TESTING_GUIDE.md

@@ -294,10 +294,16 @@ To run all E2E tests across all instrumentations:
 ./run-all-e2e-tests.sh
 
 # 2 tests in parallel
-./run-all-e2e-tests.sh 2
+./run-all-e2e-tests.sh -c 2
 
 # All tests in parallel (unlimited)
-./run-all-e2e-tests.sh 0
+./run-all-e2e-tests.sh -c 0
+
+# Run only single-instrumentation e2e tests
+./run-all-e2e-tests.sh --instrumentation-only
+
+# Run only stack tests
+./run-all-e2e-tests.sh --stack-only
 ```
 
 ## Quick Reference Commands

drift/instrumentation/README-e2e-tests.md

@@ -59,8 +59,6 @@ cd drift/instrumentation/flask/e2e-tests
 ./run.sh
 ```
 
----
-
 ### 2. fastapi
 
 **Purpose:** Test FastAPI ASGI instrumentation
@@ -81,8 +79,6 @@ cd drift/instrumentation/fastapi/e2e-tests
 ./run.sh
 ```
 
----
-
 ### 3. django
 
 **Purpose:** Test Django middleware instrumentation
@@ -103,8 +99,6 @@ cd drift/instrumentation/django/e2e-tests
 ./run.sh
 ```
 
----
-
 ### 4. redis
 
 **Purpose:** Test Redis instrumentation
@@ -124,8 +118,6 @@ cd drift/instrumentation/redis/e2e-tests
 ./run.sh
 ```
 
----
-
 ### 5. psycopg
 
 **Purpose:** Test Psycopg (v3) PostgreSQL instrumentation
@@ -146,8 +138,6 @@ cd drift/instrumentation/psycopg/e2e-tests
 ./run.sh
 ```
 
----
-
 ### 6. psycopg2
 
 **Purpose:** Test Psycopg2 (legacy) PostgreSQL instrumentation
@@ -166,8 +156,6 @@ cd drift/instrumentation/psycopg2/e2e-tests
 ./run.sh
 ```
 
----
-
 ## How E2E Tests Work
 
 ### Architecture
@@ -225,8 +213,6 @@ The e2e tests follow a **Docker entrypoint-driven architecture** where the Pytho
 - Python uses Docker entrypoint for orchestration
 - Both approaches work, but Python approach is more maintainable
 
----
-
 ## Prerequisites
 
 ### 1. Build Base Image
@@ -247,8 +233,6 @@ This image contains:
 
 All tests require Docker and Docker Compose.
 
----
-
 ## Running Tests
 
 ### Single Test
@@ -272,10 +256,16 @@ cd drift/instrumentation/flask/e2e-tests
 ./run-all-e2e-tests.sh
 
 # 2 tests in parallel
-./run-all-e2e-tests.sh 2
+./run-all-e2e-tests.sh -c 2
 
 # All tests in parallel (unlimited)
-./run-all-e2e-tests.sh 0
+./run-all-e2e-tests.sh -c 0
+
+# Run only single-instrumentation e2e tests
+./run-all-e2e-tests.sh --instrumentation-only
+
+# Run only stack tests
+./run-all-e2e-tests.sh --stack-only
 ```
 
 ### All Tests (Manual Parallel)
@@ -289,8 +279,6 @@ wait
 
 Each test uses a unique Docker Compose project name based on the port, so they don't conflict.
 
----
-
 ## Understanding Test Output
 
 ### Successful Test
@@ -341,8 +329,6 @@ If the test fails, you'll see:
 
 Traces are preserved in `.tusk/traces/` for inspection.
 
----
-
 ## Debugging Tests
 
 ### View Traces
@@ -389,8 +375,6 @@ docker compose build
 docker compose run --rm app /bin/bash
 ```
 
----
-
 ## CI Integration
 
 ### GitHub Actions Example
@@ -436,8 +420,6 @@ jobs:
           path: drift/instrumentation/${{ matrix.test }}/e2e-tests/.tusk/traces/
 ```
 
----
-
 ## Adding New Tests
 
 To add a new e2e test:
@@ -464,8 +446,6 @@ To add a new e2e test:
 
 5. **Add to CI**: Update GitHub Actions workflow
 
----
-
 ## Troubleshooting
 
 ### "python-e2e-base:latest not found"
@@ -565,8 +545,6 @@ The e2e test runner checks for these warnings after running tests. If found, the
 
 This is equivalent to the Node.js SDK's `check_tcp_instrumentation_warning` check.
 
----
-
 ## Comparison with Node.js E2E Tests
 
 | Aspect | Node.js | Python |
@@ -581,17 +559,13 @@ This is equivalent to the Node.js SDK's `check_tcp_instrumentation_warning` chec
 
 Both approaches achieve the same goal, but Python's entrypoint-driven design is simpler and more maintainable.
 
----
-
 ## Related Documentation
 
 - [Base Dockerfile](./e2e_common/Dockerfile.base) - Python e2e base Docker image
 - [Base Runner](./e2e_common/base_runner.py) - Shared e2e test runner class
 - [Python SDK README](../../README.md) - Main Python SDK documentation
 - [CONTRIBUTING.md](../../CONTRIBUTING.md) - Contribution guidelines with e2e test instructions
 
----
-
 ## Maintaining These Tests
 
 ### Updating Tusk CLI Version
@@ -623,8 +597,6 @@ Update `requirements.txt` in specific test directory:
 Flask>=3.2.0  # Update version
 ```
 
----
-
 ## Success Criteria
 
 All tests should: