diff --git a/commands/__init__.py b/commands/__init__.py index a92ff05a..f3b22463 100644 --- a/commands/__init__.py +++ b/commands/__init__.py @@ -16,6 +16,7 @@ 'nav', 'run', 'update', + 'export', ] ###################################################################### @@ -28,7 +29,7 @@ class HelpAction(argparse.Action): """ def __call__(self, parser, namespace, values, option_string=None): raise exceptions.UsageError("TradeDangerous help", parser.format_help()) - + def addArguments(group, options, required, topGroup=None): """ diff --git a/commands/export_cmd.py b/commands/export_cmd.py new file mode 100644 index 00000000..969227c5 --- /dev/null +++ b/commands/export_cmd.py @@ -0,0 +1,198 @@ +from __future__ import absolute_import, with_statement, print_function, division, unicode_literals +from commands.parsing import MutuallyExclusiveGroup, ParseArgument +from commands.exceptions import CommandLineError +from pathlib import Path +import sqlite3 +import csv + +###################################################################### +# TradeDangerous :: Commands :: Export +# +# Generate the CSV files for the master data of the database. +# +# Note: This script makes some assumptions about the structure +# of the database: +# * The column name of an foreign key reference must be the same +# * The referenced table must have a column named "name" +# which is UNIQUE +# * One column primary keys will be handled by the database engine +# +###################################################################### +# CAUTION: If the database structure gets changed this script might +# need some corrections. +###################################################################### + +###################################################################### +# Default values + +defaultPath = Path("./data") + +# for some tables the first two columns will be reversed +reverseList = [ 'AltItemNames', + 'Item', + 'ShipVendor', + 'Station', + 'StationBuying', + 'UpgradeVendor', +] + +# some tables are ignored +ignoreList = [ 'StationLink', +] + +###################################################################### +# Parser config + +help='CSV exporter for TradeDangerous database.' +name='export' +epilog=None +wantsTradeDB=False +arguments = [ +] +switches = [ + ParseArgument('--path', + help="Specify save location of the CSV files. Default: '{}'".format(defaultPath), + type=str, + default=defaultPath + ), + ParseArgument('--tables', "-T", + help='Specify comma separated tablenames to export.', + metavar='TABLE[,TABLE,...]', + type=str, + default=None + ), +] + +###################################################################### +# Helpers + +def search_dict(list, key, val): + for row in list: + if row[key] == val: return row + +###################################################################### +# Perform query and populate result set + +def run(results, cmdenv, tdb): + from tradedb import TradeDB + + # check database exists + dbFilename = cmdenv.dbFilename or TradeDB.defaultDB + if not Path(dbFilename).is_file(): + raise CommandLineError("Database '{}' not found.".format(dbFilename)) + + # check export path exists + if not Path(cmdenv.path).is_dir(): + raise CommandLineError("Save location '{}' not found.".format(cmdenv.path)) + + # connect to the database + if not cmdenv.quiet: + print("Using database '{}'".format(dbFilename)) + conn = sqlite3.connect(dbFilename) + conn.row_factory = sqlite3.Row + conn.execute("PRAGMA foreign_keys=ON") + + # extract tables from command line + if cmdenv.tables: + bindValues = cmdenv.tables.split(',') + tableStmt = " AND name COLLATE NOCASE IN ({})".format(",".join("?" * len(bindValues))) + cmdenv.DEBUG0(tableStmt) + else: + bindValues = [] + tableStmt = '' + + tableCursor = conn.cursor() + for row in tableCursor.execute(""" + SELECT name + FROM sqlite_master + WHERE type = 'table' + AND name NOT LIKE 'sqlite_%' + {cmdTables} + ORDER BY name + """.format(cmdTables=tableStmt), + bindValues): + tableName = row['name'] + if tableName in ignoreList: + # ignore the table + if not cmdenv.quiet: + print("Ignore Table '{table}'".format(table=tableName)) + continue + + # create CSV files + exportName = Path(cmdenv.path).joinpath("{table}.csv".format(table=tableName)) + if not cmdenv.quiet: + print("Export Table '{table}' to '{file}'".format(table=tableName, file=exportName)) + + with exportName.open("w") as exportFile: + exportOut = csv.writer(exportFile, delimiter=",", quotechar="'", doublequote=True, quoting=csv.QUOTE_NONNUMERIC, lineterminator="\n") + + cur = conn.cursor() + keyList = [] + for key in cur.execute("PRAGMA foreign_key_list('%s')" % tableName): + # ignore FKs to table StationItem + if key['table'] != 'StationItem': + # only support FK joins with the same column name + if key['from'] == key['to']: + keyList += [ {'table': key['table'], 'column': key['from']} ] + + pkCount = 0 + for col in cur.execute("PRAGMA table_info('%s')" % tableName): + # count the columns of the primary key + if col['pk'] > 0: pkCount += 1 + + # initialize helper lists + csvHead = [] + stmtColumn = [] + stmtTable = [ tableName ] + stmtWhere = [] + stmtOrder = [] + + # iterate over all columns of the table + for col in cur.execute("PRAGMA table_info('%s')" % tableName): + # if there is only one PK column, ignore it + if col['pk'] > 0 and pkCount == 1: continue + + # check if the column is a foreign key + key = search_dict(keyList, 'column', col['name']) + if key: + # there must be a "name" column in the referenced table + csvHead += [ "name@{}.{}".format(key['table'], key['column']) ] + stmtColumn += [ "{}.name".format(key['table']) ] + stmtWhere += [ "{}.{} = {}.{}".format(tableName, col['name'], key['table'], key['column']) ] + stmtTable += [ key['table'] ] + stmtOrder += [ "{}.name".format(key['table']) ] + else: + # ordinary column + if col['name'] == 'name': + # name columns must be unique + csvHead += [ 'unq:name' ] + stmtOrder += [ "{}.{}".format(tableName, col['name']) ] + else: + csvHead += [ col['name'] ] + stmtColumn += [ "{}.{}".format(tableName, col['name']) ] + + # reverse the first two columns for some tables + if tableName in reverseList: + csvHead[0], csvHead[1] = csvHead[1], csvHead[0] + stmtColumn[0], stmtColumn[1] = stmtColumn[1], stmtColumn[0] + if len(stmtOrder) > 1: + stmtOrder[0], stmtOrder[1] = stmtOrder[1], stmtOrder[0] + + # build the SQL statement + sqlStmt = "SELECT {} FROM {}".format(",".join(stmtColumn),",".join(stmtTable)) + if len(stmtWhere) > 0: + sqlStmt += " WHERE {}".format(" AND ".join(stmtWhere)) + if len(stmtOrder) > 0: + sqlStmt += " ORDER BY {}".format(",".join(stmtOrder)) + cmdenv.DEBUG0("SQL: %s" % sqlStmt) + + # finally generate the csv file + lineCount = 0 + # no quotes for header line + exportFile.write("{}\n".format(",".join(csvHead))) + for line in cur.execute(sqlStmt): + lineCount += 1 + exportOut.writerow(list(line)) + cmdenv.DEBUG1("{count} {table}s exported".format(count=lineCount, table=tableName)) + + return None