def db(connection: Connection) -> Generator:
transaction = connection.begin()
session = TestSession(bind=connection)
try:
yield session
finally:
transaction.rollback()
session.close()
def test_pgsql_db(
pgsql_connection: Connection,
request) -> Tuple[SqlLmsOperations, Connection, ConnectionSettings]:
"""
Fixture that takes the set-up connection and wraps in a transaction. Transaction
will be rolled-back automatically after each test.
Returns both a plain transaction-wrapped Connection and a monkey-patched
SqlLmsOperations that uses that Connection. They may be used interchangeably.
"""
# Wrap connection in transaction
transaction: Transaction = pgsql_connection.begin()
# Rollback transaction in finalizer when test is done
request.addfinalizer(lambda: transaction.rollback())
# New version of _exec using our transaction
def replace_exec(self: SqlLmsOperations, statement: str) -> int:
result = pgsql_connection.execute(statement)
if result:
return int(result.rowcount)
return 0
# New version of insert_into_staging using our transaction
def replace_insert_into_staging(self: SqlLmsOperations, df: DataFrame,
table: str):
copy = df.copy()
copy.columns = copy.columns.str.lower()
copy.to_sql(
f"stg_{table}",
con=pgsql_connection,
schema="lms",
if_exists="append",
index=False,
method="multi",
chunksize=120,
)
# Monkey-patch SqlLmsOperations to use our transaction
SqlLmsOperations._exec = replace_exec # type:ignore
SqlLmsOperations.insert_into_staging = replace_insert_into_staging # type:ignore
# Initialize monkey-patched adapter with a dummy engine, doesn't need a real one now
adapter: SqlLmsOperations = SqlLmsOperations(MagicMock())
adapter.engine = DbEngine.POSTGRESQL
return (adapter, pgsql_connection, GetSettings(request))
def refresh_and_diff_materialized_view(
connection: Connection,
view: Table,
copy: Table,
result_columns: Iterable[Column],
) -> Tuple[List[Tuple], List[Tuple], List[Tuple]]:
"""Lock the given `view` with an advisory lock, create a temporary table
of the view, refresh the view and compute the difference.
:param connection: A valid SQLAlchemy connection
:param view: The view to refresh and diff
:param copy: A temporary table to create and diff
:param result_columns: The columns to return
:return: A 3-tuple containing three lists of tuples of the `result_columns`
of added, deleted and modified records due to the refresh.
"""
with connection.begin():
lock_table(connection, view)
create_temp_copy(connection, view, copy)
refresh_materialized_view(connection, view)
return diff_tables(connection, view, copy, result_columns)
