feat: add get_number_of_posts method to the storage class

codeforjapan · Aug 18, 2024 · c5324a1 · c5324a1
1 parent 9ab07ca
commit c5324a1
Show file tree

Hide file tree

Showing 3 changed files with 52 additions and 5 deletions.
diff --git a/api/birdxplorer_api/routers/data.py b/api/birdxplorer_api/routers/data.py
@@ -112,11 +112,18 @@ def get_posts(
                 start=created_at_start,
                 end=created_at_end,
                 search_text=search_text,
+                offset=offset,
+                limit=limit,
             )
         )
+        total_count = storage.get_number_of_posts(
+            post_ids=post_id,
+            note_ids=note_id,
+            start=created_at_start,
+            end=created_at_end,
+            search_text=search_text,
+        )
 
-        total_count = len(posts)
-        paginated_posts = posts[offset : offset + limit]
         base_url = str(request.url).split("?")[0]
         next_offset = offset + limit
         prev_offset = max(offset - limit, 0)
@@ -127,6 +134,6 @@ def get_posts(
         if offset > 0:
             prev_url = f"{base_url}?offset={prev_offset}&limit={limit}"
 
-        return PostListResponse(data=paginated_posts, meta=PaginationMeta(next=next_url, prev=prev_url))
+        return PostListResponse(data=posts, meta=PaginationMeta(next=next_url, prev=prev_url))
 
     return router
diff --git a/api/tests/conftest.py b/api/tests/conftest.py
@@ -260,8 +260,14 @@ def _get_posts(
         start: Union[TwitterTimestamp, None] = None,
         end: Union[TwitterTimestamp, None] = None,
         search_text: Union[str, None] = None,
+        offset: Union[int, None] = None,
+        limit: Union[int, None] = None,
     ) -> Generator[Post, None, None]:
-        for post in post_samples:
+        gen_count = 0
+        actual_gen_count = 0
+        for idx, post in enumerate(post_samples):
+            if limit is not None and actual_gen_count >= limit:
+                break
             if post_ids is not None and post.post_id not in post_ids:
                 continue
             if note_ids is not None and not any(
@@ -274,10 +280,25 @@ def _get_posts(
                 continue
             if search_text is not None and search_text not in post.text:
                 continue
+            gen_count += 1
+            if offset is not None and gen_count <= offset:
+                continue
+            actual_gen_count += 1
             yield post
 
     mock.get_posts.side_effect = _get_posts
 
+    def _get_number_of_posts(
+        post_ids: Union[List[PostId], None] = None,
+        note_ids: Union[List[NoteId], None] = None,
+        start: Union[TwitterTimestamp, None] = None,
+        end: Union[TwitterTimestamp, None] = None,
+        search_text: Union[str, None] = None,
+    ) -> int:
+        return len(list(_get_posts(post_ids, note_ids, start, end, search_text)))
+
+    mock.get_number_of_posts.side_effect = _get_number_of_posts
+
     yield mock
 
 

diff --git a/common/birdxplorer_common/storage.py b/common/birdxplorer_common/storage.py
@@ -268,7 +268,16 @@ def get_notes(
                     created_at=note_record.created_at,
                 )
 
-    def get_posts(self) -> Generator[PostModel, None, None]:
+    def get_posts(
+        self,
+        post_ids: Union[List[PostId], None] = None,
+        note_ids: Union[List[NoteId], None] = None,
+        start: Union[TwitterTimestamp, None] = None,
+        end: Union[TwitterTimestamp, None] = None,
+        search_text: Union[str, None] = None,
+        offset: Union[int, None] = None,
+        limit: Union[int, None] = None,
+    ) -> Generator[PostModel, None, None]:
         with Session(self.engine) as sess:
             for post_record in sess.query(PostRecord).all():
                 yield self._post_record_to_model(post_record)
@@ -310,6 +319,16 @@ def search_posts_by_text(self, search_word: str) -> Generator[PostModel, None, N
             for post_record in sess.query(PostRecord).filter(PostRecord.text.like(f"%{search_word}%")):
                 yield self._post_record_to_model(post_record)
 
+    def get_number_of_posts(
+        self,
+        post_ids: Union[List[PostId], None] = None,
+        note_ids: Union[List[NoteId], None] = None,
+        start: Union[TwitterTimestamp, None] = None,
+        end: Union[TwitterTimestamp, None] = None,
+        search_text: Union[str, None] = None,
+    ) -> int:
+        raise NotImplementedError
+
 
 def gen_storage(settings: GlobalSettings) -> Storage:
     engine = create_engine(settings.storage_settings.sqlalchemy_database_url)