コード例 #1
0
    # i.e. during the combination playbook).
    # Otherwise create the loader script here.
    # If we replace the output file the loader script is something
    # that has to be generated separately.
    if args.replace_input:
        logger.info('Replacing input (%s -> %s)...',
                    augmented_file_path, args.input_nodes)
        # Remove the original input file and repl;ace it with the output
        os.remove(args.input_nodes)
        os.rename(augmented_file_path, args.input_nodes)
    else:
        # Before we finish,
        # write a convenient loader script
        # for all the files we generated...
        logger.info('Writing load script...')
        write_load_script(args.output, generated_files)

    # Finish by writing the expected edges header file...
    edges_header_file = open(os.path.join(args.output, EDGES_HDR_FILENAME), 'wt')
    edges_header_file.write(':START_ID(F2),:END_ID(F2),label,:TYPE\n')
    edges_header_file.close()

    # Now complete we write a "done" file to the output.
    # Processing may be time-consuming
    # so this file helps us avoid unnecessary re-processing on failure.
    # This can be used by the automation (ansible) framework to
    # decide whether processing was completed successfully.
    # If there's a 'done' file we can safely assume that processing
    # is complete.
    open(os.path.join(args.output, 'done'), 'a').close()
# directory content. We add every file that contains the word
# 'nodes' or 'edges' and ends in '.gz'.
#
# If the file is 'nodes.csv.gz' or 'edges.csv.gz' we expect a header
# and instead of just adding the file to the list we also add the header.
# So, if we find 'nodes.csv.gz' we insert 'nodes-header.csv,nodes.csv.gz'
# into the corresponding list.
generated_files = {'nodes': [], 'edges': []}
possible_files = os.listdir(args.path)
for possible_file in possible_files:
    if possible_file.endswith('.gz') and \
            os.path.isfile(os.path.join(args.path, possible_file)):
        if 'edges' in possible_file:
            actual_file = possible_file
            if actual_file == 'edges.csv.gz':
                actual_file = 'edges-header.csv,edges.csv.gz'
            generated_files['edges'].append(actual_file)
        elif 'nodes' in possible_file:
            actual_file = possible_file
            if actual_file == 'nodes.csv.gz':
                actual_file = 'nodes-header.csv,nodes.csv.gz'
            generated_files['nodes'].append(actual_file)

if not generated_files['nodes']:
    error('Could not find any nodes files.')
if not generated_files['edges']:
    error('Could not find any edges files.')

logger.info('Writing load script...')
write_load_script(args.path, generated_files)