fix(edge3): Resolve edge worker 405 error by deferring table creation to FastAPI startup

providers/edge3/src/airflow/providers/edge3/plugins/edge_executor_plugin.py

@@ -17,38 +17,176 @@
 
 from __future__ import annotations
 
+import logging
+import random
 import sys
+import time
+from contextlib import asynccontextmanager
 from typing import TYPE_CHECKING, Any
 
+from sqlalchemy import inspect
+from sqlalchemy.exc import OperationalError
+
 from airflow.configuration import conf
 from airflow.exceptions import AirflowConfigException
 from airflow.providers.common.compat.sdk import AirflowPlugin
 from airflow.providers.edge3.version_compat import AIRFLOW_V_3_1_PLUS
 from airflow.utils.session import NEW_SESSION, provide_session
+from airflow.utils.sqlalchemy import is_lock_not_available_error
 
 if TYPE_CHECKING:
+    from fastapi import FastAPI
+    from sqlalchemy.engine import Engine
     from sqlalchemy.orm import Session
 
 from airflow.utils.db import DBLocks, create_global_lock
 
+log = logging.getLogger(__name__)
+
+# Retry configuration for lock acquisition during table creation
+MAX_LOCK_RETRIES = 5
+LOCK_RETRY_DELAY_BASE = 1.0  # Base delay in seconds for exponential backoff
+
+# Required edge tables
+EDGE_TABLES = ("edge_job", "edge_logs", "edge_worker")
+
+
+def _tables_exist(engine: Engine) -> bool:
+    """
+    Check if all required edge tables already exist in the database.
+
+    This is a fast path to avoid acquiring a global lock when tables
+    are already present (normal operation after initial setup).
+    """
+    try:
+        inspector = inspect(engine)
+        existing_tables = set(inspector.get_table_names())
+        return all(table in existing_tables for table in EDGE_TABLES)
+    except Exception:
+        # If we can't check, assume tables don't exist and proceed with creation
+        return False
+
 
 @provide_session
-def _get_api_endpoint(session: Session = NEW_SESSION) -> dict[str, Any]:
-    # Ensure all required DB modeals are created before starting the API
-    with create_global_lock(session=session, lock=DBLocks.MIGRATIONS):
-        engine = session.get_bind().engine
-        from airflow.providers.edge3.models.edge_job import EdgeJobModel
-        from airflow.providers.edge3.models.edge_logs import EdgeLogsModel
-        from airflow.providers.edge3.models.edge_worker import EdgeWorkerModel
-
-        EdgeJobModel.metadata.create_all(engine)
-        EdgeLogsModel.metadata.create_all(engine)
-        EdgeWorkerModel.metadata.create_all(engine)
+def _ensure_tables_created(session: Session = NEW_SESSION) -> None:
+    """
+    Ensure all required DB models are created with retry logic.
+
+    This is called lazily on FastAPI app startup, not at plugin import time,
+    to avoid blocking all API server processes on a global database lock
+    during concurrent startup.
+
+    Uses exponential backoff with jitter to handle lock contention when
+    multiple API server processes start simultaneously.
+
+    Fast path: If tables already exist, skips lock acquisition entirely.
+    """
+    from airflow.providers.edge3.models.edge_job import EdgeJobModel
+    from airflow.providers.edge3.models.edge_logs import EdgeLogsModel
+    from airflow.providers.edge3.models.edge_worker import EdgeWorkerModel
 
+    engine = session.get_bind().engine
+
+    # Fast path: skip lock acquisition if tables already exist
+    if _tables_exist(engine):
+        log.debug("Edge tables already exist, skipping creation.")
+        return
+
+    last_error: OperationalError | RuntimeError | None = None
+
+    for attempt in range(MAX_LOCK_RETRIES):
+        try:
+            log.debug("Ensuring edge tables exist (attempt %d/%d)...", attempt + 1, MAX_LOCK_RETRIES)
+            with create_global_lock(session=session, lock=DBLocks.MIGRATIONS):
+                # Double-check after acquiring lock (another process may have created them)
+                if _tables_exist(engine):
+                    log.debug("Edge tables were created by another process.")
+                    return
+                EdgeJobModel.metadata.create_all(engine)
+                EdgeLogsModel.metadata.create_all(engine)
+                EdgeWorkerModel.metadata.create_all(engine)
+            log.debug("Edge tables created successfully.")
+            return
+        except OperationalError as e:
+            # Use the existing function that handles PostgreSQL (55P03) and MySQL (1205, 3572) error codes
+            if is_lock_not_available_error(e):
+                last_error = e
+                if attempt < MAX_LOCK_RETRIES - 1:
+                    # Exponential backoff with jitter to avoid thundering herd
+                    delay = LOCK_RETRY_DELAY_BASE * (2**attempt) + random.uniform(0, 1)
+                    log.warning(
+                        "Could not acquire migration lock for edge tables (attempt %d/%d), "
+                        "retrying in %.1f seconds...",
+                        attempt + 1,
+                        MAX_LOCK_RETRIES,
+                        delay,
+                    )
+                    time.sleep(delay)
+                    # Get a fresh session for retry to avoid stale connection issues
+                    session.rollback()
+                    # Check if tables were created while we were waiting
+                    if _tables_exist(engine):
+                        log.debug("Edge tables were created by another process during retry wait.")
+                        return
+            else:
+                # Re-raise non-lock-related errors immediately
+                raise
+        except RuntimeError as e:
+            # MySQL's create_global_lock raises RuntimeError when GET_LOCK times out
+            # Check if it's a lock-related error by examining the error message
+            error_str = str(e).lower()
+            if "lock" in error_str and ("could not acquire" in error_str or "timeout" in error_str):
+                last_error = e
+                if attempt < MAX_LOCK_RETRIES - 1:
+                    # Exponential backoff with jitter to avoid thundering herd
+                    delay = LOCK_RETRY_DELAY_BASE * (2**attempt) + random.uniform(0, 1)
+                    log.warning(
+                        "Could not acquire migration lock for edge tables (attempt %d/%d), "
+                        "retrying in %.1f seconds...",
+                        attempt + 1,
+                        MAX_LOCK_RETRIES,
+                        delay,
+                    )
+                    time.sleep(delay)
+                    # Get a fresh session for retry to avoid stale connection issues
+                    session.rollback()
+                    # Check if tables were created while we were waiting
+                    if _tables_exist(engine):
+                        log.debug("Edge tables were created by another process during retry wait.")
+                        return
+            else:
+                # Re-raise non-lock-related RuntimeErrors immediately
+                raise
+
+    # All retries exhausted
+    log.error(
+        "Failed to acquire migration lock for edge tables after %d attempts. "
+        "Edge worker API may not function correctly.",
+        MAX_LOCK_RETRIES,
+    )
+    if last_error:
+        raise last_error
+
+
+def _get_api_endpoint() -> dict[str, Any]:
+    """
+    Get API endpoint configuration.
+
+    Table creation is deferred to FastAPI startup event to avoid
+    blocking plugin import on database lock acquisition.
+    """
     from airflow.providers.edge3.worker_api.app import create_edge_worker_api_app
 
+    @asynccontextmanager
+    async def lifespan(app: FastAPI):  # type: ignore[type-arg]
+        """Create edge tables on app startup instead of at import time."""
+        _ensure_tables_created()
+        yield
+
+    app = create_edge_worker_api_app(lifespan=lifespan)
+
     return {
-        "app": create_edge_worker_api_app(),
+        "app": app,
         "url_prefix": "/edge_worker",
         "name": "Airflow Edge Worker",
     }

providers/edge3/src/airflow/providers/edge3/worker_api/app.py

@@ -17,6 +17,7 @@
 from __future__ import annotations
 
 from pathlib import Path
+from typing import Any
 
 from fastapi import FastAPI
 from fastapi.staticfiles import StaticFiles
@@ -28,7 +29,9 @@
 from airflow.providers.edge3.worker_api.routes.worker import worker_router
 
 
-def create_edge_worker_api_app() -> FastAPI:
+def create_edge_worker_api_app(
+    lifespan: Any = None,
+) -> FastAPI:
     """Create FastAPI app for edge worker API."""
     edge_worker_api_app = FastAPI(
         title="Airflow Edge Worker API",
@@ -40,6 +43,7 @@ def create_edge_worker_api_app() -> FastAPI:
             "All endpoints under ``/edge_worker/ui`` are used by UI and can be accessed with normal authentication. "
             "Please assume UI endpoints to change and not be stable."
         ),
+        lifespan=lifespan,
     )
 
     edge_worker_api_app.include_router(jobs_router, prefix="/v1")