From 01e4d286d2786a8466a9bc63ef84e1610b85f263 Mon Sep 17 00:00:00 2001 From: Chris Guidry Date: Thu, 8 Feb 2024 16:25:26 -0500 Subject: [PATCH] Type safety in tests between SQLAlchemy 1 and 2 (#11946) --- .github/workflows/python-tests.yaml | 159 ++++++++++++++++++ src/prefect/server/database/configurations.py | 25 ++- 2 files changed, 176 insertions(+), 8 deletions(-) diff --git a/.github/workflows/python-tests.yaml b/.github/workflows/python-tests.yaml index 1c7b72a78932..d76d2f3584c6 100644 --- a/.github/workflows/python-tests.yaml +++ b/.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() diff --git a/src/prefect/server/database/configurations.py b/src/prefect/server/database/configurations.py index 5364394cbc0d..e1a2169e5838 100644 --- a/src/prefect/server/database/configurations.py +++ b/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,8 +62,8 @@ 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() @@ -62,8 +71,8 @@ def on_connect( 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]