From dcaec24ea4287287de25c819e9c7e3a79590f5cd Mon Sep 17 00:00:00 2001
From: Benjamin Renard <brenard@easter-eggs.com>
Date: Fri, 15 Dec 2023 11:35:43 +0100
Subject: [PATCH] PgDB / MyDB / OracleDB: add limit parameter to select()
 method

---
 mylib/db.py          | 32 +++++++++++++++++++++++++++++++-
 tests/test_mysql.py  |  8 ++++++--
 tests/test_oracle.py |  6 ++++--
 tests/test_pgsql.py  |  8 ++++++--
 4 files changed, 47 insertions(+), 7 deletions(-)

diff --git a/mylib/db.py b/mylib/db.py
index 5a5aafd..142603a 100644
--- a/mylib/db.py
+++ b/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
diff --git a/tests/test_mysql.py b/tests/test_mysql.py
index 7289d98..2cd17db 100644
--- a/tests/test_mysql.py
+++ b/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):
diff --git a/tests/test_oracle.py b/tests/test_oracle.py
index 362f9cc..99081ec 100644
--- a/tests/test_oracle.py
+++ b/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
     )
 
 
diff --git a/tests/test_pgsql.py b/tests/test_pgsql.py
index 552c552..dd6112c 100644
--- a/tests/test_pgsql.py
+++ b/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):