def test_get_nb_rows(engine, schema):
# config
table_name = TableNames.WITH_YIELD
nb_rows, chunksize = 20, 3
nb_last_chunk = nb_rows % chunksize
nb_chunks = math.ceil(nb_rows / chunksize)
# MySQL does not want flexible text length in indices/PK
dtype = {
'profileid': VARCHAR(10)
} if 'mysql' in engine.dialect.dialect_description else None
df = _TestsExampleTable.create_example_df(nb_rows=nb_rows)
# iterate over upsert results
# make sure we can extract the number of updated rows and that it is correct
iterator = upsert(con=engine,
df=df,
table_name=table_name,
if_row_exists='update',
schema=schema,
chunksize=chunksize,
dtype=dtype,
yield_chunks=True)
for ix, result in enumerate(iterator):
assert result.rowcount == (chunksize
if ix != nb_chunks - 1 else nb_last_chunk)
# verify the inserted data is as expected
# we sort the index for MySQL
df_db = read_example_table_from_db(engine=engine,
schema=schema,
table_name=table_name)
pd.testing.assert_frame_equal(df.sort_index(), df_db.sort_index())
def test_methods_and_attributes(engine, schema):
"""
Makes tests in chain for PandasSpecialEngine.
This makes code easier with a small inconvenient
of not separating tests thematically.
"""
# dtype for index for MySQL... (can't have flexible text length)
dtype = {
'profileid': VARCHAR(5)
} if 'mysql' in engine.dialect.dialect_description else None
table_name = 'test_pandas_special_engine'
default_args = {
'engine': engine,
'schema': schema,
'dtype': dtype,
'table_name': table_name
}
drop_table_if_exists(engine=engine, schema=schema, table_name=table_name)
# TEST INIT
df = _TestsExampleTable.create_example_df(nb_rows=10)
pse = PandasSpecialEngine(df=df, **default_args)
# TEST ATTRIBUTE pse.table
expected_cols = list(df.index.names) + df.columns.tolist()
assert all((col in pse.table.columns for col in expected_cols))
# TEST TABLE AND SCHEMA CREATION
pse.create_schema_if_not_exists()
pse.create_table_if_not_exists()
# TEST ADD NEW COLUMNS
# don't try to add JSON columns!
# It's not supported by sqlalchemy compilers :(
df = df.assign(
new_text_col='test',
new_int_col=0,
new_float_col=1.1,
new_bool_col=False,
new_dt_col=pd.Timestamp('2020-01-01'),
# create this col for later
empty_col=None)
# recreate pse then add columns
pse = PandasSpecialEngine(df=df, **default_args)
pse.add_new_columns()
# TEST CHANGE COLUMN TYPE (not for SQlite)
if not pse._db_type == 'sqlite':
# don't try to alter from any type to JSON!
# It's not supported by sqlalchemy compilers :(
# also the order is very specific!!! we have to cast types
# even though the column is empty (to avoid losing column information such as constraint)
# so e.g. casting from BOOLEAN to BIGINT is not possible
# actually BOOLEAN is not even here because it breaks the chain
alterations = (1, 1.1, "abc", pd.Timestamp("2020-01-01", tz='UTC'))
for i in alterations:
# change empty_col
df['empty_col'] = df['empty_col'].map(
lambda x: i) # this will work for lists or dicts as well
# recreate pse then change column type
pse = PandasSpecialEngine(df=df, **default_args)
pse.adapt_dtype_of_empty_db_columns()
def test_table_attr(engine, schema):
# generate a somewhat complex table model via the _TestsExampleTable class
df = _TestsExampleTable.create_example_df(nb_rows=10)
table_name = TableNames.NO_TABLE
with sync_async_connect_switch(engine) as connection:
pse = PandasSpecialEngine(connection=connection,
schema=schema,
table_name=table_name,
df=df)
# make sure columns and table name match
expected_cols = list(df.index.names) + df.columns.tolist()
assert all((col in pse.table.columns for col in expected_cols))
assert pse.table.name == table_name
async def test_table_creation_async(engine, schema):
dtype = {
'profileid': VARCHAR(5)
} if 'mysql' in engine.dialect.dialect_description else None
df = _TestsExampleTable.create_example_df(nb_rows=10)
async with engine.connect() as connection:
pse = PandasSpecialEngine(connection=connection,
schema=schema,
dtype=dtype,
table_name=TableNames.TABLE_CREATION,
df=df)
assert not await pse.atable_exists()
await pse.acreate_table_if_not_exists()
await connection.commit()
assert await pse.atable_exists()
def insert_chunks(engine, schema, chunksize, nb_rows):
df = _TestsExampleTable.create_example_df(nb_rows=nb_rows)
# MySQL does not want flexible text length in indices/PK
dtype = {
'profileid': VARCHAR(10)
} if 'mysql' in engine.dialect.dialect_description else None
upsert_or_aupsert(schema=schema,
table_name=TableNames.VARIOUS_CHUNKSIZES,
df=df,
chunksize=chunksize,
con=engine,
if_row_exists='update',
dtype=dtype)
df_db = read_example_table_from_db(
engine=engine, schema=schema, table_name=TableNames.VARIOUS_CHUNKSIZES)
# sort index (for MySQL...)
pd.testing.assert_frame_equal(df.sort_index(), df_db.sort_index())
def insert_chunks(engine, schema, chunksize, nb_rows):
df = _TestsExampleTable.create_example_df(nb_rows=nb_rows)
table_name = f'test_insert_chunksize_{chunksize}'
drop_table_if_exists(engine=engine, schema=schema, table_name=table_name)
upsert(
schema=schema,
table_name=table_name,
df=df,
chunksize=chunksize,
engine=engine,
if_row_exists='update',
# MySQL does not want flexible text length in indices/PK
dtype={'profileid': VARCHAR(10)}
if 'mysql' in engine.dialect.dialect_description else None)
df_db = read_example_table_from_db(engine=engine,
schema=schema,
table_name=table_name)
# sort index (for MySQL...)
pd.testing.assert_frame_equal(df.sort_index(), df_db.sort_index())
def test_upsert_speed(engine, schema, benchmark, library, nb_rows, rounds, iterations, if_row_exists):
assert library == 'pangres' # in case pandas changes and we forget to update the tests
# skip async engines with pandas
if is_async_sqla_obj(engine) and library == 'pandas':
pytest.skip('async engines will not work with pandas')
# get a df
df = _TestsExampleTable.create_example_df(nb_rows=nb_rows).drop(columns=['favorite_colors'])
# setup for test (create table with no rows)
def setup():
create_or_upsert_with_pangres(engine=engine, schema=schema, if_row_exists=if_row_exists,
df=df.head(0), chunksize=nb_rows)
# test func
# insert update/ignore with `create_table=False` to maximise speed
func = lambda: create_or_upsert_with_pangres(engine=engine, schema=schema, if_row_exists=if_row_exists,
df=df, chunksize=nb_rows, create_table=False)
benchmark.pedantic(func, setup=setup, rounds=rounds, iterations=iterations)
def test_create_and_insert_speed(engine, schema, benchmark, library, nb_rows, rounds, iterations):
# skip async engines with pandas
if is_async_sqla_obj(engine) and library == 'pandas':
pytest.skip('async engines will not work with pandas')
# get a df
# we don't test JSON as this is problematic with pandas
df = _TestsExampleTable.create_example_df(nb_rows=nb_rows).drop(columns=['favorite_colors'])
# prepare funcs for benchmark and then do the benchmark
switch = {'pangres':lambda: create_or_upsert_with_pangres(engine=engine, schema=schema, if_row_exists='update',
df=df, chunksize=nb_rows),
'pandas':lambda: create_with_pandas(engine=engine, schema=schema, df=df)}
try:
benchmark.pedantic(switch[library], setup=lambda: drop_table(engine=engine, schema=schema,
table_name=TableNames.BENCHMARK),
rounds=rounds, iterations=iterations)
except NotImplementedError as e:
if 'not implemented for SQLAlchemy 2' in str(e):
pytest.skip('in Python 3.6 there is some kind of problem with engines created with '
'`future=True` flag and pandas')
from pangres.tests.conftest import read_example_table_from_db, drop_table_if_exists
# # Config
table_name = 'test_upsert'
default_args = {'table_name':table_name,
'create_schema':True,
'add_new_columns':True,
'adapt_dtype_of_empty_db_columns':False}
# # Test data
# +
df = _TestsExampleTable.create_example_df(nb_rows=5)
# test for NULL values except for boolean column
df.iloc[0,[ix for ix, col in enumerate(df.columns) if col != 'likes_pizza']] = None
# test for update
df2 = _TestsExampleTable.create_example_df(nb_rows=6)
# test for ignore
df3 = _TestsExampleTable.create_example_df(nb_rows=6)
# -
# # Tests
# ORDER MATTERS!
async def test_add_new_columns_async(engine, schema, axis=None):
# store arguments we will use for multiple PandasSpecialEngine instances
table_name = TableNames.ADD_NEW_COLUMN
common_kwargs = dict(schema=schema, table_name=table_name)
common_kwargs['dtype'] = {
'profileid': VARCHAR(5)
} if 'mysql' in engine.dialect.dialect_description else None
# create our example table
df = _TestsExampleTable.create_example_df(nb_rows=10)
async with engine.connect() as connection:
pse = PandasSpecialEngine(connection=connection,
df=df,
**common_kwargs)
await pse.acreate_table_if_not_exists()
await connection.commit()
assert await pse.atable_exists()
# we need to recreate an instance of PandasSpecialEngine
# so that a new table model with the new columns is created then add columns
async with engine.connect() as connection:
# error message if we get unexpected values for "axis"
# or we make a typo in our if/elif statements
err_msg = f'Expected axis to be one of index, columns. Got {axis}'
# add a new index level or new columns (no JSON ones,
# it's not supported by sqlalchemy compilers :( )
if axis == 'index':
df['new_index_col'] = 'foo'
df.set_index('new_index_col', append=True, inplace=True)
elif axis == 'column':
df = df.assign(
new_text_col='test',
new_int_col=0,
new_float_col=1.1,
new_bool_col=False,
new_dt_col=pd.Timestamp('2020-01-01'),
# create this col for later
empty_col=None)
else:
raise AssertionError(err_msg)
# recreate PandasSpecialEngine
pse = PandasSpecialEngine(connection=connection,
df=df,
**common_kwargs)
# check if we get an error when trying to add an index level
if axis == 'index':
with pytest.raises(MissingIndexLevelInSqlException) as exc_info:
await pse.aadd_new_columns()
assert 'Cannot add' in str(exc_info.value)
return
elif axis == 'column':
await pse.aadd_new_columns()
await connection.commit()
else:
raise AssertionError(err_msg)
# check if the columns were correctly added
# since we issued a return for 'index' earlier
# the axis must now be 'columns'
assert axis == 'column'
# check the columns where added
sync_engine = async_engine_to_sync_engine(engine)
with sync_engine.connect() as connection:
ns = get_table_namespace(schema=schema, table_name=table_name)
df_db = pd.read_sql(text(f'SELECT * FROM {ns} LIMIT 0;'),
con=connection,
index_col='profileid')
assert set(df.columns) == set(df_db.columns)