Using aiobotocore with context managers - python-3.x

I'm trying to use the aiobotocore library with context managers, but I'm having a hard time trying to configure my credentials.
I need to create a class that configure my AWS client so I can use the put, read and delete functions in this library.
The following code is being used to this:
from contextlib import AsyncExitStack
from aiobotocore.session import AioSession
from credentials import aws_access_key_id, aws_secret_access_key
class AWSConnectionManager:
def __init__(self, aws_acces_key_id, aws_secret_access_key):
self.aws_acces_key_id = aws_acces_key_id
self.aws_secret_access_key = aws_secret_access_key
self._exit_stack = AsyncExitStack()
self._client = None
print('__init__')
async def __aenter__(self):
session = AioSession
self._client = await self._exit_stack.enter_async_context(session.create_client('s3'))
print('__aenter__')
async def __aexit__(self, exc_type, exc_val, exc_tb):
await self._exit_stack.__aexit__(exc_type, exc_val, exc_tb)
print('__aexit__')
res = AWSConnectionManager(aws_acces_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key)
But, it doesn't pass through aenter and aexit method.
With the code above I have the following output:
__init__
<__main__.AWSConnectionManager object at 0x7f03921ac640>
Does anyone know what can be wrong with my code?

First: you need to fix 'session = AioSession' => 'session = AioSession()' + add return + pass credentials
async def __aenter__(self):
session = AioSession()
self._client = await self._exit_stack.enter_async_context(
session.create_client(
's3',
aws_secret_access_key=self.aws_secret_access_key,
aws_access_key_id=self.aws_access_key_id,
)
)
return self
Second: you need to write/add proxy calls for put_object/get_object or make _client public by rename _client => client
async def save_file(self, content, s3_filename: str):
await self._client.put_object(Bucket=self.bucket, Body=content, Key=f'{s3_filename}')
async def load_file(self, name):
obj = await self._client.get_object(Bucket=self.bucket, Key=f'{name}')
return obj['Body'].read()
now you can use like
async with SkyFileStorageProxy() as storage:
await storage.load_file(name='test.txt')

Related

How to properly inherit class method

I have a database connection class that creates a connection pool. Now as the application grows and I'm adding different types of database writers, I want to move database connections to a separate class and inherit from it. So far I have this:
class ServiceDB:
#classmethod
async def init(cls, settings):
self = ServiceDB()
self.pool = await asyncpg.create_pool(
database=settings["POSTGRES_DB"],
user=settings["POSTGRES_USER"],
password=settings["POSTGRES_PASSWORD"],
host=settings["DB_HOST"],
port=settings["DB_PORT"],
)
return self
class ChildWriter(ServiceDB):
async def write_db(self, query):
# Write to specific table
pass
if __name__ == "__main__":
settings = {'info': "some connection settings"}
query = "SELECT * FROM 'table'"
connection = await ChildWriter().init(settings)
await connection.write_db(msg, query)
When I run this I get AttributeError: 'ServiceDB' object has no attribute 'write_db'. How do I properly extend ServiceDB with the write_db method?
Classmethods receive the "current class" as the first argument. Instantiate this cls, not the fixed baseclass.
class ServiceDB:
#classmethod
async def init(cls, settings):
self = cls() # cls is the *current* class, not just ServiceDB
self.pool = await asyncpg.create_pool(
database=settings["POSTGRES_DB"],
user=settings["POSTGRES_USER"],
password=settings["POSTGRES_PASSWORD"],
host=settings["DB_HOST"],
port=settings["DB_PORT"],
)
return self
Note that ideally, all attributes are set via __init__ instead of a separate classmethod constructor. The separate constructor should just pass on any attributes constructed externally.
class ServiceDB:
def __init__(self, pool):
self.pool = pool
#classmethod
async def init(cls, settings, **kwargs):
pool = await asyncpg.create_pool(
database=settings["POSTGRES_DB"],
user=settings["POSTGRES_USER"],
password=settings["POSTGRES_PASSWORD"],
host=settings["DB_HOST"],
port=settings["DB_PORT"],
)
return cls(pool=pool, **kwargs)
class ChildWriter(ServiceDB):
async def write_db(self, query): ...
if __name__ == "__main__":
settings = {'info': "some connection settings"}
query = "SELECT * FROM 'table'"
# call classmethod on class V
connection = await ChildWriter.init(settings)
await connection.write_db(msg, query)

How create class Python 3.x (singleton) for asyncpg?

I wanted to organize a connection pool when initializing the class with the method below
import asyncio
import asyncpg
class DBCommands:
def __init__(self, uri: str) -> None:
loop = asyncio.get_event_loop()
self.pool: asyncpg.pool.Pool = loop.run_until_complete(asyncpg.create_pool(dsn=uri))
async def get_id_admins(self) -> list:
async with self.pool.acquire():
result = await self.pool.fetch("SELECT chat_id FROM users WHERE role_user = 'admin'")
admins_id = [row[0] for row in result]
return admins_id
Since the pool should be one, with the above implementation, this will not work. I decided to use singleton, but I don’t understand how to implement this. Below is the version that I came up with. Tell me how best to solve this problem. In addition, I do not understand how best and where to close connections. I'm new to using patterns and just starting to study OOP.
import asyncio
import asyncpg
class Singleton(type):
_instances = {}
def __call__(cls, *args, **kwargs):
if cls not in cls._instances:
cls._instances[cls] = super(Singleton, cls).__call__(*args, **kwargs)
return cls._instances[cls]
class DBManager(metaclass=Singleton):
#classmethod
def connect(cls, uri):
loop = asyncio.get_event_loop()
return loop.run_until_complete(asyncpg.create_pool(dsn=uri))
class DBCommands:
def __init__(self, uri) -> None:
self.uri = uri
self.pool = DBManager.connect(uri)
async def get_id_admins(self) -> list:
async with self.pool.acquire():
result = await self.pool.fetch("SELECT chat_id FROM users WHERE role_user = 'admin'")
admins_id = [row[0] for row in result]
return admins_id
I have an assumption that opening and closing a pool can be added to __aenter__ and __aexit__
You can use a class attribute and create the pool the first time it's needed in an async function:
class Database:
self.pool = None
...
async def get_id_admins(self)
if self.pool is None:
self.pool = await asyncpg.create_pool(dsn=...`).
I generally use a regular class and create a single instance attached to global object (like the aiohttp application for web applications) as in:
class Database:
def __init__(self, dsn):
self.dsn = dsn
self.pool = None
async def connect(self):
"""Initialize asyncpg Pool"""
self.pool = await asyncpg.create_pool(dsn=self.dsn, min_size=2, max_size=4)
logging.info("successfully initialized database pool")
async def get_id_admins(self):
...
And use it like:
async def startup(app):
await app.database.connect()
async def shutdown(app):
await app.database.pool.close()
def main():
app = web.Application()
app.database = Database(app.config.DSN)
app.on_startup.append(startup)
app.on_shutdown.append(shutdown)

RuntimeError when running coroutine from __init__

Here's a sample code.
class Foo:
def __init__(self):
self._run_coro()
def _run_coro(self):
async def init():
bar = #some I/O op
self.bar = bar
loop = asyncio.get_event_loop()
loop.run_until_complete(init())
async def spam(self):
return await #I/O op
async def main():
foo = Foo()
await foo.spam()
loop = asyncio.get_event_loop()
loop.run_until_complete(main())
When I run this code, I get following exception:
RuntimeError: This event loop is already running
If I initialize Foo outside main, the code runs without any exception. I want to initialize Foo such that during initialization it runs a coroutine which creates a class attribute bar.
I am unable to figure how to do it correctly. How can I run a coroutine from __init__.
Any help would be highly appreciated.
class Foo:
def __init__(self):
self.session = requests.Session()
self.async_session = None
#I guess this can be done to initialize it.
s = self.init_async_session()
try:
s.send(None)
except StopIteration:
pass
finally:
s.close()
async def init_async_session(self):
#ClientSession should be created inside a coroutine.
self.async_session = aiohttp.ClientSession()
What would be the right way to initialize self.async_session
If some method uses something asynchronous it should be explicitly defined as asynchronous either. This is a core idea behind asyncio: make you write code a way you always know if some arbitrary method may do something asynchronous.
In your snippet you want to do async thing (bar I/O) inside sync method __init__ and asyncio prohibits it. You should make _run_coro async and initialize Foo asynchronously, for example, using __await__ method:
import asyncio
class Foo:
def __await__(self):
return self._run_coro().__await__()
async def _run_coro(self): # real async initializer
async def init():
await asyncio.sleep(1) # bar I/O
self.bar = 123
await init()
return self
async def spam(self):
return await asyncio.sleep(1) # I/O op
async def main():
foo = await Foo()
await foo.spam()
asyncio.run(main()) # instead of two lines in Python 3.7+
You may be interested in reading this answer to understand better how asyncio works and how to handle it.
Upd:
s = self.init_async_session()
try:
s.send(None)
Don't do such things: generator's method are only details of implementation in regard of coroutines. You can predict how coroutine will react on calling .send() method and you can rely on this behavior.
If you want to execute coroutine use await, if you want to start it "in background" use task or other functions from asyncio doc.
What would be the right way to initialize self.async_session
When it comes to aiohttp.ClientSession it should not only be created, but properly closed also. Best way to do it is to use async context manager as shown in aiohttp doc.
If you want to hide this operation inside Foo you can make it async manager either. Complete example:
import aiohttp
class Foo:
async def __aenter__(self):
self._session = aiohttp.ClientSession()
await self._session.__aenter__()
return self
async def __aexit__(self, *args):
await self._session.__aexit__(*args)
async def spam(self):
url = 'http://httpbin.org/delay/1'
resp = await self._session.get(url)
text = await resp.text()
print(text)
async def main():
async with Foo() as foo:
await foo.spam()
asyncio.run(main())
Upd2:
You can combine ways to init/close object from above to achive result you like. As long as you keep in mind both operations are asynchronous and thus should be awaited, everything should be fine.
One more possible way:
import asyncio
import aiohttp
class Foo:
def __await__(self):
return self._init().__await__()
async def _init(self):
self._session = aiohttp.ClientSession()
await self._session.__aenter__()
return self
async def close(self):
await self._session.__aexit__(None, None, None)
async def spam(self):
url = 'http://httpbin.org/delay/1'
resp = await self._session.get(url)
text = await resp.text()
print(text)
async def main():
foo = await Foo()
try:
await foo.spam()
finally:
await foo.close()
asyncio.run(main())
Here's my solution.
class Session:
def __init__(self, headers):
self._headers = headers
self._session = requests.Session()
self._async_session = None
async def _init(self):
self._session = aiohttp.ClientSession(headers=headers)
async def async_request(self, url):
while True:
try:
async with self._async_session.get(url) as resp:
resp.raise_for_status()
return await resp.text()
except aiohttp.client_exceptions.ClientError:
#retry or raise
except AttributeError:
if isinstance(self._async_session, aiohttp.ClientSession):
raise
await self._init()
def request(self, url):
return self._session.get(url).text
async def close(self):
if isinstance(self._async_session, aiohttp.ClientSession):
await self._session.close()
async def main():
session = Session({})
print(await session.async_request('https://httpstat.us/200')
await session.close()
asyncio.run(main())
I can initialize the Session class and make synchronous as well as asynchronous requests. I do not have to explicitly call await session._init() to initialize self._async_session as when session._async_request is called and self._async_session is None, then await session._init() will be called and the request will be retried.

Awaiting a coroutine inside of a Class while event_loop is already running

I have an Issue with asyncio I can't really get me head around.
take this working example (with Python 3.6+ because of string interpolation)
import asyncio
import aiohttp
import async_timeout
import json
async def fetch(session, url):
async with async_timeout.timeout(10):
async with session.get(url) as response:
return await response.text()
async def get_bittrex_marketsummary(currency_pair):
url = f'https://bittrex.com/api/v1.1/public/getmarketsummary?market={currency_pair}'
async with aiohttp.ClientSession() as session:
response = await fetch(session, url)
return json.loads(response)
class MyCryptoCurrency:
def __init__(self):
self.currency = "BTC-ETH"
self.last_price = None
asyncio.ensure_future(self.get_last_price())
async def get_last_price(self):
self.last_price = await get_bittrex_marketsummary(self.currency)
async def main():
eth = MyCryptoCurrency()
print(eth.last_price)
loop = asyncio.get_event_loop()
loop.run_until_complete(main())
while this runs and doesn't throw any Exceptions, it doesn't get the result from the api request and so ... doesn't work :P
If I try to use f.e. loop.run_until_complete(get_bittrex_marketsummary()) I get "event loop is already running" error - which kind of makes sense.
Any hints how to solve this properly?
Thx in advance!
ok, after talking about this in #python channel on freenode I got the answer "don't do async I/O in __init__", so here is the working version:
import asyncio
import aiohttp
import async_timeout
import json
async def fetch(session, url):
async with async_timeout.timeout(10):
async with session.get(url) as response:
return await response.text()
async def get_bittrex_marketsummary(currency_pair):
url = f'https://bittrex.com/api/v1.1/public/getmarketsummary?market={currency_pair}'
async with aiohttp.ClientSession() as session:
response = await fetch(session, url)
return json.loads(response)
class MyCryptoCurrency:
def __init__(self):
self.currency = "BTC-ETH"
self.last_price = None
async def get_last_price(self):
self.last_price = await get_bittrex_marketsummary(self.currency)
async def main():
eth = MyCryptoCurrency()
await eth.get_last_price()
print(eth.last_price)
loop = asyncio.get_event_loop()
loop.run_until_complete(main())

How to mock "async with" statements?

I'm trying to write tests for a method that uses "async with" statements (in this case, aioredis's connection pool), i want to mock the connection to redis, but i'm having trouble figuring out how.
Here's what i have so far:
from asyncio import Future
from unittest.mock import MagicMock
import pytest
# The thing i'm trying to test
async def set_value(redis, value):
# Do things
async with redis.get() as conn:
await conn.set("key", value)
#My Mock classes
class MockRedis():
def get(self):
return MockAsyncPool()
class MockAsyncPool(MagicMock):
async def __aenter__(self):
conn = MagicMock()
f = Future()
f.set_result(True)
conn.set = MagicMock(return_value=f)
return conn
def __aexit__(self, exc_type, exc_val, exc_tb):
pass
# The actual test
#pytest.mark.asyncio
async def test_get_token():
redis = MockRedis()
token = await set_value(redis, 'something')
assert token is not None
I run it with:
py.test path/to/file.py
And i'm getting this error:
> await conn.set("key", value)
E TypeError: object NoneType can't be used in 'await' expression
__aexit__ needs to also be asyncronous (needs to return an awaitable):
async def __aexit__(self, exc_type, exc_val, exc_tb):
pass
Without being async it is returning None instead of a coroutine so it raises an error, as for the very misleading error message I have created this issue to point out that the error message could be greatly improved.

Resources