PgDB / MyDB / OracleDB: add limit parameter to select() method

mylib/db.py

@@ -77,6 +77,19 @@ class DBInvalidOrderByClause(DBException, TypeError):
         )
 
 
+class DBInvalidLimitClause(DBException, TypeError):
+    """
+    Raised when trying to select on table with invalid
+    LIMIT clause provided
+    """
+
+    def __init__(self, limit):
+        super().__init__(
+            "Invalid LIMIT clause: {limit}. Must be a non-zero positive integer.",
+            limit=limit,
+        )
+
+
 class DB:
     """Database client"""
 
@@ -340,7 +353,14 @@ class DB:
         return True
 
     def select(
-        self, table, where_clauses=None, fields=None, where_op="AND", order_by=None, just_try=False
+        self,
+        table,
+        where_clauses=None,
+        fields=None,
+        where_op="AND",
+        order_by=None,
+        limit=None,
+        just_try=False,
     ):
         """Run SELECT SQL query"""
         sql = "SELECT "
@@ -374,6 +394,16 @@ class DB:
             else:
                 raise DBInvalidOrderByClause(order_by)
 
+        if limit:
+            if not isinstance(limit, int):
+                try:
+                    limit = int(limit)
+                except ValueError as err:
+                    raise DBInvalidLimitClause(limit) from err
+            if limit <= 0:
+                raise DBInvalidLimitClause(limit)
+            sql += f" LIMIT {limit}"
+
         if just_try:
             log.debug("Just-try mode: execute SELECT query : %s", sql)
             return just_try

tests/test_mysql.py

@@ -374,17 +374,21 @@ def test_select(mocker, test_mydb):
         {"field1": 2, "field2": 3},
     ]
     order_by = "field1, DESC"
+    limit = 10
     mocker.patch(
         "mylib.mysql.MyDB.doSelect",
         generate_mock_doSQL(
             "SELECT `field1`, `field2` FROM `mytable` WHERE `test3` = %(test3)s AND `test4` ="
-            " %(test4)s ORDER BY " + order_by,  # nosec: B608
+            " %(test4)s ORDER BY " + order_by + " LIMIT " + str(limit),  # nosec: B608
             where_clauses,
             expected_return,
         ),
     )
 
-    assert test_mydb.select("mytable", where_clauses, fields, order_by=order_by) == expected_return
+    assert (
+        test_mydb.select("mytable", where_clauses, fields, order_by=order_by, limit=limit)
+        == expected_return
+    )
 
 
 def test_select_without_field_and_order_by(mocker, test_mydb):

tests/test_oracle.py

@@ -368,18 +368,20 @@ def test_select(mocker, test_oracledb):
         {"field1": 2, "field2": 3},
     ]
     order_by = "field1, DESC"
+    limit = 10
     mocker.patch(
         "mylib.oracle.OracleDB.doSelect",
         generate_mock_doSQL(
             'SELECT "field1", "field2" FROM "mytable" WHERE "test3" = :test3 AND "test4" = :test4'
-            " ORDER BY " + order_by,  # nosec: B608
+            " ORDER BY " + order_by + " LIMIT " + str(limit),  # nosec: B608
             where_clauses,
             expected_return,
         ),
     )
 
     assert (
-        test_oracledb.select("mytable", where_clauses, fields, order_by=order_by) == expected_return
+        test_oracledb.select("mytable", where_clauses, fields, order_by=order_by, limit=limit)
+        == expected_return
     )
 
 

tests/test_pgsql.py

@@ -363,17 +363,21 @@ def test_select(mocker, test_pgdb):
         {"field1": 2, "field2": 3},
     ]
     order_by = "field1, DESC"
+    limit = 10
     mocker.patch(
         "mylib.pgsql.PgDB.doSelect",
         generate_mock_doSQL(
             'SELECT "field1", "field2" FROM "mytable" WHERE "test3" = %(test3)s AND "test4" ='
-            " %(test4)s ORDER BY " + order_by,  # nosec: B608
+            " %(test4)s ORDER BY " + order_by + " LIMIT " + str(limit),  # nosec: B608
             where_clauses,
             expected_return,
         ),
     )
 
-    assert test_pgdb.select("mytable", where_clauses, fields, order_by=order_by) == expected_return
+    assert (
+        test_pgdb.select("mytable", where_clauses, fields, order_by=order_by, limit=limit)
+        == expected_return
+    )
 
 
 def test_select_without_field_and_order_by(mocker, test_pgdb):