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 and rolling back the `session`, without closing it. `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)