fix stream_pinecone_updates session issue

StampyAI · Aug 6, 2023 · 4a1c0ab · 4a1c0ab
1 parent f0a53a8
commit 4a1c0ab
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/align_data/db/session.py b/align_data/db/session.py
@@ -20,9 +20,10 @@ def make_session(auto_commit=False):
             session.commit()
 
 
-def stream_pinecone_updates(custom_sources: List[str]):
+def stream_pinecone_updates(session, custom_sources: List[str]):
     """Yield Pinecone entries that require an update."""
-    with make_session() as session:
-        query = session.query(Article).filter(Article.pinecone_update_required.is_(True))
-        query = query.filter(Article.source.in_(custom_sources))
-        yield from query.yield_per(1000)
+    yield from session.query(Article).filter(
+        Article.pinecone_update_required.is_(True)
+    ).filter(
+        Article.source.in_(custom_sources)  
+    ).yield_per(1000)
diff --git a/align_data/pinecone/update_pinecone.py b/align_data/pinecone/update_pinecone.py
@@ -84,7 +84,7 @@ def update(self, custom_sources: List[str]):
         :param custom_sources: List of sources to update.
         """
         with make_session() as session:
-            entries_stream = stream_pinecone_updates(custom_sources)
+            entries_stream = stream_pinecone_updates(session, custom_sources)
             for article, pinecone_entry in self.process_entries(entries_stream):
                 self.pinecone_db.upsert_entry(pinecone_entry.dict())
                 article.pinecone_update_required = False