Feature pg8000 graceful shutdown

legal-api/src/legal_api/__init__.py

@@ -35,14 +35,13 @@
 
 This module is the API for the Legal Entity system.
 """
-import logging
 import os
 from typing import Optional
 
+from cloud_sql_connector import setup_pg8000_close_event_listener
 from flask import Flask, jsonify
 from registry_schemas import __version__ as registry_schemas_version
 from registry_schemas.flask import SchemaServices
-from sqlalchemy import event
 
 from legal_api import config, models
 from legal_api.models import db
@@ -63,29 +62,6 @@
 setup_logging(os.path.join(os.path.abspath(os.path.dirname(__file__)), "logging.conf"))  # important to do this first
 
 
-def _setup_pg8000_graceful_shutdown(engine) -> None:
-    """Suppress pg8000 InterfaceError on connection close during Cloud Run scale-down."""
-    try:
-        from pg8000.exceptions import InterfaceError as _interface_error  # noqa: N813
-    except ImportError:
-        _interface_error = None
-
-    @event.listens_for(engine, "connect")
-    def on_connect(dbapi_conn, _connection_record):
-        orig_close = dbapi_conn.close
-
-        def safe_close():
-            try:
-                orig_close()
-            except Exception as exc:
-                if _interface_error and isinstance(exc, _interface_error):
-                    logging.getLogger(__name__).debug("Suppressed pg8000 InterfaceError on teardown.")
-                else:
-                    raise
-
-        dbapi_conn.close = safe_close
-
-
 def create_app(run_mode: Optional[str] = None, **kwargs) -> Flask:
     """Return a configured Flask App using the Factory method."""
     run_mode = run_mode or os.getenv("RUN_MODE", "production")
@@ -108,7 +84,7 @@ def create_app(run_mode: Optional[str] = None, **kwargs) -> Flask:
     with app.app_context():  # db require app context
         digital_credentials.init_app(app)
         if app.config.get("CLOUDSQL_INSTANCE_CONNECTION_NAME"):  # pragma: no cover
-            _setup_pg8000_graceful_shutdown(db.engine)
+            setup_pg8000_close_event_listener(db.engine)
 
     cache.init_app(app)