def test_add_new_column(engine, schema):
# config
table_name = TableNames.ADD_NEW_COLUMN
dtype = {
'id': VARCHAR(5)
} if 'mysql' in engine.dialect.dialect_description else None
df = pd.DataFrame({'id': ['foo']}).set_index('id')
# common kwargs for all the times we use upsert_or_aupsert
common_kwargs = dict(con=engine,
schema=schema,
table_name=table_name,
if_row_exists='update',
dtype=dtype)
# 1. create table
upsert_or_aupsert(df=df, **common_kwargs)
# 2. add a new column and repeat upsert_or_aupsert
df['new_column'] = 'bar'
upsert_or_aupsert(df=df, add_new_columns=True, **common_kwargs)
# verify content matches
df_db = select_table(engine=engine,
schema=schema,
table_name=table_name,
index_col='id').sort_index()
pd.testing.assert_frame_equal(df, df_db)
def test_mysql_pk_not_auto_incremented(engine, schema):
if 'mysql' not in engine.dialect.dialect_description:
pytest.skip('This test is only relevant for MySQL')
table_name = TableNames.PK_MYSQL
# upsert first df using pangres which creates the table automatically
df1 = pd.DataFrame({'id': [0, 1], 'name': ['foo', 'bar']}).set_index('id')
upsert_or_aupsert(con=engine,
df=df1,
table_name=table_name,
if_row_exists='update')
# upsert second df
df2 = pd.DataFrame({
'id': [100, 200],
'name': ['baz', 'qux']
}).set_index('id')
upsert_or_aupsert(con=engine,
df=df2,
table_name=table_name,
if_row_exists='update')
# read df back
df_db = select_table(engine=engine,
schema=schema,
table_name=table_name,
index_col='id')
# check mysql got that correctly
pd.testing.assert_frame_equal(df_db.sort_index(),
pd.concat((df1, df2)).sort_index())
def test_commit_as_you_go(engine, schema):
df = pd.DataFrame(index=pd.Index(['foo'], name='ix'))
table_name = TableNames.COMMIT_AS_YOU_GO
# common keyword arguments for multiple upsert operations below
common_kwargs = dict(schema=schema,
table_name=table_name,
if_row_exists='update',
dtype={'ix': VARCHAR(3)})
with engine.connect() as con:
# skip for sqlalchemy < 2.0 or when future=True flag is not passed
# during engine creation (commit-as-you-go is a new feature)
# when this is the case there is no attribute commit or rollback for
# the connection
if not hasattr(con, 'commit'):
pytest.skip(
'test not possible because there is no attribute "commit" (most likely sqlalchemy < 2)'
)
# do some random upsert operation and commit
upsert(con=con, df=df, **common_kwargs)
con.commit()
# do some other operation that requires commit and then rollback
upsert(con=con, df=df.rename(index={'foo': 'bar'}), **common_kwargs)
con.rollback()
# the table in the db should be equal to the initial df as the second
# operation was rolled back
df_db = select_table(engine=engine,
schema=schema,
table_name=table_name,
index_col='ix')
pd.testing.assert_frame_equal(df_db, df)
async def test_transaction_async(engine, schema, trans_op=None):
df = pd.DataFrame(index=pd.Index(['foo'], name='ix'))
table_name = TableNames.COMMIT_OR_ROLLBACK_TRANS
# common keyword arguments for multiple upsert operations below
common_kwargs = dict(schema=schema,
table_name=table_name,
if_row_exists='update',
dtype={'ix': VARCHAR(3)})
async with engine.connect() as con:
trans = await con.begin()
try:
# do some random upsert operation
await aupsert(con=con, df=df, **common_kwargs)
# do some other operation that requires commit
await aupsert(con=con,
df=df.rename(index={'foo': 'bar'}),
**common_kwargs)
coro = getattr(trans, trans_op) # commit or rollback
await coro()
finally:
await trans.close()
# if trans_op=='commit': make sure we have "bar" and "foo" in the index
# elif trans_op=='rollback': make sure we don't have any data
# or that the table was not even created (what is rolled back
# depends on the database type and other factors)
if trans_op == 'commit':
df_db = select_table(engine=engine,
schema=schema,
table_name=table_name,
index_col='ix')
pd.testing.assert_frame_equal(
df_db.sort_index(),
pd.DataFrame(index=pd.Index(['bar', 'foo'], name='ix')))
elif trans_op == 'rollback':
df_db = select_table(engine=engine,
schema=schema,
table_name=table_name,
error_if_missing=False)
# no table or an empty table
assert df_db is None or len(df_db) == 0
def test_only_index(engine, schema, if_row_exists):
# config
table_name = TableNames.INDEX_ONLY_INSERT
# upsert df with only index
df = pd.DataFrame({'ix': [1]}).set_index('ix')
upsert_or_aupsert(con=engine,
schema=schema,
df=df,
table_name=table_name,
if_row_exists=if_row_exists)
# check data integrity
df_db = select_table(engine=engine, schema=schema, table_name=table_name)
assert 'ix' in df_db.columns
assert len(df_db) > 0
assert df_db['ix'].iloc[0] == 1
def read_from_db():
df_db = select_table(engine=engine,
schema=schema,
table_name=table_name)
df_db['ix3'] = pd.to_datetime(df_db['ix3'])
return df_db.set_index(index_col)
def read_from_db():
return select_table(engine=engine, schema=schema,
table_name=table_name, index_col='row_id')