Skip to content

Commit

Permalink
Disable sqlalchemy 2 workaround for pandas 2 (DM-38553)
Browse files Browse the repository at this point in the history
Pandas 2.0 has a better support for sqlalchemy 2. The workaround that
was needed for pandas 1 is no longer necessary and is conditionally
disabled. When we migrate to pandas2 completely, that workaround can be
removed.
  • Loading branch information
andy-slac committed Apr 3, 2023
1 parent 0340944 commit f394890
Show file tree
Hide file tree
Showing 2 changed files with 39 additions and 29 deletions.
62 changes: 36 additions & 26 deletions python/lsst/dax/apdb/apdbSql.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,42 +50,52 @@
_LOG = logging.getLogger(__name__)


class _ConnectionHackSA2(sqlalchemy.engine.Connectable):
"""Terrible hack to workaround Pandas incomplete support for sqlalchemy 2.
if pandas.__version__.partition(".")[0] == "1":

We need to pass a Connection instance to pandas method, but in SA 2 the
Connection class lost ``connect`` method which is used by Pandas.
"""
class _ConnectionHackSA2(sqlalchemy.engine.Connectable):
"""Terrible hack to workaround Pandas 1 incomplete support for
sqlalchemy 2.
def __init__(self, connection: sqlalchemy.engine.Connection):
self._connection = connection
We need to pass a Connection instance to pandas method, but in SA 2 the
Connection class lost ``connect`` method which is used by Pandas.
"""

def connect(self) -> Any:
return self
def __init__(self, connection: sqlalchemy.engine.Connection):
self._connection = connection

@property
def execute(self) -> Callable:
return self._connection.execute
def connect(self, **kwargs: Any) -> Any:
return self

@property
def execution_options(self) -> Callable:
return self._connection.execution_options
@property
def execute(self) -> Callable:
return self._connection.execute

@property
def connection(self) -> Any:
return self._connection.connection
@property
def execution_options(self) -> Callable:
return self._connection.execution_options

def __enter__(self) -> sqlalchemy.engine.Connection:
return self._connection
@property
def connection(self) -> Any:
return self._connection.connection

def __exit__(self, type_: Any, value: Any, traceback: Any) -> None:
# Do not close connection here
pass
def __enter__(self) -> sqlalchemy.engine.Connection:
return self._connection

def __exit__(self, type_: Any, value: Any, traceback: Any) -> None:
# Do not close connection here
pass

@inspection._inspects(_ConnectionHackSA2)
def _connection_insp(conn: _ConnectionHackSA2) -> Inspector:
return Inspector._construct(Inspector._init_connection, conn._connection)

else:

@inspection._inspects(_ConnectionHackSA2)
def _connection_insp(conn: _ConnectionHackSA2) -> Inspector:
return Inspector._construct(Inspector._init_connection, conn._connection)
# Pandas 2.0 supports SQLAlchemy 2 correctly.
def _ConnectionHackSA2( # type: ignore[no-redef]
conn: sqlalchemy.engine.Connectable,
) -> sqlalchemy.engine.Connectable:
return conn


def _coerce_uint64(df: pandas.DataFrame) -> pandas.DataFrame:
Expand Down
6 changes: 3 additions & 3 deletions python/lsst/dax/apdb/apdbSqlSchema.py
Original file line number Diff line number Diff line change
Expand Up @@ -65,10 +65,10 @@ class GUID(sqlalchemy.TypeDecorator):

impl = sqlalchemy.CHAR

cache_ok = True
cache_ok = True # type: ignore[assignment]

def load_dialect_impl(self, dialect: sqlalchemy.engine.Dialect) -> sqlalchemy.types.TypeEngine:
if dialect.name == "postgresql":
if dialect.name == "postgresql": # type: ignore[attr-defined]
return dialect.type_descriptor(UUID())
else:
return dialect.type_descriptor(sqlalchemy.CHAR(32))
Expand All @@ -89,7 +89,7 @@ def process_bind_param(self, value: Any, dialect: sqlalchemy.engine.Dialect) ->
elif not isinstance(value, uuid.UUID):
raise TypeError(f"Unexpected type of a bind value: {type(value)}")

if dialect.name == "postgresql":
if dialect.name == "postgresql": # type: ignore[attr-defined]
return str(value)
else:
return "%.32x" % value.int
Expand Down

0 comments on commit f394890

Please sign in to comment.