mirror of https://github.com/ekimekim/wubloader
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
76 lines
2.6 KiB
Python
76 lines
2.6 KiB
Python
|
|
"""
|
|
Code shared between components that touch the database.
|
|
Note that this code requires psycopg2 and psycogreen, but the common module
|
|
as a whole does not to avoid needing to install them for components that don't need it.
|
|
"""
|
|
|
|
from contextlib import contextmanager
|
|
|
|
import psycopg2
|
|
import psycopg2.extensions
|
|
import psycopg2.extras
|
|
from psycogreen.gevent import patch_psycopg
|
|
|
|
|
|
class DBManager(object):
|
|
"""Patches psycopg2 before any connections are created. Stores connect info
|
|
for easy creation of new connections, and sets some defaults before
|
|
returning them.
|
|
|
|
It has the ability to serve as a primitive connection pool, as getting a
|
|
new conn will return existing conns it knows about first, but this mainly
|
|
just exists to re-use the initial conn used to test the connection, and you
|
|
should use a real conn pool for any non-trivial use.
|
|
|
|
Returned conns are set to seralizable isolation level, autocommit, and use
|
|
NamedTupleCursor cursors."""
|
|
def __init__(self, **connect_kwargs):
|
|
patch_psycopg()
|
|
self.conns = []
|
|
self.connect_kwargs = connect_kwargs
|
|
# get a connection to test whether connection is working.
|
|
conn = self.get_conn()
|
|
self.put_conn(conn)
|
|
|
|
def put_conn(self, conn):
|
|
self.conns.append(conn)
|
|
|
|
def get_conn(self):
|
|
if self.conns:
|
|
return self.conns.pop(0)
|
|
conn = psycopg2.connect(cursor_factory=psycopg2.extras.NamedTupleCursor, **self.connect_kwargs)
|
|
# We use serializable because it means less issues to think about,
|
|
# we don't care about the performance concerns and everything we do is easily retryable.
|
|
# This shouldn't matter in practice anyway since everything we're doing is either read-only
|
|
# searches or targetted single-row updates.
|
|
conn.isolation_level = psycopg2.extensions.ISOLATION_LEVEL_SERIALIZABLE
|
|
conn.autocommit = True
|
|
return conn
|
|
|
|
|
|
@contextmanager
|
|
def transaction(conn):
|
|
"""Helper context manager that runs the code block as a single database transaction
|
|
instead of in autocommit mode. The only difference between this and "with conn" is
|
|
that we explicitly disable then re-enable autocommit."""
|
|
old_autocommit = conn.autocommit
|
|
conn.autocommit = False
|
|
try:
|
|
with conn:
|
|
yield
|
|
finally:
|
|
conn.autocommit = old_autocommit
|
|
|
|
|
|
def query(conn, query, *args, **kwargs):
|
|
"""Helper that takes a conn, creates a cursor and executes query against it,
|
|
then returns the cursor.
|
|
Variables may be given as positional or keyword args (but not both), corresponding
|
|
to %s vs %(key)s placeholder forms."""
|
|
if args and kwargs:
|
|
raise TypeError("Cannot give both args and kwargs")
|
|
cur = conn.cursor()
|
|
cur.execute(query, args or kwargs or None)
|
|
return cur
|