Ejemplo n.º 1
0
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])
Ejemplo n.º 2
0
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)