From 3e15b28ec12c0386d227593c019a17e4a539c2ec Mon Sep 17 00:00:00 2001
From: Zoe Roux <zoe.roux@zoriya.dev>
Date: Fri, 9 May 2025 22:49:19 +0200
Subject: [PATCH] Create a psycopg pool

---
 scanner/requirements.txt        |  2 +-
 scanner/scanner/__init__.py     | 19 ++++++++++++++++++-
 scanner/scanner/fsscan.py       |  2 +-
 scanner/scanner/models/extra.py |  4 ++--
 scanner/scanner/models/genre.py |  4 ++--
 scanner/scanner/models/movie.py |  4 ++--
 scanner/scanner/models/serie.py |  4 ++--
 scanner/scanner/models/staff.py |  4 ++--
 8 files changed, 30 insertions(+), 13 deletions(-)

diff --git a/scanner/requirements.txt b/scanner/requirements.txt
index bda26c4e..30e94a49 100644
--- a/scanner/requirements.txt
+++ b/scanner/requirements.txt
@@ -4,5 +4,5 @@ guessit@git+https://github.com/zoriya/guessit
 aiohttp
 watchfiles
 langcodes
-psycopg[binary]
+psycopg[binary,pool]
 
diff --git a/scanner/scanner/__init__.py b/scanner/scanner/__init__.py
index 9612de4b..7798f60e 100644
--- a/scanner/scanner/__init__.py
+++ b/scanner/scanner/__init__.py
@@ -1,18 +1,35 @@
 import logging
+from contextlib import asynccontextmanager
 
 from fastapi import FastAPI
+from psycopg import AsyncConnection
+from psycopg_pool import AsyncConnectionPool
 
 logging.basicConfig(level=logging.INFO)
 logging.getLogger("watchfiles").setLevel(logging.WARNING)
 logging.getLogger("rebulk").setLevel(logging.WARNING)
 
+pool = AsyncConnectionPool(open=False, kwargs={"autocommit": True})
+
+
+@asynccontextmanager
+async def lifetime():
+	await pool.open()
+	yield
+	await pool.close()
+
+
+async def get_db() -> AsyncConnection:
+	async with pool.connection() as ret:
+		yield ret
+
 
 app = FastAPI(
 	title="Scanner",
 	description="API to control the long running scanner or interacting with external databases (themoviedb, tvdb...)\n\n"
 	+ "Most of those APIs are for admins only.",
 	root_path="/scanner",
-	# lifetime=smth
+	lifetime=lifetime,
 )
 
 
diff --git a/scanner/scanner/fsscan.py b/scanner/scanner/fsscan.py
index 4fd67017..8bb18ac0 100644
--- a/scanner/scanner/fsscan.py
+++ b/scanner/scanner/fsscan.py
@@ -94,7 +94,7 @@ class Scanner:
 				logger.error("Couldn't identify %s.", path, exc_info=e)
 		created = await self._client.create_videos(vids)
 
-		await enqueue(
+		await self._requests.enqueue(
 			[
 				Request(
 					kind=x.guess.kind,
diff --git a/scanner/scanner/models/extra.py b/scanner/scanner/models/extra.py
index 0b09985f..e495ecff 100644
--- a/scanner/scanner/models/extra.py
+++ b/scanner/scanner/models/extra.py
@@ -1,9 +1,9 @@
-from enum import Enum
+from enum import StrEnum
 
 from ..utils import Model
 
 
-class ExtraKind(str, Enum):
+class ExtraKind(StrEnum):
 	OTHER = "other"
 	TRAILER = "trailer"
 	INTERVIEW = "interview"
diff --git a/scanner/scanner/models/genre.py b/scanner/scanner/models/genre.py
index 1ba8eda8..c6651f81 100644
--- a/scanner/scanner/models/genre.py
+++ b/scanner/scanner/models/genre.py
@@ -1,7 +1,7 @@
-from enum import Enum
+from enum import StrEnum
 
 
-class Genre(str, Enum):
+class Genre(StrEnum):
 	ACTION = "action"
 	ADVENTURE = "adventure"
 	ANIMATION = "animation"
diff --git a/scanner/scanner/models/movie.py b/scanner/scanner/models/movie.py
index 4504139f..25776d36 100644
--- a/scanner/scanner/models/movie.py
+++ b/scanner/scanner/models/movie.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from datetime import date
-from enum import Enum
+from enum import StrEnum
 
 from langcodes import Language
 
@@ -13,7 +13,7 @@ from .staff import Staff
 from .studio import Studio
 
 
-class MovieStatus(str, Enum):
+class MovieStatus(StrEnum):
 	UNKNOWN = "unknown"
 	FINISHED = "finished"
 	PLANNED = "planned"
diff --git a/scanner/scanner/models/serie.py b/scanner/scanner/models/serie.py
index 18b10604..05058a3f 100644
--- a/scanner/scanner/models/serie.py
+++ b/scanner/scanner/models/serie.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from datetime import date
-from enum import Enum
+from enum import StrEnum
 
 from langcodes import Language
 
@@ -16,7 +16,7 @@ from .staff import Staff
 from .studio import Studio
 
 
-class SerieStatus(str, Enum):
+class SerieStatus(StrEnum):
 	UNKNOWN = "unknown"
 	FINISHED = "finished"
 	AIRING = "airing"
diff --git a/scanner/scanner/models/staff.py b/scanner/scanner/models/staff.py
index 759a48da..352b4ca9 100644
--- a/scanner/scanner/models/staff.py
+++ b/scanner/scanner/models/staff.py
@@ -1,12 +1,12 @@
 from __future__ import annotations
 
-from enum import Enum
+from enum import StrEnum
 
 from ..utils import Model
 from .metadataid import MetadataId
 
 
-class Role(str, Enum):
+class Role(StrEnum):
 	ACTOR = "actor"
 	DIRECTOR = "director"
 	WRITTER = "writter"