trunk/psycopg2/examples/cursor.py

   1 # cursor.py - how to subclass the cursor type
   2 #
   3 # Copyright (C) 2004 Federico Di Gregorio  <fog@debian.org>
   4 #
   5 # This program is free software; you can redistribute it and/or modify
   6 # it under the terms of the GNU General Public License as published by the
   7 # Free Software Foundation; either version 2, or (at your option) any later
   8 # version.
   9 #
  10 # This program is distributed in the hope that it will be useful, but
  11 # WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTIBILITY
  12 # or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
  13 # for more details.
  14
  15 ## put in DSN your DSN string
  16
  17 DSN = 'dbname=test'
  18
  19 ## don't modify anything below this line (except for experimenting)
  20
  21 import sys
  22 import psycopg2
  23 import psycopg2.extensions
  24
  25 if len(sys.argv) > 1:
  26     DSN = sys.argv[1]
  27
  28 print "Opening connection using dsn:", DSN
  29 conn = psycopg2.connect(DSN)
  30 print "Encoding for this connection is", conn.encoding
  31
  32
  33 class NoDataError(psycopg2.ProgrammingError):
  34     """Exception that will be raised by our cursor."""
  35     pass
  36
  37 class Cursor(psycopg2.extensions.cursor):
  38     """A custom cursor."""
  39
  40     def fetchone(self):
  41         """Like fetchone but raise an exception if no data is available.
  42
  43         Note that to have .fetchmany() and .fetchall() to raise the same
  44         exception we'll have to override them too; even if internally psycopg
  45         uses the same function to fetch rows, the code path from Python is
  46         different.
  47         """
  48         d = psycopg2.extensions.cursor.fetchone(self)
  49         if d is None:
  50             raise NoDataError("no more data")
  51         return d
  52
  53 curs = conn.cursor(cursor_factory=Cursor)
  54 curs.execute("SELECT 1 AS foo")
  55 print "Result of fetchone():", curs.fetchone()
  56
  57 # now let's raise the exception
  58 try:
  59     curs.fetchone()
  60 except NoDataError, err:
  61     print "Exception caugth:", err
  62
  63 conn.rollback()