Merge pull request #19 from tan-z-tan/tanji/add_methods

Add all() to_a()

pyfireconsole/models/pyfire_model.py

@@ -18,6 +18,7 @@ class PyfireCollection(Generic[ModelType]):
  _collection: Optional[Iterable[dict]] = None
  _where_cond: Optional[WhereCondition] = None
  _order_cond: Optional[OrderCondition] = None
+ _limit: int = 1000 # default limit to prevent loading too large collections
 
  def __init__(self, model_class: Type[ModelType]):
  self.model_class = model_class
@@ -61,6 +62,9 @@ def __iter__(self):
 
  if self._order_cond is not None:
  query = query.order(self._order_cond.field, self._order_cond.direction)
+
+ self._collection = query.limit(self._limit).iter()
+
  self._collection = query.iter()
 
  for doc in self._collection:
@@ -132,6 +136,31 @@ def order(self, field: str, direction: str = 'ASCENDING') -> 'PyfireCollection[M
 
  return coll
 
+ def all(self) -> 'PyfireCollection[ModelType]':
+ """
+ Get all documents in the collection.
+
+ Returns:
+ PyfireCollection[ModelType]: A new PyfireCollection instance containing all documents.
+ """
+ coll = PyfireCollection(self.model_class)
+ coll._where_cond = self._where_cond
+ coll._order_cond = self._order_cond
+ coll._limit = self._limit
+ if self._parent_model is not None:
+ coll.set_parent(self._parent_model)
+
+ return coll
+
+ def to_a(self, limit: int = 1000) -> list[ModelType]:
+ """
+ Convert the collection to a list.
+
+ Returns:
+ list[ModelType]: The collection as a list.
+ """
+ return list(self)
+
  def add(self, entity: ModelType) -> ModelType:
  """
  Add a new document to the collection.

pyfireconsole/queries/query_runner.py

@@ -34,6 +34,17 @@ def all(self) -> 'QueryRunner':
  self.query = AllQuery(self.query or self.collection_key).set_conn(self.conn).exec()
  return self
 
+ def limit(self, limit: int) -> 'QueryRunner':
+ """
+ Apply a limit to the query.
+ Args:
+ limit (int): The maximum number of documents to retrieve.
+ Returns:
+ QueryRunner: The current QueryRunner instance.
+ """
+ self.query = self.query.limit(limit) if self.query else self.conn.collection(self.collection_key).limit(limit)
+ return self
+
  def save(self, id: str, data: dict) -> str | None:
  return SaveQuery(self.collection_key, id, data).set_conn(self.conn).exec()
 
@@ -43,7 +54,10 @@ def create(self, data: dict) -> str | None:
  def delete(self, id: str) -> None:
  return DeleteQuery(self.collection_key, id).set_conn(self.conn).exec()
 
- def iter(self) -> Generator[Dict[str, Any], None, None]:
+ def iter(self, limit: int = 1000) -> Generator[Dict[str, Any], None, None]:
  docs = self.query.stream()
+
+ docs = self.query.limit(limit).stream()
+
  for doc in docs:
  yield dict(_doc_to_dict(doc) or {}, id=doc.id)

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pyfireconsole"
-version = "0.1.0"
+version = "0.1.1"
 description = ""
 authors = ["Makoto Tanji <tanji.makoto@gmail.com>"]
 readme = "README.md"

setup.py

@@ -2,7 +2,7 @@
 
 setup(
  name='pyfireconsole',
- version='0.1.0',
+ version='0.1.1',
  author='Makoto Tanji',
  author_email='tanji.makoto@gmail.com',
  description='An interactive console for Firestore based on Python ORM',

tests/test_usecase.py

@@ -205,6 +205,71 @@ def test_first(mock_db):
  assert user.id in [user1.id, user2.id]
 
 
+def test_class_all(mock_db):
+ mock_db = MockFirestore()
+ mock_db.reset()
+
+ user1 = User.new(
+ name="John",
+ email="",
+ ).save()
+ user2 = User.new(
+ name="Mary",
+ email="",
+ ).save()
+
+ users = User.all()
+
+ assert len([u for u in users]) == 2
+ assert user1.id in [u.id for u in users]
+ assert user2.id in [u.id for u in users]
+
+
+def test_collection_all(mock_db):
+ mock_db = MockFirestore()
+ mock_db.reset()
+
+ user1 = User.new(
+ name="John",
+ email="",
+ ).save()
+ user2 = User.new(
+ name="John",
+ email="",
+ ).save()
+ _ = User.new(
+ name="Mary",
+ email="",
+ ).save()
+
+ users = User.where("name", "==", "John").all()
+
+ assert len([u for u in users]) == 2
+ assert user1.id in [u.id for u in users]
+ assert user2.id in [u.id for u in users]
+
+
+def test_to_a(mock_db):
+ mock_db = MockFirestore()
+ mock_db.reset()
+
+ user1 = User.new(
+ name="John",
+ email="",
+ ).save()
+ user2 = User.new(
+ name="Mary",
+ email="",
+ ).save()
+
+ users = User.all().to_a()
+
+ assert len(users) == 2
+ assert type(users) is list
+ assert user1.id in [users[0].id, users[1].id]
+ assert user2.id in [users[0].id, users[1].id]
+
+
 def test_as_json(mock_db):
  book = Book.new(
  title="Math",