Type safety in tests between SQLAlchemy 1 and 2 (#11946)

.github/workflows/python-tests.yaml

@@ -381,6 +381,165 @@ jobs:
           && docker container logs postgres
           || echo "Ignoring bad exit code"
 
+  run-sqlalchemy-v1-tests:
+    runs-on:
+      group: oss-larger-runners
+    name: sqlalchemy v1, python:${{ matrix.python-version }}
+    strategy:
+      matrix:
+        database:
+          - "postgres:14"
+        python-version:
+          - "3.8"
+          - "3.12"
+
+      fail-fast: false
+
+    timeout-minutes: 45
+
+    steps:
+      - name: Display current test matrix
+        run: echo '${{ toJSON(matrix) }}'
+
+      - uses: actions/checkout@v4
+        with:
+          persist-credentials: false
+          fetch-depth: 0
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+          cache: "pip"
+          cache-dependency-path: "requirements*.txt"
+
+      - name: Get image tag
+        id: get_image_tag
+        run: |
+          SHORT_SHA=$(git rev-parse --short=7 HEAD)
+          tmp="sha-$SHORT_SHA-python${{ matrix.python-version }}"
+          echo "image_tag=${tmp}" >> $GITHUB_OUTPUT
+
+      - name: Force install sqlalchemy in docker images
+        run: |
+          sed -i 's/RUN prefect version/RUN pip install "sqlalchemy[asyncio]<2"\nRUN prefect version/' Dockerfile
+
+      - name: Build test image
+        uses: docker/build-push-action@v5
+        with:
+          context: .
+          # TODO: We do not need the UI in these tests and we may want to add a build-arg to disable building it
+          #       so that CI test runs are faster
+          build-args: |
+            PYTHON_VERSION=${{ matrix.python-version }}
+            PREFECT_EXTRAS=[dev]
+          tags: prefecthq/prefect-dev:${{ steps.get_image_tag.outputs.image_tag }}
+          outputs: type=docker,dest=/tmp/image.tar
+          cache-from: type=gha
+          cache-to: type=gha,mode=max,ignore-error=true
+
+      - name: Test Docker image
+        run: |
+          docker load --input /tmp/image.tar
+          docker run --rm prefecthq/prefect-dev:${{ steps.get_image_tag.outputs.image_tag }} prefect version
+
+      - name: Build Conda flavored test image
+        uses: docker/build-push-action@v5
+        with:
+          context: .
+          build-args: |
+            PYTHON_VERSION=${{ matrix.python-version }}
+            BASE_IMAGE=prefect-conda
+            PREFECT_EXTRAS=[dev]
+          tags: prefecthq/prefect-dev:${{ steps.get_image_tag.outputs.image_tag }}-conda
+          outputs: type=docker,dest=/tmp/image-conda.tar
+          cache-from: type=gha
+          # We do not cache Conda image layers because they very big and slow to upload
+          # cache-to: type=gha,mode=max
+
+      - name: Test Conda flavored Docker image
+        run: |
+          docker load --input /tmp/image-conda.tar
+          docker run --rm prefecthq/prefect-dev:${{ steps.get_image_tag.outputs.image_tag }}-conda prefect version
+          docker run --rm prefecthq/prefect-dev:${{ steps.get_image_tag.outputs.image_tag }}-conda conda --version
+
+      - name: Install packages
+        run: |
+          python -m pip install -U pip
+          pip install --upgrade --upgrade-strategy eager -e .[dev]
+
+      - name: Force install sqlalchemy for unit tests
+        run: pip install "sqlalchemy[asyncio]<2"
+
+      - name: Start database container
+        if: ${{ startsWith(matrix.database, 'postgres') }}
+        run: >
+          docker run
+          --name "postgres"
+          --detach
+          --health-cmd pg_isready
+          --health-interval 10s
+          --health-timeout 5s
+          --health-retries 5
+          --publish 5432:5432
+          --tmpfs /var/lib/postgresql/data
+          --env POSTGRES_USER="prefect"
+          --env POSTGRES_PASSWORD="prefect"
+          --env POSTGRES_DB="prefect"
+          --env LANG="C.UTF-8"
+          --env LANGUAGE="C.UTF-8"
+          --env LC_ALL="C.UTF-8"
+          --env LC_COLLATE="C.UTF-8"
+          --env LC_CTYPE="C.UTF-8"
+          ${{ matrix.database }}
+          -c max_connections=250
+
+          ./scripts/wait-for-healthy-container.sh postgres 30
+
+          echo "PREFECT_API_DATABASE_CONNECTION_URL=postgresql+asyncpg://prefect:prefect@localhost/prefect" >> $GITHUB_ENV
+
+        # Parallelize tests by scope to reduce expensive service fixture duplication
+        # Do not allow the test suite to build images, as we want the prebuilt images to be tested
+        # Do not run Kubernetes service tests, we do not have a cluster available
+        # maxprocesses 6 is based on empirical testing; higher than 6 sees diminishing returns
+      - name: Run tests
+        run: >
+          pytest tests
+          --numprocesses auto
+          --maxprocesses 6
+          --dist loadscope
+          --disable-docker-image-builds
+          --exclude-service kubernetes
+          --durations 26
+          --no-cov
+
+      - name: Create and Upload failure flag
+        if: ${{ failure() }}
+        id: create_failure_flag
+        run: |
+          sanitized_name="sqlalchemy-v1-${{ matrix.os }}-${{ matrix.python-version }}-${{ matrix.database }}-${{ matrix.pytest-options }}"
+          sanitized_name="${sanitized_name//:/-}"
+          echo "Failure in $sanitized_name" > "${sanitized_name}-failure.txt"
+          echo "artifact_name=${sanitized_name}-failure" >> $GITHUB_OUTPUT
+
+      - name: Upload failure flag
+        if: ${{ failure() }}
+        uses: actions/upload-artifact@v4
+        with:
+          name: ${{ steps.create_failure_flag.outputs.artifact_name }}
+          path: "${{ steps.create_failure_flag.outputs.artifact_name }}.txt"
+
+      - name: Check database container
+        # Only applicable for Postgres, but we want this to run even when tests fail
+        if: always()
+        run: >
+          docker container inspect postgres
+          && docker container logs postgres
+          || echo "Ignoring bad exit code"
+
   notify-tests-failing-on-main:
     needs: run-tests
     if: github.ref == 'refs/heads/main' && failure()

src/prefect/server/database/configurations.py

@@ -8,6 +8,15 @@
 from typing import Dict, Hashable, Optional, Tuple
 
 import sqlalchemy as sa
+
+try:
+    from sqlalchemy import AdaptedConnection
+    from sqlalchemy.pool import ConnectionPoolEntry
+except ImportError:
+    # SQLAlchemy 1.4 equivalents
+    from sqlalchemy.pool import _ConnectionFairy as AdaptedConnection
+    from sqlalchemy.pool.base import _ConnectionRecord as ConnectionPoolEntry
+
 from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine
 from typing_extensions import Literal
 
@@ -31,9 +40,9 @@ class ConnectionTracker:
     """A test utility which tracks the connections given out by a connection pool, to
     make it easy to see which connections are currently checked out and open."""
 
-    all_connections: Dict[sa.AdaptedConnection, str]
-    open_connections: Dict[sa.AdaptedConnection, str]
-    left_field_closes: Dict[sa.AdaptedConnection, str]
+    all_connections: Dict[AdaptedConnection, str]
+    open_connections: Dict[AdaptedConnection, str]
+    left_field_closes: Dict[AdaptedConnection, str]
     connects: int
     closes: int
     active: bool
@@ -53,17 +62,17 @@ def track_pool(self, pool: sa.pool.Pool):
 
     def on_connect(
         self,
-        adapted_connection: sa.AdaptedConnection,
-        connection_record: sa.pool.ConnectionPoolEntry,
+        adapted_connection: AdaptedConnection,
+        connection_record: ConnectionPoolEntry,
     ):
         self.all_connections[adapted_connection] = traceback.format_stack()
         self.open_connections[adapted_connection] = traceback.format_stack()
         self.connects += 1
 
     def on_close(
         self,
-        adapted_connection: sa.AdaptedConnection,
-        connection_record: sa.pool.ConnectionPoolEntry,
+        adapted_connection: AdaptedConnection,
+        connection_record: ConnectionPoolEntry,
     ):
         try:
             del self.open_connections[adapted_connection]
@@ -73,7 +82,7 @@ def on_close(
 
     def on_close_detached(
         self,
-        adapted_connection: sa.AdaptedConnection,
+        adapted_connection: AdaptedConnection,
     ):
         try:
             del self.open_connections[adapted_connection]