gh9084_get_schema_index_parameter

trbs · Apr 24, 2016 · ad62533 · ad62533
1 parent 2ea0601
commit ad62533
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 8 deletions.
diff --git a/pandas/io/sql.py b/pandas/io/sql.py
@@ -1299,9 +1299,10 @@ def drop_table(self, table_name, schema=None):
             self.get_table(table_name, schema).drop()
             self.meta.clear()
 
-    def _create_sql_schema(self, frame, table_name, keys=None, dtype=None):
-        table = SQLTable(table_name, self, frame=frame, index=False, keys=keys,
-                         dtype=dtype)
+    def _create_sql_schema(self, frame, table_name, keys=None, dtype=None,
+                           index=False, index_label=None):
+        table = SQLTable(table_name, self, frame=frame, index=index, keys=keys,
+                         dtype=dtype, index_label=index_label)
         return str(table.sql_schema())
 
 
@@ -1683,13 +1684,14 @@ def drop_table(self, name, schema=None):
         drop_sql = "DROP TABLE %s" % escape(name)
         self.execute(drop_sql)
 
-    def _create_sql_schema(self, frame, table_name, keys=None, dtype=None):
-        table = SQLiteTable(table_name, self, frame=frame, index=False,
-                            keys=keys, dtype=dtype)
+    def _create_sql_schema(self, frame, table_name, keys=None, dtype=None,
+                           index=False, index_label=None):
+        table = SQLiteTable(table_name, self, frame=frame, index=index,
+                            keys=keys, dtype=dtype, index_label=index_label)
         return str(table.sql_schema())
 
 
-def get_schema(frame, name, flavor='sqlite', keys=None, con=None, dtype=None):
+def get_schema(frame, name, flavor='sqlite', keys=None, con=None, dtype=None, index=False, index_label=None):
     """
     Get the SQL db table schema for the given frame.
 
@@ -1711,8 +1713,17 @@ def get_schema(frame, name, flavor='sqlite', keys=None, con=None, dtype=None):
     dtype : dict of column name to SQL type, default None
         Optional specifying the datatype for columns. The SQL type should
         be a SQLAlchemy type, or a string for sqlite3 fallback connection.
+    index : boolean, default False (for backwards compatibility)
+        Write DataFrame index as a column.
+        .. versionadded:: 0.18.1
 
+    index_label : string or sequence, default None
+        Column label for index column(s). If None is given (default) and
+        `index` is True, then the index names are used.
+        A sequence should be given if the DataFrame uses MultiIndex.
+        .. versionadded:: 0.18.1
     """
 
     pandas_sql = pandasSQL_builder(con=con, flavor=flavor)
-    return pandas_sql._create_sql_schema(frame, name, keys=keys, dtype=dtype)
+    return pandas_sql._create_sql_schema(frame, name, keys=keys, dtype=dtype,
+                                         index=index, index_label=index_label)
diff --git a/pandas/io/tests/test_sql.py b/pandas/io/tests/test_sql.py
@@ -788,6 +788,21 @@ def test_get_schema_keys(self):
         constraint_sentence = 'CONSTRAINT test_pk PRIMARY KEY ("A", "B")'
         self.assertTrue(constraint_sentence in create_sql)
 
+    def test_get_schema_index(self):
+        # support index=True (GH9084)
+        frame = DataFrame({'col':[1,2,3]}, index=pd.date_range('2012-01-01', periods=3))
+        create_sql = sql.get_schema(frame, 'test_schema', index=False)
+        self.assertFalse('index' in create_sql)
+        self.assertFalse('CREATE INDEX' in create_sql)
+
+        create_sql = sql.get_schema(frame, 'test_schema', index=True)
+        self.assertTrue('index' in create_sql)
+        self.assertTrue('CREATE INDEX "ix_test_schema_index"ON "test_schema" ("index")' in create_sql)
+
+        create_sql = sql.get_schema(frame, 'test_schema', index=True, index_label="idx")
+        self.assertTrue('idx' in create_sql)
+        self.assertTrue('CREATE INDEX "ix_test_schema_idx"ON "test_schema" ("idx")' in create_sql)
+
     def test_chunksize_read(self):
         df = DataFrame(np.random.randn(22, 5), columns=list('abcde'))
         df.to_sql('test_chunksize', self.conn, index=False)