Skip to content
Merged
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
dbConnectionParam and moveTableToSchema replaced by single function
dbConnectionParam has been deleted because now all parameters are
obtained by a single function which is called in the function
handleTable.

moveTableToSchema has been deleted because tables are now created in
the specified schema. Therefore, there is no need to move the table
after their creation.
  • Loading branch information
rdrg109 committed Aug 21, 2021
commit 5afa6714f26b9c0c75144262e81c32bc0725373d
47 changes: 5 additions & 42 deletions load_into_pg.py
Original file line number Diff line number Diff line change
Expand Up @@ -180,7 +180,7 @@ def _getTableKeys(table):
return keys


def handleTable(table, insertJson, createFk, mbDbFile, dbConnectionParam):
def handleTable(table, insertJson, createFk, mbDbFile):
"""Handle the table including the post/pre processing."""
keys = _getTableKeys(table)
dbFile = mbDbFile if mbDbFile is not None else table + ".xml"
Expand All @@ -199,7 +199,7 @@ def handleTable(table, insertJson, createFk, mbDbFile, dbConnectionParam):
sys.exit(-1)

try:
with pg.connect(dbConnectionParam) as conn:
with pg.connect(**getConnectionParameters()) as conn:
with conn.cursor() as cur:
try:
with open(dbFile, "rb") as xml:
Expand Down Expand Up @@ -279,29 +279,8 @@ def handleTable(table, insertJson, createFk, mbDbFile, dbConnectionParam):
six.print_("Warning from the database.", file=sys.stderr)
six.print_("pg.Warning: {0}".format(str(w)), file=sys.stderr)


def moveTableToSchema(table, schemaName, dbConnectionParam):
try:
with pg.connect(dbConnectionParam) as conn:
with conn.cursor() as cur:
# create the schema
cur.execute("CREATE SCHEMA IF NOT EXISTS " + schemaName + ";")
conn.commit()
# move the table to the right schema
cur.execute("ALTER TABLE " + table + " SET SCHEMA " + schemaName + ";")
conn.commit()
except pg.Error as e:
six.print_("Error in dealing with the database.", file=sys.stderr)
six.print_("pg.Error ({0}): {1}".format(e.pgcode, e.pgerror), file=sys.stderr)
six.print_(str(e), file=sys.stderr)
except pg.Warning as w:
six.print_("Warning from the database.", file=sys.stderr)
six.print_("pg.Warning: {0}".format(str(w)), file=sys.stderr)


#############################################################


parser = argparse.ArgumentParser()
parser.add_argument(
"-t",
Expand Down Expand Up @@ -390,10 +369,6 @@ def moveTableToSchema(table, schemaName, dbConnectionParam):
except NameError:
pass

dbConnectionParam = buildConnectionString(
args.dbname, args.host, args.port, args.username, args.password
)

# load given file in table
if args.file and args.table:
table = args.table
Expand All @@ -404,19 +379,10 @@ def moveTableToSchema(table, schemaName, dbConnectionParam):
specialRules[("Posts", "Body")] = "NULL"

choice = input("This will drop the {} table. Are you sure [y/n]?".format(table))

if len(choice) > 0 and choice[0].lower() == "y":
handleTable(
table, args.insert_json, args.foreign_keys, args.file, dbConnectionParam
)

# Once "handleTable" has created the table in the "public"
# schema, move the table to the given schema if and only if
# the user specified an schema as an argument for the script
# and the schema name is different than "public".

if args.schema_name != "public":
moveTableToSchema(table, args.schema_name, dbConnectionParam)

table, args.insert_json, args.foreign_keys, args.file)
else:
six.print_("Cancelled.")

Expand Down Expand Up @@ -467,7 +433,7 @@ def moveTableToSchema(table, schemaName, dbConnectionParam):

for table in tables:
six.print_("Load {0}.xml file".format(table))
handleTable(table, args.insert_json, args.foreign_keys, None, dbConnectionParam)
handleTable(table, args.insert_json, args.foreign_keys, None)
# remove file
os.remove(table + ".xml")

Expand All @@ -479,9 +445,6 @@ def moveTableToSchema(table, schemaName, dbConnectionParam):
else:
six.print_("Archive '{0}' deleted".format(filepath))

if args.schema_name != "public":
for table in tables:
moveTableToSchema(table, args.schema_name, dbConnectionParam)
exit(0)

else:
Expand Down