Make load.py more idiomatic.
[zzz-pokedex.git] / pokedex / db / __init__.py
1 from sqlalchemy import MetaData, Table, engine_from_config, orm
2
3 from ..defaults import get_default_db_uri
4 from .tables import metadata
5 from .multilang import MultilangSession, MultilangScopedSession
6
7
8 def connect(uri=None, session_args={}, engine_args={}, engine_prefix=''):
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 # If we didn't get a uri, fall back to the default
18 if uri is None:
19 uri = engine_args.get(engine_prefix + 'url', None)
20 if uri is None:
21 uri = get_default_db_uri()
22
23 ### Do some fixery for MySQL
24 if uri.startswith('mysql:'):
25 # MySQL uses latin1 for connections by default even if the server is
26 # otherwise oozing with utf8; charset fixes this
27 if 'charset' not in uri:
28 uri += '?charset=utf8'
29
30 # Tables should be InnoDB, in the event that we're creating them, and
31 # use UTF-8 goddammit!
32 for table in metadata.tables.values():
33 table.kwargs['mysql_engine'] = 'InnoDB'
34 table.kwargs['mysql_charset'] = 'utf8'
35
36 ### Connect
37 engine_args[engine_prefix + 'url'] = uri
38 engine = engine_from_config(engine_args, prefix=engine_prefix)
39 conn = engine.connect()
40 metadata.bind = engine
41
42 all_session_args = dict(autoflush=True, autocommit=False, bind=engine)
43 all_session_args.update(session_args)
44 sm = orm.sessionmaker(class_=MultilangSession, **all_session_args)
45 session = MultilangScopedSession(sm)
46
47 return session