feat: implements RawlBase.count()

mikeshultz · Sep 27, 2023 · 6f8e131 · 6f8e131
1 parent 3ce74fa
commit 6f8e131
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 0 deletions.
diff --git a/rawl/__init__.py b/rawl/__init__.py
@@ -598,6 +598,14 @@ def all(self) -> List[RawlResult]:
 
         return self.select("SELECT {0} FROM " + self.table + ";", self.columns)
 
+    def count(self) -> int:
+        """
+        Get a count of all records in the table.
+        :returns:       List of results
+        """
+
+        return self.query("SELECT COUNT(*) FROM " + self.table + ";")[0][0]
+
     def start_transaction(self) -> Connection[Any]:
         """
         Initiate a connection  to use as a transaction

diff --git a/tests/test_rawl.py b/tests/test_rawl.py
@@ -120,6 +120,14 @@ def delete_rawl_without_commit(self, rawl_id: int) -> List[RawlResult]:
 
 
 class TestRawl(object):
+    @pytest.mark.dependency()
+    def test_count(self, pgdb: Connection[Any]) -> None:
+        """Test out count() statement"""
+
+        mod = TheModel(RAWL_DSN)
+
+        assert mod.count() == 4
+
     @pytest.mark.dependency()
     def test_all(self, pgdb: Connection[Any]) -> None:
         """Test out a basic SELECT statement"""