def _import_helper(repo: Dolt, table_name: str,
write_import_file: Callable[[str], None],
primary_keys: List[str], import_mode: str) -> None:
import_modes = IMPORT_MODES_TO_FLAGS.keys()
if import_mode is not None:
assert import_mode in import_modes, 'update_mode must be one of: {}'.format(
import_modes)
else:
if table_name in [table.name for table in repo.ls()]:
logger.info(
'No import mode specified, table exists, using "{}"'.format(
UPDATE))
import_mode = UPDATE
else:
logger.info(
'No import mode specified, table exists, using "{}"'.format(
CREATE))
import_mode = CREATE
import_flags = IMPORT_MODES_TO_FLAGS[import_mode]
logger.info(
'Importing to table {} in dolt directory located in {}, import mode {}'
.format(table_name, repo.repo_dir(), import_mode))
fp = tempfile.NamedTemporaryFile(suffix='.csv')
write_import_file(fp.name)
args = [
'table', 'import', table_name, '--pk={}'.format(','.join(primary_keys))
] + import_flags
repo.execute(args + [fp.name])
def read_table(repo: Dolt, table_name: str, delimiter: str = ',') -> pd.DataFrame:
"""
Reads the contents of a table and returns it as a Pandas `DataFrame`. Under the hood this uses export and the
filesystem, in short order we are likley to replace this with use of the MySQL Server.
:param repo:
:param table_name:
:param delimiter:
:return:
"""
fp = tempfile.NamedTemporaryFile(suffix='.csv')
repo.execute(['table', 'export', table_name, fp.name, '-f'])
result = pd.read_csv(fp.name, delimiter=delimiter)
return result
def _import_helper(repo: Dolt, table_name: str,
write_import_file: Callable[[str], None],
primary_keys: List[str], import_mode: str) -> None:
import_modes = IMPORT_MODES_TO_FLAGS.keys()
if import_mode is not None:
assert import_mode in import_modes, 'update_mode must be one of: {}'.format(
import_modes)
else:
if table_name in [table.name for table in repo.ls()]:
logger.info(
'No import mode specified, table exists, using "{}"'.format(
UPDATE))
import_mode = UPDATE
else:
logger.info(
'No import mode specified, table exists, using "{}"'.format(
CREATE))
import_mode = CREATE
if import_mode == CREATE and primary_keys is None:
raise ValueError(
'Import mode CREATE requires a primary key to be specified')
import_flags = IMPORT_MODES_TO_FLAGS[import_mode]
logger.info(
'Importing to table {} in dolt directory located in {}, import mode {}'
.format(table_name, repo.repo_dir(), import_mode))
fname = tempfile.mktemp(suffix='.csv')
try:
write_import_file(fname)
args = ['table', 'import', table_name] + import_flags
if import_mode == CREATE:
args += ['--pk={}'.format(','.join(primary_keys))]
repo.execute(args + [fname])
finally:
if os.path.exists(fname):
os.remove(fname)
