Mercurial > hg > sqlex
annotate sqlex/model.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 | 22fbe50d92e8 |
| children |
| rev | line source |
|---|---|
| 6 | 1 import csv |
| 2 | 2 import sqlite3 |
| 5 | 3 from collections import OrderedDict |
| 4 | |
| 5 class NoTableException(Exception): | |
| 6 def __init__(self, table): | |
| 7 Exception.__init__(self, "Table '{}' does not exist".format(table)) | |
| 8 | |
| 2 | 9 |
| 10 class SQLEx(object): | |
| 11 """ | |
| 12 sqlite model | |
| 13 https://docs.python.org/2/library/sqlite3.html | |
| 14 """ | |
| 15 | |
| 16 def __init__(self, db): | |
| 17 self.db = db | |
| 18 self.conn = sqlite3.connect(self.db) | |
| 19 | |
| 20 def __call__(self, sql, *args): | |
| 3 | 21 c = self.conn.cursor() |
| 2 | 22 c.execute(sql, args) |
| 3 | 23 self.conn.commit() |
| 24 try: | |
| 25 return c.fetchall() | |
| 26 except Exception as e: | |
| 27 raise e | |
| 2 | 28 |
| 29 def __del__(self): | |
| 30 self.conn.close() | |
| 31 | |
| 32 def tables(self): | |
| 33 """ | |
| 34 returns table names in database | |
| 35 """ | |
| 36 # http://stackoverflow.com/questions/82875/how-to-list-the-tables-in-an-sqlite-database-file-that-was-opened-with-attach | |
| 37 | |
| 3 | 38 sql = "SELECT name FROM sqlite_master WHERE type='table';" |
| 2 | 39 tables = self(sql) |
| 3 | 40 return sum([list(item) for item in tables], []) |
| 5 | 41 |
| 42 def ensure_table(self, table): | |
| 43 """ | |
| 44 ensure that `table` exists; | |
| 45 if not, raise 1NoTableException | |
| 46 """ | |
| 47 | |
| 48 if table not in self.tables(): | |
| 49 raise NoTableException(table) | |
| 50 | |
| 51 def columns(self, table): | |
| 52 """ | |
| 53 returns columns for `table` | |
| 54 """ | |
| 55 | |
| 56 self.ensure_table(table) | |
| 57 | |
| 58 sql = "PRAGMA table_info({})".format(table) | |
| 59 data = self(sql) | |
| 60 # (Pdb) pp(columns) | |
| 61 # [(0, u'ROWID', u'INTEGER', 0, None, 1), | |
| 62 # (1, u'address', u'TEXT', 0, None, 0), | |
| 63 # (2, u'date', u'INTEGER', 0, None, 0), | |
| 64 NAME_INDEX = 1 | |
| 65 TYPE_INDEX = 2 | |
| 66 return OrderedDict([(row[NAME_INDEX], row[TYPE_INDEX]) | |
| 67 for row in data]) | |
| 68 | |
| 6 | 69 |
|
9
834b920ae345
allow output of headers in csv
Jeff Hammel <k0scist@gmail.com>
parents:
6
diff
changeset
|
70 def table2csv(self, table, fp, header=False): |
| 6 | 71 """ |
| 72 export `table` to `fp` file object in CSV format | |
| 73 """ | |
| 74 # TODO: option to add column headers | |
| 75 | |
| 76 # sanity | |
| 77 self.ensure_table(table) | |
| 78 | |
| 79 # get whole table | |
| 80 sql = 'select * from {table}'.format(table=table) | |
| 81 rows = self(sql) | |
| 82 | |
|
9
834b920ae345
allow output of headers in csv
Jeff Hammel <k0scist@gmail.com>
parents:
6
diff
changeset
|
83 if header: |
|
834b920ae345
allow output of headers in csv
Jeff Hammel <k0scist@gmail.com>
parents:
6
diff
changeset
|
84 # export header as first row, if specified |
|
834b920ae345
allow output of headers in csv
Jeff Hammel <k0scist@gmail.com>
parents:
6
diff
changeset
|
85 _header = self.columns(table).keys() |
|
834b920ae345
allow output of headers in csv
Jeff Hammel <k0scist@gmail.com>
parents:
6
diff
changeset
|
86 if _header: |
|
834b920ae345
allow output of headers in csv
Jeff Hammel <k0scist@gmail.com>
parents:
6
diff
changeset
|
87 _header[0] = '#{}'.format(_header[0]) |
|
834b920ae345
allow output of headers in csv
Jeff Hammel <k0scist@gmail.com>
parents:
6
diff
changeset
|
88 rows.insert(0, _header) |
|
834b920ae345
allow output of headers in csv
Jeff Hammel <k0scist@gmail.com>
parents:
6
diff
changeset
|
89 |
| 6 | 90 # decode unicde because the CSV module won't |
| 91 # http://stackoverflow.com/questions/22733642/how-to-write-a-unicode-csv-in-python-2-7 | |
| 92 rows = [[unicode(s).encode("utf-8") for s in row] | |
| 93 for row in rows] | |
| 94 | |
| 95 # write | |
| 96 writer = csv.writer(fp) | |
| 97 writer.writerows(rows) | |
| 98 |