ENH: primary key for sqlite schema

Author: changhiskhan

pandas/io/sql.py

@@ -169,10 +169,10 @@ def write_frame(frame, name=None, con=None, flavor='sqlite'):
     data = [tuple(x) for x in frame.values]
     con.executemany(insert_sql, data)
 
-def get_sqlite_schema(frame, name, dtypes=None):
+def get_sqlite_schema(frame, name, dtypes=None, keys=None):
     template = """
 CREATE TABLE %(name)s (
-  %(columns)s
+  %(columns)s%(keystr)s
 );"""
 
     column_types = []
@@ -195,7 +195,12 @@ def get_sqlite_schema(frame, name, dtypes=None):
 
     columns = ',\n  '.join('%s %s' % x for x in column_types)
 
-    return template % {'name' : name, 'columns' : columns}
+    keystr = ''
+    if keys is not None:
+        if isinstance(keys, basestring):
+            keys = (keys,)
+        keystr = ', PRIMARY KEY (%s)' % ','.join(keys)
+    return template % {'name' : name, 'columns' : columns, 'keystr' : keystr}
 
 
 

pandas/io/tests/test_sql.py

@@ -60,6 +60,12 @@ def test_schema(self):
             if len(tokens) == 2 and tokens[0] == 'A':
                 self.assert_(tokens[1] == 'DATETIME')
 
+        frame = tm.makeTimeDataFrame()
+        create_sql = sql.get_sqlite_schema(frame, 'test', keys=['A', 'B'])
+        lines = create_sql.splitlines()
+        self.assert_('PRIMARY KEY (A,B)' in create_sql)
+        self.db.execute(create_sql)
+
     def test_execute_fail(self):
         create_sql = """
         CREATE TABLE test