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