Merge pull request #21 from codemation/model-filter-counting

Feature -  DataBaseModel Count and filter count

.github/workflows/main.yaml

@@ -57,6 +57,7 @@ jobs:
             pytest tests/test_query_caching.py;
             pytest tests/test_model_filtering_operators.py;
             pytest tests/test_model_limit_offset.py;
+            pytest tests/test_model_counting.py;
             pytest tests/test_querying.py;
             pytest tests/test_query_no_caching.py;
             pytest tests/test_integration_fastapi.py;
@@ -118,6 +119,7 @@ jobs:
             pytest tests/test_query_caching.py;
             pytest tests/test_model_filtering_operators.py;
             pytest tests/test_model_limit_offset.py;
+            pytest tests/test_model_counting.py;
             pytest tests/test_querying.py;
             pytest tests/test_query_no_caching.py;
             pytest tests/test_integration_fastapi.py;
@@ -172,6 +174,7 @@ jobs:
             pytest tests/test_query_caching.py;
             pytest tests/test_model_filtering_operators.py;
             pytest tests/test_model_limit_offset.py;
+            pytest tests/test_model_counting.py;
             pytest tests/test_querying.py;
             pytest tests/test_query_no_caching.py;
             pytest tests/test_integration_fastapi.py;

.github/workflows/package.yaml

@@ -58,6 +58,7 @@ jobs:
             pytest tests/test_query_caching.py;
             pytest tests/test_model_filtering_operators.py;
             pytest tests/test_model_limit_offset.py;
+            pytest tests/test_model_limit_offset.py;
             pytest tests/test_querying.py;
             pytest tests/test_query_no_caching.py;
             pytest tests/test_integration_fastapi.py;
@@ -113,6 +114,7 @@ jobs:
             pytest tests/test_query_caching.py;
             pytest tests/test_model_filtering_operators.py;
             pytest tests/test_model_limit_offset.py;
+            pytest tests/test_model_limit_offset.py;
             pytest tests/test_querying.py;
             pytest tests/test_query_no_caching.py;
             pytest tests/test_integration_fastapi.py;
@@ -179,6 +181,7 @@ jobs:
             pytest tests/test_query_caching.py;
             pytest tests/test_model_filtering_operators.py;
             pytest tests/test_model_limit_offset.py;
+            pytest tests/test_model_limit_offset.py;
             pytest tests/test_querying.py;
             pytest tests/test_query_no_caching.py;
             pytest tests/test_integration_fastapi.py;

docs/model-usage.md

@@ -169,7 +169,17 @@ latest_employees = await Employees.filter(
     offset=175
 )
 ```
+#### Counting
+`DataBaseModel` objects can be counted by calling the `.count()` method. Filtered `DataBaseModel` objects can use `.filter(.., count_rows=True)` to return a total count of objects matching a given filter.
 
+```python
+employee_count = await Employees.count()
+
+employed_count = await Employees.filter(
+      is_employed=True,
+      count_rows=True
+)
+```
 
 ### Model Usage - Updating
 Updates to `DataBaseModel` objects must be done directly via an object instance, related `DataBaseModel` field objects must be updated by calling the related fields object's `.save()` or `.update()` method. 

pydbantic/core.py

@@ -1,8 +1,9 @@
+from sqlalchemy.sql.functions import count
 from pydantic import BaseModel, Field
 import typing
 from typing import Optional, Union, List
 import sqlalchemy
-from sqlalchemy import select
+from sqlalchemy import select, func
 from pickle import dumps, loads
 
 class _Generic(BaseModel):
@@ -408,12 +409,13 @@ async def select(cls,
         alias: Optional[dict] = None,
         limit: Optional[int] = None,
         offset: Optional[int] = 0,
-        order_by = None, 
+        order_by = None
     ) -> List[dict]:
         if alias is None:
             alias = {}
 
         table = cls.get_table()
+        database = cls.__metadata__.database
 
         if selection[0] == '*':
             selection = [k for k in cls.__metadata__.tables[cls.__name__]['column_map']]
@@ -444,8 +446,6 @@ async def select(cls,
             sel = sel.order_by(order_by)
 
         decoded_results = []
-
-        database = cls.__metadata__.database
 
         results = await database.fetch(sel, cls.__name__, values)
 
@@ -531,30 +531,43 @@ async def all(
 
         return await cls.select('*', **parameters)
 
+    @classmethod
+    async def count(cls):
+        table = cls.get_table()
+        database = cls.__metadata__.database
+        sel = select([func.count()]).select_from(table)
+        results = await database.fetch(sel, cls.__name__)
+        return results[0][0] if results else 0
+
     @classmethod
     async def filter(
         cls, 
         *conditions, 
         limit: int = None, 
         offset: int = 0,
         order_by = None, 
+        count_rows: bool = False,
         **column_filters
     ):
         table = cls.get_table()
+        database = cls.__metadata__.database
+
         columns = [k for k in cls.__fields__]
         if not column_filters and not conditions:
             raise Exception(f"{cls.__name__}.filter() expects keyword arguments for columns: {columns} or conditions")
-        sel = table.select()
+        sel = table.select() if not count_rows else select([func.count()]).select_from(table)
 
         sel, values = cls.where(sel, column_filters, *conditions)
 
         sel, values = cls.check_limit_offset(sel, values, limit, offset)
+
+        if count_rows:
+            row_count = await database.fetch(sel, cls.__name__)
+            return row_count[0][0] if row_count else 0
 
         if not order_by is None:
             sel = sel.order_by(order_by)
 
-        database = cls.__metadata__.database
-
         results = await database.fetch(sel, cls.__name__, values)
         rows = []
         for result in cls.normalize(results):

tests/test_model_counting.py

@@ -0,0 +1,31 @@
+import pytest
+from tests.models import EmployeeInfo
+
+@pytest.mark.asyncio
+async def test_model_counting(loaded_database_and_model):
+    db, Employees = loaded_database_and_model
+
+    all_employees = await Employees.all()
+    employee_count = await Employees.count()
+
+    print(f"Number of Employees is ", employee_count)
+
+    assert employee_count == len(all_employees)
+
+    employed = await Employees.filter(
+        is_employed=True
+    )
+
+    employed_count = await Employees.filter(
+        is_employed=True,
+        count_rows=True,
+    )
+
+    assert len(employed) == employed_count
+
+    un_employed = await Employees.filter(
+        is_employed=False,
+        count_rows=True
+    )
+    assert un_employed == 0
+