async def perform_operation(pool: asyncpg.pool,
operation: str,
statement: str,
timeout: float = 5) -> OperationResults:
global inserted_count
"""Performs a simple transaction with the provided pool. """
success, conn_timer, trans_timer = True, Timer(), Timer()
# Run the operation without letting it exceed the timeout given
deadline = time.monotonic() + timeout
try:
with conn_timer: # Start the connection timer
time_left = deadline - conn_timer.start
async with pool.acquire(timeout=time_left) as conn:
with trans_timer: # Start transaction timer
time_left = deadline - conn_timer.start
stmt = await conn.fetch(statement, timeout=time_left)
except Exception as ex:
success = False
logger.warning(f"Statement: {statement}")
logger.warning("Transaction failed with exception: %s", ex)
if inserted_count > 1 and re.search("INSERT", statement):
inserted_count = inserted_count - 1
return (
operation,
success,
conn_timer.start,
conn_timer.stop,
trans_timer.start if trans_timer.start else conn_timer.stop,
trans_timer.stop if trans_timer.stop else conn_timer.stop,
)
async def get_user(pool: asyncpg.pool, username: str) -> Maybe[User]:
async with pool.acquire() as conn:
row = await conn.fetchrow(
f'SELECT * FROM {USERS_TABLE_NAME} WHERE username = $1', username)
if not row:
return ObjectNotFound()
return User(**dict(row))
async def update_user(pool: asyncpg.pool, user: User) -> Maybe[User]:
async with pool.acquire() as conn:
await conn.execute(
f'''
UPDATE {USERS_TABLE_NAME} SET password = $2, scope = $3 WHERE username = $1
''', user.username, passlib.hash.pbkdf2_sha256.hash(user.password),
user.scope)
return user
async def delete_user(pool: asyncpg.pool, username: str) -> Maybe[str]:
async with pool.acquire() as conn:
res = await conn.execute(
f'''
DELETE FROM {USERS_TABLE_NAME} WHERE username = $1
''', username)
if res == 'DELETE 0':
return ObjectNotFound()
return username
async def update_password(pool: asyncpg.pool, username: str,
password: str) -> Maybe[str]:
async with pool.acquire() as connection:
res = await connection.execute(
f'''
UPDATE {USERS_TABLE_NAME} SET password = $2 WHERE username = $1
''', username, passlib.hash.pbkdf2_sha256.hash(password))
if res != 'UPDATE 1':
return Exception()
return password
async def create_user(pool: asyncpg.pool, user: User) -> Maybe[User]:
async with pool.acquire() as conn:
try:
await conn.execute(
f'''
INSERT INTO {USERS_TABLE_NAME}(username, password, scope) VALUES($1, $2, $3)
''', user.username, passlib.hash.pbkdf2_sha256.hash(user.password),
user.scope)
except asyncpg.exceptions.UniqueViolationError:
return ObjectAlreadyExists()
return user
async def _check_table_exists(pool: asyncpg.pool, table: str) -> bool:
async with pool.acquire() as conn:
r = await conn.fetchrow(
f'''
SELECT EXISTS (
SELECT 1
FROM pg_tables
WHERE schemaname = 'public'
AND tablename = $1
);
''', table)
return r['exists']
async def init_db(pool: asyncpg.pool) -> None:
if await _check_table_exists(pool, USERS_TABLE_NAME):
logger.info('Already initializated.')
return
logger.info(f'Create table: {USERS_TABLE_NAME}')
async with pool.acquire() as conn:
await conn.execute(f'''
CREATE TABLE {USERS_TABLE_NAME}(
username text PRIMARY KEY,
password text,
scope text
)
''')
logger.info(f'Create admin user')
await create_user(pool, User(ADMIN_USER, ADMIN_PASSWORD, 'admin'))