Mercurial > hg > sqlex
annotate 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 |
| rev | line source |
|---|---|
| 0 | 1 #!/usr/bin/env python |
| 2 # -*- coding: utf-8 -*- | |
| 3 | |
| 4 """ | |
| 5 sql(ite) explorer/exporter | |
| 6 """ | |
| 7 | |
| 8 # imports | |
| 9 import argparse | |
| 5 | 10 import csv |
| 0 | 11 import os |
| 12 import sys | |
| 3 | 13 from .model import SQLEx |
| 14 | |
| 0 | 15 |
| 16 def ensure_dir(directory): | |
| 17 """ensure a directory exists""" | |
| 18 if os.path.exists(directory): | |
| 19 if not os.path.isdir(directory): | |
| 20 raise OSError("Not a directory: '{}'".format(directory)) | |
| 21 return directory | |
| 22 os.makedirs(directory) | |
| 23 return directory | |
| 24 | |
| 25 | |
| 26 class SQLExParser(argparse.ArgumentParser): | |
| 27 """CLI option parser""" | |
| 28 | |
| 29 def __init__(self, **kwargs): | |
| 30 kwargs.setdefault('formatter_class', argparse.RawTextHelpFormatter) | |
| 31 kwargs.setdefault('description', __doc__) | |
| 32 argparse.ArgumentParser.__init__(self, **kwargs) | |
| 1 | 33 self.add_argument('db', |
| 34 help="sqlite `.db` file") | |
|
4
b440206930ac
stub table optional positional argument
Jeff Hammel <k0scist@gmail.com>
parents:
3
diff
changeset
|
35 self.add_argument('table', nargs='?', |
|
b440206930ac
stub table optional positional argument
Jeff Hammel <k0scist@gmail.com>
parents:
3
diff
changeset
|
36 help="table to operate on") |
| 3 | 37 self.add_argument('--tables', '--list-tables', dest='list_tables', |
| 38 action='store_true', default=False, | |
| 39 help="list tables and exit") | |
| 5 | 40 self.add_argument('--columns', '--list-columns', dest='list_columns', |
| 41 action='store_true', default=False, | |
| 42 help="list columns in `table` and exit") | |
| 6 | 43 self.add_argument('-o', '--output', |
| 44 help="output to directory (if `table` not given), or filename or stdout by default") | |
|
9
834b920ae345
allow output of headers in csv
Jeff Hammel <k0scist@gmail.com>
parents:
8
diff
changeset
|
45 self.add_argument('--header', dest='header', |
|
834b920ae345
allow output of headers in csv
Jeff Hammel <k0scist@gmail.com>
parents:
8
diff
changeset
|
46 action='store_true', default=False, |
|
834b920ae345
allow output of headers in csv
Jeff Hammel <k0scist@gmail.com>
parents:
8
diff
changeset
|
47 help="export header as first row") |
| 0 | 48 self.options = None |
| 49 | |
| 50 def parse_args(self, *args, **kw): | |
| 51 options = argparse.ArgumentParser.parse_args(self, *args, **kw) | |
| 52 self.validate(options) | |
| 53 self.options = options | |
| 54 return options | |
| 55 | |
| 56 def validate(self, options): | |
| 57 """validate options""" | |
| 58 | |
| 1 | 59 try: |
| 3 | 60 open(options.db).close() |
| 1 | 61 except Exception as e: |
| 62 self.error("Could not open '{}': {}".format(options.db, e)) | |
| 63 | |
| 6 | 64 if not any((options.output, |
| 65 options.list_tables, | |
| 66 options.table)): | |
| 67 self.error("`--output` directory must be specified to output entire database") | |
| 68 | |
| 5 | 69 if options.list_columns and not options.table: |
| 70 self.error("`--list-columns` requires `table`") | |
| 2 | 71 |
| 0 | 72 def main(args=sys.argv[1:]): |
| 73 """CLI""" | |
| 74 | |
| 75 # parse command line options | |
| 76 parser = SQLExParser() | |
| 77 options = parser.parse_args(args) | |
| 78 | |
| 3 | 79 # connect to database |
| 80 db = SQLEx(options.db) | |
| 81 | |
| 82 if options.list_tables: | |
|
4
b440206930ac
stub table optional positional argument
Jeff Hammel <k0scist@gmail.com>
parents:
3
diff
changeset
|
83 # list tables and return |
|
b440206930ac
stub table optional positional argument
Jeff Hammel <k0scist@gmail.com>
parents:
3
diff
changeset
|
84 # if `table` argument is provided, exit 1 |
|
b440206930ac
stub table optional positional argument
Jeff Hammel <k0scist@gmail.com>
parents:
3
diff
changeset
|
85 # if not available. Otherwise exit 0 |
| 3 | 86 tables = db.tables() |
| 87 print ('\n'.join(tables)) | |
|
4
b440206930ac
stub table optional positional argument
Jeff Hammel <k0scist@gmail.com>
parents:
3
diff
changeset
|
88 retval = 0 |
|
b440206930ac
stub table optional positional argument
Jeff Hammel <k0scist@gmail.com>
parents:
3
diff
changeset
|
89 if options.table: |
|
b440206930ac
stub table optional positional argument
Jeff Hammel <k0scist@gmail.com>
parents:
3
diff
changeset
|
90 retval = int(options.table not in tables) |
|
b440206930ac
stub table optional positional argument
Jeff Hammel <k0scist@gmail.com>
parents:
3
diff
changeset
|
91 return retval |
| 3 | 92 |
| 5 | 93 if options.table: |
| 94 # ensure selected table exists | |
| 95 if options.table not in db.tables(): | |
| 96 parser.error("No table '{}' in {} tables:\n{}".format(options.table, options.db, ', '.join(db.tables()))) | |
| 97 | |
| 98 if options.list_columns: | |
| 99 # list columns and return | |
| 100 print ('\n'.join(db.columns(options.table).keys())) | |
| 101 return | |
| 102 | |
| 6 | 103 if options.table: |
| 104 # output table | |
| 105 | |
| 106 if options.output: | |
| 107 with open(options.output, 'w') as f: | |
|
9
834b920ae345
allow output of headers in csv
Jeff Hammel <k0scist@gmail.com>
parents:
8
diff
changeset
|
108 db.table2csv(options.table, f, header=options.header) |
| 6 | 109 else: |
|
9
834b920ae345
allow output of headers in csv
Jeff Hammel <k0scist@gmail.com>
parents:
8
diff
changeset
|
110 db.table2csv(options.table, sys.stdout, header=options.header) |
| 6 | 111 sys.stdout.flush() |
| 112 else: | |
| 113 # output entire db to CSV files in directory | |
|
8
adf056d67c01
ability to export entire db to directory
Jeff Hammel <k0scist@gmail.com>
parents:
6
diff
changeset
|
114 |
|
adf056d67c01
ability to export entire db to directory
Jeff Hammel <k0scist@gmail.com>
parents:
6
diff
changeset
|
115 # ensure directory exists |
|
adf056d67c01
ability to export entire db to directory
Jeff Hammel <k0scist@gmail.com>
parents:
6
diff
changeset
|
116 ensure_dir(options.output) |
|
adf056d67c01
ability to export entire db to directory
Jeff Hammel <k0scist@gmail.com>
parents:
6
diff
changeset
|
117 |
|
adf056d67c01
ability to export entire db to directory
Jeff Hammel <k0scist@gmail.com>
parents:
6
diff
changeset
|
118 for table in db.tables(): |
|
adf056d67c01
ability to export entire db to directory
Jeff Hammel <k0scist@gmail.com>
parents:
6
diff
changeset
|
119 # export each table |
|
adf056d67c01
ability to export entire db to directory
Jeff Hammel <k0scist@gmail.com>
parents:
6
diff
changeset
|
120 path = os.path.join(options.output, '{}.csv'.format(table)) |
|
adf056d67c01
ability to export entire db to directory
Jeff Hammel <k0scist@gmail.com>
parents:
6
diff
changeset
|
121 with open(path, 'w') as f: |
|
9
834b920ae345
allow output of headers in csv
Jeff Hammel <k0scist@gmail.com>
parents:
8
diff
changeset
|
122 db.table2csv(table, f, header=options.header) |
|
8
adf056d67c01
ability to export entire db to directory
Jeff Hammel <k0scist@gmail.com>
parents:
6
diff
changeset
|
123 |
| 6 | 124 |
| 0 | 125 if __name__ == '__main__': |
| 126 main() |