'SQLAlchemy - Closing Session Blocks Code Execution

I have a query that runs as a batch using SQLAlchemy yield_per. This is a long duration query and I wanted to have a time limit for how long I run this query for. When the time limit is crossed I close the SQLAlchemy Session and break out of the for. The code gets stuck at the session.close statement after raising the following warning -

/Users/user/opt/miniconda3/envs/project/lib/python3.6/site-packages/pymysql/connections.py:799: UserWarning: Previous unbuffered result was left incomplete warnings.warn("Previous unbuffered result was left incomplete")

Below you can find a simple recreation of the code.

class QueryExec:
    def __init__(self, uri, connect_timeout, statement_timeout, pool_size, max_overflow) -> None:
        try:
            self.engine = create_engine(
                uri, #PyMySQL Driver is used.
                connect_args={
                    "connect_timeout": connect_timeout,
                    "read_timeout": statement_timeout,
                    "write_timeout": statement_timeout,
                },
                pool_size=pool_size,
                max_overflow=max_overflow,
            )
        except Exception as e:
            exit("Exiting as we couldn't connect to database")
        
    def execute(self, batch_size):
        start = datetime.now()
        
        #Create Session
        session_factory = sessionmaker(bind=self.engine)
        Session = scoped_session(session_factory)
        session = Session()

        query = (
            session.query(Table)
            .with_entities(
                Table.Row1,
                Table.Row2,
                Table.Row3,
            )
            .order_by(Table.Row3.asc())
            .execution_options(yield_per=batch_size)
        )
        
        for query_results in session.execute(query).partitions(batch_size):
            # Perform Operations on query_results for every batch

            # After every batch check -
            time_since_start = datetime.now() - start
            if time_since_start>timedelta(hours=1):
                session.close() # Gets stuck here
                break

        return

I am still figuring out how to work with SQLAlchemy, could someone help me with how to continue code execution after session close. Thanks!



Sources

This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.

Source: Stack Overflow

Solution Source