pokedex/db/load.py

   1 """CSV to database or vice versa."""
   2 import csv
   3 import fnmatch
   4 import os.path
   5 import sys
   6
   7 from sqlalchemy.orm.attributes import instrumentation_registry
   8 import sqlalchemy.sql.util
   9 import sqlalchemy.types
  10
  11 import pokedex
  12 from pokedex.db import metadata, tables, translations
  13 from pokedex.defaults import get_default_csv_dir
  14 from pokedex.db.dependencies import find_dependent_tables
  15
  16
  17 def _get_table_names(metadata, patterns):
  18     """Returns a list of table names from the given metadata.  If `patterns`
  19     exists, only tables matching one of the patterns will be returned.
  20     """
  21     if patterns:
  22         table_names = set()
  23         for pattern in patterns:
  24             if '.' in pattern or '/' in pattern:
  25                 # If it looks like a filename, pull out just the table name
  26                 _, filename = os.path.split(pattern)
  27                 table_name, _ = os.path.splitext(filename)
  28                 pattern = table_name
  29
  30             table_names.update(fnmatch.filter(metadata.tables.keys(), pattern))
  31     else:
  32         table_names = metadata.tables.keys()
  33
  34     return list(table_names)
  35
  36 def _get_verbose_prints(verbose):
  37     """If `verbose` is true, returns three functions: one for printing a
  38     starting message, one for printing an interim status update, and one for
  39     printing a success or failure message when finished.
  40
  41     If `verbose` is false, returns no-op functions.
  42     """
  43
  44     if not verbose:
  45         # Return dummies
  46         def dummy(*args, **kwargs):
  47             pass
  48
  49         return dummy, dummy, dummy
  50
  51     ### Okay, verbose == True; print stuff
  52
  53     def print_start(thing):
  54         # Truncate to 66 characters, leaving 10 characters for a success
  55         # or failure message
  56         truncated_thing = thing[:66]
  57
  58         # Also, space-pad to keep the cursor in a known column
  59         num_spaces = 66 - len(truncated_thing)
  60
  61         print "%s...%s" % (truncated_thing, ' ' * num_spaces),
  62         sys.stdout.flush()
  63
  64     if sys.stdout.isatty():
  65         # stdout is a terminal; stupid backspace tricks are OK.
  66         # Don't use print, because it always adds magical spaces, which
  67         # makes backspace accounting harder
  68
  69         backspaces = [0]
  70         def print_status(msg):
  71             # Overwrite any status text with spaces before printing
  72             sys.stdout.write('\b' * backspaces[0])
  73             sys.stdout.write(' ' * backspaces[0])
  74             sys.stdout.write('\b' * backspaces[0])
  75             sys.stdout.write(msg)
  76             sys.stdout.flush()
  77             backspaces[0] = len(msg)
  78
  79         def print_done(msg='ok'):
  80             # Overwrite any status text with spaces before printing
  81             sys.stdout.write('\b' * backspaces[0])
  82             sys.stdout.write(' ' * backspaces[0])
  83             sys.stdout.write('\b' * backspaces[0])
  84             sys.stdout.write(msg + "\n")
  85             sys.stdout.flush()
  86             backspaces[0] = 0
  87
  88     else:
  89         # stdout is a file (or something); don't bother with status at all
  90         def print_status(msg):
  91             pass
  92
  93         def print_done(msg='ok'):
  94             print msg
  95
  96     return print_start, print_status, print_done
  97
  98
  99 def load(session, tables=[], directory=None, drop_tables=False, verbose=False, safe=True, recursive=True, langs=None):
 100     """Load data from CSV files into the given database session.
 101
 102     Tables are created automatically.
 103
 104     `session`
 105         SQLAlchemy session to use.
 106
 107     `tables`
 108         List of tables to load.  If omitted, all tables are loaded.
 109
 110     `directory`
 111         Directory the CSV files reside in.  Defaults to the `pokedex` data
 112         directory.
 113
 114     `drop_tables`
 115         If set to True, existing `pokedex`-related tables will be dropped.
 116
 117     `verbose`
 118         If set to True, status messages will be printed to stdout.
 119
 120     `safe`
 121         If set to False, load can be faster, but can corrupt the database if
 122         it crashes or is interrupted.
 123
 124     `recursive`
 125         If set to True, load all dependent tables too.
 126
 127     `langs`
 128         List of identifiers of extra language to load, or None to load them all
 129     """
 130
 131     # First take care of verbosity
 132     print_start, print_status, print_done = _get_verbose_prints(verbose)
 133
 134
 135     if directory is None:
 136         directory = get_default_csv_dir()
 137
 138     # XXX why isn't this done in command_load
 139     table_names = _get_table_names(metadata, tables)
 140     table_objs = [metadata.tables[name] for name in table_names]
 141
 142     if recursive:
 143         table_objs.extend(find_dependent_tables(table_objs))
 144
 145     table_objs = sqlalchemy.sql.util.sort_tables(table_objs)
 146
 147     # SQLite speed tweaks
 148     if not safe and session.connection().dialect.name == 'sqlite':
 149         session.connection().execute("PRAGMA synchronous=OFF")
 150         session.connection().execute("PRAGMA journal_mode=OFF")
 151
 152     # Drop all tables if requested
 153     if drop_tables:
 154         print_start('Dropping tables')
 155         for n, table in enumerate(reversed(table_objs)):
 156             table.drop(checkfirst=True)
 157             print_status('%s/%s' % (n, len(table_objs)))
 158         print_done()
 159
 160     print_start('Creating tables')
 161     for n, table in enumerate(table_objs):
 162         table.create()
 163         print_status('%s/%s' % (n, len(table_objs)))
 164     print_done()
 165     connection = session.connection()
 166
 167     # Okay, run through the tables and actually load the data now
 168     for table_obj in table_objs:
 169         table_name = table_obj.name
 170         insert_stmt = table_obj.insert()
 171
 172         print_start(table_name)
 173
 174         try:
 175             csvpath = "%s/%s.csv" % (directory, table_name)
 176             csvfile = open(csvpath, 'rb')
 177         except IOError:
 178             # File doesn't exist; don't load anything!
 179             print_done('missing?')
 180             continue
 181
 182         csvsize = os.stat(csvpath).st_size
 183
 184         reader = csv.reader(csvfile, lineterminator='\n')
 185         column_names = [unicode(column) for column in reader.next()]
 186
 187         if not safe and session.connection().dialect.name == 'postgresql':
 188             """
 189             Postgres' CSV dialect works with our data, if we mark the not-null
 190             columns with FORCE NOT NULL.
 191             COPY is only allowed for DB superusers. If you're not one, use safe
 192             loading (pokedex load -S).
 193             """
 194             session.commit()
 195             not_null_cols = [c for c in column_names if not table_obj.c[c].nullable]
 196             if not_null_cols:
 197                 force_not_null = 'FORCE NOT NULL ' + ','.join('"%s"' % c for c in not_null_cols)
 198             else:
 199                 force_not_null = ''
 200             command = "COPY %(table_name)s (%(columns)s) FROM '%(csvpath)s' CSV HEADER %(force_not_null)s"
 201             session.connection().execute(
 202                 command % dict(
 203                     table_name=table_name,
 204                     csvpath=csvpath,
 205                     columns=','.join('"%s"' % c for c in column_names),
 206                     force_not_null=force_not_null,
 207                 )
 208             )
 209             session.commit()
 210             print_done()
 211             continue
 212
 213         # Self-referential tables may contain rows with foreign keys of other
 214         # rows in the same table that do not yet exist.  Pull these out and add
 215         # them to the session last
 216         # ASSUMPTION: Self-referential tables have a single PK called "id"
 217         deferred_rows = []  # ( row referring to id, [foreign ids we need] )
 218         seen_ids = set()    # primary keys we've seen
 219
 220         # Fetch foreign key columns that point at this table, if any
 221         self_ref_columns = []
 222         for column in table_obj.c:
 223             if any(x.references(table_obj) for x in column.foreign_keys):
 224                 self_ref_columns.append(column)
 225
 226         new_rows = []
 227         def insert_and_commit():
 228             if not new_rows:
 229                 return
 230             session.connection().execute(insert_stmt, new_rows)
 231             session.commit()
 232             new_rows[:] = []
 233
 234             progress = "%d%%" % (100 * csvfile.tell() // csvsize)
 235             print_status(progress)
 236
 237         for csvs in reader:
 238             row_data = {}
 239
 240             for column_name, value in zip(column_names, csvs):
 241                 column = table_obj.c[column_name]
 242                 if column.nullable and value == '':
 243                     # Empty string in a nullable column really means NULL
 244                     value = None
 245                 elif isinstance(column.type, sqlalchemy.types.Boolean):
 246                     # Boolean values are stored as string values 0/1, but both
 247                     # of those evaluate as true; SQLA wants True/False
 248                     if value == '0':
 249                         value = False
 250                     else:
 251                         value = True
 252                 else:
 253                     # Otherwise, unflatten from bytes
 254                     value = value.decode('utf-8')
 255
 256                 # nb: Dictionaries flattened with ** have to have string keys
 257                 row_data[ str(column_name) ] = value
 258
 259             # May need to stash this row and add it later if it refers to a
 260             # later row in this table
 261             if self_ref_columns:
 262                 foreign_ids = set(row_data[x.name] for x in self_ref_columns)
 263                 foreign_ids.discard(None)  # remove NULL ids
 264
 265                 if not foreign_ids:
 266                     # NULL key.  Remember this row and add as usual.
 267                     seen_ids.add(row_data['id'])
 268
 269                 elif foreign_ids.issubset(seen_ids):
 270                     # Non-NULL key we've already seen.  Remember it and commit
 271                     # so we know the old row exists when we add the new one
 272                     insert_and_commit()
 273                     seen_ids.add(row_data['id'])
 274
 275                 else:
 276                     # Non-NULL future id.  Save this and insert it later!
 277                     deferred_rows.append((row_data, foreign_ids))
 278                     continue
 279
 280             # Insert row!
 281             new_rows.append(row_data)
 282
 283             # Remembering some zillion rows in the session consumes a lot of
 284             # RAM.  Let's not do that.  Commit every 1000 rows
 285             if len(new_rows) >= 1000:
 286                 insert_and_commit()
 287
 288         insert_and_commit()
 289
 290         # Attempt to add any spare rows we've collected
 291         for row_data, foreign_ids in deferred_rows:
 292             if not foreign_ids.issubset(seen_ids):
 293                 # Could happen if row A refers to B which refers to C.
 294                 # This is ridiculous and doesn't happen in my data so far
 295                 raise ValueError("Too many levels of self-reference!  "
 296                                  "Row was: " + str(row))
 297
 298             session.connection().execute(
 299                 insert_stmt.values(**row_data)
 300             )
 301             seen_ids.add(row_data['id'])
 302         session.commit()
 303
 304         print_done()
 305
 306
 307     print_start('Translations')
 308     transl = translations.Translations(csv_directory=directory)
 309
 310     new_row_count = 0
 311     for translation_class, rows in transl.get_load_data(langs):
 312         table_obj = translation_class.__table__
 313         if table_obj in table_objs:
 314             insert_stmt = table_obj.insert()
 315             session.connection().execute(insert_stmt, rows)
 316             session.commit()
 317             # We don't have a total, but at least show some increasing number
 318             new_row_count += len(rows)
 319             print_status(str(new_row_count))
 320
 321     print_done()
 322
 323     # SQLite check
 324     if session.connection().dialect.name == 'sqlite':
 325         session.connection().execute("PRAGMA integrity_check")
 326
 327
 328
 329 def dump(session, tables=[], directory=None, verbose=False, langs=['en']):
 330     """Dumps the contents of a database to a set of CSV files.  Probably not
 331     useful to anyone besides a developer.
 332
 333     `session`
 334         SQLAlchemy session to use.
 335
 336     `tables`
 337         List of tables to dump.  If omitted, all tables are dumped.
 338
 339     `directory`
 340         Directory the CSV files should be put in.  Defaults to the `pokedex`
 341         data directory.
 342
 343     `verbose`
 344         If set to True, status messages will be printed to stdout.
 345
 346     `langs`
 347         List of identifiers of languages to dump unofficial texts for
 348     """
 349
 350     # First take care of verbosity
 351     print_start, print_status, print_done = _get_verbose_prints(verbose)
 352
 353     languages = dict((l.id, l) for l in session.query(pokedex.db.tables.Language))
 354
 355     if not directory:
 356         directory = get_default_csv_dir()
 357
 358     table_names = _get_table_names(metadata, tables)
 359     table_names.sort()
 360
 361
 362     for table_name in table_names:
 363         print_start(table_name)
 364         table = metadata.tables[table_name]
 365
 366         writer = csv.writer(open("%s/%s.csv" % (directory, table_name), 'wb'),
 367                             lineterminator='\n')
 368         columns = [col.name for col in table.columns]
 369
 370         # For name tables, dump rows for official languages, as well as
 371         # for those in `langs`.
 372         # For other translation tables, only dump rows for languages in `langs`
 373         # For non-translation tables, dump all rows.
 374         if 'local_language_id' in columns:
 375             if any(col.info.get('official') for col in table.columns):
 376                 def include_row(row):
 377                     return (languages[row.local_language_id].official or
 378                             languages[row.local_language_id].identifier in langs)
 379             else:
 380                 def include_row(row):
 381                     return languages[row.local_language_id].identifier in langs
 382         else:
 383             def include_row(row):
 384                 return True
 385
 386         writer.writerow(columns)
 387
 388         primary_key = table.primary_key
 389         for row in session.query(table).order_by(*primary_key).all():
 390             if include_row(row):
 391                 csvs = []
 392                 for col in columns:
 393                     # Convert Pythony values to something more universal
 394                     val = getattr(row, col)
 395                     if val == None:
 396                         val = ''
 397                     elif val == True:
 398                         val = '1'
 399                     elif val == False:
 400                         val = '0'
 401                     else:
 402                         val = unicode(val).encode('utf-8')
 403
 404                     csvs.append(val)
 405
 406                 writer.writerow(csvs)
 407
 408         print_done()