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 from pokedex.db import metadata
  12 import pokedex.db.tables as tables
  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=False):
 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
 128     # First take care of verbosity
 129     print_start, print_status, print_done = _get_verbose_prints(verbose)
 130
 131
 132     if directory is None:
 133         directory = get_default_csv_dir()
 134
 135     # XXX why isn't this done in command_load
 136     table_names = _get_table_names(metadata, tables)
 137     table_objs = [metadata.tables[name] for name in table_names]
 138
 139     if recursive:
 140         table_objs.extend(find_dependent_tables(table_objs))
 141
 142     table_objs = sqlalchemy.sql.util.sort_tables(table_objs)
 143
 144     # SQLite speed tweaks
 145     if not safe and session.connection().dialect.name == 'sqlite':
 146         session.connection().execute("PRAGMA synchronous=OFF")
 147         session.connection().execute("PRAGMA journal_mode=OFF")
 148
 149     # Drop all tables if requested
 150     if drop_tables:
 151         print_start('Dropping tables')
 152         for n, table in enumerate(reversed(table_objs)):
 153             table.drop(checkfirst=True)
 154             print_status('%s/%s' % (n, len(table_objs)))
 155         print_done()
 156
 157     print_start('Creating tables')
 158     for n, table in enumerate(table_objs):
 159         table.create()
 160         print_status('%s/%s' % (n, len(table_objs)))
 161     print_done()
 162     connection = session.connection()
 163
 164     # Okay, run through the tables and actually load the data now
 165     for table_obj in table_objs:
 166         table_name = table_obj.name
 167         insert_stmt = table_obj.insert()
 168
 169         print_start(table_name)
 170
 171         try:
 172             csvpath = "%s/%s.csv" % (directory, table_name)
 173             csvfile = open(csvpath, 'rb')
 174         except IOError:
 175             # File doesn't exist; don't load anything!
 176             print_done('missing?')
 177             continue
 178
 179         csvsize = os.stat(csvpath).st_size
 180
 181         reader = csv.reader(csvfile, lineterminator='\n')
 182         column_names = [unicode(column) for column in reader.next()]
 183
 184         if not safe and session.connection().dialect.name == 'postgresql':
 185             """
 186             Postgres' CSV dialect works with our data, if we mark the not-null
 187             columns with FORCE NOT NULL.
 188             COPY is only allowed for DB superusers. If you're not one, use safe
 189             loading (pokedex load -S).
 190             """
 191             session.commit()
 192             not_null_cols = [c for c in column_names if not table_obj.c[c].nullable]
 193             if not_null_cols:
 194                 force_not_null = 'FORCE NOT NULL ' + ','.join('"%s"' % c for c in not_null_cols)
 195             else:
 196                 force_not_null = ''
 197             command = "COPY %(table_name)s (%(columns)s) FROM '%(csvpath)s' CSV HEADER %(force_not_null)s"
 198             session.connection().execute(
 199                 command % dict(
 200                     table_name=table_name,
 201                     csvpath=csvpath,
 202                     columns=','.join('"%s"' % c for c in column_names),
 203                     force_not_null=force_not_null,
 204                 )
 205             )
 206             session.commit()
 207             print_done()
 208             continue
 209
 210         # Self-referential tables may contain rows with foreign keys of other
 211         # rows in the same table that do not yet exist.  Pull these out and add
 212         # them to the session last
 213         # ASSUMPTION: Self-referential tables have a single PK called "id"
 214         deferred_rows = []  # ( row referring to id, [foreign ids we need] )
 215         seen_ids = set()    # primary keys we've seen
 216
 217         # Fetch foreign key columns that point at this table, if any
 218         self_ref_columns = []
 219         for column in table_obj.c:
 220             if any(x.references(table_obj) for x in column.foreign_keys):
 221                 self_ref_columns.append(column)
 222
 223         new_rows = []
 224         def insert_and_commit():
 225             if not new_rows:
 226                 return
 227             session.connection().execute(insert_stmt, new_rows)
 228             session.commit()
 229             new_rows[:] = []
 230
 231             progress = "%d%%" % (100 * csvfile.tell() // csvsize)
 232             print_status(progress)
 233
 234         for csvs in reader:
 235             row_data = {}
 236
 237             for column_name, value in zip(column_names, csvs):
 238                 column = table_obj.c[column_name]
 239                 if column.nullable and value == '':
 240                     # Empty string in a nullable column really means NULL
 241                     value = None
 242                 elif isinstance(column.type, sqlalchemy.types.Boolean):
 243                     # Boolean values are stored as string values 0/1, but both
 244                     # of those evaluate as true; SQLA wants True/False
 245                     if value == '0':
 246                         value = False
 247                     else:
 248                         value = True
 249                 else:
 250                     # Otherwise, unflatten from bytes
 251                     value = value.decode('utf-8')
 252
 253                 # nb: Dictionaries flattened with ** have to have string keys
 254                 row_data[ str(column_name) ] = value
 255
 256             # May need to stash this row and add it later if it refers to a
 257             # later row in this table
 258             if self_ref_columns:
 259                 foreign_ids = set(row_data[x.name] for x in self_ref_columns)
 260                 foreign_ids.discard(None)  # remove NULL ids
 261
 262                 if not foreign_ids:
 263                     # NULL key.  Remember this row and add as usual.
 264                     seen_ids.add(row_data['id'])
 265
 266                 elif foreign_ids.issubset(seen_ids):
 267                     # Non-NULL key we've already seen.  Remember it and commit
 268                     # so we know the old row exists when we add the new one
 269                     insert_and_commit()
 270                     seen_ids.add(row_data['id'])
 271
 272                 else:
 273                     # Non-NULL future id.  Save this and insert it later!
 274                     deferred_rows.append((row_data, foreign_ids))
 275                     continue
 276
 277             # Insert row!
 278             new_rows.append(row_data)
 279
 280             # Remembering some zillion rows in the session consumes a lot of
 281             # RAM.  Let's not do that.  Commit every 1000 rows
 282             if len(new_rows) >= 1000:
 283                 insert_and_commit()
 284
 285         insert_and_commit()
 286
 287         # Attempt to add any spare rows we've collected
 288         for row_data, foreign_ids in deferred_rows:
 289             if not foreign_ids.issubset(seen_ids):
 290                 # Could happen if row A refers to B which refers to C.
 291                 # This is ridiculous and doesn't happen in my data so far
 292                 raise ValueError("Too many levels of self-reference!  "
 293                                  "Row was: " + str(row))
 294
 295             session.connection().execute(
 296                 insert_stmt.values(**row_data)
 297             )
 298             seen_ids.add(row_data['id'])
 299         session.commit()
 300
 301         print_done()
 302
 303     # SQLite check
 304     if session.connection().dialect.name == 'sqlite':
 305         session.connection().execute("PRAGMA integrity_check")
 306
 307
 308
 309 def dump(session, tables=[], directory=None, verbose=False):
 310     """Dumps the contents of a database to a set of CSV files.  Probably not
 311     useful to anyone besides a developer.
 312
 313     `session`
 314         SQLAlchemy session to use.
 315
 316     `tables`
 317         List of tables to dump.  If omitted, all tables are dumped.
 318
 319     `directory`
 320         Directory the CSV files should be put in.  Defaults to the `pokedex`
 321         data directory.
 322
 323     `verbose`
 324         If set to True, status messages will be printed to stdout.
 325     """
 326
 327     # First take care of verbosity
 328     print_start, print_status, print_done = _get_verbose_prints(verbose)
 329
 330
 331     if not directory:
 332         directory = get_default_csv_dir()
 333
 334     table_names = _get_table_names(metadata, tables)
 335     table_names.sort()
 336
 337
 338     for table_name in table_names:
 339         print_start(table_name)
 340         table = metadata.tables[table_name]
 341
 342         writer = csv.writer(open("%s/%s.csv" % (directory, table_name), 'wb'),
 343                             lineterminator='\n')
 344         columns = [col.name for col in table.columns]
 345         writer.writerow(columns)
 346
 347         primary_key = table.primary_key
 348         for row in session.query(table).order_by(*primary_key).all():
 349             csvs = []
 350             for col in columns:
 351                 # Convert Pythony values to something more universal
 352                 val = getattr(row, col)
 353                 if val == None:
 354                     val = ''
 355                 elif val == True:
 356                     val = '1'
 357                 elif val == False:
 358                     val = '0'
 359                 else:
 360                     val = unicode(val).encode('utf-8')
 361
 362                 csvs.append(val)
 363
 364             writer.writerow(csvs)
 365
 366         print_done()