handle disconnect exceptions raised from dbapi_connection or driver_connection operations in Connection context handler

[leo-b]

Currently SQLAlchemy doesn't handle SQL dialect exceptions that are raised during execution of direct dbapi_connection or driver_connection operations.
This may result in undetected disconnect errors and dead connections being returned to the connection pool.

However when using the engine.connect() context handler, a SQL dialect exception that has not already been processed and wrapped by SQLAlchemy could easily be detected in its __exit__ handler.

I am currently using subclasses of Connection and AsyncConnection that do the cleanup job well in my environment:

from typing import Any
import sqlalchemy
import sqlalchemy.ext.asyncio
from sqlalchemy.util.concurrency import greenlet_spawn


class EngineConnectionCleanup(sqlalchemy.engine.Connection):
    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        if exc_val is not None \
            and self._check_disconnect_exception(exc_val):
            self._handle_dbapi_exception(exc_val, None, None, None, None)
            # doesn't return: _handle_dbapi_exception() always raises an exception
            # the exception raised in _handle_dbapi_exception() will replace
            # the original one
            raise RuntimeError("shouldn't be here")
        return super().__exit__(exc_type, exc_val, exc_tb)

    def _check_disconnect_exception(self, e):
        if not isinstance(e, self.dialect.loaded_dbapi.Error) \
           or isinstance(e, sqlalchemy.exc.DBAPIError):
            # not a driver db error or already wrapped
            return False
        if self.invalidated:
            return False
        is_disconnect = (
            not self.closed
            and self.dialect.is_disconnect(
                e,
                self._dbapi_connection if not self.invalidated else None,
                None,
            )
        )
        if is_disconnect:
            return True
        else:
            return False


class AsyncEngineConnectionCleanup(sqlalchemy.ext.asyncio.AsyncConnection):
    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        if exc_val is not None \
           and self.sync_connection._check_disconnect_exception(exc_val):
            await greenlet_spawn(self._proxied._handle_dbapi_exception,
                                 exc_val, None, None, None, None,
                                 _require_await=False,
            )
            raise RuntimeError("shouldn't be here")
        return await super().__aexit__(exc_type, exc_val, exc_tb)

Is there a chance that this could be accepted upstream?

Cheers,
--leo

[zzzeek]

people dont necessarily use engine.connect() as a context handler, so putting this in there as an implicit behavior would be surprising and confusing. additionally, handling DBAPI exceptions has to be done immediately, because if people have things like try: / finally: going on where there's DB logic inside the finally, it turns into a catastrophe very quickly.

handle_dbapi_exception() is a good method to use here but it does the "disconnect" check internally, so we would not want to do a separate disconnect check outside and then call it conditionally, applications that use DBAPI connections should explicitly opt in to this error handling.

so i would propose a single context manager, perhaps like this:

conn = engine.connect()
try:
    dbapi_connection = conn.driver_connection
    with conn.handle_dbapi_exceptions():
        cursor = dbapi_connection.cursor()
        cursor.execute("select 1")
        cursor.close()
finally:
    conn.close()

or maybe like this (I like this better):

conn = engine.connect()
try:
    with conn.use_dbapi_connection() as dbapi_connection:
        cursor = dbapi_connection.cursor()
        cursor.execute("select 1")
        cursor.close()
finally:
    conn.close()

having folks opt in to the block explicitly with explicit behaviors is the way to go here. @CaselIT any thoughts?

[zzzeek]

note that this block would do the traditional "handle_dbapi_exception" logic unconditionally for all errors raised in the block

[CaselIT]

in the async connection we have get_raw_connection.

maybe we could call it raw_connection_context?. It's implied by the context that it manages errors.

conn = engine.connect()
try:
    with conn.raw_connection_context() as dbapi_connection:
        cursor = dbapi_connection.cursor()
        cursor.execute("select 1")
        cursor.close()
finally:
    conn.close()
# async
conn = await engine.connect()
try:
    async with aconn.raw_connection_context() as dbapi_connection:
        cursor = dbapi_connection.cursor()
        # this likely explodes since there is not greenlet, but it's but fine for this example
        cursor.execute("select 1")
        cursor.close()
finally:
    conn.close()

[zzzeek]

for async this is tricky because we have our "dbapi" version of the connection but what they should actually be using for this use case is the driver connection

i think these ctx mangers should be focused on the "driver" connection, not the "dbapi". becuase that's the use case for using the driver connection - you want to work with the raw driver's capabilities

[CaselIT]

that's a good point, so in async we would not return like raw_connection an internal proxy but directly the driver, so asyncpg when using asyncpg, psycopg when using it etc.

so probably something like driver_connection_context() or similar. Can't use dbapi in the name since it's ambiguous in the async connection

[leo-b]

handle_dbapi_exception() is a good method to use here but it does the "disconnect" check internally, so we would not want to do a separate disconnect check outside and then call it conditionally

In my case the separate check is needed to verify if the exception has already been handled (isinstance(e, sqlalchemy.exc.DBAPIError)), e.g. because engine.Connection methods are also used inside the context handler.

[zzzeek]

are you referring to your actual use case or the proposed context handlers in your example? handle_dbapi_exception is aware of being called in a re-entrant fashion

[leo-b]

handle_dbapi_exception is aware of being called in a re-entrant fashion

Ah - good to know. I had not been aware of that.
So if the exception has already been processed and wrapped, should_wrap and newraise will be False and the previous exception will be raised again?

[zzzeek]

I would hope, but with this area of study everything has to be demonstrated and vetted. suffice to say if handle_dbapi_exception has a case that it doesnt recognize what's going on , it has to be improved so that it does (im pretty sure it handles all this right now)

[manmayaray]

Hi Everyone, in my case, I am still getting these errors and I am lost ... we have added retry , so all query of the DB is handled..whereever retry is not there we are fcing the issue
We have the following exception handling

 try:
                db.session.commit()
                LOG.info(f"File_upload: Updated header for claim ID: {header_id}. Documents remaining: {docs_remaining}")
            except SQLAlchemyError as e:
                LOG.error(f"File_upload: Error committing to the database, error: {str(e)}", exc_info=True)
                LOG.error("Full traceback: " + traceback.format_exc())
            
            Al_Number=''
            if docs_remaining == 0:
                # producer = connect_kafka_producer()
                LOG.info(f"File_upload: All documents received for claim ID: {header_id}. Starting processing.")
                # try:
                #     claims_queue_active = get_queue_implementation_status()
                # except SQLAlchemyError as e:
                #     LOG.error(f"File_upload: Database error after all retries for querying claim queue status for header id: {header_id}, error: {str(e)}", exc_info=True)
                #     LOG.error("Full traceback: " + traceback.format_exc())
                
                claims_queue_active = "inactive"
                claim_status = "processing" if str(claims_queue_active.lower()) != 'active' else "pending_in_queue"
                keys = ["claim_id", "mode", "from_email","hospital_id"]
                values = [args["iAssist_Case_id"], "API", header_filename or "", hospital_id]
                Al_Number=header.al_number
                claims_ = CustMedClaimHdr(claim_header_id=args["iAssist_Case_id"],
                                        hospital_name="None",
                                        payment_mode="Cashless",
                                        request_type="Cashless",
                                        created_date=datetime.now(),
                                        updated_date=datetime.now(),
                                        claim_status=claim_status,
                                        attributes_1=header_filename or Al_Number or '')
                try:
                    commit_claims(claims_)
                except SQLAlchemyError as e:
                    LOG.error(f"File_upload: All retry attempts failed. Error: {str(e)}", exc_info=True)
                    LOG.error("Full traceback: " + traceback.format_exc())
                   
                    try:
                        with open(APP.config["MOUNT_DIR"], "fileupload_failed.json") as file:
                            data = file.read()
                            claim_id = args["iAssist_Case_id"]
                            data[claim_id] = claim_id
                    except Exception as e:
                        LOG.error(f"File_upload: Error cases, error: {str(e)}", exc_info=True)
                        LOG.error("Full traceback: " + traceback.format_exc())

                db.session.close()
                
                update_timestamp_at_level(args["iAssist_Case_id"], "claim_received_time")

                if str(claims_queue_active.lower()) != 'active':
                    print("cla",claims_queue_active)
                    if str(hospital_id) in ["1102394",'1271351', '1523','1294063','3166', '1102394','1007982', '1012345', '1010604']:
                        threading.Thread(target=
                                        process_classification, args=(APP,dict(zip(keys, values)))).start()
                    else:
                        threading.Thread(target=process_docs, args=(APP, dict(zip(keys, values)))).start()
                    print(f"email_dict is :: {dict(zip(keys, values))}")

                # produce(producer, APP.config["TOPIC"], str(keys), str(values))
            return {"success": True,
                    "data": {"iAssist_case_id": header_id,
                            "Al_Number":Al_Number,
                            "document_no": doc,
                            "no_of_document_remaining": docs_remaining},
                    "errors": ""}


            # return {"success": False, "data": {"iAssist_case_id": header_id}, "errors": "Unknown error"}
        except UnprocessableEntity as e:
            LOG.error(f"File_upload: Client error in file upload for claim ID: {header_id}, error: {str(e)}", exc_info=True)
            LOG.error("Full traceback: " + traceback.format_exc())
            return {"success": False, "data": {"iAssist_case_id": header_id}, "errors": str(e)}
        except NotAcceptable as e:
            LOG.error(f"File_upload: Not Acceptable error in file upload for claim ID: {header_id}, error: {str(e)}", exc_info=True)
            LOG.error("Full traceback: " + traceback.format_exc())
            return {"success": False, "data": {"iAssist_case_id": header_id}, "errors": str(e)}
        except Exception as e:
            LOG.error(f"File_upload: Unexpected error in file upload for claim ID: {header_id}, error: {str(e)}", exc_info=True)
            LOG.error("Full traceback: " + traceback.format_exc())
            try:
                with open(APP.config["MOUNT_DIR"], "fileupload_failed.json") as file:
                    data = file.read()
                    claim_id = args["iAssist_Case_id"]
                    data[claim_id] = claim_id
            except Exception as e:
                LOG.error(f"File_upload: Error cases, error: {str(e)}", exc_info=True)
                LOG.error("Full traceback: " + traceback.format_exc())
            return {"success": False, "data": {"iAssist_case_id": header_id}, "errors": str(e)}