Fix database initialization

Author: FreeFlyingSheep

compose.yaml

@@ -10,6 +10,8 @@ services:
         condition: service_healthy
       api:
         condition: service_healthy
+      grafana-probe:
+        condition: service_healthy
     healthcheck:
       test: ["CMD", "wget", "-qO-", "http://localhost/"]
       interval: 30s

scripts/create_db.py

@@ -41,26 +41,43 @@ def create_database(settings: Settings) -> None:
                     sql.Identifier(settings.database_db)
                 )
             )
-            cur.execute(sql.SQL("CREATE EXTENSION IF NOT EXISTS vector"))
-            cur.execute(sql.SQL("CREATE EXTENSION IF NOT EXISTS pg_textsearch"))
-            cur.execute(sql.SQL("CREATE EXTENSION IF NOT EXISTS zhparser"))
-            cur.execute(
-                sql.SQL("DROP TEXT SEARCH CONFIGURATION IF EXISTS chinese CASCADE")
-            )
-            cur.execute(
-                sql.SQL("CREATE TEXT SEARCH CONFIGURATION chinese (parser = zhparser)")
-            )
-            cur.execute(
-                sql.SQL(
-                    "ALTER TEXT SEARCH CONFIGURATION chinese "
-                    "ADD MAPPING FOR n,v,a,i,e,l WITH simple"
-                )
-            )
             logger.info("Created database: %s", settings.database_db)
         else:
             logger.info("Database already exists: %s", settings.database_db)
 
 
+def enable_extensions(settings: Settings) -> None:
+    """Enable the pgvector extension in the database.
+
+    Args:
+        settings: Application settings containing database connection details.
+    """
+    with (
+        psycopg.connect(
+            dbname=settings.database_db,
+            user=settings.database_user,
+            password=settings.database_password.get_secret_value(),
+            host=settings.database_host,
+            port=settings.database_port,
+            autocommit=True,
+        ) as conn,
+        conn.cursor() as cur,
+    ):
+        cur.execute(sql.SQL("CREATE EXTENSION IF NOT EXISTS vector"))
+        cur.execute(sql.SQL("CREATE EXTENSION IF NOT EXISTS pg_textsearch"))
+        cur.execute(sql.SQL("CREATE EXTENSION IF NOT EXISTS zhparser"))
+        cur.execute(sql.SQL("DROP TEXT SEARCH CONFIGURATION IF EXISTS chinese CASCADE"))
+        cur.execute(
+            sql.SQL("CREATE TEXT SEARCH CONFIGURATION chinese (parser = zhparser)")
+        )
+        cur.execute(
+            sql.SQL(
+                "ALTER TEXT SEARCH CONFIGURATION chinese "
+                "ADD MAPPING FOR n,v,a,i,e,l WITH simple"
+            )
+        )
+
+
 def main() -> None:
     """Create the database.
 
@@ -70,6 +87,7 @@ def main() -> None:
     settings: Settings = get_settings()
     logging.basicConfig(level=settings.log_level)
     create_database(settings)
+    enable_extensions(settings)
 
 
 if __name__ == "__main__":

scripts/import_csv.py

@@ -28,7 +28,7 @@
 logger: logging.Logger = logging.getLogger(__name__)
 
 
-def read_csv(csv_path: str | os.PathLike) -> pd.DataFrame:
+async def read_csv(csv_path: str | os.PathLike) -> pd.DataFrame:
     """Read CSV file into DataFrame.
 
     Args:
@@ -42,11 +42,13 @@ def read_csv(csv_path: str | os.PathLike) -> pd.DataFrame:
         ValueError: If CSV format is invalid or headers don't match expected format.
     """
     csv_file = Path(csv_path)
-    if not csv_file.exists():
+    if not await csv_file.exists():
         msg: str = f"CSV file not found: {csv_path}"
         raise FileNotFoundError(msg)
 
-    df: pd.DataFrame = pd.read_csv(csv_file, encoding="utf-8", dtype=str)
+    df: pd.DataFrame = await asyncio.to_thread(
+        pd.read_csv, csv_file, encoding="utf-8", dtype=str
+    )
 
     expected_headers: list[str] = [
         "上市公司代码",
@@ -93,11 +95,11 @@ async def import_stocks_from_csv(db: AsyncSession, csv_path: str | os.PathLike)
         ValueError: If CSV format is invalid or headers don't match expected format.
     """
     csv_file = Path(csv_path)
-    if not csv_file.exists():
+    if not await csv_file.exists():
         msg: str = f"CSV file not found: {csv_path}"
         raise FileNotFoundError(msg)
 
-    df: pd.DataFrame = await asyncio.to_thread(read_csv, csv_file)
+    df: pd.DataFrame = await read_csv(csv_file)
     valid_records: list[dict[str, str]] = []
     for r in df.to_dict(orient="records"):
         try:

src/stock_analysis/agent/ingest.py

@@ -589,7 +589,7 @@ async def _process_pdf(
         finally:
             if tmp_pdf:
                 tmp_pdf_path: Path = Path(tmp_pdf)
-                if tmp_pdf_path.exists():
+                if await tmp_pdf_path.exists():
                     await tmp_pdf_path.unlink()
 
     async def ingest(self, object_key: str) -> None: