refactor(db): use asyncpg.Pool for fetching data (#19)

* feat(db): refacto

* fix: retours

* fix: retours

Author: maxgfr

srdt_analysis/__main__.py

@@ -8,10 +8,10 @@
 
 
 def main():
-    data = get_data()
+    data = get_data(["information"])
     exploiter = PageInfosExploiter()
     result = exploiter.process_documents(
-        [data[3][0]], "page_infos.csv", "cdtn_page_infos"
+        [data["information"][0]], "page_infos.csv", "cdtn_page_infos"
     )
     collections = Collections()
     res = collections.search(

srdt_analysis/database_manager.py

@@ -1,88 +1,66 @@
 import asyncio
 import os
-from typing import Tuple
+from contextlib import asynccontextmanager
+from typing import Literal, Optional, Sequence
 
 import asyncpg
 
 from srdt_analysis.models import Document, DocumentsList
 
+CollectionName = Literal[
+    "code_du_travail",
+    "fiches_service_public",
+    "page_fiche_ministere_travail",
+    "contributions",
+    "information",
+]
+
 
 class DatabaseManager:
     def __init__(self):
-        self.conn
+        self.pool: Optional[asyncpg.Pool] = None
 
     async def connect(self):
-        self.conn = await asyncpg.connect(
+        self.pool = await asyncpg.create_pool(
             user=os.getenv("POSTGRES_USER"),
             password=os.getenv("POSTGRES_PASSWORD"),
             database=os.getenv("POSTGRES_DATABASE_NAME"),
             host=os.getenv("POSTGRES_DATABASE_URL"),
         )
 
     async def close(self):
-        if self.conn:
-            await self.conn.close()
-
-    async def fetch_articles_code_du_travail(self) -> DocumentsList:
-        results = await self.conn.fetch(
-            "SELECT * from public.documents WHERE source = 'code_du_travail'"
-        )
-        return [Document.from_record(r) for r in results]
-
-    async def fetch_fiches_mt(self) -> DocumentsList:
-        result = await self.conn.fetch(
-            "SELECT * from public.documents WHERE source = 'page_fiche_ministere_travail'"
-        )
-        return [Document.from_record(r) for r in result]
-
-    async def fetch_fiches_sp(self) -> DocumentsList:
-        result = await self.conn.fetch(
-            "SELECT * from public.documents WHERE source = 'fiches_service_public'"
-        )
-        return [Document.from_record(r) for r in result]
-
-    async def fetch_page_infos(self) -> DocumentsList:
-        result = await self.conn.fetch(
-            "SELECT * from public.documents WHERE source = 'information'"
-        )
-        return [Document.from_record(r) for r in result]
-
-    async def fetch_page_contribs(self) -> DocumentsList:
-        result = await self.conn.fetch(
-            "SELECT * from public.documents WHERE source = 'contributions'"
-        )
-        return [Document.from_record(r) for r in result]
-
-    async def fetch_all(
-        self,
-    ) -> Tuple[
-        DocumentsList,
-        DocumentsList,
-        DocumentsList,
-        DocumentsList,
-        DocumentsList,
-    ]:
-        await self.connect()
-
-        result1 = await self.fetch_articles_code_du_travail()
-        result2 = await self.fetch_fiches_mt()
-        result3 = await self.fetch_fiches_sp()
-        result4 = await self.fetch_page_infos()
-        result5 = await self.fetch_page_contribs()
-
-        await self.close()
-
-        return (result1, result2, result3, result4, result5)
-
-
-def get_data() -> (
-    Tuple[
-        DocumentsList,
-        DocumentsList,
-        DocumentsList,
-        DocumentsList,
-        DocumentsList,
-    ]
-):
+        if self.pool:
+            await self.pool.close()
+
+    @asynccontextmanager
+    async def get_connection(self):
+        if not self.pool:
+            await self.connect()
+        if self.pool is None:
+            raise ValueError("Pool is not initialized")
+        async with self.pool.acquire() as conn:
+            yield conn
+
+    async def fetch_documents_by_source(self, source: str) -> DocumentsList:
+        async with self.get_connection() as conn:
+            result = await conn.fetch(
+                "SELECT * from public.documents WHERE source = $1", source
+            )
+            return [Document.from_record(r) for r in result]
+
+    async def fetch_sources(
+        self, sources: Sequence[CollectionName]
+    ) -> dict[CollectionName, DocumentsList]:
+        try:
+            tasks = [self.fetch_documents_by_source(source) for source in sources]
+            results = await asyncio.gather(*tasks)
+            return {source: result for source, result in zip(sources, results)}
+        finally:
+            await self.close()
+
+
+def get_data(
+    sources: Sequence[CollectionName],
+) -> dict[CollectionName, DocumentsList]:
     db = DatabaseManager()
-    return asyncio.run(db.fetch_all())
+    return asyncio.run(db.fetch_sources(sources))