From 6a7250d1842849fd852bf68610963d866684b476 Mon Sep 17 00:00:00 2001
From: pablodanswer <pablo@danswer.ai>
Date: Sat, 5 Oct 2024 15:39:48 -0700
Subject: [PATCH] use new kv store

---
 backend/danswer/key_value_store/store.py |  9 +++------
 backend/danswer/setup.py                 | 12 ++++++------
 2 files changed, 9 insertions(+), 12 deletions(-)

diff --git a/backend/danswer/key_value_store/store.py b/backend/danswer/key_value_store/store.py
index 450056c40b1..4306743f875 100644
--- a/backend/danswer/key_value_store/store.py
+++ b/backend/danswer/key_value_store/store.py
@@ -5,7 +5,7 @@
 
 from sqlalchemy.orm import Session
 
-from danswer.db.engine import get_session_factory
+from danswer.db.engine import get_sqlalchemy_engine
 from danswer.db.models import KVStore
 from danswer.key_value_store.interface import JSON_ro
 from danswer.key_value_store.interface import KeyValueStore
@@ -26,12 +26,9 @@ def __init__(self) -> None:
 
     @contextmanager
     def get_session(self) -> Iterator[Session]:
-        factory = get_session_factory()
-        session: Session = factory()
-        try:
+        engine = get_sqlalchemy_engine()
+        with Session(engine, expire_on_commit=False) as session:
             yield session
-        finally:
-            session.close()
 
     def store(self, key: str, val: JSON_ro, encrypt: bool = False) -> None:
         # Not encrypted in Redis, but encrypted in Postgres
diff --git a/backend/danswer/setup.py b/backend/danswer/setup.py
index 9d6a2ca8884..2baeda4a811 100644
--- a/backend/danswer/setup.py
+++ b/backend/danswer/setup.py
@@ -29,9 +29,9 @@
 from danswer.db.swap_index import check_index_swap
 from danswer.document_index.factory import get_default_document_index
 from danswer.document_index.interfaces import DocumentIndex
-from danswer.dynamic_configs.factory import get_dynamic_config_store
-from danswer.dynamic_configs.interface import ConfigNotFoundError
 from danswer.indexing.models import IndexingSetting
+from danswer.key_value_store.factory import get_kv_store
+from danswer.key_value_store.interface import KvKeyNotFoundError
 from danswer.natural_language_processing.search_nlp_models import EmbeddingModel
 from danswer.natural_language_processing.search_nlp_models import warm_up_bi_encoder
 from danswer.natural_language_processing.search_nlp_models import warm_up_cross_encoder
@@ -134,7 +134,7 @@ def setup_danswer(db_session: Session) -> None:
 
 
 def translate_saved_search_settings(db_session: Session) -> None:
-    kv_store = get_dynamic_config_store()
+    kv_store = get_kv_store()
 
     try:
         search_settings_dict = kv_store.load(KV_SEARCH_SETTINGS)
@@ -172,17 +172,17 @@ def translate_saved_search_settings(db_session: Session) -> None:
             logger.notice("Search settings updated and KV store entry deleted.")
         else:
             logger.notice("KV store search settings is empty.")
-    except ConfigNotFoundError:
+    except KvKeyNotFoundError:
         logger.notice("No search config found in KV store.")
 
 
 def mark_reindex_flag(db_session: Session) -> None:
-    kv_store = get_dynamic_config_store()
+    kv_store = get_kv_store()
     try:
         value = kv_store.load(KV_REINDEX_KEY)
         logger.debug(f"Re-indexing flag has value {value}")
         return
-    except ConfigNotFoundError:
+    except KvKeyNotFoundError:
         # Only need to update the flag if it hasn't been set
         pass