def test_sql_tracker_store_with_login_db_db_already_exists( postgres_login_db_connection: sa.engine.Connection, ): postgres_login_db_connection.execution_options( isolation_level="AUTOCOMMIT" ).execute(f"CREATE DATABASE {POSTGRES_TRACKER_STORE_DB}") tracker_store = SQLTrackerStore( dialect="postgresql", host=POSTGRES_HOST, port=POSTGRES_PORT, username=POSTGRES_USER, password=POSTGRES_PASSWORD, db=POSTGRES_TRACKER_STORE_DB, login_db=POSTGRES_LOGIN_DB, ) matching_rows = ( postgres_login_db_connection.execution_options(isolation_level="AUTOCOMMIT") .execute( sa.text( "SELECT 1 FROM pg_catalog.pg_database WHERE datname = :database_name" ), database_name=POSTGRES_TRACKER_STORE_DB, ) .rowcount ) assert matching_rows == 1 tracker_store.engine.dispose()
def test_sql_tracker_store_with_login_db_race_condition( postgres_login_db_connection: sa.engine.Connection, caplog: LogCaptureFixture, monkeypatch: MonkeyPatch, ): original_execute = sa.engine.Connection.execute def mock_execute(self, *args, **kwargs): # this simulates a race condition if kwargs == {"database_name": POSTGRES_TRACKER_STORE_DB}: original_execute( self.execution_options(isolation_level="AUTOCOMMIT"), f"CREATE DATABASE {POSTGRES_TRACKER_STORE_DB}", ) return Mock(rowcount=0) else: return original_execute(self, *args, **kwargs) with monkeypatch.context() as mp: mp.setattr(sa.engine.Connection, "execute", mock_execute) with caplog.at_level(logging.ERROR): tracker_store = SQLTrackerStore( dialect="postgresql", host=POSTGRES_HOST, port=POSTGRES_PORT, username=POSTGRES_USER, password=POSTGRES_PASSWORD, db=POSTGRES_TRACKER_STORE_DB, login_db=POSTGRES_LOGIN_DB, ) # IntegrityError has been caught and we log the error assert any( [ f"Could not create database '{POSTGRES_TRACKER_STORE_DB}'" in record.message for record in caplog.records ] ) matching_rows = ( postgres_login_db_connection.execution_options(isolation_level="AUTOCOMMIT") .execute( sa.text( "SELECT 1 FROM pg_catalog.pg_database WHERE datname = :database_name" ), database_name=POSTGRES_TRACKER_STORE_DB, ) .rowcount ) assert matching_rows == 1 tracker_store.engine.dispose()
def _drop_db(connection: sa.engine.Connection, database_name: Text) -> None: connection.execution_options(isolation_level="AUTOCOMMIT").execute( f"DROP DATABASE IF EXISTS {database_name}")
def _create_login_db(connection: sa.engine.Connection) -> None: connection.execution_options(isolation_level="AUTOCOMMIT").execute( f"CREATE DATABASE {POSTGRES_LOGIN_DB}")