def for_prod_data_client(
cls,
database_key: SQLAlchemyDatabaseKey,
ssl_cert_path: str,
*,
autocommit: bool = True,
) -> Iterator[Session]:
"""Implements a context manager for db sessions for use in prod-data-client."""
engine = SQLAlchemyEngineManager.get_engine_for_database_with_ssl_certs(
database_key=database_key, ssl_cert_path=ssl_cert_path)
if engine is None:
raise ValueError(f"No engine set for key [{database_key}]")
try:
session = Session(bind=engine)
cls._alter_session_variables(session)
cls._apply_session_listener_for_schema_base(
database_key.declarative_meta, session)
yield session
if autocommit:
try:
session.commit()
except Exception as e:
session.rollback()
raise e
finally:
session.close()
def _execute_statement(self, statement: str) -> None:
session = Session(bind=self.postgres_engine)
try:
session.execute(statement)
session.commit()
except Exception as e:
logging.warning("Failed to cleanup: %s", e)
session.rollback()
finally:
session.close()
def for_schema_base(cls, schema_base: DeclarativeMeta) -> Session:
engine = SQLAlchemyEngineManager.get_engine_for_schema_base(
schema_base)
if engine is None:
raise ValueError(
f"No engine set for base [{schema_base.__name__}]")
session = Session(bind=engine)
cls._apply_session_listener_for_schema_base(schema_base, session)
return session
def _for_database(cls, database_key: SQLAlchemyDatabaseKey) -> Session:
# TODO(#8046): When the above method is deleted, move this into `using_database`
# directly.
engine = SQLAlchemyEngineManager.get_engine_for_database(
database_key=database_key)
if engine is None:
raise ValueError(f"No engine set for key [{database_key}]")
session = Session(bind=engine)
cls._alter_session_variables(session)
cls._apply_session_listener_for_schema_base(
database_key.declarative_meta, session)
return session