class KOKOROEngine: __slots__ = ( '_engine', '_worker', ) def __init__(self, pool, dialect, u, single_worker=True, **kwargs): if single_worker: worker = ExecutorThread() else: worker = None self._worker = worker self._engine = Engine(pool, dialect, u, **kwargs) @property def uses_single_worker(self): """ Returns whether the engine uses a single worker. Returns ------- uses_single_worker : `bool` """ return (self._worker is not None) @property def dialect(self): return self._engine.dialect @property def _has_events(self): return self._engine._has_events @property def logger(self): return self._engine.logger @property def _execution_options(self): return self._engine._execution_options def _should_log_info(self): return self._engine._should_log_info() def connect(self): return ConnectionContextManager(self._connect()) async def _connect(self): executor = self._worker if executor is None: executor = current_thread().claim_executor() connection = await executor.execute(self._engine.connect) return AsyncConnection(connection, executor) def begin(self, close_with_result=False): executor = self._worker if executor is None: executor = current_thread().claim_executor() return EngineTransactionContextManager(self, close_with_result, executor) async def execute(self, *args, **kwargs): executor = self._worker if executor is None: executor = current_thread().claim_executor() result_proxy = await executor.execute( alchemy_incendiary(self._engine.execute, args, kwargs)) return AsyncResultProxy(result_proxy, executor) async def scalar(self, *args, **kwargs): executor = self._worker if executor is None: executor = current_thread().claim_executor() result_proxy = await executor.execute( alchemy_incendiary(self._engine.execute, args, kwargs)) async_result_proxy = AsyncResultProxy(result_proxy, executor) return await async_result_proxy.scalar() async def has_table(self, table_name, schema=None): return await current_thread().run_in_executor( alchemy_incendiary(self._engine.has_table, (table_name, schema))) async def table_names(self, schema=None, connection=None): task = alchemy_incendiary( self._engine.table_names, (schema, None if connection else connection._connection)) executor = self._worker if executor is None: return await current_thread().run_in_executor(task) else: return await executor.execute(task) def __del__(self): worker = self._worker if (worker is not None): self._worker = None worker.cancel()
class AsyncioEngine: """Mostly like :class:`sqlalchemy.engine.Engine` except some of the methods are coroutines.""" def __init__(self, pool, dialect, url, logging_name=None, echo=None, execution_options=None, loop=None, **kwargs): self._engine = Engine(pool, dialect, url, logging_name=logging_name, echo=echo, execution_options=execution_options, **kwargs) self._loop = loop max_workers = None # https://www.python.org/dev/peps/pep-0249/#threadsafety if dialect.dbapi.threadsafety < 2: # This might seem overly-restrictive, but when we instantiate an # AsyncioResultProxy from AsyncioEngine.execute, subsequent # fetchone calls could be in different threads. Let's limit to one. max_workers = 1 self._engine_executor = ThreadPoolExecutor(max_workers=max_workers) async def _run_in_thread(_self, _func, *args, **kwargs): return await _run_in_executor(_self._engine_executor, _func, _self._loop, *args, **kwargs) @property def dialect(self): return self._engine.dialect @property def _has_events(self): return self._engine._has_events @property def logger(self): return self._engine.logger @property def _execution_options(self): return self._engine._execution_options def _should_log_info(self): return self._engine._should_log_info() def connect(self): """Like :meth:`Engine.connect <sqlalchemy.engine.Engine.connect>`, but returns an awaitable that can also be used as an asynchronous context manager. Examples: .. code-block:: python conn = await engine.connect() await conn.execute(...) await conn.close() .. code-block:: python async with engine.connect() as conn: await conn.execute(...) """ return _ConnectionContextManager(self._connect()) async def _connect(self): executor = ThreadPoolExecutor(max_workers=1) connection = await _run_in_executor(executor, self._engine.connect, self._loop) return AsyncioConnection(connection, executor, self._loop) def begin(self, close_with_result=False): """Like :meth:`Engine.begin <sqlalchemy.engine.Engine.begin>`, but returns an asynchronous context manager. Example: .. code-block:: python async with engine.begin(): await engine.execute(...) """ return _EngineTransactionContextManager(self, close_with_result) async def execute(self, *args, **kwargs): """Like :meth:`Engine.execute <sqlalchemy.engine.Engine.execute>`, but is a coroutine that returns an :class:`AsyncioResultProxy`. Example: .. code-block:: python result = await engine.execute(...) data = await result.fetchall() .. warning:: Make sure to explicitly call :meth:`AsyncioResultProxy.close` if the :class:`~sqlalchemy.engine.ResultProxy` has pending rows remaining otherwise it will be closed during garbage collection. With SQLite, this will raise an exception since the DBAPI connection was created in a different thread. """ rp = await self._run_in_thread(self._engine.execute, *args, **kwargs) return AsyncioResultProxy(rp, self._run_in_thread) async def scalar(self, *args, **kwargs): """Like :meth:`Connection.close <sqlalchemy.engine.Engine.scalar>`, but is a coroutine. """ rp = await self.execute(*args, **kwargs) return await rp.scalar() async def has_table(self, table_name, schema=None): """Like :meth:`Engine.has_table <sqlalchemy.engine.Engine.has_table>`, but is a coroutine. """ return await self._run_in_thread(self._engine.has_table, table_name, schema) async def table_names(self, schema=None, connection: 'AsyncioConnection' = None): """Like :meth:`Engine.table_names <sqlalchemy.engine.Engine.table_names>`, but is a coroutine. """ run_in_thread = self._run_in_thread if connection is not None: run_in_thread = connection._run_in_thread connection = connection._connection return await run_in_thread(self._engine.table_names, schema, connection) def __repr__(self): r = ReprHelper(self) r.parantheses = ('<', '>') r.positional_from_attr('_engine') return str(r)
class AsyncEngine(ABC): def __init__(self, pool, dialect, url, logging_name=None, echo=None, execution_options=None, **kwargs): self._engine = Engine(pool, dialect, url, logging_name=logging_name, echo=echo, execution_options=execution_options, **kwargs) self._engine_worker = None @abstractmethod def _make_worker(self): raise NotImplementedError async def _run_in_thread(_self, _func, *args, **kwargs): if _self._engine_worker is None: _self._engine_worker = _self._make_worker() return await _self._engine_worker.run(_func, *args, **kwargs) @property def dialect(self): return self._engine.dialect @property def _has_events(self): return self._engine._has_events @property def logger(self): return self._engine.logger @property def _execution_options(self): return self._engine._execution_options def _should_log_info(self): return self._engine._should_log_info() @property def sync_engine(self): """Public property of the underlying SQLAlchemy engine.""" return self._engine def connect(self): """Like :meth:`Engine.connect <sqlalchemy.engine.Engine.connect>`, but returns an awaitable that can also be used as an asynchronous context manager. Examples: .. code-block:: python conn = await engine.connect() await conn.execute(...) await conn.close() .. code-block:: python async with engine.connect() as conn: await conn.execute(...) """ return _ConnectionContextManager(self._make_async_connection()) async def _make_async_connection(self): worker = self._make_worker() connection = await worker.run(self._engine.connect) return AsyncConnection(connection, worker, self) def begin(self, close_with_result=False): """Like :meth:`Engine.begin <sqlalchemy.engine.Engine.begin>`, but returns an asynchronous context manager. Example: .. code-block:: python async with engine.begin(): await engine.execute(...) """ return _EngineTransactionContextManager(self, close_with_result) async def execute(self, *args, **kwargs): """Like :meth:`Engine.execute <sqlalchemy.engine.Engine.execute>`, but is a coroutine that returns an :class:`AsyncioResultProxy`. Example: .. code-block:: python result = await engine.execute(...) data = await result.fetchall() .. warning:: Make sure to explicitly call :meth:`AsyncioResultProxy.close` if the :class:`~sqlalchemy.engine.ResultProxy` has pending rows remaining otherwise it will be closed during garbage collection. With SQLite, this will raise an exception since the DBAPI connection was created in a different thread. """ rp = await self._run_in_thread(self._engine.execute, *args, **kwargs) return AsyncResultProxy(rp, self._run_in_thread) async def scalar(self, *args, **kwargs): """Like :meth:`Connection.scalar <sqlalchemy.engine.Engine.scalar>`, but is a coroutine. """ rp = await self.execute(*args, **kwargs) return await rp.scalar() async def has_table(self, table_name, schema=None): """Like :meth:`Engine.has_table <sqlalchemy.engine.Engine.has_table>`, but is a coroutine. """ return await self._run_in_thread(self._engine.has_table, table_name, schema) async def table_names(self, schema=None, connection: 'AsyncConnection' = None): """Like :meth:`Engine.table_names <sqlalchemy.engine.Engine.table_names>`, but is a coroutine. """ run_in_thread = self._run_in_thread if connection is not None: run_in_thread = connection._worker.run connection = connection._connection return await run_in_thread(self._engine.table_names, schema, connection) def run_callable(self, callable_, *args, **kwargs): """Like :meth:`Engine.run_callable\ <sqlalchemy.engine.Engine.run_callable>`. .. warning:: This method blocks. It exists so that we can warn the user if they try to use an async engine for table reflection: .. code-block:: python Table(..., autoload_with=engine) """ warnings.warn( 'The AsyncEngine has been called in a blocking fashion, e.g. with ' 'Table(..., autoload_with=engine). You may wish to run it in a ' 'separate thread to avoid blocking the event loop. You can use ' 'Table(..., autoload_with=engine.sync_engine) to opt out of the ' 'warning for this blocking behaviour.', BlockingWarning) self._engine.run_callable(callable_, *args, **kwargs) def __repr__(self): r = ReprHelper(self) r.parantheses = ('<', '>') r.positional_from_attr('_engine') return str(r) def __getattr__(self, item): msg = '{!r} object has no attribute {!r}.'.format( self.__class__.__name__, item) if item == '_run_visitor': raise AttributeError( msg + ' Did you try to use Table.create(engine) or similar? ' 'You must use Table.create(engine.sync_engine) instead, which' 'is a blocking function. Consider using sqlalchemy.schema.' 'CreateTable instead') raise AttributeError(msg)