Merge pull request #219 from fantix/t14

Fixed #14, implemented prepared statement

Author: fantix

gino/crud.py

@@ -413,7 +413,8 @@ def _init_table(cls, sub_cls):
  return rv
 
  @classmethod
- async def _create_without_instance(cls, bind=None, timeout=DEFAULT, **values):
+ async def _create_without_instance(cls, bind=None, timeout=DEFAULT,
+ **values):
  return await cls(**values)._create(bind=bind, timeout=timeout)
 
  async def _create(self, bind=None, timeout=DEFAULT):

gino/dialects/asyncpg.py

@@ -99,8 +99,8 @@ async def forward(self, n, *, timeout=base.DEFAULT):
 
 
 class PreparedStatement(base.PreparedStatement):
- def __init__(self, prepared):
- super().__init__()
+ def __init__(self, prepared, clause=None):
+ super().__init__(clause)
  self._prepared = prepared
 
  def _get_iterator(self, *params, **kwargs):
@@ -111,27 +111,41 @@ async def _get_cursor(self, *params, **kwargs):
  iterator = await self._prepared.cursor(*params, **kwargs)
  return AsyncpgCursor(self.context, iterator)
 
+ async def _execute(self, params, one):
+ if one:
+ rv = await self._prepared.fetchrow(*params)
+ if rv is None:
+ rv = []
+ else:
+ rv = [rv]
+ else:
+ rv = await self._prepared.fetch(*params)
+ return self._prepared.get_statusmsg(), rv
+
 
 class DBAPICursor(base.DBAPICursor):
  def __init__(self, dbapi_conn):
  self._conn = dbapi_conn
  self._attributes = None
  self._status = None
 
- async def prepare(self, query, timeout):
+ async def prepare(self, context, clause=None):
+ timeout = context.timeout
  if timeout is None:
  conn = await self._conn.acquire(timeout=timeout)
  else:
  before = time.monotonic()
  conn = await self._conn.acquire(timeout=timeout)
  after = time.monotonic()
  timeout -= after - before
- prepared = await conn.prepare(query, timeout=timeout)
+ prepared = await conn.prepare(context.statement, timeout=timeout)
  try:
  self._attributes = prepared.get_attributes()
  except TypeError: # asyncpg <= 0.12.0
  self._attributes = []
- return PreparedStatement(prepared)
+ rv = PreparedStatement(prepared, clause)
+ rv.context = context
+ return rv
 
  async def async_execute(self, query, timeout, args, limit=0, many=False):
  if timeout is None:

gino/dialects/base.py

@@ -26,7 +26,7 @@ def executemany(self, statement, parameters):
  def description(self):
  raise NotImplementedError
 
- async def prepare(self, query, timeout):
+ async def prepare(self, context, clause=None):
  raise NotImplementedError
 
  async def async_execute(self, query, timeout, args, limit=0, many=False):
@@ -67,18 +67,62 @@ async def rollback(self):
 
 
 class PreparedStatement:
- def __init__(self):
+ def __init__(self, clause=None):
  self.context = None
+ self.clause = clause
 
  def iterate(self, *params, **kwargs):
  return _PreparedIterableCursor(self, params, kwargs)
 
+ async def _do_execute(self, multiparams, params, one=False,
+ return_model=True, status=False):
+ ctx = self.context.connection.execute(
+ self.clause, *multiparams, **params).context
+ if ctx.executemany:
+ raise ValueError('PreparedStatement does not support multiple '
+ 'parameters.')
+ assert ctx.statement == self.context.statement, (
+ 'Prepared statement generated different SQL with parameters')
+ params = []
+ for val in ctx.parameters[0]:
+ params.append(val)
+ msg, rows = await self._execute(params, one)
+ if status:
+ return msg
+ item = self.context.process_rows(rows, return_model=return_model)
+ if one:
+ if item:
+ item = item[0]
+ else:
+ item = None
+ return item
+
+ async def all(self, *multiparams, **params):
+ return await self._do_execute(multiparams, params)
+
+ async def first(self, *multiparams, **params):
+ return await self._do_execute(multiparams, params, one=True)
+
+ async def scalar(self, *multiparams, **params):
+ rv = await self._do_execute(multiparams, params, one=True,
+ return_model=False)
+ if rv:
+ return rv[0]
+ else:
+ return None
+
+ async def status(self, *multiparams, **params):
+ return await self._do_execute(multiparams, params, status=True)
+
  def _get_iterator(self, *params, **kwargs):
  raise NotImplementedError
 
  async def _get_cursor(self, *params, **kwargs):
  raise NotImplementedError
 
+ async def _execute(self, params, one):
+ raise NotImplementedError
+
 
 class _PreparedIterableCursor:
  def __init__(self, prepared, params, kwargs):
@@ -100,9 +144,7 @@ def __init__(self, context):
  self._context = context
 
  async def _iterate(self):
- prepared = await self._context.cursor.prepare(self._context.statement,
- self._context.timeout)
- prepared.context = self._context
+ prepared = await self._context.cursor.prepare(self._context)
  return prepared.iterate(*self._context.parameters[0],
  timeout=self._context.timeout)
 
@@ -173,6 +215,9 @@ def iterate(self):
  raise ValueError('too many multiparams')
  return _IterableCursor(self._context)
 
+ async def prepare(self, clause):
+ return await self._context.cursor.prepare(self._context, clause)
+
  def _soft_close(self):
  pass
 
@@ -237,6 +282,88 @@ def process_rows(self, rows, return_model=True):
  def get_result_proxy(self):
  return _ResultProxy(self)
 
+ @classmethod
+ def _init_compiled_prepared(cls, dialect, connection, dbapi_connection,
+ compiled, parameters):
+ self = cls.__new__(cls)
+ self.root_connection = connection
+ self._dbapi_connection = dbapi_connection
+ self.dialect = connection.dialect
+
+ self.compiled = compiled
+
+ # this should be caught in the engine before
+ # we get here
+ assert compiled.can_execute
+
+ self.execution_options = compiled.execution_options.union(
+ connection._execution_options)
+
+ self.result_column_struct = (
+ compiled._result_columns, compiled._ordered_columns,
+ compiled._textual_ordered_columns)
+
+ self.unicode_statement = util.text_type(compiled)
+ if not dialect.supports_unicode_statements:
+ self.statement = self.unicode_statement.encode(
+ self.dialect.encoding)
+ else:
+ self.statement = self.unicode_statement
+
+ self.isinsert = compiled.isinsert
+ self.isupdate = compiled.isupdate
+ self.isdelete = compiled.isdelete
+ self.is_text = compiled.isplaintext
+
+ self.executemany = False
+
+ self.cursor = self.create_cursor()
+
+ if self.isinsert or self.isupdate or self.isdelete:
+ self.is_crud = True
+ self._is_explicit_returning = bool(compiled.statement._returning)
+ self._is_implicit_returning = bool(
+ compiled.returning and not compiled.statement._returning)
+
+ if self.dialect.positional:
+ self.parameters = [dialect.execute_sequence_format()]
+ else:
+ self.parameters = [{}]
+ self.compiled_parameters = [{}]
+
+ return self
+
+ @classmethod
+ def _init_statement_prepared(cls, dialect, connection, dbapi_connection,
+ statement, parameters):
+ """Initialize execution context for a string SQL statement."""
+
+ self = cls.__new__(cls)
+ self.root_connection = connection
+ self._dbapi_connection = dbapi_connection
+ self.dialect = connection.dialect
+ self.is_text = True
+
+ # plain text statement
+ self.execution_options = connection._execution_options
+
+ if self.dialect.positional:
+ self.parameters = [dialect.execute_sequence_format()]
+ else:
+ self.parameters = [{}]
+
+ self.executemany = False
+
+ if not dialect.supports_unicode_statements and \
+ isinstance(statement, util.text_type):
+ self.unicode_statement = statement
+ self.statement = dialect._encoder(statement)[0]
+ else:
+ self.statement = self.unicode_statement = statement
+
+ self.cursor = self.create_cursor()
+ return self
+
 
 class AsyncDialectMixin:
  cursor_cls = DBAPICursor

gino/engine.py

@@ -129,9 +129,26 @@ async def _release(self):
  pass
 
 
+# noinspection PyPep8Naming,PyMethodMayBeStatic
+class _bypass_no_param:
+ def keys(self):
+ return []
+
+
+_bypass_no_param = _bypass_no_param()
+
+
 # noinspection PyAbstractClass
 class _SAConnection(Connection):
- pass
+ def _execute_context(self, dialect, constructor,
+ statement, parameters,
+ *args):
+ if parameters == [_bypass_no_param]:
+ constructor = getattr(self.dialect.execution_ctx_cls,
+ constructor.__name__ + '_prepared',
+ constructor)
+ return super()._execute_context(dialect, constructor, statement,
+ parameters, *args)
 
 
 # noinspection PyAbstractClass
@@ -503,6 +520,10 @@ async def _run_visitor(self, visitorcallable, element, **kwargs):
  await visitorcallable(self.dialect, self,
  **kwargs).traverse_single(element)
 
+ async def prepare(self, clause):
+ return await self._execute(
+ clause, (_bypass_no_param,), {}).prepare(clause)
+
 
 class GinoEngine:
  """

tests/test_engine.py

@@ -170,6 +170,7 @@ async def test_scalar_return_none(bind):
 
 
 async def test_asyncpg_0120(bind, mocker):
+ # for asyncpg 0.12.0
  assert await bind.first('rollback') is None
 
  orig = getattr(asyncpg.Connection, '_do_execute')

tests/test_prepared_stmt.py

@@ -0,0 +1,46 @@
+from datetime import datetime
+
+import pytest
+from .models import db, User
+
+pytestmark = pytest.mark.asyncio
+
+
+async def test_compiled_and_bindparam(bind):
+ async with db.acquire() as conn:
+ # noinspection PyArgumentList
+ ins = await conn.prepare(User.insert().returning(
+ *User).execution_options(loader=User))
+ users = {}
+ for name in '12345':
+ u = await ins.first(nickname=name)
+ assert u.nickname == name
+ users[u.id] = u
+ get = await conn.prepare(
+ User.query.where(User.id == db.bindparam('uid')))
+ for key in users:
+ u = await get.first(uid=key)
+ assert u.nickname == users[key].nickname
+ assert (await get.all(uid=key))[0].nickname == u.nickname
+
+ assert await get.scalar(uid=-1) is None
+
+ with pytest.raises(ValueError, match='does not support multiple'):
+ await get.all([dict(uid=1), dict(uid=2)])
+
+ delete = await conn.prepare(
+ User.delete.where(User.nickname == db.bindparam('name')))
+ for name in '12345':
+ msg = await delete.status(name=name)
+ assert msg == 'DELETE 1'
+
+
+async def test_statement(engine):
+ async with engine.acquire() as conn:
+ stmt = await conn.prepare('SELECT now()')
+ last = None
+ for i in range(5):
+ now = await stmt.scalar()
+ assert isinstance(now, datetime)
+ assert last != now
+ last = now