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 retry_transaction(
session: Session,
measurements: MeasurementMap,
txn_body: Callable[[Session], bool],
max_retries: Optional[int],
) -> bool:
"""Retries the transaction if a serialization failure occurs.
Handles management of committing, rolling back, and closing the `session`. `txn_body` can return False to force the
transaction to be aborted, otherwise return True.
Returns:
True, if the transaction succeeded.
False, if the transaction was aborted by `txn_body`.
"""
num_retries = 0
try:
while True:
try:
should_continue = txn_body(session)
if not should_continue:
session.rollback()
return should_continue
session.commit()
return True
except sqlalchemy.exc.DBAPIError as e:
session.rollback()
if max_retries and num_retries >= max_retries:
raise
if (isinstance(e.orig, psycopg2.OperationalError)
and e.orig.pgcode == SERIALIZATION_FAILURE):
logging.info(
"Retrying transaction due to serialization failure: %s",
e)
num_retries += 1
continue
raise
except Exception:
session.rollback()
raise
finally:
measurements.measure_int_put(m_retries, num_retries)
session.close()