-from __future__ import unicode_literals
+from __future__ import absolute_import
+from __future__ import division
from __future__ import print_function
+from __future__ import unicode_literals
+
+from future_builtins import ascii, filter, hex, map, oct, zip
import inspect
import os
import psycopg2.pool
psycopg2.extensions.register_type(psycopg2.extensions.UNICODE)
+psycopg2.extensions.register_type(psycopg2.extensions.UNICODEARRAY)
-def one(values):
- if values == None or len(values) == 0:
- return None
- else:
- assert len(values) == 1
- return values[0]
+class ConnectionError():
+ pass
class connect(object):
def __init__(self, dsn):
+ options = dsn.copy()
+ if 'cache' in options:
+ del options['cache']
+
+ if 'cache' in dsn:
+ cached = True
+ cache = dsn['cache']
+ else:
+ cached = False
+ cache = {
+ 'hstore': None,
+ }
+
attempt = 0
while True:
try:
- self.driver = psycopg2.connect(**dsn)
+ self.driver = psycopg2.connect(**options)
break
except psycopg2.OperationalError, e:
+ if e.message.startswith('could not connect to server: '):
+ raise ConnectionError()
if attempt == 2:
- raise e
+ raise
attempt = attempt + 1
- try:
- self.driver.set_client_encoding('UNICODE')
- self.driver.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
- except:
- self.driver.close()
+ self.driver.autocommit = True
+
+ # XXX: all of my databases default to this...
+ #try:
+ # self.driver.set_client_encoding('UNICODE')
+ #except:
+ # self.driver.close()
+ # raise
+
+ hstore = cache['hstore']
+ if hstore == None:
+ hstore = psycopg2.extras.HstoreAdapter.get_oids(self.driver)
+ if hstore != None:
+ hstore = hstore[0]
+ cache['hstore'] = hstore
+
+ if hstore != None:
+ try:
+ psycopg2.extras.register_hstore(self.driver, globally=False, unicode=True, oid=hstore)
+ except psycopg2.ProgrammingError, e:
+ pass
+
+ if not cached:
+ dsn['cache'] = cache
def close(self):
self.driver.close()
def __exit__(self, type, value, traceback):
self.close()
+ def begin(self):
+ self.driver.autocommit = False
+
+ def commit(self):
+ self.driver.commit()
+
+ def rollback(self):
+ self.driver.rollback()
+
@contextmanager
def cursor(self):
cursor = self.driver.cursor(cursor_factory=psycopg2.extras.DictCursor)
@contextmanager
def execute(self, statement, depth=0, context=None):
+ # two frames, accounting for execute() and @contextmanager
+ frame = inspect.currentframe(depth + 2)
+
with self.cursor() as cursor:
- # two frames, accounting for execute() and @contextmanager
- locals = inspect.currentframe(depth + 2).f_locals
- try:
- if context == None:
- context = locals
- cursor.execute(statement.format(**locals), context)
- finally:
- del locals
+ f_globals = None
+ f_locals = frame.f_locals
+
+ if context == None:
+ context = dict(**f_locals)
+
+ start = 0
+ while True:
+ percent = statement.find('%', start)
+ if percent == -1:
+ break
+
+ next = statement[percent + 1]
+ if next == '(':
+ start = statement.index(')', percent + 2) + 2
+ assert statement[start - 1] == 's'
+ elif next == '{':
+ start = statement.index('}', percent + 2)
+ assert statement[start + 1] == 's'
+ code = statement[percent + 2:start]
+
+ if f_globals == None:
+ f_globals = frame.f_globals
+
+ key = '__cyql__%i' % (percent,)
+ # XXX: compile() in the frame's context
+ context[key] = eval(code, f_globals, f_locals)
+
+ statement = '%s%%(%s)%s' % (statement[0:percent], key, statement[start + 1:])
+ start = percent + len(key) + 4
+ elif next in ('%', 's'):
+ start = percent + 2
+ else:
+ assert False
+
+ cursor.execute(statement, context)
+
+ del context
+ del f_locals
+ del f_globals
+
yield cursor
@contextmanager
def transact(self, synchronous_commit=True):
- self.driver.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_READ_COMMITTED)
+ self.driver.autocommit = False
try:
with self.cursor() as cursor:
if not synchronous_commit:
self.driver.rollback()
raise
finally:
- self.driver.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
+ self.driver.autocommit = True
- def one_(self, statement):
- with self.execute(statement, 2) as cursor:
+ def one_(self, statement, context=None):
+ with self.execute(statement, 2, context) as cursor:
one = cursor.fetchone()
if one == None:
return None
with self.execute(statement, 1) as cursor:
return cursor.callproc(procedure, *parameters)
- def run(self, statement, locals=None):
- with self.execute(statement, 1, locals) as cursor:
+ def run(self, statement, context=None):
+ with self.execute(statement, 1, context) as cursor:
return cursor.rowcount
+ def gen(self, statement):
+ with self.execute(statement, 1) as cursor:
+ while True:
+ fetch = cursor.fetchone()
+ if fetch == None:
+ break
+ yield fetch
+
@contextmanager
def set(self, statement):
- with self.execute(statement, 1) as cursor:
+ with self.execute(statement, 2) as cursor:
yield cursor
- def all(self, statement):
- with self.execute(statement, 1) as cursor:
+ def all(self, statement, context=None):
+ with self.execute(statement, 1, context) as cursor:
return cursor.fetchall()
- def one(self, statement):
- return self.one_(statement)
+ def one(self, statement, context=None):
+ return self.one_(statement, context)
def has(self, statement):
- return one(self.one_('select exists(%s)' % (statement,)))
+ exists, = self.one_('select exists(%s)' % (statement,))
+ return exists
def connected(dsn):
def wrapped(method):
projects / cyql.git / blobdiff