This is an automated email from the ASF dual-hosted git repository. humbedooh pushed a commit to branch master in repository https://gitbox.apache.org/repos/asf/incubator-ponymail-foal.git
commit a847810a6829a78c86d1c340d00898238c35d4e6 Author: Daniel Gruno <[email protected]> AuthorDate: Sun Mar 28 22:16:01 2021 +0200 async_scan breaks type tests, refactor into our own database class Some auto-linting also... --- server/plugins/database.py | 57 ++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 47 insertions(+), 10 deletions(-) diff --git a/server/plugins/database.py b/server/plugins/database.py index ae79252..1a7b651 100644 --- a/server/plugins/database.py +++ b/server/plugins/database.py @@ -20,7 +20,7 @@ This is the Database library stub for Pony Mail codename Foal """ import uuid - +import typing import elasticsearch import plugins.configuration @@ -50,17 +50,17 @@ class Database: self.uuid = str(uuid.uuid4()) self.dbs = DBNames(config.db_prefix) if self.config.dburl: - self.client = elasticsearch.AsyncElasticsearch([self.config.dburl,]) + self.client = elasticsearch.AsyncElasticsearch([self.config.dburl, ]) else: self.client = elasticsearch.AsyncElasticsearch( - [ - { - "host": config.hostname, - "port": config.port, - "url_prefix": config.url_prefix or "", - "use_ssl": config.secure, - }, - ] + [ + { + "host": config.hostname, + "port": config.port, + "url_prefix": config.url_prefix or "", + "use_ssl": config.secure, + }, + ] ) async def search(self, index="", **kwargs): @@ -86,3 +86,40 @@ class Database: index = self.dbs.session res = await self.client.index(index=index, **kwargs) return res + + async def scan(self, + query=None, + scroll="5m", + preserve_order=False, + size=1000, + request_timeout=None, + clear_scroll=True, + scroll_kwargs=None, + **kwargs) -> typing.AsyncIterator[dict]: + + scroll_kwargs = scroll_kwargs or {} + + if not preserve_order: + query = query.copy() if query else {} + query["sort"] = "_doc" + + # Do the search + resp = await self.search( + body=query, scroll=scroll, size=size, request_timeout=request_timeout, **kwargs + ) + scroll_id = resp.get("_scroll_id") + + # While we can scroll, fetch a page + try: + while scroll_id and resp["hits"]["hits"]: + for hit in resp["hits"]["hits"]: + yield hit + resp = await self.client.scroll( + body={"scroll_id": scroll_id, "scroll": scroll}, **scroll_kwargs + ) + scroll_id = resp.get("_scroll_id") + + # Shut down and clear scroll once done + finally: + if scroll_id and clear_scroll: + await self.client.clear_scroll(body={"scroll_id": [scroll_id]}, ignore=(404,))
