diff sqlex/model.py @ 5:3a7f515571dc

can list column names
author Jeff Hammel <k0scist@gmail.com>
date Sat, 01 Apr 2017 11:11:28 -0700
parents 5f1e1ac96aa7
children 22fbe50d92e8
line wrap: on
line diff
--- a/sqlex/model.py	Sat Apr 01 09:28:49 2017 -0700
+++ b/sqlex/model.py	Sat Apr 01 11:11:28 2017 -0700
@@ -1,4 +1,10 @@
 import sqlite3
+from collections import OrderedDict
+
+class NoTableException(Exception):
+    def __init__(self, table):
+        Exception.__init__(self, "Table '{}' does not exist".format(table))
+
 
 class SQLEx(object):
     """
@@ -31,3 +37,31 @@
         sql = "SELECT name FROM sqlite_master WHERE type='table';"
         tables = self(sql)
         return sum([list(item) for item in tables], [])
+
+    def ensure_table(self, table):
+        """
+        ensure that `table` exists;
+        if not, raise 1NoTableException
+        """
+
+        if table not in self.tables():
+            raise NoTableException(table)
+
+    def columns(self, table):
+        """
+        returns columns for `table`
+        """
+
+        self.ensure_table(table)
+
+        sql = "PRAGMA table_info({})".format(table)
+        data = self(sql)
+        # (Pdb) pp(columns)
+        # [(0, u'ROWID', u'INTEGER', 0, None, 1),
+        #  (1, u'address', u'TEXT', 0, None, 0),
+        #  (2, u'date', u'INTEGER', 0, None, 0),
+        NAME_INDEX = 1
+        TYPE_INDEX = 2
+        return OrderedDict([(row[NAME_INDEX], row[TYPE_INDEX])
+                            for row in data])
+