fix: prevent race condition in vectordb tables creation

Signed-off-by: kyteinsky <kyteinsky@gmail.com>

Author: kyteinsky

context_chat_backend/vectordb/pgvector.py

@@ -134,6 +134,13 @@ def __init__(self, embedding: Embeddings | None = None, **kwargs):
 			# tried to create the tables but it was already created in another process
 			# init the client again to detect it already exists, and continue from there
 			self.client = PGVector(embedding, collection_name=COLLECTION_NAME, **kwargs)
+		except sa.exc.ProgrammingError as pe:  # pyright: ignore[reportAttributeAccessIssue]
+			if not isinstance(pe.orig, psycopg.errors.DuplicateTable):
+				raise
+
+			# multiple processes raced to CREATE TABLE, all but one lost so tables now exist, retry
+			# may not be necessary now since initial tables creation happens in main.py now
+			self.client = PGVector(embedding, collection_name=COLLECTION_NAME, **kwargs)
 
 	def get_instance(self) -> VectorStore:
 		return self.client

main.py

@@ -16,6 +16,7 @@
 from context_chat_backend.controller import app  # isort: skip
 from context_chat_backend.logger import get_logging_config, setup_logging  # isort: skip
 from context_chat_backend.utils import is_k8s_env, redact_config  # isort: skip
+from context_chat_backend.dyn_loader import VectorDBLoader  # isort: skip
 
 LOGGER_CONFIG_NAME = 'logger_config.yaml'
 LOGGER_K8S_CONFIG_NAME = 'logger_config.k8s.yaml'
@@ -64,6 +65,10 @@ def _setup_log_levels(debug: bool):
 		'sqlalchemy',
 	])
 
+	# init the vectordb in the main process before workers are forked so that
+	# no two worker processes race to CREATE TABLE at the same time
+	VectorDBLoader(app_config).load()
+
 	print(f'CPU count: {cpu_count()}, Memory: {psutil.virtual_memory()}')
 	print('App config:\n' + redact_config(app_config).model_dump_json(indent=2), flush=True)