Mercurial > hg > sqlex
view sqlex/main.py @ 9:834b920ae345 default tip
allow output of headers in csv
author | Jeff Hammel <k0scist@gmail.com> |
---|---|
date | Sat, 01 Apr 2017 15:11:34 -0700 |
parents | adf056d67c01 |
children |
line wrap: on
line source
#!/usr/bin/env python # -*- coding: utf-8 -*- """ sql(ite) explorer/exporter """ # imports import argparse import csv import os import sys from .model import SQLEx def ensure_dir(directory): """ensure a directory exists""" if os.path.exists(directory): if not os.path.isdir(directory): raise OSError("Not a directory: '{}'".format(directory)) return directory os.makedirs(directory) return directory class SQLExParser(argparse.ArgumentParser): """CLI option parser""" def __init__(self, **kwargs): kwargs.setdefault('formatter_class', argparse.RawTextHelpFormatter) kwargs.setdefault('description', __doc__) argparse.ArgumentParser.__init__(self, **kwargs) self.add_argument('db', help="sqlite `.db` file") self.add_argument('table', nargs='?', help="table to operate on") self.add_argument('--tables', '--list-tables', dest='list_tables', action='store_true', default=False, help="list tables and exit") self.add_argument('--columns', '--list-columns', dest='list_columns', action='store_true', default=False, help="list columns in `table` and exit") self.add_argument('-o', '--output', help="output to directory (if `table` not given), or filename or stdout by default") self.add_argument('--header', dest='header', action='store_true', default=False, help="export header as first row") self.options = None def parse_args(self, *args, **kw): options = argparse.ArgumentParser.parse_args(self, *args, **kw) self.validate(options) self.options = options return options def validate(self, options): """validate options""" try: open(options.db).close() except Exception as e: self.error("Could not open '{}': {}".format(options.db, e)) if not any((options.output, options.list_tables, options.table)): self.error("`--output` directory must be specified to output entire database") if options.list_columns and not options.table: self.error("`--list-columns` requires `table`") def main(args=sys.argv[1:]): """CLI""" # parse command line options parser = SQLExParser() options = parser.parse_args(args) # connect to database db = SQLEx(options.db) if options.list_tables: # list tables and return # if `table` argument is provided, exit 1 # if not available. Otherwise exit 0 tables = db.tables() print ('\n'.join(tables)) retval = 0 if options.table: retval = int(options.table not in tables) return retval if options.table: # ensure selected table exists if options.table not in db.tables(): parser.error("No table '{}' in {} tables:\n{}".format(options.table, options.db, ', '.join(db.tables()))) if options.list_columns: # list columns and return print ('\n'.join(db.columns(options.table).keys())) return if options.table: # output table if options.output: with open(options.output, 'w') as f: db.table2csv(options.table, f, header=options.header) else: db.table2csv(options.table, sys.stdout, header=options.header) sys.stdout.flush() else: # output entire db to CSV files in directory # ensure directory exists ensure_dir(options.output) for table in db.tables(): # export each table path = os.path.join(options.output, '{}.csv'.format(table)) with open(path, 'w') as f: db.table2csv(table, f, header=options.header) if __name__ == '__main__': main()