-class _TestingCursor(psycopg2.extras.DictCursor):
- """ Extension to the DictCursor class that provides execution
- short-cuts that simplify writing assertions.
- """
-
- def scalar(self, sql, params=None):
- """ Execute a query with a single return value and return this value.
- Raises an assertion when not exactly one row is returned.
- """
- self.execute(sql, params)
- assert self.rowcount == 1
- return self.fetchone()[0]
-
- def row_set(self, sql, params=None):
- """ Execute a query and return the result as a set of tuples.
- """
- self.execute(sql, params)
- if self.rowcount == 1:
- return set(tuple(self.fetchone()))
-
- return set((tuple(row) for row in self))
-
- def table_exists(self, table):
- """ Check that a table with the given name exists in the database.
- """
- num = self.scalar("""SELECT count(*) FROM pg_tables
- WHERE tablename = %s""", (table, ))
- return num == 1
-
- def table_rows(self, table):
- """ Return the number of rows in the given table.
- """
- return self.scalar('SELECT count(*) FROM ' + table)