Skip to content

Commit

Permalink
Added schema kwarg and test
Browse files Browse the repository at this point in the history
Added what's new cleaned up code so that it follows flake8

Cleaned up after commit

Removed print statement
  • Loading branch information
gregorylivschitz committed Apr 4, 2020
1 parent a44ac34 commit 6f8baf2
Show file tree
Hide file tree
Showing 3 changed files with 32 additions and 7 deletions.
1 change: 1 addition & 0 deletions doc/source/whatsnew/v1.1.0.rst
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,7 @@ Other enhancements
- :class:`Series.str` now has a `fullmatch` method that matches a regular expression against the entire string in each row of the series, similar to `re.fullmatch` (:issue:`32806`).
- :meth:`DataFrame.sample` will now also allow array-like and BitGenerator objects to be passed to ``random_state`` as seeds (:issue:`32503`)
- :meth:`MultiIndex.union` will now raise `RuntimeWarning` if the object inside are unsortable, pass `sort=False` to suppress this warning (:issue:`33015`)
- :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 @@ -1372,9 +1372,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 @@ -1505,9 +1511,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 @@ -1762,14 +1772,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 @@ -1790,4 +1806,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
)
6 changes: 6 additions & 0 deletions pandas/tests/io/test_sql.py
Original file line number Diff line number Diff line change
Expand Up @@ -855,6 +855,12 @@ 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"
)
assert "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 6f8baf2

Please sign in to comment.