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