def write_to_table(repo: Dolt,
table: Table,
data: List[dict],
commit: bool = False,
message: str = None):
"""
Given a repo, table, and data, will try and use the repo's MySQL Server instance to write the provided data to the
table. Since Dolt does not yet support ON DUPLICATE KEY clause to INSERT statements we also have to separate
updates from inserts and run sets of statements.
:param repo:
:param table:
:param data:
:param commit:
:param message:
:return:
"""
coerced_data = list(clean_types(data))
inserts, updates = get_inserts_and_updates(repo.engine, table,
coerced_data)
if inserts:
logger.info('Inserting {} rows'.format(len(inserts)))
with repo.engine.connect() as conn:
conn.execute(table.insert(), inserts)
# We need to prefix the columns with "_" in order to use bindparam properly
from copy import deepcopy
_updates = deepcopy(updates)
for dic in _updates:
for col in list(dic.keys()):
dic['_{}'.format(col)] = dic.pop(col)
if _updates:
logger.info('Updating {} rows'.format(len(_updates)))
with repo.engine.connect() as conn:
statement = table.update()
for pk_col in [
col.name for col in table.columns if col.primary_key
]:
statement = statement.where(
table.c[pk_col] == bindparam('_{}'.format(pk_col)))
non_pk_cols = [
col.name for col in table.columns if not col.primary_key
]
statement = statement.values(
{col: bindparam('_{}'.format(col))
for col in non_pk_cols})
conn.execute(statement, _updates)
if commit:
repo.add(str(table.name))
message = message or 'Inserting {} records at '.format(
len(data), datetime.now())
repo.commit(message)
def inner(repo: Dolt):
current_branch, current_branch_list = repo.branch()
original_branch = current_branch.name
if branch != original_branch and not commit:
raise ValueError(
'If writes are to another branch, and commit is not True, writes will be lost'
)
if current_branch.name != branch:
logger.info('Current branch is {}, checking out {}'.format(
current_branch.name, branch))
if branch not in [b.name for b in current_branch_list]:
logger.info('{} does not exist, creating'.format(branch))
repo.branch(branch_name=branch)
repo.checkout(branch)
if transaction_mode:
raise NotImplementedError(
'transaction_mode is not yet implemented')
tables_updated = [writer(repo) for writer in writers]
if commit:
if not repo.status().is_clean:
logger.info(
'Committing to repo located in {} for tables:\n{}'.format(
repo.repo_dir, tables_updated))
for table in tables_updated:
repo.add(table)
repo.commit(message)
else:
logger.warning('No changes to repo in:\n{}'.format(
repo.repo_dir))
current_branch, branches = repo.branch()
if original_branch != current_branch.name:
logger.info(
'Checked out {} from {}, checking out {} to restore state'.
format([b.name for b in branches], original_branch,
original_branch))
repo.checkout(original_branch)
return branch
def _create_table_from_schema_import_helper(
repo: Dolt,
table: str,
pks: List[str],
path: str,
transformers: List[DataframeTransformer] = None,
commit: bool = True,
commit_message: str = None):
if transformers:
fp = tempfile.NamedTemporaryFile(suffix='.csv')
temp = pd.read_csv(path)
transformed = _apply_df_transformers(temp, transformers)
transformed.to_csv(fp.name, index=False)
path = fp.name
repo.schema_import(table=table, pks=pks, filename=path, create=True)
if commit:
message = commit_message or 'Creating table {}'.format(table)
repo.add(table)
repo.commit(message)
