def test_same_name():
v1 = ContextVar('v')
v2 = ContextVar('v')
v1.set(1)
v2.set(2)
assert v1.get() == 1
assert v2.get() == 2
def __init__(self, dialect, pool, loop,
logging_name=None, echo=None, execution_options=None):
self._sa_engine = _SAEngine(
dialect,
logging_name=logging_name, echo=echo,
execution_options=execution_options)
self._dialect = dialect
self._pool = pool
self._loop = loop
self._ctx = ContextVar('gino', default=None)
def __init__(self,
connection_name: str,
fetch_inserted: bool = True,
**kwargs) -> None:
self.log = logging.getLogger("db_client")
self.connection_name = connection_name
self.fetch_inserted = fetch_inserted
self._current_transaction = ContextVar(self.connection_name,
default=self) # Type: dict
async def test_max_count_cancel(loop):
t_exec = 0
delay_exec = 0.1
event = Event()
executions = 0
leeway = 100
max_count = 5
arg = ContextVar("arg")
tasks = []
executing_task = None
@aggregate(leeway * 1000, max_count)
async def pow(*args: float, power: float = 2) -> List[float]:
nonlocal executions, executing_task, t_exec, delay_exec
await asyncio.sleep(0.03)
t_exec = time.time()
executions += 1
executing_task = tasks[arg.get()]
event.set()
await asyncio.sleep(delay_exec)
return [math.pow(num, power) for num in args]
async def pho(num: int):
arg.set(num)
return await pow(float(num))
t = time.time()
tasks = []
for i in range(5):
tasks.append(loop.create_task(pho(i)))
# Execution must have started
await event.wait()
event.clear()
elapsed = t_exec - t
assert 0 < elapsed < leeway
assert all(not task.done() for task in tasks)
assert executions == 1
first_executing_task = executing_task
first_executing_task.cancel()
# Another task must have tried to execute
await event.wait()
assert executions == 2
assert first_executing_task.cancelled()
assert all(
not task.done() for task in tasks
if task is not first_executing_task
)
# Must have finished
await wait(tasks)
assert first_executing_task.cancelled()
for i, task in enumerate(tasks):
if task is first_executing_task:
continue
assert task.done()
assert task.result() == math.pow(i, 2)
class BaseDBAsyncClient:
query_class = Query
executor_class = BaseExecutor
schema_generator = BaseSchemaGenerator
capabilities = Capabilities("")
def __init__(self,
connection_name: str,
fetch_inserted: bool = True,
**kwargs) -> None:
self.log = logging.getLogger("db_client")
self.connection_name = connection_name
self.fetch_inserted = fetch_inserted
self._current_transaction = ContextVar(self.connection_name,
default=self) # Type: dict
async def create_connection(self, with_db: bool) -> None:
raise NotImplementedError() # pragma: nocoverage
async def close(self) -> None:
raise NotImplementedError() # pragma: nocoverage
async def db_create(self) -> None:
raise NotImplementedError() # pragma: nocoverage
async def db_delete(self) -> None:
raise NotImplementedError() # pragma: nocoverage
def acquire_connection(self):
raise NotImplementedError() # pragma: nocoverage
def _in_transaction(self) -> "BaseTransactionWrapper":
raise NotImplementedError() # pragma: nocoverage
async def execute_insert(self, query: str, values: list) -> Any:
raise NotImplementedError() # pragma: nocoverage
async def execute_query(self, query: str) -> Sequence[dict]:
raise NotImplementedError() # pragma: nocoverage
async def execute_script(self, query: str) -> None:
raise NotImplementedError() # pragma: nocoverage
def get_current_transaction(self) -> "BaseDBAsyncClient":
return self._current_transaction.get()
class GinoEngine:
"""
Connects a :class:`~.dialects.base.Pool` and
:class:`~sqlalchemy.engine.interfaces.Dialect` together to provide a source
of database connectivity and behavior.
A :class:`.GinoEngine` object is instantiated publicly using the
:func:`gino.create_engine` function or
:func:`db.set_bind() <gino.api.Gino.set_bind>` method.
.. seealso::
:doc:`/engine`
"""
connection_cls = GinoConnection
"""Customizes the connection class to use, default is
:class:`.GinoConnection`."""
def __init__(self, dialect, pool, loop,
logging_name=None, echo=None, execution_options=None):
self._sa_engine = _SAEngine(
dialect,
logging_name=logging_name, echo=echo,
execution_options=execution_options)
self._dialect = dialect
self._pool = pool
self._loop = loop
self._ctx = ContextVar('gino', default=None)
@property
def dialect(self):
"""
Read-only property for the
:class:`~sqlalchemy.engine.interfaces.Dialect` of this engine.
"""
return self._dialect
@property
def raw_pool(self):
"""
Read-only access to the underlying database connection pool instance.
This depends on the actual dialect in use, :class:`~asyncpg.pool.Pool`
of asyncpg for example.
"""
return self._pool.raw_pool
def acquire(self, *, timeout=None, reuse=False, lazy=False, reusable=True):
"""
Acquire a connection from the pool.
There are two ways using this method - as an asynchronous context
manager::
async with engine.acquire() as conn:
# play with the connection
which will guarantee the connection is returned to the pool when
leaving the ``async with`` block; or as a coroutine::
conn = await engine.acquire()
try:
# play with the connection
finally:
await conn.release()
where the connection should be manually returned to the pool with
:meth:`conn.release() <.GinoConnection.release>`.
Within the same context (usually the same :class:`~asyncio.Task`, see
also :doc:`/transaction`), a nesting acquire by default re
:param timeout: Block up to ``timeout`` seconds until there is one free
connection in the pool. Default is ``None`` - block forever until
succeeded. This has no effect when ``lazy=True``, and depends on the
actual situation when ``reuse=True``.
:param reuse: Reuse the latest reusable acquired connection (before
it's returned to the pool) in current context if there is one, or
borrow a new one if none present. Default is ``False`` for always
borrow a new one. This is useful when you are in a nested method call
series, wishing to use the same connection without passing it around
as parameters. See also: :doc:`/transaction`. A reusing connection is
not reusable even if ``reusable=True``. If the reused connection
happened to be a lazy one, then the reusing connection is lazy too.
:param lazy: Don't acquire the actual underlying connection yet - do it
only when needed. Default is ``False`` for always do it immediately.
This is useful before entering a code block which may or may not make
use of a given connection object. Feeding in a lazy connection will
save the borrow-return job if the connection is never used. If
setting ``reuse=True`` at the same time, then the reused connection -
if any - applies the same laziness. For example, reusing a lazy
connection with ``lazy=False`` will cause the reused connection to
acquire an underlying connection immediately.
:param reusable: Mark this connection as reusable or otherwise. This
has no effect if it is a reusing connection. All reusable connections
are placed in a stack, any reusing acquire operation will always
reuse the top (latest) reusable connection. One reusable connection
may be reused by several reusing connections - they all share one
same underlying connection. Acquiring a connection with
``reusable=False`` and ``reusing=False`` makes it a cleanly isolated
connection which is only referenced once here.
:return: A :class:`.GinoConnection` object.
"""
return _AcquireContext(functools.partial(
self._acquire, timeout, reuse, lazy, reusable))
async def _acquire(self, timeout, reuse, lazy, reusable):
stack = _ContextualStack(self._ctx)
if reuse and stack:
dbapi_conn = _ReusingDBAPIConnection(self._dialect.cursor_cls,
stack.top)
reusable = False
else:
dbapi_conn = _DBAPIConnection(self._dialect.cursor_cls, self._pool)
rv = self.connection_cls(self._dialect,
_SAConnection(self._sa_engine, dbapi_conn),
stack if reusable else None)
dbapi_conn.gino_conn = rv
if not lazy:
await dbapi_conn.acquire(timeout=timeout)
if reusable:
stack.push(dbapi_conn)
return rv
@property
def current_connection(self):
"""
Gets the most recently acquired reusable connection in the context.
``None`` if there is no such connection.
:return: :class:`.GinoConnection`
"""
stack = self._ctx.get()
if stack:
return stack[-1].gino_conn
async def close(self):
"""
Close the engine, by closing the underlying pool.
"""
await self._pool.close()
async def all(self, clause, *multiparams, **params):
"""
Acquires a connection with ``reuse=True`` and runs
:meth:`~.GinoConnection.all` on it. ``reuse=True`` means you can safely
do this without borrowing more than one underlying connection::
async with engine.acquire():
await engine.all('SELECT ...')
The same applies for other query methods.
"""
async with self.acquire(reuse=True) as conn:
return await conn.all(clause, *multiparams, **params)
async def first(self, clause, *multiparams, **params):
"""
Runs :meth:`~.GinoConnection.first`, See :meth:`.all`.
"""
async with self.acquire(reuse=True) as conn:
return await conn.first(clause, *multiparams, **params)
async def scalar(self, clause, *multiparams, **params):
"""
Runs :meth:`~.GinoConnection.scalar`, See :meth:`.all`.
"""
async with self.acquire(reuse=True) as conn:
return await conn.scalar(clause, *multiparams, **params)
async def status(self, clause, *multiparams, **params):
"""
Runs :meth:`~.GinoConnection.status`. See also :meth:`.all`.
"""
async with self.acquire(reuse=True) as conn:
return await conn.status(clause, *multiparams, **params)
def compile(self, clause, *multiparams, **params):
"""
A shortcut for :meth:`~gino.dialects.base.AsyncDialectMixin.compile` on
the dialect, returns raw SQL string and parameters according to the
rules of the dialect.
"""
return self._dialect.compile(clause, *multiparams, **params)
def transaction(self, *args, timeout=None, reuse=True, reusable=True,
**kwargs):
"""
Borrows a new connection and starts a transaction with it.
Different to :meth:`.GinoConnection.transaction`, transaction on engine
level supports only managed usage::
async with engine.transaction() as tx:
# play with transaction here
Where the implicitly acquired connection is available as
:attr:`tx.connection <gino.transaction.GinoTransaction.connection>`.
By default, :meth:`.transaction` acquires connection with
``reuse=True`` and ``reusable=True``, that means it by default tries to
create a nested transaction instead of a new transaction on a new
connection. You can change the default behavior by setting these two
arguments.
The other arguments are the same as
:meth:`~.GinoConnection.transaction` on connection.
.. seealso::
:meth:`.GinoEngine.acquire`
:meth:`.GinoConnection.transaction`
:class:`~gino.transaction.GinoTransaction`
:return: A asynchronous context manager that yields a
:class:`~gino.transaction.GinoTransaction`
"""
return _TransactionContext(self.acquire(
timeout=timeout, reuse=reuse, reusable=reusable), (args, kwargs))
def iterate(self, clause, *multiparams, **params):
"""
Creates a server-side cursor in database for large query results.
This requires that there is a reusable connection in the current
context, and an active transaction is present. Then its
:meth:`.GinoConnection.iterate` is executed and returned.
"""
connection = self.current_connection
if connection is None:
raise ValueError(
'No Connection in context, please provide one')
return connection.iterate(clause, *multiparams, **params)
def update_execution_options(self, **opt):
"""Update the default execution_options dictionary
of this :class:`.GinoEngine`.
.. seealso::
:meth:`sqlalchemy.engine.Engine.update_execution_options`
:meth:`.GinoConnection.execution_options`
"""
self._sa_engine.update_execution_options(**opt)
async def _run_visitor(self, *args, **kwargs):
async with self.acquire(reuse=True) as conn:
await getattr(conn, '_run_visitor')(*args, **kwargs)
import asyncio
import pytest
from aiocontextvars import ContextVar
v = ContextVar('v')
@pytest.mark.asyncio
async def parallel(x):
v.set(x)
await asyncio.sleep(0.1)
assert v.get() == x
@pytest.mark.asyncio
async def test_parallel():
await asyncio.gather(*map(parallel, range(16)))
# noinspection PyUnusedLocal
@pytest.mark.asyncio
async def test_parallel_with_inherit():
await asyncio.gather(*map(parallel, range(16)))
# noinspection PyUnusedLocal
@pytest.mark.asyncio
async def test_inherit():
v.set('initial')
async def sub():
def test_get():
v = ContextVar('v')
with pytest.raises(LookupError):
v.get()
assert v.get(456) == 456
return v
def test_get_with_default():
v = ContextVar('v', default=123)
assert v.get() == 123
assert v.get(456) == 456
return v
import asyncio
import textwrap
import gettext
from aiocontextvars import ContextVar
import discord
from discord.abc import Messageable
from discord.ext import commands
from meowth.core import checks
from meowth.utils import convert_to_bool, make_embed, bold
cvar = ContextVar('bot')
def ctx_setup(loop):
import builtins
builtins.__dict__['_'] = use_current_gettext
builtins.__dict__['get_ctx'] = cvar.get
builtins.__dict__['__cvar__'] = cvar
def use_current_gettext(*args, **kwargs):
return cvar.get().get_text(*args, **kwargs)
class Context(commands.Context):
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.dbi = self.bot.dbi
@property
@abstractmethod
def user(self) -> User:
"""Get the current request user"""
@property
@abstractmethod
def roles(self) -> List[str]:
"""Get current user roles"""
@abstractmethod
def validate_roles(self, required_roles: List[str] = None):
"""Check if a user is authenticated"""
user_var = ContextVar('user', default=None)
class StandardAuthProvider(AuthProvider):
def setup(self, user: User) -> None:
user_var.set(user)
@property
def user(self) -> User:
if not user_var.get():
raise AuthenticationError("Not authenticated.")
return user_var.get()
@property
def roles(self) -> List[str]:
from abc import ABC, abstractmethod
from aiocontextvars import ContextVar
from .tenant import Tenant
class TenantProvider(ABC):
"""Tenant service."""
@abstractmethod
def setup(self, tenant: Tenant) -> None:
"Setup current tenant method to be implemented."
@property
@abstractmethod
def tenant(self) -> Tenant:
"""Get the current tenant"""
tenant_var = ContextVar('tenant', default=None)
class StandardTenantProvider(TenantProvider):
def setup(self, tenant: Tenant) -> None:
tenant_var.set(tenant)
@property
def tenant(self) -> Tenant:
if not tenant_var.get():
raise ValueError('No tenant has been set.')
return tenant_var.get()
async def test_max_count_multiple_batches_cancel(loop):
t_exec = 0
delay_exec = 0.1
event = Event()
executions = 0
leeway = 0.1
max_count = 5
arg = ContextVar('arg')
tasks = []
executing_task = None
@aggregate(leeway * 1000, max_count)
async def pow(*args: float, power: float = 2) -> List[float]:
nonlocal executions, executing_task, t_exec, delay_exec
t_exec = monotonic()
executions += 1
executing_task = tasks[arg.get()]
event.set()
await asyncio.sleep(delay_exec)
return [math.pow(num, power) for num in args]
async def pho(num: int):
arg.set(num)
return await pow(float(num))
t = monotonic()
tasks = []
for i in range(9):
tasks.append(loop.create_task(pho(i)))
# Execution of the first batch must have started
await event.wait()
event.clear()
assert all(not task.done() for task in tasks)
assert executions == 1
first_executing_task = executing_task
first_executing_task.cancel()
# Another task must have tried to execute
await event.wait()
event.clear()
assert executions == 2
assert first_executing_task.cancelled()
assert all(
not task.done() for task in tasks
if task is not first_executing_task
)
await wait(tasks[:5])
# First batch must have finished
assert first_executing_task.cancelled()
for i, task in enumerate(tasks[:5]):
if task is first_executing_task:
continue
assert task.done()
assert task.result() == math.pow(i, 2)
# Second batch must have started execution
await event.wait()
elapsed = t_exec - t
assert leeway * 0.9 < elapsed < leeway * 2
assert all(not task.done() for task in tasks[5:])
assert executions == 3
# Second batch mast have finished
await wait(tasks[5:])
for i, task in enumerate(tasks[5:], start=5):
assert task.done()
assert task.result() == math.pow(i, 2)
from .base import state_machine
from .base import super_check, call_with_super_check
from .base.state_machine import StateEntryFailed, StateMachine, event
from .base import TransitionFailed
from . import events
from . import persistence
from . import process_comms
from . import process_states
from . import utils
# pylint: disable=too-many-lines
__all__ = ['Process', 'ProcessSpec', 'BundleKeys', 'TransitionFailed']
_LOGGER = logging.getLogger(__name__)
PROCESS_STACK = ContextVar('process stack', default=[])
class BundleKeys:
"""
String keys used by the process to save its state in the state bundle.
See :func:`save_instance_state` and :func:`load_instance_state`.
"""
# pylint: disable=too-few-public-methods
INPUTS_RAW = 'INPUTS_RAW'
INPUTS_PARSED = 'INPUTS_PARSED'
OUTPUTS = 'OUTPUTS'
class ProcessStateMachineMeta(abc.ABCMeta, state_machine.StateMachineMeta):
Except:
- we do it on top of asyncio
- we do not pretend implementing an equivalent of Trio; just the philosophy!
- no pretensions; if you want more Trio-like, just switch to Trio!
"""
import asyncio
import logging
import sys
import time
from typing import Awaitable, Optional
from aiocontextvars import ContextVar
from traio.task import NamedFuture, TaskWrapper
SCOPE = ContextVar('traio_scope')
DEFAULT_LOGGER = logging.getLogger('traio')
DEFAULT_LOGGER.setLevel(logging.CRITICAL)
PY37 = sys.version_info >= (3, 7)
def current_task(loop=None):
"""
Return current task. Wraps difference between before/after py37
"""
if PY37:
return asyncio.current_task(loop)
return asyncio.Task.current_task(loop)
class CancelError(Exception):
try:
return "%s.%s.%s" % (
func.im_class.__module__, # type: ignore
func.im_class.__name__, # type: ignore
func.__name__,
)
except Exception:
pass
func_qualname = (
getattr(func, "__qualname__", None) or getattr(func, "__name__", None) or None
) # type: Optional[str]
if not func_qualname:
# No idea what it is
return None
# Methods in Python 3
# Functions
# Classes
try:
return "%s.%s" % (func.__module__, func_qualname)
except Exception:
pass
# Possibly a lambda
return func_qualname
disable_capture_event = ContextVar("disable_capture_event")
from aiocontextvars import ContextVar
from asyncpg import Connection, create_pool
from asyncpg.pool import Pool
class ConnectionManager(ABC):
@abstractmethod
async def get(self, pool: str = "") -> Connection:
"""Get a database connection."""
@abstractmethod
async def put(self, connection: Connection, pool: str = "") -> None:
"""Put a database connection back in its pool."""
connections_var = ContextVar('connections', default=None)
class DefaultConnectionManager(ConnectionManager):
def __init__(self, settings: List[Dict[str, Any]],
default: str = "") -> None:
self.settings = settings
self.pools: Dict[str, Pool] = {}
self.default = default or self.settings[0]['name']
async def get(self, pool: str = "") -> Connection:
connection: Connection = connections_var.get(None)
if connection is None:
if self.pools == {}:
await self._setup()
