public
Last active

An example of psycopg2 cursor supporting prepared statements

  • Download Gist
prepare.py
Python
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229
#!/usr/bin/env python
"""An example of cursor dealing with prepared statements.
 
A cursor can be used as a regular one, but has also a prepare() statement. If
prepare() is called, execute() and executemany() can be used without query: in
this case the parameters are passed to the prepared statement. The functions
also execute the prepared statement if the query is the same prepared before.
 
Prepared statements aren't automatically deallocated when the cursor is
deleted, but are when the cursor is closed. For long-running sessions creating
an unbound number of cursors you should make sure to deallocate the prepared
statements (calling close() or deallocate() on the cursor).
 
"""
 
# Copyright (C) 2012 Daniele Varrazzo <daniele.varrazzo@gmail.com>
 
import os
import re
from threading import Lock
 
import psycopg2
import psycopg2.extensions as ext
 
class PreparingCursor(ext.cursor):
_lock = Lock()
_ncur = 0
def __init__(self, *args, **kwargs):
super(PreparingCursor, self).__init__(*args, **kwargs)
# create a distinct name for the statements prepared by this cursor
self._lock.acquire()
self._prepname = "psyco_%x" % self._ncur
PreparingCursor._ncur += 1
self._lock.release()
 
self._prepared = None
self._execstmt = None
 
_re_replargs = re.compile(r'(%s)|(%\([^)]+\)s)')
 
def prepare(self, stmt):
"""Prepare a query for execution.
 
TODO: handle literal %s and $s in the string.
"""
# replace the python placeholders with postgres placeholders
parlist = []
parmap = {}
parord = []
 
def repl(m):
par = m.group(1)
if par is not None:
parlist.append(par)
return "$%d" % len(parlist)
else:
par = m.group(2)
assert par
idx = parmap.get(par)
if idx is None:
idx = parmap[par] = "$%d" % (len(parmap) + 1)
parord.append(par)
 
return idx
 
pgstmt = self._re_replargs.sub(repl, stmt)
 
if parlist and parmap:
raise psycopg2.ProgrammingError(
"you can't mix positional and named placeholders")
 
self.deallocate()
self.execute("prepare %s as %s" % (self._prepname, pgstmt))
 
if parlist:
self._execstmt = "execute %s (%s)" % (
self._prepname, ','.join(parlist))
elif parmap:
self._execstmt = "execute %s (%s)" % (
self._prepname, ','.join(parord))
else:
self._execstmt = "execute %s" % (self._prepname)
 
self._prepared = stmt
 
@property
def prepared(self):
"""The query currently prepared."""
return self._prepared
 
def deallocate(self):
"""Deallocate the currently prepared statement."""
if self._prepared is not None:
self.execute("deallocate " + self._prepname)
self._prepared = None
self._execstmt = None
 
def execute(self, stmt=None, args=None):
if stmt is None or stmt == self._prepared:
stmt = self._execstmt
elif not isinstance(stmt, basestring):
args = stmt
stmt = self._execstmt
 
if stmt is None:
raise psycopg2.ProgrammingError(
"execute() with no query called without prepare")
 
return super(PreparingCursor, self).execute(stmt, args)
 
def executemany(self, stmt, args=None):
if args is None:
args = stmt
stmt = self._execstmt
 
if stmt is None:
raise psycopg2.ProgrammingError(
"executemany() with no query called without prepare")
else:
if stmt != self._prepared:
self.prepare(stmt)
 
return super(PreparingCursor, self).executemany(self._execstmt, args)
 
def close(self):
if not self.closed and not self.connection.closed and self._prepared:
self.deallocate()
 
return super(PreparingCursor, self).close()
 
 
 
import unittest
 
class PreparingCursorTestCase(unittest.TestCase):
def setUp(self):
self.conn = psycopg2.connect(os.environ['TEST_DSN'])
 
def tearDown(self):
self.conn.close()
 
def cursor(self):
return self.conn.cursor(cursor_factory=PreparingCursor)
 
def test_prepare_noargs(self):
cur = self.cursor()
self.assert_(not cur.prepared)
cur.prepare("select 1")
self.assertEqual(cur.prepared, "select 1")
cur.execute()
self.assertEqual(cur.fetchone(), (1,))
 
def test_deallocate(self):
cur = self.cursor()
cur.execute("select * from pg_prepared_statements where name = %s",
(cur._prepname,))
self.assert_(not cur.fetchone())
 
cur.prepare("select 1")
self.assert_(cur.prepared)
cur.execute("select * from pg_prepared_statements where name = %s",
(cur._prepname,))
self.assert_(cur.fetchone())
 
cur.deallocate()
self.assert_(not cur.prepared)
cur.execute("select * from pg_prepared_statements where name = %s",
(cur._prepname,))
self.assert_(not cur.fetchone())
 
def test_prepare_posargs(self):
cur = self.cursor()
cur.prepare("select 1 + %s::integer, %s::text")
cur.execute((1,'foo'))
self.assertEqual(cur.fetchone(), (2, 'foo'))
cur.execute(None, [2,'bar'])
self.assertEqual(cur.fetchone(), (3, 'bar'))
 
def test_prepare_kwargs(self):
cur = self.cursor()
cur.prepare(
"select %(foo)s::integer, %(bar)s::integer, %(foo)s::integer")
cur.execute({'foo': 10, 'bar': 20})
self.assertEqual(cur.fetchone(), (10, 20, 10))
cur.execute(None, {'foo': 10, 'bar': 20})
self.assertEqual(cur.fetchone(), (10, 20, 10))
 
def test_executemany(self):
cur = self.cursor()
cur.execute("create table testem(id int, data text)")
cur.executemany("insert into testem values (%s, %s)",
[(i, i) for i in range(10)])
cur.execute("select min(id), max(data), count(*) from testem")
self.assertEqual(cur.fetchone(), (0, "9", 10))
 
def test_prepare_executemany(self):
cur = self.cursor()
cur.execute("create table testem(id int, data text)")
cur.prepare("""
insert into testem values
(%(foo)s::integer, %(foo)s::text)""")
cur.executemany([{'foo': i, 'bar': 2*i} for i in range(10)])
cur.execute("select min(id), max(data), count(*) from testem")
self.assertEqual(cur.fetchone(), (0, "9", 10))
 
def test_nomix(self):
cur = self.cursor()
self.assertRaises(psycopg2.ProgrammingError,
cur.prepare, "select 1 + %s::integer, %(foo)s::text")
 
def test_many(self):
cur1 = self.cursor()
cur2 = self.cursor()
cur1.prepare("select 1")
cur2.prepare("select 2")
cur1.execute()
cur2.execute()
self.assertEqual(cur1.fetchone(), (1,))
self.assertEqual(cur2.fetchone(), (2,))
 
def test_execute_prepared(self):
cur = self.cursor()
cur.prepare("select %s::integer")
cur.execute("select %s::integer", (10,))
self.assert_(cur.query.startswith("execute"))
 
 
if __name__ == '__main__':
unittest.main()

Please sign in to comment on this gist.

Something went wrong with that request. Please try again.