diff --git a/mylib/pgsql.py b/mylib/pgsql.py index 7c0f847..77cf5fd 100644 --- a/mylib/pgsql.py +++ b/mylib/pgsql.py @@ -50,6 +50,19 @@ class PgDBUnsupportedWHEREClauses(PgDBException, TypeError): ) +class PgDBInvalidOrderByClause(PgDBException, TypeError): + """ + Raised when trying to select on table with invalid + ORDER BY clause provided + """ + + def __init__(self, order_by): + super().__init__( + "Invalid ORDER BY clause: {order_by}. Must be a string or a list of two values (ordering field name and direction)", + order_by=order_by + ) + + class PgDB: """ PostgreSQL client """ @@ -367,7 +380,17 @@ class PgDB: return False if order_by: - sql += ' ORDER BY {0}'.format(order_by) + if isinstance(order_by, str): + sql += ' ORDER BY {0}'.format(order_by) + elif ( + isinstance(order_by, (list, tuple)) and len(order_by) == 2 and + isinstance(order_by[0], str) and + isinstance(order_by[1], str) and + order_by[1].upper() in ('ASC', 'UPPER') + ): + sql += ' ORDER BY "{0}" {1}'.format(order_by[0], order_by[1].upper()) + else: + raise PgDBInvalidOrderByClause(order_by) if just_try: log.debug("Just-try mode: execute SELECT query : %s", sql)