AttributeError: 'psycopg2.extensions.connection' object has no attribute 'staleness' with 1.0.0

[hadim]

I am using this SQLAlchemy to interact with 2 different databases: Spanner and PostgreSQL.

The code used to work fine using sqlalchemy-spanner ==0.1.0 but after upgrading to sqlalchemy-spanner ==1.0.0, I observe the following error:

Exception during reset or similar
Traceback (most recent call last):
  File "/home/hadim/conda-env/lib/python3.9/site-packages/sqlalchemy/pool/base.py", line 697, in _finalize_fairy
    fairy._reset(pool)
  File "/home/hadim/conda-env/lib/python3.9/site-packages/sqlalchemy/pool/base.py", line 874, in _reset
    pool.dispatch.reset(self, self._connection_record)
  File "/home/hadim/conda-env/lib/python3.9/site-packages/sqlalchemy/event/attr.py", line 261, in __call__
    fn(*args, **kw)
  File "/home/hadim/conda-env/lib/python3.9/site-packages/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py", line 51, in reset_connection
    dbapi_conn.connection.staleness = None
AttributeError: 'psycopg2.extensions.connection' object has no attribute 'staleness'

The code is too complex to be posted here, but the error happens when I tried to interact with the non Spanner database: task = session.query(models.Task).get(task_uid) (nothing special here).

See below some lib versions that could be relevant:

python 3.9
psycopg2 2.9.3
sqlalchemy 1.3.23

and the piece of code I used to instantiate the two alchemy engine:

import sqlalchemy as sa

from .settings import get_settings

# Regular DB
engine = sa.create_engine(get_settings().backend_database_url, connect_args=connect_args)
SessionLocal = sa.orm.sessionmaker(autocommit=False, autoflush=False, bind=engine)

# Spanner DB
spanner_engine = sa.create_engine(get_settings().spanner_db_url)
SpannerSessionLocal = sa.orm.sessionmaker(autocommit=False, autoflush=False, bind=spanner_engine)

Everything used to work well with 0.1.0.

[hadim]

Please confirm it seems to be related to #186 and fix by #188 ?

[aderk]

Until a new release is cut, you can hotfix with the following code after importing sqlalchemy but before executing any DB queries:

import sqlalchemy as sa
import google.cloud.sqlalchemy_spanner.sqlalchemy_spanner as sa_spanner
# Remove listener and re-add in sqlalchemy to include fix for:
# https://github.com/googleapis/python-spanner-sqlalchemy/issues/186
# TODO: Remove after upgrading sqlalchemy-spanner > 1.0.0
sa.event.remove(sa.pool.Pool, 'reset', sa_spanner.reset_connection)
@sa.event.listens_for(sa.pool.Pool, 'reset')
def reset_connection(dbapi_conn, connection_record):  # type: ignore
    """An event of returning a connection back to a pool."""
    if getattr(dbapi_conn.connection, 'staleness', None) is not None:
        dbapi_conn.connection.staleness = None

Edit: I am not a maintainer, but yes, this is #186 and #188 fixes it.

[hadim]

Thanks for the snippet.

I'll close here then.