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

This commit is contained in:
Benjamin Renard 2023-12-15 11:35:43 +01:00
parent 2736fc30ae
commit dcaec24ea4
Signed by: bn8
GPG key ID: 3E2E1CE1907115BC
4 changed files with 47 additions and 7 deletions

View file

@ -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: class DB:
"""Database client""" """Database client"""
@ -340,7 +353,14 @@ class DB:
return True return True
def select( 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""" """Run SELECT SQL query"""
sql = "SELECT " sql = "SELECT "
@ -374,6 +394,16 @@ class DB:
else: else:
raise DBInvalidOrderByClause(order_by) 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: if just_try:
log.debug("Just-try mode: execute SELECT query : %s", sql) log.debug("Just-try mode: execute SELECT query : %s", sql)
return just_try return just_try

View file

@ -374,17 +374,21 @@ def test_select(mocker, test_mydb):
{"field1": 2, "field2": 3}, {"field1": 2, "field2": 3},
] ]
order_by = "field1, DESC" order_by = "field1, DESC"
limit = 10
mocker.patch( mocker.patch(
"mylib.mysql.MyDB.doSelect", "mylib.mysql.MyDB.doSelect",
generate_mock_doSQL( generate_mock_doSQL(
"SELECT `field1`, `field2` FROM `mytable` WHERE `test3` = %(test3)s AND `test4` =" "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, where_clauses,
expected_return, 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): def test_select_without_field_and_order_by(mocker, test_mydb):

View file

@ -368,18 +368,20 @@ def test_select(mocker, test_oracledb):
{"field1": 2, "field2": 3}, {"field1": 2, "field2": 3},
] ]
order_by = "field1, DESC" order_by = "field1, DESC"
limit = 10
mocker.patch( mocker.patch(
"mylib.oracle.OracleDB.doSelect", "mylib.oracle.OracleDB.doSelect",
generate_mock_doSQL( generate_mock_doSQL(
'SELECT "field1", "field2" FROM "mytable" WHERE "test3" = :test3 AND "test4" = :test4' '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, where_clauses,
expected_return, expected_return,
), ),
) )
assert ( 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
) )

View file

@ -363,17 +363,21 @@ def test_select(mocker, test_pgdb):
{"field1": 2, "field2": 3}, {"field1": 2, "field2": 3},
] ]
order_by = "field1, DESC" order_by = "field1, DESC"
limit = 10
mocker.patch( mocker.patch(
"mylib.pgsql.PgDB.doSelect", "mylib.pgsql.PgDB.doSelect",
generate_mock_doSQL( generate_mock_doSQL(
'SELECT "field1", "field2" FROM "mytable" WHERE "test3" = %(test3)s AND "test4" =' '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, where_clauses,
expected_return, 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): def test_select_without_field_and_order_by(mocker, test_pgdb):