X-Git-Url: http://git.veekun.com/zzz-pokedex.git/blobdiff_plain/831dfe6f33ba078ebe4b1a27e13f5d267d8688f5..d17c1c221b4c536c80094e46561904149a84b181:/pokedex/__init__.py?ds=sidebyside diff --git a/pokedex/__init__.py b/pokedex/__init__.py index c24a5db..1e6c6d1 100644 --- a/pokedex/__init__.py +++ b/pokedex/__init__.py @@ -1,9 +1,11 @@ # encoding: utf8 import sys +from sqlalchemy.exc import IntegrityError import sqlalchemy.types from .db import connect, metadata, tables as tables_module +from pokedex.lookup import lookup as pokedex_lookup def main(): if len(sys.argv) <= 1: @@ -13,14 +15,14 @@ def main(): args = sys.argv[2:] # Find the command as a function in this file - func = globals().get(command, None) - if func and callable(func) and command != 'main': + func = globals().get("command_%s" % command, None) + if func: func(*args) else: - help() + command_help() -def csvimport(engine_uri, directory='.'): +def command_csvimport(engine_uri, directory='.'): import csv from sqlalchemy.orm.attributes import instrumentation_registry @@ -29,17 +31,11 @@ def csvimport(engine_uri, directory='.'): metadata.create_all() - # Oh, mysql-chan. - # TODO try to insert data in preorder so we don't need this hack and won't - # break similarly on other engines - if 'mysql' in engine_uri: - session.execute('SET FOREIGN_KEY_CHECKS = 0') - # SQLAlchemy is retarded and there is no way for me to get a list of ORM # classes besides to inspect the module they all happen to live in for # things that look right. table_base = tables_module.TableBase - orm_classes = {} + orm_classes = {} # table object => table class for name in dir(tables_module): # dir() returns strings! How /convenient/. @@ -56,12 +52,16 @@ def csvimport(engine_uri, directory='.'): continue # thingy is definitely a table class! Hallelujah. - orm_classes[thingy.__table__.name] = thingy + orm_classes[thingy.__table__] = thingy # Okay, run through the tables and actually load the data now - for table_name, table in sorted(orm_classes.items()): + for table_obj in metadata.sorted_tables: + table_class = orm_classes[table_obj] + table_name = table_obj.name + # Print the table name but leave the cursor in a fixed column print table_name + '...', ' ' * (40 - len(table_name)), + sys.stdout.flush() try: csvfile = open("%s/%s.csv" % (directory, table_name), 'rb') @@ -73,11 +73,24 @@ def csvimport(engine_uri, directory='.'): reader = csv.reader(csvfile, lineterminator='\n') column_names = [unicode(column) for column in reader.next()] + # Self-referential tables may contain rows with foreign keys of other + # rows in the same table that do not yet exist. Pull these out and add + # them to the session last + # ASSUMPTION: Self-referential tables have a single PK called "id" + deferred_rows = [] # ( row referring to id, [foreign ids we need] ) + seen_ids = {} # primary key we've seen => 1 + + # Fetch foreign key columns that point at this table, if any + self_ref_columns = [] + for column in table_obj.c: + if any(_.references(table_obj) for _ in column.foreign_keys): + self_ref_columns.append(column) + for csvs in reader: - row = table() + row = table_class() for column_name, value in zip(column_names, csvs): - column = table.__table__.c[column_name] + column = table_obj.c[column_name] if column.nullable and value == '': # Empty string in a nullable column really means NULL value = None @@ -94,18 +107,46 @@ def csvimport(engine_uri, directory='.'): setattr(row, column_name, value) + # May need to stash this row and add it later if it refers to a + # later row in this table + if self_ref_columns: + foreign_ids = [getattr(row, _.name) for _ in self_ref_columns] + foreign_ids = [_ for _ in foreign_ids if _] # remove NULL ids + + if not foreign_ids: + # NULL key. Remember this row and add as usual. + seen_ids[row.id] = 1 + + elif all(_ in seen_ids for _ in foreign_ids): + # Non-NULL key we've already seen. Remember it and commit + # so we know the old row exists when we add the new one + session.commit() + seen_ids[row.id] = 1 + + else: + # Non-NULL future id. Save this and insert it later! + deferred_rows.append((row, foreign_ids)) + continue + session.add(row) session.commit() - print 'loaded' - # Shouldn't matter since this is usually the end of the program and thus - # the connection too, but let's change this back just in case - if 'mysql' in engine_uri: - session.execute('SET FOREIGN_KEY_CHECKS = 1') + # Attempt to add any spare rows we've collected + for row, foreign_ids in deferred_rows: + if not all(_ in seen_ids for _ in foreign_ids): + # Could happen if row A refers to B which refers to C. + # This is ridiculous and doesn't happen in my data so far + raise ValueError("Too many levels of self-reference! " + "Row was: " + str(row.__dict__)) + session.add(row) + seen_ids[row.id] = 1 + session.commit() -def csvexport(engine_uri, directory='.'): + print 'loaded' + +def command_csvexport(engine_uri, directory='.'): import csv session = connect(engine_uri) @@ -118,7 +159,8 @@ def csvexport(engine_uri, directory='.'): columns = [col.name for col in table.columns] writer.writerow(columns) - for row in session.query(table).all(): + primary_key = table.primary_key + for row in session.query(table).order_by(*primary_key).all(): csvs = [] for col in columns: # Convert Pythony values to something more universal @@ -136,11 +178,25 @@ def csvexport(engine_uri, directory='.'): writer.writerow(csvs) +def command_lookup(engine_uri, name): + # XXX don't require uri! somehow + session = connect(engine_uri) + + results, exact = pokedex_lookup(session, name) + if exact: + print "Matched:" + else: + print "Fuzzy-matched:" + + for object in results: + print object.__tablename__, object.name + -def help(): +def command_help(): print u"""pokedex -- a command-line Pokédex interface help Displays this message. + lookup {uri} [name] Look up something in the Pokédex. These commands are only useful for developers: csvimport {uri} [dir] Import data from a set of CSVs to the database @@ -148,6 +204,6 @@ def help(): csvexport {uri} [dir] Export data from the database given by the URI to a set of CSVs. Directory defaults to cwd. -""" +""".encode(sys.getdefaultencoding(), 'replace') sys.exit(0)