Skip to content

Commit

Permalink
Added schema kwarg and test
Browse files Browse the repository at this point in the history
  • Loading branch information
gregorylivschitz committed Apr 15, 2020
1 parent 3d4f9dc commit 59e756a
Show file tree
Hide file tree
Showing 3 changed files with 34 additions and 7 deletions.
2 changes: 2 additions & 0 deletions doc/source/whatsnew/v1.1.0.rst
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,8 @@ Other enhancements
This can be used to set a custom compression level, e.g.,
``df.to_csv(path, compression={'method': 'gzip', 'compresslevel': 1}``
(:issue:`33196`)
- :meth:`get_schema` will now allow a schema kwarg that will add a schema into the create table statement
-

.. ---------------------------------------------------------------------------
Expand Down
32 changes: 25 additions & 7 deletions pandas/io/sql.py
Original file line number Diff line number Diff line change
Expand Up @@ -1360,9 +1360,15 @@ 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):
def _create_sql_schema(self, frame, table_name, keys=None, dtype=None, schema=None):
table = SQLTable(
table_name, self, frame=frame, index=False, keys=keys, dtype=dtype
table_name,
self,
frame=frame,
index=False,
keys=keys,
dtype=dtype,
schema=schema,
)
return str(table.sql_schema())

Expand Down Expand Up @@ -1493,9 +1499,13 @@ def _create_table_setup(self):
create_tbl_stmts.append(
f"CONSTRAINT {self.name}_pk PRIMARY KEY ({cnames_br})"
)

if self.schema:
schema_to_add = self.schema + "."
else:
schema_to_add = ""
create_stmts = [
"CREATE TABLE "
+ schema_to_add
+ escape(self.name)
+ " (\n"
+ ",\n ".join(create_tbl_stmts)
Expand Down Expand Up @@ -1750,14 +1760,20 @@ def drop_table(self, name, schema=None):
drop_sql = f"DROP TABLE {_get_valid_sqlite_name(name)}"
self.execute(drop_sql)

def _create_sql_schema(self, frame, table_name, keys=None, dtype=None):
def _create_sql_schema(self, frame, table_name, keys=None, dtype=None, schema=None):
table = SQLiteTable(
table_name, self, frame=frame, index=False, keys=keys, dtype=dtype
table_name,
self,
frame=frame,
index=False,
keys=keys,
dtype=dtype,
schema=schema,
)
return str(table.sql_schema())


def get_schema(frame, name, keys=None, con=None, dtype=None):
def get_schema(frame, name, keys=None, con=None, dtype=None, schema=None):
"""
Get the SQL db table schema for the given frame.
Expand All @@ -1778,4 +1794,6 @@ def get_schema(frame, name, keys=None, con=None, dtype=None):
"""
pandas_sql = pandasSQL_builder(con=con)
return pandas_sql._create_sql_schema(frame, name, keys=keys, dtype=dtype)
return pandas_sql._create_sql_schema(
frame, name, keys=keys, dtype=dtype, schema=schema
)
7 changes: 7 additions & 0 deletions pandas/tests/io/test_sql.py
Original file line number Diff line number Diff line change
Expand Up @@ -855,6 +855,13 @@ def test_get_schema(self):
create_sql = sql.get_schema(self.test_frame1, "test", con=self.conn)
assert "CREATE" in create_sql

def test_get_schema_with_schema(self):
create_sql = sql.get_schema(
self.test_frame1, "test", con=self.conn, schema="pypi"
)
print(create_sql)
assert "CREATE TABLE pypi." in create_sql

def test_get_schema_dtypes(self):
float_frame = DataFrame({"a": [1.1, 1.2], "b": [2.1, 2.2]})
dtype = sqlalchemy.Integer if self.mode == "sqlalchemy" else "INTEGER"
Expand Down

0 comments on commit 59e756a

Please sign in to comment.