Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
from asyncpgsa import compile_query
from asyncpg.pool import Pool
from .mockconnection import MockConnection
from .mocktransactionmanager import MockTransactionManager
class MockSAPool(Pool):
def __init__(self, connection=None):
super().__init__(min_size=1,
max_size=1,
max_queries=1,
max_inactive_connection_lifetime=300.0,
setup=None,
loop=None,
init=None,
connection_class=MockConnection)
self.connection = connection
if not self.connection:
self.connection = MockConnection()
def __getattr__(self, item):
from asyncpg.pool import Pool
from asyncpg.connection import Connection
from asyncio.queues import Queue
from time import perf_counter
from sqlalchemy.dialects.postgresql import JSONB
import asyncpg
import yo.db
from yo.db import metadata
from yo.schema import TransportType
from yo.json import loads
logger = structlog.getLogger(__name__, source='YoDB')
PoolOrConn = TypeVar('PoolOrConn', Pool, Connection)
QItemId = int
NotificationId = int
class QueueStorage:
'''A mixin class to preserve compatability with asyncio.Queue which
calls len(self._queue)
'''
def __init__(self, loop, pool):
self.loop = loop
self.pool = pool
from pylru import WriteThroughCacheManager
from sqlalchemy.dialects.postgresql import JSONB
from yo.rpc_client import get_user_data
from ..schema import NOTIFICATION_TYPES
from ..schema import NotificationType
from yo.db import metadata
class UserNotFoundError(Exception):
pass
logger = structlog.getLogger(__name__, source='users')
PoolOrConn = TypeVar('PoolOrConn', Pool, Connection)
DEFAULT_SENTINEL = 'default_transports'
DEFAULT_USER_TRANSPORT_SETTINGS = {
"desktop": {
"notification_types": NOTIFICATION_TYPES,
"data": None
}
}
DEFAULT_USER_TRANSPORT_SETTINGS_STRING = ujson.dumps(DEFAULT_USER_TRANSPORT_SETTINGS)
CREATE_USER_STMT = '''INSERT INTO users(username,transports, created, updated) VALUES($1,$2,NOW(),NOW()) ON CONFLICT DO NOTHING RETURNING username'''
CREATE_USER_WITHOUT_TRANSPORTS_STMT = '''INSERT INTO users(username) VALUES($1) ON CONFLICT DO NOTHING RETURNING username'''
UPDATE_USER_TRANSPORTS_STMT = '''UPDATE users SET transports = $1 WHERE username = $2 RETURNING username'''
GET_USER_TRANSPORTS_STMT = '''SELECT transports FROM users WHERE username = $1'''
async def __aenter__(self):
if self._used:
raise RuntimeError('GinoAcquireContext is entered twice')
self._used = True
local = get_local()
if self._reuse and local:
stack = local.get('connection_stack')
if stack:
conn = stack[-1]
if not self._lazy:
conn = await conn.get()
return conn
if isinstance(self._bind, Pool):
self._lazy_conn = conn = LazyConnection(self._bind, self._timeout)
if local is not None:
local.setdefault('connection_stack', deque()).append(conn)
self._pop = True
if not self._lazy:
conn = await conn.get()
return conn
else:
return self._bind
from typing import TypeVar
import sqlalchemy as sa
import structlog
from asyncpg.pool import Pool
from asyncpg.connection import Connection
from ..schema import ActionStatus
from ..schema import TransportType
from yo.db import metadata
logger = structlog.getLogger(__name__, source='YoDB')
PoolOrConn = TypeVar('PoolOrConn', Pool, Connection)
ActionId = int
NotificationId = int
PERMANENT_FAIL_COUNT = 3
actions_table = sa.Table(
'actions',
metadata,
sa.Column('aid', sa.BigInteger, primary_key=True),
sa.Column('nid', sa.BigInteger),
sa.Column('username', sa.Text(), index=True, nullable=False),
sa.Column('transport', sa.Integer, nullable=False, index=True),
sa.Column('status',sa.Integer,nullable=False,index=True),
sa.Column('created', sa.DateTime, index=True)
)
async def fetchrow_b(self, query_template, *, _timeout: float = None, print_=False, **kwargs):
query, args = render(query_template, **kwargs)
self._print_query(print_, query, args)
return await self.fetchrow(query, *args, timeout=_timeout)
class BuildPgConnection(_BuildPgMixin, Connection): # noqa
pass
async def connect_b(*args, **kwargs):
kwargs.setdefault('connection_class', BuildPgConnection)
return await connect(*args, **kwargs) # noqa
class BuildPgPool(_BuildPgMixin, Pool):
pass
def create_pool_b(
dsn=None,
*,
min_size=10,
max_size=10,
max_queries=50000,
max_inactive_connection_lifetime=300.0,
setup=None,
init=None,
loop=None,
connection_class=BuildPgConnection,
**connect_kwargs,
):