Skip to content

Commit

Permalink
Add methods .list_table_names() and .alter_table_schema() to the …
Browse files Browse the repository at this point in the history 
…class `PostgreSQL` in sql.py
  • Loading branch information
@mikeqfu
mikeqfu committed Jan 28, 2021
1 parent 4fdcb4a commit 9589424
Showing 1 changed file with 160 additions and 41 deletions.
201 changes: 160 additions & 41 deletions pyhelpers/sql.py
View file
Expand Up @@ -108,17 +108,15 @@ class PostgreSQL:
>>> from pyhelpers.sql import PostgreSQL
>>> # Connect the default database 'postgres'
>>> postgres = PostgreSQL(host='localhost', port=5432, username='postgres',
... database_name='postgres')
>>> postgres = PostgreSQL('localhost', 5432, username='postgres', database_name='postgres')
Password (postgres@localhost:5432): ***
Connecting postgres:***@localhost:5432/postgres ... Successfully.
>>> print(postgres.address)
postgres:***@localhost:5432/postgres
>>> # Connect a database 'testdb' (which will be created if it does not exist)
>>> testdb = PostgreSQL(host='localhost', port=5432, username='postgres',
... database_name='testdb')
>>> testdb = PostgreSQL('localhost', 5432, username='postgres', database_name='testdb')
Password (postgres@localhost:5432): ***
Connecting postgres:***@localhost:5432/testdb ... Successfully.
Expand Down Expand Up @@ -213,8 +211,7 @@ def database_exists(self, database_name=None):
>>> from pyhelpers.sql import PostgreSQL
>>> testdb = PostgreSQL(host='localhost', port=5432, username='postgres',
... database_name='testdb')
>>> testdb = PostgreSQL('localhost', 5432, username='postgres', database_name='testdb')
Password (postgres@localhost:5432): ***
Connecting postgres:***@localhost:5432/testdb ... Successfully.
Expand Down Expand Up @@ -253,8 +250,7 @@ def connect_database(self, database_name=None, verbose=False):
>>> from pyhelpers.sql import PostgreSQL
>>> testdb = PostgreSQL(host='localhost', port=5432, username='postgres',
... database_name='testdb')
>>> testdb = PostgreSQL('localhost', 5432, username='postgres', database_name='testdb')
Password (postgres@localhost:5432): ***
Connecting postgres:***@localhost:5432/testdb ... Successfully.
Expand Down Expand Up @@ -306,8 +302,7 @@ def create_database(self, database_name, verbose=False):
>>> from pyhelpers.sql import PostgreSQL
>>> testdb = PostgreSQL(host='localhost', port=5432, username='postgres',
... database_name='testdb')
>>> testdb = PostgreSQL('localhost', 5432, username='postgres', database_name='testdb')
Password (postgres@localhost:5432): ***
Connecting postgres:***@localhost:5432/testdb ... Successfully.
Expand Down Expand Up @@ -352,8 +347,7 @@ def get_database_size(self, database_name=None):
>>> from pyhelpers.sql import PostgreSQL
>>> testdb = PostgreSQL(host='localhost', port=5432, username='postgres',
... database_name='testdb')
>>> testdb = PostgreSQL('localhost', 5432, username='postgres', database_name='testdb')
Password (postgres@localhost:5432): ***
Connecting postgres:***@localhost:5432/testdb ... Successfully.
Expand Down Expand Up @@ -385,8 +379,7 @@ def disconnect_database(self, database_name=None, verbose=False):
>>> from pyhelpers.sql import PostgreSQL
>>> testdb = PostgreSQL(host='localhost', port=5432, username='postgres',
... database_name='testdb')
>>> testdb = PostgreSQL('localhost', 5432, username='postgres', database_name='testdb')
Password (postgres@localhost:5432): ***
Connecting postgres:***@localhost:5432/testdb ... Successfully.
Expand Down Expand Up @@ -426,8 +419,7 @@ def disconnect_other_databases(self):
>>> from pyhelpers.sql import PostgreSQL
>>> testdb = PostgreSQL(host='localhost', port=5432, username='postgres',
... database_name='testdb')
>>> testdb = PostgreSQL('localhost', 5432, username='postgres', database_name='testdb')
Password (postgres@localhost:5432): ***
Connecting postgres:***@localhost:5432/testdb ... Successfully.
Expand Down Expand Up @@ -463,8 +455,7 @@ def drop_database(self, database_name=None, confirmation_required=True, verbose=
>>> from pyhelpers.sql import PostgreSQL
>>> testdb = PostgreSQL(host='localhost', port=5432, username='postgres',
... database_name='testdb')
>>> testdb = PostgreSQL('localhost', 5432, username='postgres', database_name='testdb')
Password (postgres@localhost:5432): ***
Connecting postgres:***@localhost:5432/testdb ... Successfully.
Expand Down Expand Up @@ -523,8 +514,7 @@ def schema_exists(self, schema_name):
>>> from pyhelpers.sql import PostgreSQL
>>> testdb = PostgreSQL(host='localhost', port=5432, username='postgres',
... database_name='testdb')
>>> testdb = PostgreSQL('localhost', 5432, username='postgres', database_name='testdb')
Password (postgres@localhost:5432): ***
Connecting postgres:***@localhost:5432/testdb ... Successfully.
Expand Down Expand Up @@ -557,8 +547,7 @@ def create_schema(self, schema_name, verbose=False):
>>> from pyhelpers.sql import PostgreSQL
>>> testdb = PostgreSQL(host='localhost', port=5432, username='postgres',
... database_name='testdb')
>>> testdb = PostgreSQL('localhost', 5432, username='postgres', database_name='testdb')
Password (postgres@localhost:5432): ***
Connecting postgres:***@localhost:5432/testdb ... Successfully.
Expand Down Expand Up @@ -607,8 +596,7 @@ def msg_for_multi_schemas(self, item_names, desc='schema'):
>>> from pyhelpers.sql import PostgreSQL
>>> testdb = PostgreSQL(host='localhost', port=5432, username='postgres',
... database_name='testdb')
>>> testdb = PostgreSQL('localhost', 5432, username='postgres', database_name='testdb')
Password (postgres@localhost:5432): ***
Connecting postgres:***@localhost:5432/testdb ... Successfully.
Expand Down Expand Up @@ -667,8 +655,7 @@ def drop_schema(self, schema_names, confirmation_required=True, verbose=False):
>>> from pyhelpers.sql import PostgreSQL
>>> testdb = PostgreSQL(host='localhost', port=5432, username='postgres',
... database_name='testdb')
>>> testdb = PostgreSQL('localhost', 5432, username='postgres', database_name='testdb')
Password (postgres@localhost:5432): ***
Connecting postgres:***@localhost:5432/testdb ... Successfully.
Expand All @@ -687,7 +674,7 @@ def drop_schema(self, schema_names, confirmation_required=True, verbose=False):
"polygons"
"test_schema"
from postgres:***@localhost:5432/testdb
? [No]|Yes: >? yes
? [No]|Yes: yes
Dropping ...
"points" ... Done.
"lines" ... Done.
Expand Down Expand Up @@ -748,8 +735,7 @@ def table_exists(self, table_name, schema_name='public'):
>>> from pyhelpers.sql import PostgreSQL
>>> testdb = PostgreSQL(host='localhost', port=5432, username='postgres',
... database_name='testdb')
>>> testdb = PostgreSQL('localhost', 5432, username='postgres', database_name='testdb')
Password (postgres@localhost:5432): ***
Connecting postgres:***@localhost:5432/testdb ... Successfully.
Expand Down Expand Up @@ -788,8 +774,7 @@ def create_table(self, table_name, column_specs, schema_name='public', verbose=F
>>> from pyhelpers.sql import PostgreSQL
>>> testdb = PostgreSQL(host='localhost', port=5432, username='postgres',
... database_name='testdb')
>>> testdb = PostgreSQL('localhost', 5432, username='postgres', database_name='testdb')
Password (postgres@localhost:5432): ***
Connecting postgres:***@localhost:5432/testdb ... Successfully.
Expand Down Expand Up @@ -853,7 +838,7 @@ def create_table(self, table_name, column_specs, schema_name='public', verbose=F
>>> # Drop the table
>>> testdb.drop_table(tbl_name, verbose=True)
To drop the table "public"."test_table" from postgres:***@localhost:5432/testdb
? [No]|Yes: >? yes
? [No]|Yes: yes
Dropping "public"."test_table" ... Done.
"""

Expand Down Expand Up @@ -907,8 +892,7 @@ def get_column_info(self, table_name, schema_name='public', as_dict=True):

return info_tbl

def drop_table(self, table_name, schema_name='public', confirmation_required=True,
verbose=False):
def drop_table(self, table_name, schema_name='public', confirmation_required=True, verbose=False):
"""
Remove a table from the database being connected.
Expand All @@ -926,21 +910,21 @@ def drop_table(self, table_name, schema_name='public', confirmation_required=Tru
See the example for the method :py:meth:`.create_table()<pyhelpers.sql.PostgreSQL.create_table>`.
"""

table = '\"{}\".\"{}\"'.format(schema_name, table_name)
table_name_ = '"{}"."{}"'.format(schema_name, table_name)

if not self.table_exists(table_name=table_name, schema_name=schema_name):
if verbose:
print("The table {} does not exist.".format(table))
print("The table {} does not exist.".format(table_name_))

else:
if confirmed("To drop the table {} from {}\n?".format(table, self.address),
if confirmed("To drop the table {} from {}\n?".format(table_name_, self.address),
confirmation_required=confirmation_required):

if verbose:
if confirmation_required:
log_msg = "Dropping {}".format(table)
log_msg = "Dropping {}".format(table_name_)
else:
log_msg = "Dropping the table {} from {}".format(table, self.address)
log_msg = "Dropping the table {} from {}".format(table_name_, self.address)
print(log_msg, end=" ... ")

try:
Expand All @@ -949,6 +933,142 @@ def drop_table(self, table_name, schema_name='public', confirmation_required=Tru
except Exception as e:
print("Failed. {}".format(e))

def list_table_names(self, schema_name='public', verbose=False):
"""
List the names of all tables in a schema.
:param schema_name: name of a schema, defaults to ``'public'``
:type schema_name: str
:param verbose: whether to print relevant information in console as the function runs,
defaults to ``False``
:type verbose: bool
:return: a list of table names
:rtype: list or None
**Examples**::
>>> from pyhelpers.sql import PostgreSQL
>>> testdb = PostgreSQL('localhost', 5432, username='postgres', database_name='testdb')
Password (postgres@localhost:5432): ***
Connecting postgres:***@localhost:5432/testdb ... Successfully.
>>> lst_tbl_names = testdb.list_table_names()
>>> print(lst_tbl_names)
[]
>>> lst_tbl_names = testdb.list_table_names(schema_name='testdb', verbose=True)
The schema "testdb" does not exist.
>>> # Create a new table named "test_table" in the schema "testdb"
>>> new_table_name = 'test_table'
>>> column_specs = 'col_name_1 INT, col_name_2 TEXT'
>>> testdb.create_table(new_table_name, column_specs, verbose=True)
Creating a table: "public"."test_table" ... Done.
>>> lst_tbl_names = testdb.list_table_names(schema_name='public')
>>> print(lst_tbl_names)
['test_table']
>>> # Delete the database "testdb"
>>> testdb.drop_database(verbose=True)
To drop the database "testdb" from postgres:***@localhost:5432
? [No]|Yes: yes
Dropping "testdb" ... Done.
"""

if not self.schema_exists(schema_name=schema_name):
if verbose:
print("The schema \"{}\" does not exist.".format(schema_name))

else:
res = self.engine.execute(
"SELECT table_name FROM information_schema.tables "
"WHERE table_schema='{}' AND table_type='BASE TABLE';".format(schema_name))

temp_list = res.fetchall()

table_list = [x[0] for x in temp_list]

return table_list

def alter_table_schema(self, table_name, schema_name, new_schema_name, confirmation_required=True,
verbose=False):
"""
Move a table from one schema to another within the database being connected.
:param table_name: name of a table
:type table_name: str
:param schema_name: name of a schema
:type schema_name: str
:param new_schema_name: name of a new schema
:param confirmation_required: whether to prompt a message for confirmation to proceed,
defaults to ``True``
:type confirmation_required: bool
:param verbose: whether to print relevant information in console as the function runs,
defaults to ``False``
:type verbose: bool
**Examples**::
>>> from pyhelpers.sql import PostgreSQL
>>> testdb = PostgreSQL('localhost', 5432, username='postgres', database_name='testdb')
Password (postgres@localhost:5432): ***
Connecting postgres:***@localhost:5432/testdb ... Successfully.
>>> # Create a new table named "test_table" in the schema "testdb"
>>> new_table_name = 'test_table'
>>> column_specs = 'col_name_1 INT, col_name_2 TEXT'
>>> testdb.create_table(new_table_name, column_specs, verbose=True)
Creating a table: "public"."test_table" ... Done.
>>> # Create a new schema "test_schema"
>>> testdb.create_schema(schema_name='test_schema', verbose=True)
Creating a schema: "test_schema" ... Done.
>>> # Move the table "public"."test_table" to the schema "test_schema"
>>> testdb.alter_table_schema('test_table', 'public', 'test_schema', verbose=True)
To move the table "test_table" from the schema "public" to "test_schema"
? [No]|Yes: yes
Moving "public"."test_table" to "test_schema" ... Done.
>>> lst_tbl_names = testdb.list_table_names(schema_name='test_schema')
>>> print(lst_tbl_names)
['test_table']
>>> testdb.drop_database(verbose=True)
To drop the database "testdb" from postgres:***@localhost:5432
? [No]|Yes: yes
Dropping "testdb" ... Done.
"""

table_name_ = '"{}"."{}"'.format(schema_name, table_name)

if confirmed("To move the table \"{}\" from the schema \"{}\" to \"{}\"\n?".format(
table_name, schema_name, new_schema_name),
confirmation_required=confirmation_required):

if not self.schema_exists(schema_name=new_schema_name):
self.create_schema(schema_name=new_schema_name, verbose=verbose)

if verbose:
if confirmation_required:
log_msg = "Moving {} to \"{}\"".format(table_name_, new_schema_name)
else:
log_msg = "Moving the table \"{}\" from \"{}\" to \"{}\"".format(
table_name, schema_name, new_schema_name)
print(log_msg, end=" ... ")

try:
self.engine.execute(
'ALTER TABLE {} SET SCHEMA "{}";'.format(table_name_, new_schema_name))

print("Done.") if verbose else ""

except Exception as e:
print("Failed. {}".format(e))

@staticmethod
def psql_insert_copy(sql_table, sql_db_engine, column_name_list, data_iter):
"""
Expand Down Expand Up @@ -1188,8 +1308,7 @@ def read_sql_query(self, sql_query, method='tempfile', tempfile_mode='w+b', max_
>>> import pandas
>>> from pyhelpers.sql import PostgreSQL
>>> testdb = PostgreSQL(host='localhost', port=5432, username='postgres',
... database_name='testdb')
>>> testdb = PostgreSQL('localhost', 5432, username='postgres', database_name='testdb')
Password (postgres@localhost:5432): ***
Connecting postgres:***@localhost:5432/testdb ... Successfully.
Expand Down

0 comments on commit 9589424

Please sign in to comment.