Make search_companies async

leftmove · May 29, 2024 · bee0c18 · bee0c18
1 parent 3740c98
commit bee0c18
Show file tree

Hide file tree

Showing 5 changed files with 119 additions and 45 deletions.
diff --git a/backend/poetry.lock b/backend/poetry.lock
diff --git a/backend/pyproject.toml b/backend/pyproject.toml
@@ -23,7 +23,7 @@ gunicorn = "^21.2.0"
 h11 = "^0.14.0"
 idna = "^3.6"
 lxml = "^5.1.0"
-meilisearch = "^0.31.0"
+meilisearch-python-sdk = "^2.9.0"
 pydantic = "^2.6.4"
 pymongo = "^4.6.2"
 python-dotenv = "^1.0.1"

diff --git a/backend/routers/filer.py b/backend/routers/filer.py
@@ -300,12 +300,7 @@ async def rollback_filer(cik: str, password: str):
 @cache(24)
 @router.get("/search", tags=["filers"], status_code=200)
 async def search_filers(q: str, limit: int = 4):
-    options = {"limit": limit, "filter": "thirteen_f = true"}
-    hits = search_companies(q, options)
-
-    results = []
-    for result in hits:
-        results.append(result)
+    hits = await search_companies(q, limit=limit, filter="thirteen_f = true")
 
     return {"description": "Successfully queried 13F filers.", "results": hits}
 

diff --git a/backend/routers/lib/search.py b/backend/routers/lib/search.py
@@ -1,7 +1,6 @@
-import meilisearch
+import meilisearch_python_sdk
 
 import os
-import time
 
 ENVIRONMENT = os.environ.get("ENVIRONMENT", "development")
 production_environment = True if ENVIRONMENT == "production" else False
@@ -13,25 +12,32 @@
 MEILI_SERVER_URL = os.environ["MEILI_SERVER_URL"]
 MEILI_MASTER_KEY = os.environ["MEILI_MASTER_KEY"]
 
-search = meilisearch.Client(MEILI_SERVER_URL, MEILI_MASTER_KEY)
-if "companies" not in [index.uid for index in search.get_indexes()["results"]]:
-    search.create_index("companies", {"primaryKey": "cik"})
-    time.sleep(3)
-    search = meilisearch.Client(MEILI_SERVER_URL, MEILI_MASTER_KEY)
+
+def _prepare_meilisearch():
+    client = meilisearch_python_sdk.Client(MEILI_SERVER_URL, MEILI_MASTER_KEY)
+    companies_index = client.index("companies")
+    indexes = client.get_indexes()
+    if not indexes or "companies" not in [index.uid for index in indexes]:
+        task = client.create_index("companies", {"primaryKey": "cik"})
+        client.wait_for_task(task.task_uid, timeout=None)
+    companies_index.update_displayed_attributes(
+        [
+            "name",
+            "cik",
+            "tickers",
+        ]
+    )
+    companies_index.update_searchable_attributes(["name", "tickers", "cik"])
+    companies_index.update_filterable_attributes(["thirteen_f"])
+
+
+_prepare_meilisearch()
+search = meilisearch_python_sdk.AsyncClient(MEILI_SERVER_URL, MEILI_MASTER_KEY)
 companies_index = search.index("companies")
-companies_index.update_displayed_attributes(
-    [
-        "name",
-        "cik",
-        "tickers",
-    ]
-)
-companies_index.update_searchable_attributes(["name", "tickers", "cik"])
-companies_index.update_filterable_attributes(["thirteen_f"])
-
-
-def search_companies(query, options={}):
-    result = companies_index.search(query, options)
-    hits = result["hits"]
+
+
+async def search_companies(query, limit, filter):
+    result = await companies_index.search(query, limit=limit, filter=filter)
+    hits = result.hits
 
     return hits
diff --git a/backend/routers/utils.py b/backend/routers/utils.py
@@ -8,7 +8,7 @@
 from dotenv import load_dotenv
 
 import redis
-import meilisearch
+import meilisearch_python_sdk
 import pymongo
 import uvicorn
 
@@ -118,14 +118,14 @@ def initialize():
     try:
         retries = 3
         while retries:
-            search = meilisearch.Client(MEILI_SERVER_URL, MEILI_MASTER_KEY)
-            search.create_index("companies", {"primaryKey": "cik"})
+            search = meilisearch_python_sdk.Client(MEILI_SERVER_URL, MEILI_MASTER_KEY)
+            search.create_index("companies", primary_key="cik")
             companies_index = search.index("companies")
             companies_index.add_documents([{"cik": "TEST"}])
             retries -= 1
         raise RuntimeError  # @IgnoreException
     except RuntimeError:
-        search = meilisearch.Client(MEILI_SERVER_URL, MEILI_MASTER_KEY)
+        search = meilisearch_python_sdk.Client(MEILI_SERVER_URL, MEILI_MASTER_KEY)
         companies_index = search.index("companies")
 
     store = redis.Redis(