pokedex/db/__init__.py

   1 import os
   2 import pkg_resources
   3
   4 from sqlalchemy import MetaData, Table, create_engine, orm
   5
   6 from .tables import metadata
   7
   8 def connect(uri=None, session_args={}, engine_args={}):
   9     """Connects to the requested URI.  Returns a session object.
  10
  11     With the URI omitted, attempts to connect to a default SQLite database
  12     contained within the package directory.
  13
  14     Calling this function also binds the metadata object to the created engine.
  15     """
  16
  17     # Fall back to the environment, then a URI within the package
  18     if not uri:
  19         uri = os.environ.get('POKEDEX_DB_ENGINE', None)
  20
  21     if not uri:
  22         sqlite_path = pkg_resources.resource_filename('pokedex',
  23                                                       'data/pokedex.sqlite')
  24         uri = 'sqlite:///' + sqlite_path
  25
  26     ### Do some fixery for MySQL
  27     if uri[0:5] == 'mysql':
  28         # MySQL uses latin1 for connections by default even if the server is
  29         # otherwise oozing with utf8; charset fixes this
  30         if 'charset' not in uri:
  31             uri += '?charset=utf8'
  32
  33         # Tables should be InnoDB, in the event that we're creating them, and
  34         # use UTF-8 goddammit!
  35         for table in metadata.tables.values():
  36             table.kwargs['mysql_engine'] = 'InnoDB'
  37             table.kwargs['mysql_charset'] = 'utf8'
  38
  39     ### Connect
  40     engine = create_engine(uri, **engine_args)
  41     conn = engine.connect()
  42     metadata.bind = engine
  43
  44     all_session_args = dict(autoflush=True, autocommit=False, bind=engine)
  45     all_session_args.update(session_args)
  46     sm = orm.sessionmaker(**all_session_args)
  47     session = orm.scoped_session(sm)
  48
  49     return session