psycopg/psycopg2
1
1
Fork 0
mirror of https://github.com/psycopg/psycopg2.git synced 2025-08-01 02:50:08 +03:00
Code Issues Packages Projects Releases Wiki Activity

Merge e5322a71bb into cecff195fc

Browse Source
This commit is contained in:
Jon Dufresne 2020-07-13 23:52:28 +00:00 committed by GitHub
commit 5f93150707
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
32 changed files with 3315 additions and 3052 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
lib/_json.py
View File

@ -180,8 +180,7 @@ def _get_json_oids(conn_or_curs, name='json'):
from psycopg2.extensions import STATUS_IN_TRANSACTION
from psycopg2.extras import _solve_conn_curs
conn, curs = _solve_conn_curs(conn_or_curs)
with _solve_conn_curs(conn_or_curs) as (conn, curs):
# Store the transaction status of the connection to revert it after use
conn_status = conn.status

3
lib/_range.py
View File

@ -351,8 +351,7 @@ class RangeCaster(object):
"""
from psycopg2.extensions import STATUS_IN_TRANSACTION
from psycopg2.extras import _solve_conn_curs
conn, curs = _solve_conn_curs(conn_or_curs)
with _solve_conn_curs(conn_or_curs) as (conn, curs):
if conn.info.server_version < 90200:
raise ProgrammingError("range types not available in version %s"
% conn.info.server_version)

11
lib/extras.py
View File

@ -26,6 +26,7 @@ and classes until a better place in the distribution is found.
# FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public
# License for more details.
import contextlib
import os as _os
import time as _time
import re as _re
@ -787,6 +788,7 @@ def wait_select(conn):
continue
@contextlib.contextmanager
def _solve_conn_curs(conn_or_curs):
"""Return the connection and a DBAPI cursor from a connection or cursor."""
if conn_or_curs is None:
@ -799,7 +801,8 @@ def _solve_conn_curs(conn_or_curs):
conn = conn_or_curs
curs = conn.cursor(cursor_factory=_cursor)
return conn, curs
with curs:
yield conn, curs
class HstoreAdapter(object):
@ -910,8 +913,7 @@ class HstoreAdapter(object):
def get_oids(self, conn_or_curs):
"""Return the lists of OID of the hstore and hstore[] types.
"""
conn, curs = _solve_conn_curs(conn_or_curs)
with _solve_conn_curs(conn_or_curs) as (conn, curs):
# Store the transaction status of the connection to revert it after use
conn_status = conn.status
@ -1089,8 +1091,7 @@ class CompositeCaster(object):
Raise `ProgrammingError` if the type is not found.
"""
conn, curs = _solve_conn_curs(conn_or_curs)
with _solve_conn_curs(conn_or_curs) as (conn, curs):
# Store the transaction status of the connection to revert it after use
conn_status = conn.status

80
psycopg/connection_type.c
View File

@ -59,6 +59,7 @@ static PyObject *
psyco_conn_cursor(connectionObject *self, PyObject *args, PyObject *kwargs)
{
PyObject *obj = NULL;
cursorObject *curs = NULL;
PyObject *rv = NULL;
PyObject *name = Py_None;
PyObject *factory = Py_None;
@ -113,23 +114,46 @@ psyco_conn_cursor(connectionObject *self, PyObject *args, PyObject *kwargs)
goto exit;
}
if (0 > curs_withhold_set((cursorObject *)obj, withhold)) {
goto exit;
/* pass ownership from obj to curs */
curs = (cursorObject *)obj;
obj = NULL;
if (0 > curs_withhold_set(curs, withhold)) {
goto error;
}
if (0 > curs_scrollable_set((cursorObject *)obj, scrollable)) {
goto exit;
if (0 > curs_scrollable_set(curs, scrollable)) {
goto error;
}
Dprintf("psyco_conn_cursor: new cursor at %p: refcnt = "
FORMAT_CODE_PY_SSIZE_T,
obj, Py_REFCNT(obj)
curs, Py_REFCNT(curs)
);
rv = obj;
obj = NULL;
/* pass ownership from curs to rv */
rv = (PyObject *)curs;
curs = NULL;
goto exit;
error:
{
PyObject *error_type, *error_value, *error_traceback;
PyObject *close;
PyErr_Fetch(&error_type, &error_value, &error_traceback);
if (curs) {
close = PyObject_CallMethod((PyObject *)curs, "close", NULL);
if (close)
Py_DECREF(close);
else
PyErr_WriteUnraisable((PyObject *)curs);
}
PyErr_Restore(error_type, error_value, error_traceback);
}
exit:
Py_XDECREF(obj);
Py_XDECREF(curs);
return rv;
}
@ -1366,6 +1390,9 @@ connection_dealloc(PyObject* obj)
{
connectionObject *self = (connectionObject *)obj;
if (PyObject_CallFinalizerFromDealloc(obj) < 0)
return;
/* Make sure to untrack the connection before calling conn_close, which may
* allow a different thread to try and dealloc the connection again,
* resulting in a double-free segfault (ticket #166). */
@ -1405,6 +1432,31 @@ connection_dealloc(PyObject* obj)
Py_TYPE(obj)->tp_free(obj);
}
#if PY_3
static void
connection_finalize(PyObject *obj)
{
connectionObject *self = (connectionObject *)obj;
#ifdef CONN_CHECK_PID
if (self->procpid == getpid())
#endif
{
if (!self->closed) {
PyObject *error_type, *error_value, *error_traceback;
/* Save the current exception, if any. */
PyErr_Fetch(&error_type, &error_value, &error_traceback);
if (PyErr_WarnFormat(PyExc_ResourceWarning, 1, "unclosed connection %R", obj))
PyErr_WriteUnraisable(obj);
/* Restore the saved exception. */
PyErr_Restore(error_type, error_value, error_traceback);
}
}
}
#endif
static int
connection_init(PyObject *obj, PyObject *args, PyObject *kwds)
{
@ -1479,7 +1531,7 @@ PyTypeObject connectionType = {
0, /*tp_setattro*/
0, /*tp_as_buffer*/
Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC |
Py_TPFLAGS_HAVE_WEAKREFS,
Py_TPFLAGS_HAVE_WEAKREFS | Py_TPFLAGS_HAVE_FINALIZE,
/*tp_flags*/
connectionType_doc, /*tp_doc*/
(traverseproc)connection_traverse, /*tp_traverse*/
@ -1499,4 +1551,16 @@ PyTypeObject connectionType = {
connection_init, /*tp_init*/
0, /*tp_alloc*/
connection_new, /*tp_new*/
0, /* tp_free */
0, /* tp_is_gc */
0, /* tp_bases */
0, /* tp_mro */
0, /* tp_cache */
0, /* tp_subclasses */
0, /* tp_weaklist */
#if PY_3
0, /* tp_del */
0, /* tp_version_tag */
connection_finalize, /* tp_finalize */
#endif
};

42
psycopg/cursor_type.c
View File

@ -39,6 +39,7 @@
#include <stdlib.h>
#define cursor_closed(self) ((self)->closed || ((self)->conn && (self)->conn->closed))
/** DBAPI methods **/
@ -1637,7 +1638,7 @@ exit:
static PyObject *
curs_closed_get(cursorObject *self, void *closure)
{
return PyBool_FromLong(self->closed || (self->conn && self->conn->closed));
return PyBool_FromLong(cursor_closed(self));
}
/* extension: withhold - get or set "WITH HOLD" for named cursors */
@ -1945,6 +1946,9 @@ cursor_dealloc(PyObject* obj)
{
cursorObject *self = (cursorObject *)obj;
if (PyObject_CallFinalizerFromDealloc(obj) < 0)
return;
PyObject_GC_UnTrack(self);
if (self->weakreflist) {
@ -1965,6 +1969,28 @@ cursor_dealloc(PyObject* obj)
Py_TYPE(obj)->tp_free(obj);
}
#if PY_3
static void
cursor_finalize(PyObject *obj)
{
cursorObject *self = (cursorObject *)obj;
if (!cursor_closed(self)) {
PyObject *error_type, *error_value, *error_traceback;
/* Save the current exception, if any. */
PyErr_Fetch(&error_type, &error_value, &error_traceback);
if (PyErr_WarnFormat(PyExc_ResourceWarning, 1,
"unclosed cursor %R for connection %R",
obj, (PyObject *)self->conn))
PyErr_WriteUnraisable(obj);
/* Restore the saved exception. */
PyErr_Restore(error_type, error_value, error_traceback);
}
}
#endif
static int
cursor_init(PyObject *obj, PyObject *args, PyObject *kwargs)
{
@ -2056,7 +2082,7 @@ PyTypeObject cursorType = {
0, /*tp_setattro*/
0, /*tp_as_buffer*/
Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_ITER |
Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_HAVE_WEAKREFS ,
Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_HAVE_WEAKREFS | Py_TPFLAGS_HAVE_FINALIZE,
/*tp_flags*/
cursorType_doc, /*tp_doc*/
(traverseproc)cursor_traverse, /*tp_traverse*/
@ -2076,4 +2102,16 @@ PyTypeObject cursorType = {
cursor_init, /*tp_init*/
0, /*tp_alloc*/
cursor_new, /*tp_new*/
0, /* tp_free */
0, /* tp_is_gc */
0, /* tp_bases */
0, /* tp_mro */
0, /* tp_cache */
0, /* tp_subclasses */
0, /* tp_weaklist */
#if PY_3
0, /* tp_del */
0, /* tp_version_tag */
cursor_finalize, /* tp_finalize */
#endif
};

4
psycopg/python.h
View File

@ -86,6 +86,8 @@ typedef unsigned long Py_uhash_t;
#define Bytes_ConcatAndDel PyString_ConcatAndDel
#define _Bytes_Resize _PyString_Resize
#define Py_TPFLAGS_HAVE_FINALIZE 0L
#define PyDateTime_DELTA_GET_DAYS(o) (((PyDateTime_Delta*)o)->days)
#define PyDateTime_DELTA_GET_SECONDS(o) (((PyDateTime_Delta*)o)->seconds)
#define PyDateTime_DELTA_GET_MICROSECONDS(o) (((PyDateTime_Delta*)o)->microseconds)
@ -97,6 +99,8 @@ typedef unsigned long Py_uhash_t;
PyLong_FromUnsignedLong((unsigned long)(x)) : \
PyInt_FromLong((x)))
#define PyObject_CallFinalizerFromDealloc(obj) 0
#endif /* PY_2 */
#if PY_3

2
tests/__init__.py
View File

@ -56,6 +56,7 @@ from . import test_sql
from . import test_transaction
from . import test_types_basic
from . import test_types_extras
from . import test_warnings
from . import test_with
if sys.version_info[:2] < (3, 6):
@ -101,6 +102,7 @@ def test_suite():
suite.addTest(test_transaction.test_suite())
suite.addTest(test_types_basic.test_suite())
suite.addTest(test_types_extras.test_suite())
suite.addTest(test_warnings.test_suite())
suite.addTest(test_with.test_suite())
return suite

2
tests/dbapi20.py
View File

@ -232,6 +232,7 @@ class DatabaseAPI20Test(unittest.TestCase):
self.failUnless(con.InternalError is drv.InternalError)
self.failUnless(con.ProgrammingError is drv.ProgrammingError)
self.failUnless(con.NotSupportedError is drv.NotSupportedError)
con.close()
def test_commit(self):
@ -251,6 +252,7 @@ class DatabaseAPI20Test(unittest.TestCase):
con.rollback()
except self.driver.NotSupportedError:
pass
con.close()
def test_cursor(self):
con = self._connect()

3
tests/dbapi20_tpc.py
View File

@ -23,6 +23,7 @@ class TwoPhaseCommitTests(unittest.TestCase):
def test_xid(self):
con = self.connect()
try:
try:
xid = con.xid(42, "global", "bqual")
except self.driver.NotSupportedError:
@ -37,6 +38,8 @@ class TwoPhaseCommitTests(unittest.TestCase):
self.assertEquals(tuple(xid), (0, "", ""))
xid = con.xid(0x7fffffff, "a" * 64, "b" * 64)
self.assertEquals(tuple(xid), (0x7fffffff, "a" * 64, "b" * 64))
finally:
con.close()
def test_tpc_begin(self):
con = self.connect()

75
tests/test_async.py
View File

@ -61,7 +61,7 @@ class AsyncTests(ConnectingTestCase):
self.wait(self.conn)
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute('''
CREATE TEMPORARY TABLE table1 (
id int PRIMARY KEY
@ -71,6 +71,8 @@ class AsyncTests(ConnectingTestCase):
def test_connection_setup(self):
cur = self.conn.cursor()
sync_cur = self.sync_conn.cursor()
cur.close()
sync_cur.close()
del cur, sync_cur
self.assert_(self.conn.async_)
@ -90,7 +92,7 @@ class AsyncTests(ConnectingTestCase):
self.conn.cursor, "name")
def test_async_select(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
self.assertFalse(self.conn.isexecuting())
cur.execute("select 'a'")
self.assertTrue(self.conn.isexecuting())
@ -103,7 +105,7 @@ class AsyncTests(ConnectingTestCase):
@slow
@skip_before_postgres(8, 2)
def test_async_callproc(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.callproc("pg_sleep", (0.1, ))
self.assertTrue(self.conn.isexecuting())
@ -113,8 +115,9 @@ class AsyncTests(ConnectingTestCase):
@slow
def test_async_after_async(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur2 = self.conn.cursor()
cur2.close()
del cur2
cur.execute("insert into table1 values (1)")
@ -141,7 +144,7 @@ class AsyncTests(ConnectingTestCase):
self.assertEquals(cur.fetchone(), None)
def test_fetch_after_async(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("select 'a'")
# a fetch after an asynchronous query should raise an error
@ -152,16 +155,14 @@ class AsyncTests(ConnectingTestCase):
self.assertEquals(cur.fetchall()[0][0], "a")
def test_rollback_while_async(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("select 'a'")
# a rollback should not work in asynchronous mode
self.assertRaises(psycopg2.ProgrammingError, self.conn.rollback)
def test_commit_while_async(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("begin")
self.wait(cur)
@ -188,8 +189,7 @@ class AsyncTests(ConnectingTestCase):
self.assertEquals(cur.fetchone(), None)
def test_set_parameters_while_async(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("select 'c'")
self.assertTrue(self.conn.isexecuting())
@ -207,7 +207,7 @@ class AsyncTests(ConnectingTestCase):
self.conn.set_isolation_level, 1)
def test_reset_while_async(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("select 'c'")
self.assertTrue(self.conn.isexecuting())
@ -215,8 +215,7 @@ class AsyncTests(ConnectingTestCase):
self.assertRaises(psycopg2.ProgrammingError, self.conn.reset)
def test_async_iter(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("begin")
self.wait(cur)
cur.execute("""
@ -236,7 +235,7 @@ class AsyncTests(ConnectingTestCase):
self.assertFalse(self.conn.isexecuting())
def test_copy_while_async(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("select 'a'")
# copy should fail
@ -250,13 +249,13 @@ class AsyncTests(ConnectingTestCase):
self.conn.lobject)
def test_async_executemany(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
self.assertRaises(
psycopg2.ProgrammingError,
cur.executemany, "insert into table1 values (%s)", [1, 2, 3])
def test_async_scroll(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("""
insert into table1 values (1);
insert into table1 values (2);
@ -274,16 +273,16 @@ class AsyncTests(ConnectingTestCase):
cur.scroll(1)
self.assertEquals(cur.fetchall(), [(2, ), (3, )])
cur = self.conn.cursor()
with self.conn.cursor() as cur2:
cur.execute("select id from table1 order by id")
self.wait(cur)
cur2 = self.conn.cursor()
with self.conn.cursor() as cur2:
self.assertRaises(psycopg2.ProgrammingError, cur2.scroll, 1)
self.assertRaises(psycopg2.ProgrammingError, cur.scroll, 4)
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("select id from table1 order by id")
self.wait(cur)
cur.scroll(2)
@ -291,7 +290,7 @@ class AsyncTests(ConnectingTestCase):
self.assertEquals(cur.fetchall(), [(2, ), (3, )])
def test_scroll(self):
cur = self.sync_conn.cursor()
with self.sync_conn.cursor() as cur:
cur.execute("create table table1 (id int)")
cur.execute("""
insert into table1 values (1);
@ -304,7 +303,7 @@ class AsyncTests(ConnectingTestCase):
self.assertEquals(cur.fetchall(), [(2, ), (3, )])
def test_async_dont_read_all(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("select repeat('a', 10000); select repeat('b', 10000)")
# fetch the result
@ -326,7 +325,7 @@ class AsyncTests(ConnectingTestCase):
@slow
def test_flush_on_write(self):
# a very large query requires a flush loop to be sent to the backend
curs = self.conn.cursor()
with self.conn.cursor() as curs:
for mb in 1, 5, 10, 20, 50:
size = mb * 1024 * 1024
stub = PollableStub(self.conn)
@ -343,7 +342,7 @@ class AsyncTests(ConnectingTestCase):
warnings.warn("sending a large query didn't trigger block on write.")
def test_sync_poll(self):
cur = self.sync_conn.cursor()
with self.sync_conn.cursor() as cur:
cur.execute("select 1")
# polling with a sync query works
cur.connection.poll()
@ -351,9 +350,7 @@ class AsyncTests(ConnectingTestCase):
@slow
def test_notify(self):
cur = self.conn.cursor()
sync_cur = self.sync_conn.cursor()
with self.conn.cursor() as cur, self.sync_conn.cursor() as sync_cur:
sync_cur.execute("listen test_notify")
self.sync_conn.commit()
cur.execute("notify test_notify")
@ -374,8 +371,7 @@ class AsyncTests(ConnectingTestCase):
self.fail("No NOTIFY in 2.5 seconds")
def test_async_fetch_wrong_cursor(self):
cur1 = self.conn.cursor()
cur2 = self.conn.cursor()
with self.conn.cursor() as cur1, self.conn.cursor() as cur2:
cur1.execute("select 1")
self.wait(cur1)
@ -386,7 +382,7 @@ class AsyncTests(ConnectingTestCase):
self.assertEquals(cur1.fetchone()[0], 1)
def test_error(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("insert into table1 values (%s)", (1, ))
self.wait(cur)
cur.execute("insert into table1 values (%s)", (1, ))
@ -409,7 +405,7 @@ class AsyncTests(ConnectingTestCase):
self.wait(cur)
def test_stop_on_first_error(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("select 1; select x; select 1/0; select 2")
self.assertRaises(psycopg2.errors.UndefinedColumn, self.wait, cur)
@ -418,8 +414,7 @@ class AsyncTests(ConnectingTestCase):
self.assertEqual(cur.fetchone(), (1,))
def test_error_two_cursors(self):
cur = self.conn.cursor()
cur2 = self.conn.cursor()
with self.conn.cursor() as cur, self.conn.cursor() as cur2:
cur.execute("select * from no_such_table")
self.assertRaises(psycopg2.ProgrammingError, self.wait, cur)
cur2.execute("select 1")
@ -428,7 +423,7 @@ class AsyncTests(ConnectingTestCase):
def test_notices(self):
del self.conn.notices[:]
cur = self.conn.cursor()
with self.conn.cursor() as cur:
if self.conn.info.server_version >= 90300:
cur.execute("set client_min_messages=debug1")
self.wait(cur)
@ -438,31 +433,33 @@ class AsyncTests(ConnectingTestCase):
self.assert_(self.conn.notices)
def test_async_cursor_gone(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("select 42;")
del cur
gc.collect()
self.assertRaises(psycopg2.InterfaceError, self.wait, self.conn)
# The connection is still usable
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("select 42;")
self.wait(self.conn)
self.assertEqual(cur.fetchone(), (42,))
def test_async_connection_error_message(self):
try:
cnn = psycopg2.connect('dbname=thisdatabasedoesntexist', async_=True)
try:
self.wait(cnn)
except psycopg2.Error as e:
self.assertNotEqual(str(e), "asynchronous connection failed",
"connection error reason lost")
else:
self.fail("no exception raised")
finally:
cnn.close()
@skip_before_postgres(8, 2)
def test_copy_no_hang(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("copy (select 1) to stdout")
self.assertRaises(psycopg2.ProgrammingError, self.wait, self.conn)
@ -471,7 +468,7 @@ class AsyncTests(ConnectingTestCase):
def test_non_block_after_notification(self):
from select import select
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("""
select 1;
do $$

14
tests/test_async_keyword.py
View File

@ -47,7 +47,7 @@ class AsyncTests(ConnectingTestCase):
self.wait(self.conn)
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute('''
CREATE TEMPORARY TABLE table1 (
id int PRIMARY KEY
@ -57,6 +57,8 @@ class AsyncTests(ConnectingTestCase):
def test_connection_setup(self):
cur = self.conn.cursor()
sync_cur = self.sync_conn.cursor()
cur.close()
sync_cur.close()
del cur, sync_cur
self.assert_(self.conn.async)
@ -89,13 +91,15 @@ class AsyncTests(ConnectingTestCase):
"connection error reason lost")
else:
self.fail("no exception raised")
finally:
cnn.close()
class CancelTests(ConnectingTestCase):
def setUp(self):
ConnectingTestCase.setUp(self)
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute('''
CREATE TEMPORARY TABLE table1 (
id int PRIMARY KEY
@ -108,7 +112,7 @@ class CancelTests(ConnectingTestCase):
async_conn = psycopg2.connect(dsn, async=True)
self.assertRaises(psycopg2.OperationalError, async_conn.cancel)
extras.wait_select(async_conn)
cur = async_conn.cursor()
with async_conn.cursor() as cur:
cur.execute("select pg_sleep(10)")
time.sleep(1)
self.assertTrue(async_conn.isexecuting())
@ -118,6 +122,7 @@ class CancelTests(ConnectingTestCase):
cur.execute("select 1")
extras.wait_select(async_conn)
self.assertEqual(cur.fetchall(), [(1, )])
async_conn.close()
def test_async_connection_cancel(self):
async_conn = psycopg2.connect(dsn, async=True)
@ -180,8 +185,7 @@ class AsyncReplicationTest(ReplicationTestCase):
if conn is None:
return
cur = conn.cursor()
with conn.cursor() as cur:
self.create_replication_slot(cur, output_plugin='test_decoding')
self.wait(cur)

2
tests/test_bug_gc.py
View File

@ -39,7 +39,7 @@ class StolenReferenceTestCase(ConnectingTestCase):
return 42
UUID = psycopg2.extensions.new_type((2950,), "UUID", fish)
psycopg2.extensions.register_type(UUID, self.conn)
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("select 'b5219e01-19ab-4994-b71e-149225dc51e4'::uuid")
curs.fetchone()

9
tests/test_cancel.py
View File

@ -41,7 +41,7 @@ class CancelTests(ConnectingTestCase):
def setUp(self):
ConnectingTestCase.setUp(self)
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute('''
CREATE TEMPORARY TABLE table1 (
id int PRIMARY KEY
@ -57,7 +57,7 @@ class CancelTests(ConnectingTestCase):
errors = []
def neverending(conn):
cur = conn.cursor()
with conn.cursor() as cur:
try:
self.assertRaises(psycopg2.extensions.QueryCanceledError,
cur.execute, "select pg_sleep(60)")
@ -70,7 +70,7 @@ class CancelTests(ConnectingTestCase):
raise
def canceller(conn):
cur = conn.cursor()
with conn.cursor() as cur:
try:
conn.cancel()
except Exception as e:
@ -95,7 +95,7 @@ class CancelTests(ConnectingTestCase):
async_conn = psycopg2.connect(dsn, async_=True)
self.assertRaises(psycopg2.OperationalError, async_conn.cancel)
extras.wait_select(async_conn)
cur = async_conn.cursor()
with async_conn.cursor() as cur:
cur.execute("select pg_sleep(10)")
time.sleep(1)
self.assertTrue(async_conn.isexecuting())
@ -105,6 +105,7 @@ class CancelTests(ConnectingTestCase):
cur.execute("select 1")
extras.wait_select(async_conn)
self.assertEqual(cur.fetchall(), [(1, )])
async_conn.close()
def test_async_connection_cancel(self):
async_conn = psycopg2.connect(dsn, async_=True)

118
tests/test_connection.py
View File

@ -80,7 +80,7 @@ class ConnectionTests(ConnectingTestCase):
def test_cleanup_on_badconn_close(self):
# ticket #148
conn = self.conn
cur = conn.cursor()
with conn.cursor() as cur:
self.assertRaises(psycopg2.OperationalError,
cur.execute, "select pg_terminate_backend(pg_backend_pid())")
@ -113,7 +113,7 @@ class ConnectionTests(ConnectingTestCase):
def test_notices(self):
conn = self.conn
cur = conn.cursor()
with conn.cursor() as cur:
if self.conn.info.server_version >= 90300:
cur.execute("set client_min_messages=debug1")
cur.execute("create temp table chatty (id serial primary key);")
@ -122,7 +122,7 @@ class ConnectionTests(ConnectingTestCase):
def test_notices_consistent_order(self):
conn = self.conn
cur = conn.cursor()
with conn.cursor() as cur:
if self.conn.info.server_version >= 90300:
cur.execute("set client_min_messages=debug1")
cur.execute("""
@ -142,7 +142,7 @@ class ConnectionTests(ConnectingTestCase):
@slow
def test_notices_limited(self):
conn = self.conn
cur = conn.cursor()
with conn.cursor() as cur:
if self.conn.info.server_version >= 90300:
cur.execute("set client_min_messages=debug1")
for i in range(0, 100, 10):
@ -157,7 +157,7 @@ class ConnectionTests(ConnectingTestCase):
def test_notices_deque(self):
conn = self.conn
self.conn.notices = deque()
cur = conn.cursor()
with conn.cursor() as cur:
if self.conn.info.server_version >= 90300:
cur.execute("set client_min_messages=debug1")
@ -187,7 +187,7 @@ class ConnectionTests(ConnectingTestCase):
def test_notices_noappend(self):
conn = self.conn
self.conn.notices = None # will make an error swallowes ok
cur = conn.cursor()
with conn.cursor() as cur:
if self.conn.info.server_version >= 90300:
cur.execute("set client_min_messages=debug1")
@ -233,7 +233,7 @@ class ConnectionTests(ConnectingTestCase):
def test_encoding_name(self):
self.conn.set_client_encoding("EUC_JP")
# conn.encoding is 'EUCJP' now.
cur = self.conn.cursor()
with self.conn.cursor() as cur:
ext.register_type(ext.UNICODE, cur)
cur.execute("select 'foo'::text;")
self.assertEqual(cur.fetchone()[0], u'foo')
@ -281,7 +281,7 @@ class ConnectionTests(ConnectingTestCase):
while conn.notices:
notices.append((2, conn.notices.pop()))
cur = conn.cursor()
with conn.cursor() as cur:
t1 = threading.Thread(target=committer)
t1.start()
for i in range(1000):
@ -297,36 +297,36 @@ class ConnectionTests(ConnectingTestCase):
def test_connect_cursor_factory(self):
conn = self.connect(cursor_factory=psycopg2.extras.DictCursor)
cur = conn.cursor()
with conn.cursor() as cur:
cur.execute("select 1 as a")
self.assertEqual(cur.fetchone()['a'], 1)
def test_cursor_factory(self):
self.assertEqual(self.conn.cursor_factory, None)
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("select 1 as a")
self.assertRaises(TypeError, (lambda r: r['a']), cur.fetchone())
self.conn.cursor_factory = psycopg2.extras.DictCursor
self.assertEqual(self.conn.cursor_factory, psycopg2.extras.DictCursor)
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("select 1 as a")
self.assertEqual(cur.fetchone()['a'], 1)
self.conn.cursor_factory = None
self.assertEqual(self.conn.cursor_factory, None)
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("select 1 as a")
self.assertRaises(TypeError, (lambda r: r['a']), cur.fetchone())
def test_cursor_factory_none(self):
# issue #210
conn = self.connect()
cur = conn.cursor(cursor_factory=None)
with conn.cursor(cursor_factory=None) as cur:
self.assertEqual(type(cur), ext.cursor)
conn = self.connect(cursor_factory=psycopg2.extras.DictCursor)
cur = conn.cursor(cursor_factory=None)
with conn.cursor(cursor_factory=None) as cur:
self.assertEqual(type(cur), psycopg2.extras.DictCursor)
def test_failed_init_status(self):
@ -585,8 +585,7 @@ class IsolationLevelsTestCase(ConnectingTestCase):
def test_set_isolation_level(self):
conn = self.connect()
curs = conn.cursor()
with conn.cursor() as curs:
levels = [
('read uncommitted',
ext.ISOLATION_LEVEL_READ_UNCOMMITTED),
@ -617,8 +616,7 @@ class IsolationLevelsTestCase(ConnectingTestCase):
def test_set_isolation_level_autocommit(self):
conn = self.connect()
curs = conn.cursor()
with conn.cursor() as curs:
conn.set_isolation_level(ext.ISOLATION_LEVEL_AUTOCOMMIT)
self.assertEqual(conn.isolation_level, ext.ISOLATION_LEVEL_DEFAULT)
self.assert_(conn.autocommit)
@ -632,8 +630,7 @@ class IsolationLevelsTestCase(ConnectingTestCase):
def test_set_isolation_level_default(self):
conn = self.connect()
curs = conn.cursor()
with conn.cursor() as curs:
conn.autocommit = True
curs.execute("set default_transaction_isolation to 'read committed'")
@ -651,8 +648,7 @@ class IsolationLevelsTestCase(ConnectingTestCase):
def test_set_isolation_level_abort(self):
conn = self.connect()
cur = conn.cursor()
with conn.cursor() as cur:
self.assertEqual(ext.TRANSACTION_STATUS_IDLE,
conn.info.transaction_status)
cur.execute("insert into isolevel values (10);")
@ -693,12 +689,12 @@ class IsolationLevelsTestCase(ConnectingTestCase):
cnn2 = self.connect()
cnn2.set_isolation_level(ext.ISOLATION_LEVEL_AUTOCOMMIT)
cur1 = cnn1.cursor()
with cnn1.cursor() as cur1:
cur1.execute("select count(*) from isolevel;")
self.assertEqual(0, cur1.fetchone()[0])
cnn1.commit()
cur2 = cnn2.cursor()
with cnn2.cursor() as cur2:
cur2.execute("insert into isolevel values (10);")
cur1.execute("select count(*) from isolevel;")
@ -709,12 +705,12 @@ class IsolationLevelsTestCase(ConnectingTestCase):
cnn2 = self.connect()
cnn2.set_isolation_level(ext.ISOLATION_LEVEL_READ_COMMITTED)
cur1 = cnn1.cursor()
with cnn1.cursor() as cur1:
cur1.execute("select count(*) from isolevel;")
self.assertEqual(0, cur1.fetchone()[0])
cnn1.commit()
cur2 = cnn2.cursor()
with cnn2.cursor() as cur2:
cur2.execute("insert into isolevel values (10);")
cur1.execute("insert into isolevel values (20);")
@ -735,12 +731,12 @@ class IsolationLevelsTestCase(ConnectingTestCase):
cnn2 = self.connect()
cnn2.set_isolation_level(ext.ISOLATION_LEVEL_SERIALIZABLE)
cur1 = cnn1.cursor()
with cnn1.cursor() as cur1:
cur1.execute("select count(*) from isolevel;")
self.assertEqual(0, cur1.fetchone()[0])
cnn1.commit()
cur2 = cnn2.cursor()
with cnn2.cursor() as cur2:
cur2.execute("insert into isolevel values (10);")
cur1.execute("insert into isolevel values (20);")
@ -768,7 +764,7 @@ class IsolationLevelsTestCase(ConnectingTestCase):
cnn.set_isolation_level, 1)
def test_setattr_isolation_level_int(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
self.conn.isolation_level = ext.ISOLATION_LEVEL_SERIALIZABLE
self.assertEqual(self.conn.isolation_level, ext.ISOLATION_LEVEL_SERIALIZABLE)
@ -816,7 +812,7 @@ class IsolationLevelsTestCase(ConnectingTestCase):
self.assertEqual(cur.fetchone()[0], isol)
def test_setattr_isolation_level_str(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
self.conn.isolation_level = "serializable"
self.assertEqual(self.conn.isolation_level, ext.ISOLATION_LEVEL_SERIALIZABLE)
@ -950,7 +946,7 @@ class ConnectionTwoPhaseTests(ConnectingTestCase):
cnn.tpc_begin(xid)
self.assertEqual(cnn.status, ext.STATUS_BEGIN)
cur = cnn.cursor()
with cnn.cursor() as cur:
cur.execute("insert into test_tpc values ('test_tpc_commit');")
self.assertEqual(0, self.count_xacts())
self.assertEqual(0, self.count_test_records())
@ -973,7 +969,7 @@ class ConnectionTwoPhaseTests(ConnectingTestCase):
cnn.tpc_begin(xid)
self.assertEqual(cnn.status, ext.STATUS_BEGIN)
cur = cnn.cursor()
with cnn.cursor() as cur:
cur.execute("insert into test_tpc values ('test_tpc_commit_1p');")
self.assertEqual(0, self.count_xacts())
self.assertEqual(0, self.count_test_records())
@ -1017,7 +1013,7 @@ class ConnectionTwoPhaseTests(ConnectingTestCase):
cnn.tpc_begin(xid)
self.assertEqual(cnn.status, ext.STATUS_BEGIN)
cur = cnn.cursor()
with cnn.cursor() as cur:
cur.execute("insert into test_tpc values ('test_tpc_rollback');")
self.assertEqual(0, self.count_xacts())
self.assertEqual(0, self.count_test_records())
@ -1040,7 +1036,7 @@ class ConnectionTwoPhaseTests(ConnectingTestCase):
cnn.tpc_begin(xid)
self.assertEqual(cnn.status, ext.STATUS_BEGIN)
cur = cnn.cursor()
with cnn.cursor() as cur:
cur.execute("insert into test_tpc values ('test_tpc_rollback_1p');")
self.assertEqual(0, self.count_xacts())
self.assertEqual(0, self.count_test_records())
@ -1058,7 +1054,7 @@ class ConnectionTwoPhaseTests(ConnectingTestCase):
cnn.tpc_begin(xid)
self.assertEqual(cnn.status, ext.STATUS_BEGIN)
cur = cnn.cursor()
with cnn.cursor() as cur:
cur.execute("insert into test_tpc values ('test_tpc_commit_rec');")
self.assertEqual(0, self.count_xacts())
self.assertEqual(0, self.count_test_records())
@ -1082,7 +1078,7 @@ class ConnectionTwoPhaseTests(ConnectingTestCase):
cnn.tpc_recover()
self.assertEqual(ext.STATUS_READY, cnn.status)
cur = cnn.cursor()
with cnn.cursor() as cur:
cur.execute("select 1")
self.assertEqual(ext.STATUS_BEGIN, cnn.status)
cnn.tpc_recover()
@ -1092,7 +1088,7 @@ class ConnectionTwoPhaseTests(ConnectingTestCase):
# insert a few test xns
cnn = self.connect()
cnn.set_isolation_level(0)
cur = cnn.cursor()
with cnn.cursor() as cur:
cur.execute("begin; prepare transaction '1-foo';")
cur.execute("begin; prepare transaction '2-bar';")
@ -1125,7 +1121,7 @@ class ConnectionTwoPhaseTests(ConnectingTestCase):
cnn.tpc_prepare()
cnn = self.connect()
cur = cnn.cursor()
with cnn.cursor() as cur:
cur.execute("select gid from pg_prepared_xacts where database = %s;",
(dbname,))
self.assertEqual('42_Z3RyaWQ=_YnF1YWw=', cur.fetchone()[0])
@ -1252,14 +1248,14 @@ class TransactionControlTests(ConnectingTestCase):
ext.ISOLATION_LEVEL_SERIALIZABLE)
def test_not_in_transaction(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("select 1")
self.assertRaises(psycopg2.ProgrammingError,
self.conn.set_session,
ext.ISOLATION_LEVEL_SERIALIZABLE)
def test_set_isolation_level(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
self.conn.set_session(
ext.ISOLATION_LEVEL_SERIALIZABLE)
cur.execute("SHOW transaction_isolation;")
@ -1291,7 +1287,7 @@ class TransactionControlTests(ConnectingTestCase):
self.conn.rollback()
def test_set_isolation_level_str(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
self.conn.set_session("serializable")
cur.execute("SHOW transaction_isolation;")
self.assertEqual(cur.fetchone()[0], 'serializable')
@ -1326,7 +1322,7 @@ class TransactionControlTests(ConnectingTestCase):
def test_set_read_only(self):
self.assert_(self.conn.readonly is None)
cur = self.conn.cursor()
with self.conn.cursor() as cur:
self.conn.set_session(readonly=True)
self.assert_(self.conn.readonly is True)
cur.execute("SHOW transaction_read_only;")
@ -1343,7 +1339,7 @@ class TransactionControlTests(ConnectingTestCase):
self.conn.rollback()
def test_setattr_read_only(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
self.conn.readonly = True
self.assert_(self.conn.readonly is True)
cur.execute("SHOW transaction_read_only;")
@ -1356,7 +1352,7 @@ class TransactionControlTests(ConnectingTestCase):
self.assertEqual(cur.fetchone()[0], 'on')
self.conn.rollback()
cur = self.conn.cursor()
with self.conn.cursor() as cur:
self.conn.readonly = None
self.assert_(self.conn.readonly is None)
cur.execute("SHOW transaction_read_only;")
@ -1370,7 +1366,7 @@ class TransactionControlTests(ConnectingTestCase):
self.conn.rollback()
def test_set_default(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("SHOW transaction_isolation;")
isolevel = cur.fetchone()[0]
cur.execute("SHOW transaction_read_only;")
@ -1388,7 +1384,7 @@ class TransactionControlTests(ConnectingTestCase):
@skip_before_postgres(9, 1)
def test_set_deferrable(self):
self.assert_(self.conn.deferrable is None)
cur = self.conn.cursor()
with self.conn.cursor() as cur:
self.conn.set_session(readonly=True, deferrable=True)
self.assert_(self.conn.deferrable is True)
cur.execute("SHOW transaction_read_only;")
@ -1417,7 +1413,7 @@ class TransactionControlTests(ConnectingTestCase):
@skip_before_postgres(9, 1)
def test_setattr_deferrable(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
self.conn.deferrable = True
self.assert_(self.conn.deferrable is True)
cur.execute("SHOW transaction_deferrable;")
@ -1430,7 +1426,7 @@ class TransactionControlTests(ConnectingTestCase):
self.assertEqual(cur.fetchone()[0], 'on')
self.conn.rollback()
cur = self.conn.cursor()
with self.conn.cursor() as cur:
self.conn.deferrable = None
self.assert_(self.conn.deferrable is None)
cur.execute("SHOW transaction_deferrable;")
@ -1444,7 +1440,7 @@ class TransactionControlTests(ConnectingTestCase):
self.conn.rollback()
def test_mixing_session_attribs(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
self.conn.autocommit = True
self.conn.readonly = True
@ -1467,7 +1463,7 @@ class TransactionControlTests(ConnectingTestCase):
self.conn.autocommit = True
self.conn.readonly = True
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("SHOW transaction_read_only")
self.assertEqual(cur.fetchone()[0], 'on')
@ -1490,7 +1486,7 @@ class TestEncryptPassword(ConnectingTestCase):
@skip_before_libpq(10)
@skip_before_postgres(10)
def test_encrypt_server(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("SHOW password_encryption;")
server_encryption_algorithm = cur.fetchone()[0]
@ -1572,7 +1568,7 @@ class AutocommitTests(ConnectingTestCase):
self.assertEqual(self.conn.info.transaction_status,
ext.TRANSACTION_STATUS_IDLE)
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute('select 1;')
self.assertEqual(self.conn.status, ext.STATUS_BEGIN)
self.assertEqual(self.conn.info.transaction_status,
@ -1590,7 +1586,7 @@ class AutocommitTests(ConnectingTestCase):
self.assertEqual(self.conn.info.transaction_status,
ext.TRANSACTION_STATUS_IDLE)
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute('select 1;')
self.assertEqual(self.conn.status, ext.STATUS_READY)
self.assertEqual(self.conn.info.transaction_status,
@ -1608,7 +1604,7 @@ class AutocommitTests(ConnectingTestCase):
ext.TRANSACTION_STATUS_INTRANS)
def test_set_intrans_error(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute('select 1;')
self.assertRaises(psycopg2.ProgrammingError,
setattr, self.conn, 'autocommit', True)
@ -1620,7 +1616,7 @@ class AutocommitTests(ConnectingTestCase):
self.assertEqual(self.conn.info.transaction_status,
ext.TRANSACTION_STATUS_IDLE)
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute('select 1;')
self.assertEqual(self.conn.status, ext.STATUS_READY)
self.assertEqual(self.conn.info.transaction_status,
@ -1766,7 +1762,7 @@ class TestConnectionInfo(ConnectingTestCase):
self.assert_(self.bconn.info.dbname is None)
def test_user(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("select user")
self.assertEqual(self.conn.info.user, cur.fetchone()[0])
self.assert_(self.bconn.info.user is None)
@ -1805,13 +1801,13 @@ class TestConnectionInfo(ConnectingTestCase):
def test_transaction_status(self):
self.assertEqual(self.conn.info.transaction_status, 0)
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("select 1")
self.assertEqual(self.conn.info.transaction_status, 2)
self.assertEqual(self.bconn.info.transaction_status, 4)
def test_parameter_status(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
try:
cur.execute("show server_version")
except psycopg2.DatabaseError:
@ -1830,7 +1826,7 @@ class TestConnectionInfo(ConnectingTestCase):
self.assertEqual(self.bconn.info.protocol_version, 0)
def test_server_version(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
try:
cur.execute("show server_version_num")
except psycopg2.DatabaseError:
@ -1845,7 +1841,7 @@ class TestConnectionInfo(ConnectingTestCase):
self.assertIsNone(self.conn.info.error_message)
self.assertIsNotNone(self.bconn.info.error_message)
cur = self.conn.cursor()
with self.conn.cursor() as cur:
try:
cur.execute("select 1 from nosuchtable")
except psycopg2.DatabaseError:
@ -1858,7 +1854,7 @@ class TestConnectionInfo(ConnectingTestCase):
self.assert_(self.bconn.info.socket < 0)
def test_backend_pid(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
try:
cur.execute("select pg_backend_pid()")
except psycopg2.DatabaseError:

28
tests/test_copy.py
View File

@ -66,7 +66,7 @@ class CopyTests(ConnectingTestCase):
self._create_temp_table()
def _create_temp_table(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute('''
CREATE TEMPORARY TABLE tcopy (
id serial PRIMARY KEY,
@ -92,7 +92,7 @@ class CopyTests(ConnectingTestCase):
curs.close()
def test_copy_from_cols(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
f = StringIO()
for i in range(10):
f.write("%s\n" % (i,))
@ -104,7 +104,7 @@ class CopyTests(ConnectingTestCase):
self.assertEqual([(i, None) for i in range(10)], curs.fetchall())
def test_copy_from_cols_err(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
f = StringIO()
for i in range(10):
f.write("%s\n" % (i,))
@ -140,7 +140,7 @@ class CopyTests(ConnectingTestCase):
+ list(range(160, 256))).decode('latin1')
about = abin.replace('\\', '\\\\')
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute('insert into tcopy values (%s, %s)',
(42, abin))
@ -161,7 +161,7 @@ class CopyTests(ConnectingTestCase):
+ list(range(160, 255))).decode('latin1')
about = abin.replace('\\', '\\\\').encode('latin1')
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute('insert into tcopy values (%s, %s)',
(42, abin))
@ -188,7 +188,7 @@ class CopyTests(ConnectingTestCase):
f.write(about)
f.seek(0)
curs = self.conn.cursor()
with self.conn.cursor() as curs:
psycopg2.extensions.register_type(
psycopg2.extensions.UNICODE, curs)
@ -254,14 +254,14 @@ class CopyTests(ConnectingTestCase):
pass
f = Whatever()
curs = self.conn.cursor()
with self.conn.cursor() as curs:
self.assertRaises(TypeError,
curs.copy_expert, 'COPY tcopy (data) FROM STDIN', f)
def test_copy_no_column_limit(self):
cols = ["c%050d" % i for i in range(200)]
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute('CREATE TEMPORARY TABLE manycols (%s)' % ',\n'.join(
["%s int" % c for c in cols]))
curs.execute("INSERT INTO manycols DEFAULT VALUES")
@ -278,8 +278,7 @@ class CopyTests(ConnectingTestCase):
@skip_before_postgres(8, 2) # they don't send the count
def test_copy_rowcount(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.copy_from(StringIO('aaa\nbbb\nccc\n'), 'tcopy', columns=['data'])
self.assertEqual(curs.rowcount, 3)
@ -296,8 +295,7 @@ class CopyTests(ConnectingTestCase):
self.assertEqual(curs.rowcount, 6)
def test_copy_rowcount_error(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("insert into tcopy (data) values ('fff')")
self.assertEqual(curs.rowcount, 1)
@ -317,6 +315,7 @@ try:
curs.execute("copy copy_segf from stdin")
except psycopg2.ProgrammingError:
pass
curs.close()
conn.close()
""" % {'dsn': dsn})
@ -336,6 +335,7 @@ try:
curs.execute("copy copy_segf to stdout")
except psycopg2.ProgrammingError:
pass
curs.close()
conn.close()
""" % {'dsn': dsn})
@ -351,7 +351,7 @@ conn.close()
def readline(self):
return 1 / 0
curs = self.conn.cursor()
with self.conn.cursor() as curs:
# It seems we cannot do this, but now at least we propagate the error
# self.assertRaises(ZeroDivisionError,
# curs.copy_from, BrokenRead(), "tcopy")
@ -365,7 +365,7 @@ conn.close()
def write(self, data):
return 1 / 0
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("insert into tcopy values (10, 'hi')")
self.assertRaises(ZeroDivisionError,
curs.copy_to, BrokenWrite(), "tcopy")

63
tests/test_cursor.py
View File

@ -51,7 +51,7 @@ class CursorTests(ConnectingTestCase):
self.assert_(cur.closed)
def test_empty_query(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
self.assertRaises(psycopg2.ProgrammingError, cur.execute, "")
self.assertRaises(psycopg2.ProgrammingError, cur.execute, " ")
self.assertRaises(psycopg2.ProgrammingError, cur.execute, ";")
@ -70,8 +70,7 @@ class CursorTests(ConnectingTestCase):
def test_mogrify_unicode(self):
conn = self.conn
cur = conn.cursor()
with conn.cursor() as cur:
# test consistency between execute and mogrify.
# unicode query containing only ascii data
@ -108,7 +107,7 @@ class CursorTests(ConnectingTestCase):
def test_mogrify_decimal_explodes(self):
conn = self.conn
cur = conn.cursor()
with conn.cursor() as cur:
self.assertEqual(b'SELECT 10.3;',
cur.mogrify("SELECT %s;", (Decimal("10.3"),)))
@ -116,7 +115,7 @@ class CursorTests(ConnectingTestCase):
def test_mogrify_leak_on_multiple_reference(self):
# issue #81: reference leak when a parameter value is referenced
# more than once from a dict.
cur = self.conn.cursor()
with self.conn.cursor() as cur:
foo = (lambda x: x)('foo') * 10
nref1 = sys.getrefcount(foo)
cur.mogrify("select %(foo)s, %(foo)s, %(foo)s", {'foo': foo})
@ -130,7 +129,7 @@ class CursorTests(ConnectingTestCase):
self.assertEqual(sql, b"select 10")
def test_bad_placeholder(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
self.assertRaises(psycopg2.ProgrammingError,
cur.mogrify, "select %(foo", {})
self.assertRaises(psycopg2.ProgrammingError,
@ -141,8 +140,7 @@ class CursorTests(ConnectingTestCase):
cur.mogrify, "select %(foo, %(bar)", {'foo': 1, 'bar': 2})
def test_cast(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
self.assertEqual(42, curs.cast(20, '42'))
self.assertAlmostEqual(3.14, curs.cast(700, '3.14'))
@ -152,7 +150,7 @@ class CursorTests(ConnectingTestCase):
self.assertEqual("who am i?", curs.cast(705, 'who am i?')) # unknown
def test_cast_specificity(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
self.assertEqual("foo", curs.cast(705, 'foo'))
D = psycopg2.extensions.new_type((705,), "DOUBLING", lambda v, c: v * 2)
@ -163,18 +161,18 @@ class CursorTests(ConnectingTestCase):
psycopg2.extensions.register_type(T, curs)
self.assertEqual("foofoofoo", curs.cast(705, 'foo'))
curs2 = self.conn.cursor()
with self.conn.cursor() as curs2:
self.assertEqual("foofoo", curs2.cast(705, 'foo'))
def test_weakref(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
w = ref(curs)
del curs
gc.collect()
self.assert_(w() is None)
def test_null_name(self):
curs = self.conn.cursor(None)
with self.conn.cursor(None) as curs:
self.assertEqual(curs.name, None)
def test_invalid_name(self):
@ -184,7 +182,7 @@ class CursorTests(ConnectingTestCase):
curs.execute("insert into invname values (%s)", (i,))
curs.close()
curs = self.conn.cursor(r'1-2-3 \ "test"')
with self.conn.cursor(r'1-2-3 \ "test"') as curs:
curs.execute("select data from invname order by data")
self.assertEqual(curs.fetchall(), [(10,), (20,), (30,)])
@ -213,13 +211,13 @@ class CursorTests(ConnectingTestCase):
self.assertEqual(curs.fetchall(), [(10,), (20,), (30,)])
curs.close()
curs = self.conn.cursor("W", withhold=True)
with self.conn.cursor("W", withhold=True) as curs:
self.assertEqual(curs.withhold, True)
curs.execute("select data from withhold order by data")
self.conn.commit()
self.assertEqual(curs.fetchall(), [(10,), (20,), (30,)])
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("drop table withhold")
self.conn.commit()
@ -328,7 +326,7 @@ class CursorTests(ConnectingTestCase):
return self.skipTest("can't evaluate non-scrollable cursor")
curs.close()
curs = self.conn.cursor("S", scrollable=False)
with self.conn.cursor("S", scrollable=False) as curs:
self.assertEqual(curs.scrollable, False)
curs.execute("select * from scrollable")
curs.scroll(2)
@ -337,7 +335,7 @@ class CursorTests(ConnectingTestCase):
@slow
@skip_before_postgres(8, 2)
def test_iter_named_cursor_efficient(self):
curs = self.conn.cursor('tmp')
with self.conn.cursor('tmp') as curs:
# if these records are fetched in the same roundtrip their
# timestamp will not be influenced by the pause in Python world.
curs.execute("""select clock_timestamp() from generate_series(1,2)""")
@ -351,7 +349,7 @@ class CursorTests(ConnectingTestCase):
@skip_before_postgres(8, 0)
def test_iter_named_cursor_default_itersize(self):
curs = self.conn.cursor('tmp')
with self.conn.cursor('tmp') as curs:
curs.execute('select generate_series(1,50)')
rv = [(r[0], curs.rownumber) for r in curs]
# everything swallowed in one gulp
@ -359,7 +357,7 @@ class CursorTests(ConnectingTestCase):
@skip_before_postgres(8, 0)
def test_iter_named_cursor_itersize(self):
curs = self.conn.cursor('tmp')
with self.conn.cursor('tmp') as curs:
curs.itersize = 30
curs.execute('select generate_series(1,50)')
rv = [(r[0], curs.rownumber) for r in curs]
@ -368,7 +366,7 @@ class CursorTests(ConnectingTestCase):
@skip_before_postgres(8, 0)
def test_iter_named_cursor_rownumber(self):
curs = self.conn.cursor('tmp')
with self.conn.cursor('tmp') as curs:
# note: this fails if itersize < dataset: internally we check
# rownumber == rowcount to detect when to read anoter page, so we
# would need an extra attribute to have a monotonic rownumber.
@ -378,7 +376,7 @@ class CursorTests(ConnectingTestCase):
self.assertEqual(i + 1, curs.rownumber)
def test_description_attribs(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("""select
3.14::decimal(10,2) as pi,
'hello'::text as hi,
@ -413,7 +411,7 @@ class CursorTests(ConnectingTestCase):
self.assertEqual(c.scale, None)
def test_description_extra_attribs(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("""
create table testcol (
pi decimal(10,2),
@ -439,7 +437,7 @@ class CursorTests(ConnectingTestCase):
self.assertEqual(curs.description[0][0:2], ('a', 23))
def test_pickle_description(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute('SELECT 1 AS foo')
description = curs.description
@ -452,11 +450,11 @@ class CursorTests(ConnectingTestCase):
def test_named_cursor_stealing(self):
# you can use a named cursor to iterate on a refcursor created
# somewhere else
cur1 = self.conn.cursor()
with self.conn.cursor() as cur1:
cur1.execute("DECLARE test CURSOR WITHOUT HOLD "
" FOR SELECT generate_series(1,7)")
cur2 = self.conn.cursor('test')
with self.conn.cursor('test') as cur2:
# can call fetch without execute
self.assertEqual((1,), cur2.fetchone())
self.assertEqual([(2,), (3,), (4,)], cur2.fetchmany(3))
@ -469,7 +467,7 @@ class CursorTests(ConnectingTestCase):
@skip_before_postgres(8, 2)
def test_stolen_named_cursor_close(self):
cur1 = self.conn.cursor()
with self.conn.cursor() as cur1:
cur1.execute("DECLARE test CURSOR WITHOUT HOLD "
" FOR SELECT generate_series(1,7)")
cur2 = self.conn.cursor('test')
@ -482,7 +480,7 @@ class CursorTests(ConnectingTestCase):
@skip_before_postgres(8, 0)
def test_scroll(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("select generate_series(0,9)")
cur.scroll(2)
self.assertEqual(cur.fetchone(), (2,))
@ -516,7 +514,7 @@ class CursorTests(ConnectingTestCase):
@skip_before_postgres(8, 0)
def test_scroll_named(self):
cur = self.conn.cursor('tmp', scrollable=True)
with self.conn.cursor('tmp', scrollable=True) as cur:
cur.execute("select generate_series(0,9)")
cur.scroll(2)
self.assertEqual(cur.fetchone(), (2,))
@ -536,13 +534,13 @@ class CursorTests(ConnectingTestCase):
# I am stupid so not calling superclass init
pass
cur = StupidCursor()
with StupidCursor() as cur:
self.assertRaises(psycopg2.InterfaceError, cur.execute, 'select 1')
self.assertRaises(psycopg2.InterfaceError, cur.executemany,
'select 1', [])
def test_callproc_badparam(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
self.assertRaises(TypeError, cur.callproc, 'lower', 42)
# It would be inappropriate to test callproc's named parameters in the
@ -556,8 +554,7 @@ class CursorTests(ConnectingTestCase):
escaped_paramname = '"%s"' % paramname.replace('"', '""')
procname = 'pg_temp.randall'
cur = self.conn.cursor()
with self.conn.cursor() as cur:
# Set up the temporary function
cur.execute('''
CREATE FUNCTION %s(%s INT)
@ -643,7 +640,7 @@ class CursorTests(ConnectingTestCase):
@skip_before_postgres(8, 2)
def test_rowcount_on_executemany_returning(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("create table execmany(id serial primary key, data int)")
cur.executemany(
"insert into execmany (data) values (%s)",

6
tests/test_dates.py
View File

@ -369,7 +369,7 @@ class DatetimeTests(ConnectingTestCase, CommonDatetimeTestsMixin):
self.assertEqual(total_seconds(t), 1e-6)
def test_interval_overflow(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
# hack a cursor to receive values too extreme to be represented
# but still I want an error, not a random number
psycopg2.extensions.register_type(
@ -405,7 +405,7 @@ class DatetimeTests(ConnectingTestCase, CommonDatetimeTestsMixin):
def test_redshift_day(self):
# Redshift is reported returning 1 day interval as microsec (bug #558)
cur = self.conn.cursor()
with self.conn.cursor() as cur:
psycopg2.extensions.register_type(
psycopg2.extensions.new_type(
psycopg2.STRING.values, 'WAT', psycopg2.extensions.INTERVAL),
@ -426,7 +426,7 @@ class DatetimeTests(ConnectingTestCase, CommonDatetimeTestsMixin):
@skip_before_postgres(8, 4)
def test_interval_iso_8601_not_supported(self):
# We may end up supporting, but no pressure for it
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("set local intervalstyle to iso_8601")
cur.execute("select '1 day 2 hours'::interval")
self.assertRaises(psycopg2.NotSupportedError, cur.fetchone)

125
tests/test_extras_dictcursor.py
View File

@ -33,7 +33,7 @@ from .testutils import ConnectingTestCase, skip_before_postgres, \
class _DictCursorBase(ConnectingTestCase):
def setUp(self):
ConnectingTestCase.setUp(self)
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("CREATE TEMPORARY TABLE ExtrasDictCursorTests (foo text)")
curs.execute("INSERT INTO ExtrasDictCursorTests VALUES ('bar')")
self.conn.commit()
@ -62,15 +62,18 @@ class _DictCursorBase(ConnectingTestCase):
class ExtrasDictCursorTests(_DictCursorBase):
"""Test if DictCursor extension class works."""
@skip_before_postgres(8, 2)
def testDictConnCursorArgs(self):
self.conn.close()
self.conn = self.connect(connection_factory=psycopg2.extras.DictConnection)
cur = self.conn.cursor()
with self.conn.cursor() as cur:
self.assert_(isinstance(cur, psycopg2.extras.DictCursor))
self.assertEqual(cur.name, None)
# overridable
cur = self.conn.cursor('foo',
cursor_factory=psycopg2.extras.NamedTupleCursor)
with self.conn.cursor(
'foo',
cursor_factory=psycopg2.extras.NamedTupleCursor
) as cur:
self.assertEqual(cur.name, 'foo')
self.assert_(isinstance(cur, psycopg2.extras.NamedTupleCursor))
@ -100,11 +103,11 @@ class ExtrasDictCursorTests(_DictCursorBase):
@skip_before_postgres(8, 0)
def testDictCursorWithPlainCursorIterRowNumber(self):
curs = self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
with self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as curs:
self._testIterRowNumber(curs)
def _testWithPlainCursor(self, getter):
curs = self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
with self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as curs:
curs.execute("SELECT * FROM ExtrasDictCursorTests")
row = getter(curs)
self.failUnless(row['foo'] == 'bar')
@ -131,23 +134,32 @@ class ExtrasDictCursorTests(_DictCursorBase):
@skip_before_postgres(8, 2)
def testDictCursorWithNamedCursorNotGreedy(self):
curs = self.conn.cursor('tmp', cursor_factory=psycopg2.extras.DictCursor)
with self.conn.cursor(
'tmp',
cursor_factory=psycopg2.extras.DictCursor
) as curs:
self._testNamedCursorNotGreedy(curs)
@skip_before_postgres(8, 0)
def testDictCursorWithNamedCursorIterRowNumber(self):
curs = self.conn.cursor('tmp', cursor_factory=psycopg2.extras.DictCursor)
with self.conn.cursor(
'tmp',
cursor_factory=psycopg2.extras.DictCursor
) as curs:
self._testIterRowNumber(curs)
def _testWithNamedCursor(self, getter):
curs = self.conn.cursor('aname', cursor_factory=psycopg2.extras.DictCursor)
with self.conn.cursor(
'aname',
cursor_factory=psycopg2.extras.DictCursor
) as curs:
curs.execute("SELECT * FROM ExtrasDictCursorTests")
row = getter(curs)
self.failUnless(row['foo'] == 'bar')
self.failUnless(row[0] == 'bar')
def testPickleDictRow(self):
curs = self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
with self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as curs:
curs.execute("select 10 as a, 20 as b")
r = curs.fetchone()
d = pickle.dumps(r)
@ -175,7 +187,7 @@ class ExtrasDictCursorTests(_DictCursorBase):
@skip_from_python(3)
def test_iter_methods_2(self):
curs = self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
with self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as curs:
curs.execute("select 10 as a, 20 as b")
r = curs.fetchone()
self.assert_(isinstance(r.keys(), list))
@ -194,7 +206,7 @@ class ExtrasDictCursorTests(_DictCursorBase):
@skip_before_python(3)
def test_iter_methods_3(self):
curs = self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
with self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as curs:
curs.execute("select 10 as a, 20 as b")
r = curs.fetchone()
self.assert_(not isinstance(r.keys(), list))
@ -205,7 +217,7 @@ class ExtrasDictCursorTests(_DictCursorBase):
self.assertEqual(len(list(r.items())), 2)
def test_order(self):
curs = self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
with self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as curs:
curs.execute("select 5 as foo, 4 as bar, 33 as baz, 2 as qux")
r = curs.fetchone()
self.assertEqual(list(r), [5, 4, 33, 2])
@ -238,7 +250,7 @@ class ExtrasDictCursorTests(_DictCursorBase):
class ExtrasDictCursorRealTests(_DictCursorBase):
def testRealMeansReal(self):
curs = self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
with self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as curs:
curs.execute("SELECT * FROM ExtrasDictCursorTests")
row = curs.fetchone()
self.assert_(isinstance(row, dict))
@ -263,17 +275,17 @@ class ExtrasDictCursorRealTests(_DictCursorBase):
@skip_before_postgres(8, 0)
def testDictCursorWithPlainCursorRealIterRowNumber(self):
curs = self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
with self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as curs:
self._testIterRowNumber(curs)
def _testWithPlainCursorReal(self, getter):
curs = self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
with self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as curs:
curs.execute("SELECT * FROM ExtrasDictCursorTests")
row = getter(curs)
self.failUnless(row['foo'] == 'bar')
def testPickleRealDictRow(self):
curs = self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
with self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as curs:
curs.execute("select 10 as a, 20 as b")
r = curs.fetchone()
d = pickle.dumps(r)
@ -316,24 +328,32 @@ class ExtrasDictCursorRealTests(_DictCursorBase):
@skip_before_postgres(8, 2)
def testDictCursorRealWithNamedCursorNotGreedy(self):
curs = self.conn.cursor('tmp', cursor_factory=psycopg2.extras.RealDictCursor)
with self.conn.cursor(
'tmp',
cursor_factory=psycopg2.extras.RealDictCursor
) as curs:
self._testNamedCursorNotGreedy(curs)
@skip_before_postgres(8, 0)
def testDictCursorRealWithNamedCursorIterRowNumber(self):
curs = self.conn.cursor('tmp', cursor_factory=psycopg2.extras.RealDictCursor)
with self.conn.cursor(
'tmp',
cursor_factory=psycopg2.extras.RealDictCursor
) as curs:
self._testIterRowNumber(curs)
def _testWithNamedCursorReal(self, getter):
curs = self.conn.cursor('aname',
cursor_factory=psycopg2.extras.RealDictCursor)
with self.conn.cursor(
'aname',
cursor_factory=psycopg2.extras.RealDictCursor
) as curs:
curs.execute("SELECT * FROM ExtrasDictCursorTests")
row = getter(curs)
self.failUnless(row['foo'] == 'bar')
@skip_from_python(3)
def test_iter_methods_2(self):
curs = self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
with self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as curs:
curs.execute("select 10 as a, 20 as b")
r = curs.fetchone()
self.assert_(isinstance(r.keys(), list))
@ -352,7 +372,7 @@ class ExtrasDictCursorRealTests(_DictCursorBase):
@skip_before_python(3)
def test_iter_methods_3(self):
curs = self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
with self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as curs:
curs.execute("select 10 as a, 20 as b")
r = curs.fetchone()
self.assert_(not isinstance(r.keys(), list))
@ -363,7 +383,7 @@ class ExtrasDictCursorRealTests(_DictCursorBase):
self.assertEqual(len(list(r.items())), 2)
def test_order(self):
curs = self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
with self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as curs:
curs.execute("select 5 as foo, 4 as bar, 33 as baz, 2 as qux")
r = curs.fetchone()
self.assertEqual(list(r), ['foo', 'bar', 'baz', 'qux'])
@ -380,7 +400,7 @@ class ExtrasDictCursorRealTests(_DictCursorBase):
@skip_from_python(3)
def test_order_iter(self):
curs = self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
with self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as curs:
curs.execute("select 5 as foo, 4 as bar, 33 as baz, 2 as qux")
r = curs.fetchone()
self.assertEqual(list(r.iterkeys()), ['foo', 'bar', 'baz', 'qux'])
@ -394,7 +414,7 @@ class ExtrasDictCursorRealTests(_DictCursorBase):
self.assertEqual(list(r1.iteritems()), list(r.iteritems()))
def test_pop(self):
curs = self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
with self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as curs:
curs.execute("select 1 as a, 2 as b, 3 as c")
r = curs.fetchone()
self.assertEqual(r.pop('b'), 2)
@ -407,7 +427,7 @@ class ExtrasDictCursorRealTests(_DictCursorBase):
self.assertRaises(KeyError, r.pop, 'b')
def test_mod(self):
curs = self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
with self.conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as curs:
curs.execute("select 1 as a, 2 as b, 3 as c")
r = curs.fetchone()
r['d'] = 4
@ -428,20 +448,24 @@ class NamedTupleCursorTest(ConnectingTestCase):
ConnectingTestCase.setUp(self)
self.conn = self.connect(connection_factory=NamedTupleConnection)
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("CREATE TEMPORARY TABLE nttest (i int, s text)")
curs.execute("INSERT INTO nttest VALUES (1, 'foo')")
curs.execute("INSERT INTO nttest VALUES (2, 'bar')")
curs.execute("INSERT INTO nttest VALUES (3, 'baz')")
self.conn.commit()
@skip_before_postgres(8, 2)
def test_cursor_args(self):
cur = self.conn.cursor('foo', cursor_factory=psycopg2.extras.DictCursor)
with self.conn.cursor(
'foo',
cursor_factory=psycopg2.extras.DictCursor
) as cur:
self.assertEqual(cur.name, 'foo')
self.assert_(isinstance(cur, psycopg2.extras.DictCursor))
def test_fetchone(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("select * from nttest order by 1")
t = curs.fetchone()
self.assertEqual(t[0], 1)
@ -452,7 +476,7 @@ class NamedTupleCursorTest(ConnectingTestCase):
self.assertEqual(curs.rowcount, 3)
def test_fetchmany_noarg(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.arraysize = 2
curs.execute("select * from nttest order by 1")
res = curs.fetchmany()
@ -465,7 +489,7 @@ class NamedTupleCursorTest(ConnectingTestCase):
self.assertEqual(curs.rowcount, 3)
def test_fetchmany(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("select * from nttest order by 1")
res = curs.fetchmany(2)
self.assertEqual(2, len(res))
@ -477,7 +501,7 @@ class NamedTupleCursorTest(ConnectingTestCase):
self.assertEqual(curs.rowcount, 3)
def test_fetchall(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("select * from nttest order by 1")
res = curs.fetchall()
self.assertEqual(3, len(res))
@ -491,7 +515,7 @@ class NamedTupleCursorTest(ConnectingTestCase):
self.assertEqual(curs.rowcount, 3)
def test_executemany(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.executemany("delete from nttest where i = %s",
[(1,), (2,)])
curs.execute("select * from nttest order by 1")
@ -501,7 +525,7 @@ class NamedTupleCursorTest(ConnectingTestCase):
self.assertEqual(res[0].s, 'baz')
def test_iter(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("select * from nttest order by 1")
i = iter(curs)
self.assertEqual(curs.rownumber, 0)
@ -526,7 +550,7 @@ class NamedTupleCursorTest(ConnectingTestCase):
self.assertEqual(curs.rowcount, 3)
def test_record_updated(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("select 1 as foo;")
r = curs.fetchone()
self.assertEqual(r.foo, 1)
@ -537,7 +561,7 @@ class NamedTupleCursorTest(ConnectingTestCase):
self.assertRaises(AttributeError, getattr, r, 'foo')
def test_no_result_no_surprise(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("update nttest set s = s")
self.assertRaises(psycopg2.ProgrammingError, curs.fetchone)
@ -545,7 +569,7 @@ class NamedTupleCursorTest(ConnectingTestCase):
self.assertRaises(psycopg2.ProgrammingError, curs.fetchall)
def test_bad_col_names(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute('select 1 as "foo.bar_baz", 2 as "?column?", 3 as "3"')
rv = curs.fetchone()
self.assertEqual(rv.foo_bar_baz, 1)
@ -555,7 +579,7 @@ class NamedTupleCursorTest(ConnectingTestCase):
@skip_before_python(3)
@skip_before_postgres(8)
def test_nonascii_name(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute('select 1 as \xe5h\xe9')
rv = curs.fetchone()
self.assertEqual(getattr(rv, '\xe5h\xe9'), 1)
@ -572,7 +596,7 @@ class NamedTupleCursorTest(ConnectingTestCase):
NamedTupleCursor._make_nt = f_patched
try:
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("select * from nttest order by 1")
curs.fetchone()
curs.fetchone()
@ -594,7 +618,7 @@ class NamedTupleCursorTest(ConnectingTestCase):
@skip_before_postgres(8, 0)
def test_named(self):
curs = self.conn.cursor('tmp')
with self.conn.cursor('tmp') as curs:
curs.execute("""select i from generate_series(0,9) i""")
recs = []
recs.extend(curs.fetchmany(5))
@ -603,28 +627,29 @@ class NamedTupleCursorTest(ConnectingTestCase):
self.assertEqual(list(range(10)), [t.i for t in recs])
def test_named_fetchone(self):
curs = self.conn.cursor('tmp')
with self.conn.cursor('tmp') as curs:
curs.execute("""select 42 as i""")
t = curs.fetchone()
self.assertEqual(t.i, 42)
def test_named_fetchmany(self):
curs = self.conn.cursor('tmp')
with self.conn.cursor('tmp') as curs:
curs.execute("""select 42 as i""")
recs = curs.fetchmany(10)
self.assertEqual(recs[0].i, 42)
def test_named_fetchall(self):
curs = self.conn.cursor('tmp')
with self.conn.cursor('tmp') as curs:
curs.execute("""select 42 as i""")
recs = curs.fetchall()
self.assertEqual(recs[0].i, 42)
@skip_before_postgres(8, 2)
def test_not_greedy(self):
curs = self.conn.cursor('tmp')
with self.conn.cursor('tmp') as curs:
curs.itersize = 2
curs.execute("""select clock_timestamp() as ts from generate_series(1,3)""")
curs.execute(
"""select clock_timestamp() as ts from generate_series(1,3)""")
recs = []
for t in curs:
time.sleep(0.01)
@ -636,7 +661,7 @@ class NamedTupleCursorTest(ConnectingTestCase):
@skip_before_postgres(8, 0)
def test_named_rownumber(self):
curs = self.conn.cursor('tmp')
with self.conn.cursor('tmp') as curs:
# Only checking for dataset < itersize:
# see CursorTests.test_iter_named_cursor_rownumber
curs.itersize = 4
@ -647,14 +672,14 @@ class NamedTupleCursorTest(ConnectingTestCase):
def test_cache(self):
NamedTupleCursor._cached_make_nt.cache_clear()
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("select 10 as a, 20 as b")
r1 = curs.fetchone()
curs.execute("select 10 as a, 20 as c")
r2 = curs.fetchone()
# Get a new cursor to check that the cache works across multiple ones
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("select 10 as a, 30 as b")
r3 = curs.fetchone()
@ -672,7 +697,7 @@ class NamedTupleCursorTest(ConnectingTestCase):
lru_cache(8)(NamedTupleCursor._cached_make_nt.__wrapped__)
try:
recs = []
curs = self.conn.cursor()
with self.conn.cursor() as curs:
for i in range(10):
curs.execute("select 1 as f%s" % i)
recs.append(curs.fetchone())

38
tests/test_fast_executemany.py
View File

@ -46,14 +46,14 @@ class TestPaginate(unittest.TestCase):
class FastExecuteTestMixin(object):
def setUp(self):
super(FastExecuteTestMixin, self).setUp()
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("""create table testfast (
id serial primary key, date date, val int, data text)""")
class TestExecuteBatch(FastExecuteTestMixin, testutils.ConnectingTestCase):
def test_empty(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
psycopg2.extras.execute_batch(cur,
"insert into testfast (id, val) values (%s, %s)",
[])
@ -61,7 +61,7 @@ class TestExecuteBatch(FastExecuteTestMixin, testutils.ConnectingTestCase):
self.assertEqual(cur.fetchall(), [])
def test_one(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
psycopg2.extras.execute_batch(cur,
"insert into testfast (id, val) values (%s, %s)",
iter([(1, 10)]))
@ -69,7 +69,7 @@ class TestExecuteBatch(FastExecuteTestMixin, testutils.ConnectingTestCase):
self.assertEqual(cur.fetchall(), [(1, 10)])
def test_tuples(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
psycopg2.extras.execute_batch(cur,
"insert into testfast (id, date, val) values (%s, %s, %s)",
((i, date(2017, 1, i + 1), i * 10) for i in range(10)))
@ -78,7 +78,7 @@ class TestExecuteBatch(FastExecuteTestMixin, testutils.ConnectingTestCase):
[(i, date(2017, 1, i + 1), i * 10) for i in range(10)])
def test_many(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
psycopg2.extras.execute_batch(cur,
"insert into testfast (id, val) values (%s, %s)",
((i, i * 10) for i in range(1000)))
@ -86,7 +86,7 @@ class TestExecuteBatch(FastExecuteTestMixin, testutils.ConnectingTestCase):
self.assertEqual(cur.fetchall(), [(i, i * 10) for i in range(1000)])
def test_composed(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
psycopg2.extras.execute_batch(cur,
sql.SQL("insert into {0} (id, val) values (%s, %s)")
.format(sql.Identifier('testfast')),
@ -95,7 +95,7 @@ class TestExecuteBatch(FastExecuteTestMixin, testutils.ConnectingTestCase):
self.assertEqual(cur.fetchall(), [(i, i * 10) for i in range(1000)])
def test_pages(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
psycopg2.extras.execute_batch(cur,
"insert into testfast (id, val) values (%s, %s)",
((i, i * 10) for i in range(25)),
@ -109,7 +109,7 @@ class TestExecuteBatch(FastExecuteTestMixin, testutils.ConnectingTestCase):
@testutils.skip_before_postgres(8, 0)
def test_unicode(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
ext.register_type(ext.UNICODE, cur)
snowman = u"\u2603"
@ -138,7 +138,7 @@ class TestExecuteBatch(FastExecuteTestMixin, testutils.ConnectingTestCase):
@testutils.skip_before_postgres(8, 2)
class TestExecuteValues(FastExecuteTestMixin, testutils.ConnectingTestCase):
def test_empty(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
psycopg2.extras.execute_values(cur,
"insert into testfast (id, val) values %s",
[])
@ -146,7 +146,7 @@ class TestExecuteValues(FastExecuteTestMixin, testutils.ConnectingTestCase):
self.assertEqual(cur.fetchall(), [])
def test_one(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
psycopg2.extras.execute_values(cur,
"insert into testfast (id, val) values %s",
iter([(1, 10)]))
@ -154,7 +154,7 @@ class TestExecuteValues(FastExecuteTestMixin, testutils.ConnectingTestCase):
self.assertEqual(cur.fetchall(), [(1, 10)])
def test_tuples(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
psycopg2.extras.execute_values(cur,
"insert into testfast (id, date, val) values %s",
((i, date(2017, 1, i + 1), i * 10) for i in range(10)))
@ -163,7 +163,7 @@ class TestExecuteValues(FastExecuteTestMixin, testutils.ConnectingTestCase):
[(i, date(2017, 1, i + 1), i * 10) for i in range(10)])
def test_dicts(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
psycopg2.extras.execute_values(cur,
"insert into testfast (id, date, val) values %s",
(dict(id=i, date=date(2017, 1, i + 1), val=i * 10, foo="bar")
@ -174,7 +174,7 @@ class TestExecuteValues(FastExecuteTestMixin, testutils.ConnectingTestCase):
[(i, date(2017, 1, i + 1), i * 10) for i in range(10)])
def test_many(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
psycopg2.extras.execute_values(cur,
"insert into testfast (id, val) values %s",
((i, i * 10) for i in range(1000)))
@ -182,7 +182,7 @@ class TestExecuteValues(FastExecuteTestMixin, testutils.ConnectingTestCase):
self.assertEqual(cur.fetchall(), [(i, i * 10) for i in range(1000)])
def test_composed(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
psycopg2.extras.execute_values(cur,
sql.SQL("insert into {0} (id, val) values %s")
.format(sql.Identifier('testfast')),
@ -191,7 +191,7 @@ class TestExecuteValues(FastExecuteTestMixin, testutils.ConnectingTestCase):
self.assertEqual(cur.fetchall(), [(i, i * 10) for i in range(1000)])
def test_pages(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
psycopg2.extras.execute_values(cur,
"insert into testfast (id, val) values %s",
((i, i * 10) for i in range(25)),
@ -204,7 +204,7 @@ class TestExecuteValues(FastExecuteTestMixin, testutils.ConnectingTestCase):
self.assertEqual(cur.fetchall(), [(i, i * 10) for i in range(25)])
def test_unicode(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
ext.register_type(ext.UNICODE, cur)
snowman = u"\u2603"
@ -230,7 +230,7 @@ class TestExecuteValues(FastExecuteTestMixin, testutils.ConnectingTestCase):
self.assertEqual(cur.fetchone(), (3, snowman))
def test_returning(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
result = psycopg2.extras.execute_values(cur,
"insert into testfast (id, val) values %s returning id",
((i, i * 10) for i in range(25)),
@ -239,7 +239,7 @@ class TestExecuteValues(FastExecuteTestMixin, testutils.ConnectingTestCase):
self.assertEqual([r[0] for r in result], list(range(25)))
def test_invalid_sql(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
self.assertRaises(ValueError, psycopg2.extras.execute_values, cur,
"insert", [])
self.assertRaises(ValueError, psycopg2.extras.execute_values, cur,
@ -250,7 +250,7 @@ class TestExecuteValues(FastExecuteTestMixin, testutils.ConnectingTestCase):
"insert %f %s", [])
def test_percent_escape(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
psycopg2.extras.execute_values(cur,
"insert into testfast (id, data) values %s -- a%%b",
[(1, 'hi')])

15
tests/test_green.py
View File

@ -71,7 +71,7 @@ class GreenTestCase(ConnectingTestCase):
# a very large query requires a flush loop to be sent to the backend
conn = self.conn
stub = self.set_stub_wait_callback(conn)
curs = conn.cursor()
with conn.cursor() as curs:
for mb in 1, 5, 10, 20, 50:
size = mb * 1024 * 1024
del stub.polls[:]
@ -105,7 +105,7 @@ class GreenTestCase(ConnectingTestCase):
# if there is an error in a green query, don't freak out and close
# the connection
conn = self.conn
curs = conn.cursor()
with conn.cursor() as curs:
self.assertRaises(psycopg2.ProgrammingError,
curs.execute, "select the unselectable")
@ -117,7 +117,7 @@ class GreenTestCase(ConnectingTestCase):
@skip_before_postgres(8, 2)
def test_copy_no_hang(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
self.assertRaises(psycopg2.ProgrammingError,
cur.execute, "copy (select 1) to stdout")
@ -137,7 +137,7 @@ class GreenTestCase(ConnectingTestCase):
raise conn.OperationalError("bad state from poll: %s" % state)
stub = self.set_stub_wait_callback(self.conn, wait)
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("""
select 1;
do $$
@ -203,7 +203,7 @@ class CallbackErrorTestCase(ConnectingTestCase):
for i in range(100):
self.to_error = None
cnn = self.connect()
cur = cnn.cursor()
with cnn.cursor() as cur:
self.to_error = i
try:
cur.execute("select 1")
@ -220,7 +220,7 @@ class CallbackErrorTestCase(ConnectingTestCase):
for i in range(100):
self.to_error = None
cnn = self.connect()
cur = cnn.cursor('foo')
with cnn.cursor('foo') as cur:
self.to_error = i
try:
cur.execute("select 1")
@ -230,6 +230,9 @@ class CallbackErrorTestCase(ConnectingTestCase):
else:
# The query completed
return
finally:
# Don't raise an exception in the cursor context manager.
self.to_error = None
self.fail("you should have had a success or an error by now")

12
tests/test_ipaddress.py
View File

@ -33,7 +33,7 @@ except ImportError:
@unittest.skipIf(ip is None, "'ipaddress' module not available")
class NetworkingTestCase(testutils.ConnectingTestCase):
def test_inet_cast(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
psycopg2.extras.register_ipaddress(cur)
cur.execute("select null::inet")
@ -51,7 +51,7 @@ class NetworkingTestCase(testutils.ConnectingTestCase):
@testutils.skip_before_postgres(8, 2)
def test_inet_array_cast(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
psycopg2.extras.register_ipaddress(cur)
cur.execute("select '{NULL,127.0.0.1,::ffff:102:300/128}'::inet[]")
l = cur.fetchone()[0]
@ -62,7 +62,7 @@ class NetworkingTestCase(testutils.ConnectingTestCase):
self.assert_(isinstance(l[2], ip.IPv6Interface), l)
def test_inet_adapt(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
psycopg2.extras.register_ipaddress(cur)
cur.execute("select %s", [ip.ip_interface('127.0.0.1/24')])
@ -72,7 +72,7 @@ class NetworkingTestCase(testutils.ConnectingTestCase):
self.assertEquals(cur.fetchone()[0], '::ffff:102:300/128')
def test_cidr_cast(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
psycopg2.extras.register_ipaddress(cur)
cur.execute("select null::cidr")
@ -90,7 +90,7 @@ class NetworkingTestCase(testutils.ConnectingTestCase):
@testutils.skip_before_postgres(8, 2)
def test_cidr_array_cast(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
psycopg2.extras.register_ipaddress(cur)
cur.execute("select '{NULL,127.0.0.1,::ffff:102:300/128}'::cidr[]")
l = cur.fetchone()[0]
@ -101,7 +101,7 @@ class NetworkingTestCase(testutils.ConnectingTestCase):
self.assert_(isinstance(l[2], ip.IPv6Network), l)
def test_cidr_adapt(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
psycopg2.extras.register_ipaddress(cur)
cur.execute("select %s", [ip.ip_network('127.0.0.0/24')])

6
tests/test_module.py
View File

@ -154,7 +154,7 @@ class ConnectTestCase(unittest.TestCase):
class ExceptionsTestCase(ConnectingTestCase):
def test_attributes(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
try:
cur.execute("select * from nonexist")
except psycopg2.Error as exc:
@ -165,7 +165,7 @@ class ExceptionsTestCase(ConnectingTestCase):
self.assert_(e.cursor is cur)
def test_diagnostics_attributes(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
try:
cur.execute("select * from nonexist")
except psycopg2.Error as exc:
@ -195,7 +195,7 @@ class ExceptionsTestCase(ConnectingTestCase):
def test_diagnostics_life(self):
def tmp():
cur = self.conn.cursor()
with self.conn.cursor() as cur:
try:
cur.execute("select * from nonexist")
except psycopg2.Error as exc:

3
tests/test_notify.py
View File

@ -120,7 +120,8 @@ conn.close()
self.listen('foo')
pid = int(self.notify('foo').communicate()[0])
self.assertEqual(0, len(self.conn.notifies))
self.conn.cursor().execute('select 1;')
with self.conn.cursor() as cur:
cur.execute('select 1;')
self.assertEqual(1, len(self.conn.notifies))
self.assertEqual(pid, self.conn.notifies[0][0])
self.assertEqual('foo', self.conn.notifies[0][1])

17
tests/test_quote.py
View File

@ -56,7 +56,7 @@ class QuotingTestCase(ConnectingTestCase):
"""
data += "".join(map(chr, range(1, 127)))
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("SELECT %s;", (data,))
res = curs.fetchone()[0]
@ -64,7 +64,7 @@ class QuotingTestCase(ConnectingTestCase):
self.assert_(not self.conn.notices)
def test_string_null_terminator(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
data = 'abcd\x01\x00cdefg'
try:
@ -84,7 +84,7 @@ class QuotingTestCase(ConnectingTestCase):
else:
data += bytes(list(range(256)))
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("SELECT %s::bytea;", (psycopg2.Binary(data),))
if PY2:
res = str(curs.fetchone()[0])
@ -99,7 +99,7 @@ class QuotingTestCase(ConnectingTestCase):
self.assert_(not self.conn.notices)
def test_unicode(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("SHOW server_encoding")
server_encoding = curs.fetchone()[0]
if server_encoding != "UTF8":
@ -123,7 +123,7 @@ class QuotingTestCase(ConnectingTestCase):
def test_latin1(self):
self.conn.set_client_encoding('LATIN1')
curs = self.conn.cursor()
with self.conn.cursor() as curs:
if PY2:
data = ''.join(map(chr, range(32, 127) + range(160, 256)))
else:
@ -138,7 +138,8 @@ class QuotingTestCase(ConnectingTestCase):
# as unicode
if PY2:
psycopg2.extensions.register_type(psycopg2.extensions.UNICODE, self.conn)
psycopg2.extensions.register_type(
psycopg2.extensions.UNICODE, self.conn)
data = data.decode('latin1')
curs.execute("SELECT %s::text;", (data,))
@ -148,7 +149,7 @@ class QuotingTestCase(ConnectingTestCase):
def test_koi8(self):
self.conn.set_client_encoding('KOI8')
curs = self.conn.cursor()
with self.conn.cursor() as curs:
if PY2:
data = ''.join(map(chr, range(32, 127) + range(128, 256)))
else:
@ -176,7 +177,7 @@ class QuotingTestCase(ConnectingTestCase):
conn = self.connect()
conn.set_client_encoding('UNICODE')
psycopg2.extensions.register_type(psycopg2.extensions.BYTES, conn)
curs = conn.cursor()
with conn.cursor() as curs:
curs.execute("select %s::text", (snowman,))
x = curs.fetchone()[0]
self.assert_(isinstance(x, bytes))

27
tests/test_replication.py
View File

@ -73,8 +73,7 @@ class ReplicationTestCase(ConnectingTestCase):
conn = self.connect()
if conn is None:
return
cur = conn.cursor()
with conn.cursor() as cur:
try:
cur.execute("DROP TABLE dummy1")
except psycopg2.ProgrammingError:
@ -90,7 +89,7 @@ class ReplicationTest(ReplicationTestCase):
conn = self.repl_connect(connection_factory=PhysicalReplicationConnection)
if conn is None:
return
cur = conn.cursor()
with conn.cursor() as cur:
cur.execute("IDENTIFY_SYSTEM")
cur.fetchall()
@ -104,7 +103,7 @@ class ReplicationTest(ReplicationTestCase):
connection_factory=PhysicalReplicationConnection)
if conn is None:
return
cur = conn.cursor()
with conn.cursor() as cur:
cur.execute("IDENTIFY_SYSTEM")
cur.fetchall()
@ -113,7 +112,7 @@ class ReplicationTest(ReplicationTestCase):
conn = self.repl_connect(connection_factory=LogicalReplicationConnection)
if conn is None:
return
cur = conn.cursor()
with conn.cursor() as cur:
cur.execute("IDENTIFY_SYSTEM")
cur.fetchall()
@ -122,8 +121,7 @@ class ReplicationTest(ReplicationTestCase):
conn = self.repl_connect(connection_factory=PhysicalReplicationConnection)
if conn is None:
return
cur = conn.cursor()
with conn.cursor() as cur:
self.create_replication_slot(cur)
self.assertRaises(
psycopg2.ProgrammingError, self.create_replication_slot, cur)
@ -134,8 +132,7 @@ class ReplicationTest(ReplicationTestCase):
conn = self.repl_connect(connection_factory=PhysicalReplicationConnection)
if conn is None:
return
cur = conn.cursor()
with conn.cursor() as cur:
self.assertRaises(psycopg2.ProgrammingError,
cur.start_replication, self.slot)
@ -148,8 +145,7 @@ class ReplicationTest(ReplicationTestCase):
conn = self.repl_connect(connection_factory=LogicalReplicationConnection)
if conn is None:
return
cur = conn.cursor()
with conn.cursor() as cur:
self.create_replication_slot(cur, output_plugin='test_decoding')
cur.start_replication_expert(
sql.SQL("START_REPLICATION SLOT {slot} LOGICAL 0/00000000").format(
@ -161,8 +157,7 @@ class ReplicationTest(ReplicationTestCase):
conn = self.repl_connect(connection_factory=LogicalReplicationConnection)
if conn is None:
return
cur = conn.cursor()
with conn.cursor() as cur:
self.create_replication_slot(cur, output_plugin='test_decoding')
self.make_replication_events()
@ -208,8 +203,7 @@ class ReplicationTest(ReplicationTestCase):
conn = self.repl_connect(connection_factory=LogicalReplicationConnection)
if conn is None:
return
cur = conn.cursor()
with conn.cursor() as cur:
self.create_replication_slot(cur, output_plugin='test_decoding')
self.make_replication_events()
@ -230,8 +224,7 @@ class AsyncReplicationTest(ReplicationTestCase):
if conn is None:
return
cur = conn.cursor()
with conn.cursor() as cur:
self.create_replication_slot(cur, output_plugin='test_decoding')
self.wait(cur)

6
tests/test_sql.py
View File

@ -117,7 +117,7 @@ class SqlFormatTests(ConnectingTestCase):
sql.SQL("select {0};").format(sql.Literal(Foo())).as_string, self.conn)
def test_execute(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("""
create table test_compose (
id serial primary key,
@ -134,7 +134,7 @@ class SqlFormatTests(ConnectingTestCase):
self.assertEqual(cur.fetchall(), [(10, 'a', 'b', 'c')])
def test_executemany(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("""
create table test_compose (
id serial primary key,
@ -154,7 +154,7 @@ class SqlFormatTests(ConnectingTestCase):
@skip_copy_if_green
@skip_before_postgres(8, 2)
def test_copy(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("""
create table test_compose (
id serial primary key,

22
tests/test_transaction.py
View File

@ -37,7 +37,7 @@ class TransactionTests(ConnectingTestCase):
def setUp(self):
ConnectingTestCase.setUp(self)
self.conn.set_isolation_level(ISOLATION_LEVEL_SERIALIZABLE)
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute('''
CREATE TEMPORARY TABLE table1 (
id int PRIMARY KEY
@ -55,7 +55,7 @@ class TransactionTests(ConnectingTestCase):
def test_rollback(self):
# Test that rollback undoes changes
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute('INSERT INTO table2 VALUES (2, 1)')
# Rollback takes us from BEGIN state to READY state
self.assertEqual(self.conn.status, STATUS_BEGIN)
@ -66,7 +66,7 @@ class TransactionTests(ConnectingTestCase):
def test_commit(self):
# Test that commit stores changes
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute('INSERT INTO table2 VALUES (2, 1)')
# Rollback takes us from BEGIN state to READY state
self.assertEqual(self.conn.status, STATUS_BEGIN)
@ -80,7 +80,7 @@ class TransactionTests(ConnectingTestCase):
def test_failed_commit(self):
# Test that we can recover from a failed commit.
# We use a deferred constraint to cause a failure on commit.
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute('SET CONSTRAINTS table2__table1_id__fk DEFERRED')
curs.execute('INSERT INTO table2 VALUES (2, 42)')
# The commit should fail, and move the cursor back to READY state
@ -103,7 +103,7 @@ class DeadlockSerializationTests(ConnectingTestCase):
def setUp(self):
ConnectingTestCase.setUp(self)
curs = self.conn.cursor()
with self.conn.cursor() as curs:
# Drop table if it already exists
try:
curs.execute("DROP TABLE table1")
@ -126,7 +126,7 @@ class DeadlockSerializationTests(ConnectingTestCase):
self.conn.commit()
def tearDown(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("DROP TABLE table1")
curs.execute("DROP TABLE table2")
self.conn.commit()
@ -142,7 +142,7 @@ class DeadlockSerializationTests(ConnectingTestCase):
def task1():
try:
conn = self.connect()
curs = conn.cursor()
with conn.cursor() as curs:
curs.execute("LOCK table1 IN ACCESS EXCLUSIVE MODE")
step1.set()
step2.wait()
@ -155,7 +155,7 @@ class DeadlockSerializationTests(ConnectingTestCase):
def task2():
try:
conn = self.connect()
curs = conn.cursor()
with conn.cursor() as curs:
step1.wait()
curs.execute("LOCK table2 IN ACCESS EXCLUSIVE MODE")
step2.set()
@ -190,7 +190,7 @@ class DeadlockSerializationTests(ConnectingTestCase):
def task1():
try:
conn = self.connect()
curs = conn.cursor()
with conn.cursor() as curs:
curs.execute("SELECT name FROM table1 WHERE id = 1")
curs.fetchall()
step1.set()
@ -205,7 +205,7 @@ class DeadlockSerializationTests(ConnectingTestCase):
def task2():
try:
conn = self.connect()
curs = conn.cursor()
with conn.cursor() as curs:
step1.wait()
curs.execute("UPDATE table1 SET name='task2' WHERE id = 1")
conn.commit()
@ -240,7 +240,7 @@ class QueryCancellationTests(ConnectingTestCase):
@skip_before_postgres(8, 2)
def test_statement_timeout(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
# Set a low statement timeout, then sleep for a longer period.
curs.execute('SET statement_timeout TO 10')
self.assertRaises(psycopg2.extensions.QueryCanceledError,

21
tests/test_types_basic.py
View File

@ -41,7 +41,7 @@ class TypesBasicTests(ConnectingTestCase):
"""Test that all type conversions are working."""
def execute(self, *args):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute(*args)
return curs.fetchone()[0]
@ -156,7 +156,7 @@ class TypesBasicTests(ConnectingTestCase):
def testEmptyArrayRegression(self):
# ticket #42
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute(
"create table array_test "
"(id integer, col timestamp without time zone[])")
@ -164,7 +164,8 @@ class TypesBasicTests(ConnectingTestCase):
curs.execute("insert into array_test values (%s, %s)",
(1, [datetime.date(2011, 2, 14)]))
curs.execute("select col from array_test where id = 1")
self.assertEqual(curs.fetchone()[0], [datetime.datetime(2011, 2, 14, 0, 0)])
self.assertEqual(
curs.fetchone()[0], [datetime.datetime(2011, 2, 14, 0, 0)])
curs.execute("insert into array_test values (%s, %s)", (2, []))
curs.execute("select col from array_test where id = 2")
@ -173,7 +174,7 @@ class TypesBasicTests(ConnectingTestCase):
@testutils.skip_before_postgres(8, 4)
def testNestedEmptyArray(self):
# issue #788
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("select 10 = any(%s::int[])", ([[]], ))
self.assertFalse(curs.fetchone()[0])
@ -204,14 +205,14 @@ class TypesBasicTests(ConnectingTestCase):
self.failUnlessEqual(ss, r)
def testArrayMalformed(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
ss = ['', '{', '{}}', '{' * 20 + '}' * 20]
for s in ss:
self.assertRaises(psycopg2.DataError,
psycopg2.extensions.STRINGARRAY, s.encode('utf8'), curs)
def testTextArray(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("select '{a,b,c}'::text[]")
x = curs.fetchone()[0]
self.assert_(isinstance(x[0], str))
@ -220,7 +221,7 @@ class TypesBasicTests(ConnectingTestCase):
def testUnicodeArray(self):
psycopg2.extensions.register_type(
psycopg2.extensions.UNICODEARRAY, self.conn)
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("select '{a,b,c}'::text[]")
x = curs.fetchone()[0]
self.assert_(isinstance(x[0], text_type))
@ -229,7 +230,7 @@ class TypesBasicTests(ConnectingTestCase):
def testBytesArray(self):
psycopg2.extensions.register_type(
psycopg2.extensions.BYTESARRAY, self.conn)
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("select '{a,b,c}'::text[]")
x = curs.fetchone()[0]
self.assert_(isinstance(x[0], bytes))
@ -237,7 +238,7 @@ class TypesBasicTests(ConnectingTestCase):
@testutils.skip_before_postgres(8, 2)
def testArrayOfNulls(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("""
create table na (
texta text[],
@ -273,7 +274,7 @@ class TypesBasicTests(ConnectingTestCase):
@testutils.skip_before_postgres(8, 2)
def testNestedArrays(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
for a in [
[[1]],
[[None]],

178
tests/test_types_extras.py
View File

@ -36,7 +36,7 @@ from psycopg2._json import _get_json_oids
from psycopg2.extras import (
CompositeCaster, DateRange, DateTimeRange, DateTimeTZRange, HstoreAdapter,
Inet, Json, NumericRange, Range, RealDictConnection,
register_composite, register_hstore, register_range,
register_composite, register_hstore, register_range, _solve_conn_curs
)
from psycopg2.tz import FixedOffsetTimezone
@ -45,7 +45,7 @@ class TypesExtrasTests(ConnectingTestCase):
"""Test that all type conversions are working."""
def execute(self, *args):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute(*args)
return curs.fetchone()[0]
@ -231,7 +231,7 @@ class HstoreTestCase(ConnectingTestCase):
@skip_if_no_hstore
def test_register_conn(self):
register_hstore(self.conn)
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("select null::hstore, ''::hstore, 'a => b'::hstore")
t = cur.fetchone()
self.assert_(t[0] is None)
@ -240,7 +240,7 @@ class HstoreTestCase(ConnectingTestCase):
@skip_if_no_hstore
def test_register_curs(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
register_hstore(cur)
cur.execute("select null::hstore, ''::hstore, 'a => b'::hstore")
t = cur.fetchone()
@ -252,7 +252,7 @@ class HstoreTestCase(ConnectingTestCase):
@skip_from_python(3)
def test_register_unicode(self):
register_hstore(self.conn, unicode=True)
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("select null::hstore, ''::hstore, 'a => b'::hstore")
t = cur.fetchone()
self.assert_(t[0] is None)
@ -268,7 +268,7 @@ class HstoreTestCase(ConnectingTestCase):
register_hstore(self.conn, globally=True)
conn2 = self.connect()
try:
cur2 = self.conn.cursor()
with self.conn.cursor() as cur2:
cur2.execute("select 'a => b'::hstore")
r = cur2.fetchone()
self.assert_(isinstance(r[0], dict))
@ -278,8 +278,7 @@ class HstoreTestCase(ConnectingTestCase):
@skip_if_no_hstore
def test_roundtrip(self):
register_hstore(self.conn)
cur = self.conn.cursor()
with self.conn.cursor() as cur:
def ok(d):
cur.execute("select %s", (d,))
d1 = cur.fetchone()[0]
@ -299,7 +298,8 @@ class HstoreTestCase(ConnectingTestCase):
if PY2:
ab = map(chr, range(32, 127) + range(160, 255))
else:
ab = bytes(list(range(32, 127)) + list(range(160, 255))).decode('latin1')
ab = bytes(
list(range(32, 127)) + list(range(160, 255))).decode('latin1')
ok({''.join(ab): ''.join(ab)})
ok(dict(zip(ab, ab)))
@ -308,8 +308,7 @@ class HstoreTestCase(ConnectingTestCase):
@skip_from_python(3)
def test_roundtrip_unicode(self):
register_hstore(self.conn, unicode=True)
cur = self.conn.cursor()
with self.conn.cursor() as cur:
def ok(d):
cur.execute("select %s", (d,))
d1 = cur.fetchone()[0]
@ -330,7 +329,7 @@ class HstoreTestCase(ConnectingTestCase):
@skip_if_no_hstore
@restore_types
def test_oid(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("select 'hstore'::regtype::oid")
oid = cur.fetchone()[0]
@ -363,7 +362,7 @@ class HstoreTestCase(ConnectingTestCase):
ds.append({''.join(ab): ''.join(ab)})
ds.append(dict(zip(ab, ab)))
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("select %s", (ds,))
ds1 = cur.fetchone()[0]
self.assertEqual(ds, ds1)
@ -372,7 +371,7 @@ class HstoreTestCase(ConnectingTestCase):
@skip_before_postgres(8, 3)
def test_array_cast(self):
register_hstore(self.conn)
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("select array['a=>1'::hstore, 'b=>2'::hstore];")
a = cur.fetchone()[0]
self.assertEqual(a, [{'a': '1'}, {'b': '2'}])
@ -380,7 +379,7 @@ class HstoreTestCase(ConnectingTestCase):
@skip_if_no_hstore
@restore_types
def test_array_cast_oid(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("select 'hstore'::regtype::oid, 'hstore[]'::regtype::oid")
oid, aoid = cur.fetchone()
@ -399,7 +398,7 @@ class HstoreTestCase(ConnectingTestCase):
conn = self.connect(connection_factory=RealDictConnection)
try:
register_hstore(conn)
curs = conn.cursor()
with conn.cursor() as curs:
curs.execute("select ''::hstore as x")
self.assertEqual(curs.fetchone()['x'], {})
finally:
@ -407,7 +406,7 @@ class HstoreTestCase(ConnectingTestCase):
conn = self.connect(connection_factory=RealDictConnection)
try:
curs = conn.cursor()
with conn.cursor() as curs:
register_hstore(curs)
curs.execute("select ''::hstore as x")
self.assertEqual(curs.fetchone()['x'], {})
@ -431,7 +430,7 @@ def skip_if_no_composite(f):
class AdaptTypeTestCase(ConnectingTestCase):
@skip_if_no_composite
def test_none_in_record(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
s = curs.mogrify("SELECT %s;", [(42, None)])
self.assertEqual(b"SELECT (42, NULL);", s)
curs.execute("SELECT %s;", [(42, None)])
@ -448,8 +447,7 @@ class AdaptTypeTestCase(ConnectingTestCase):
def getquoted(self):
return "NOPE!"
curs = self.conn.cursor()
with self.conn.cursor() as curs:
orig_adapter = ext.adapters[type(None), ext.ISQLQuote]
try:
ext.register_adapter(type(None), WonkyAdapter)
@ -502,7 +500,7 @@ class AdaptTypeTestCase(ConnectingTestCase):
self.assertEqual(t.attnames, ['anint', 'astring', 'adate'])
self.assertEqual(t.atttypes, [23, 25, 1082])
curs = self.conn.cursor()
with self.conn.cursor() as curs:
r = (10, 'hello', date(2011, 1, 2))
curs.execute("select %s::type_isd;", (r,))
v = curs.fetchone()[0]
@ -519,7 +517,7 @@ class AdaptTypeTestCase(ConnectingTestCase):
def test_empty_string(self):
# issue #141
self._create_type("type_ss", [('s1', 'text'), ('s2', 'text')])
curs = self.conn.cursor()
with self.conn.cursor() as curs:
psycopg2.extras.register_composite("type_ss", curs)
def ok(t):
@ -548,7 +546,7 @@ class AdaptTypeTestCase(ConnectingTestCase):
psycopg2.extras.register_composite("type_r_dt", self.conn)
psycopg2.extras.register_composite("type_r_ft", self.conn)
curs = self.conn.cursor()
with self.conn.cursor() as curs:
r = (0.25, (date(2011, 1, 2), (42, "hello")))
curs.execute("select %s::type_r_ft;", (r,))
v = curs.fetchone()[0]
@ -560,8 +558,7 @@ class AdaptTypeTestCase(ConnectingTestCase):
def test_register_on_cursor(self):
self._create_type("type_ii", [("a", "integer"), ("b", "integer")])
curs1 = self.conn.cursor()
curs2 = self.conn.cursor()
with self.conn.cursor() as curs1, self.conn.cursor() as curs2:
psycopg2.extras.register_composite("type_ii", curs1)
curs1.execute("select (1,2)::type_ii")
self.assertEqual(curs1.fetchone()[0], (1, 2))
@ -576,8 +573,7 @@ class AdaptTypeTestCase(ConnectingTestCase):
conn2 = self.connect()
try:
psycopg2.extras.register_composite("type_ii", conn1)
curs1 = conn1.cursor()
curs2 = conn2.cursor()
with conn1.cursor() as curs1, conn2.cursor() as curs2:
curs1.execute("select (1,2)::type_ii")
self.assertEqual(curs1.fetchone()[0], (1, 2))
curs2.execute("select (1,2)::type_ii")
@ -595,8 +591,7 @@ class AdaptTypeTestCase(ConnectingTestCase):
conn2 = self.connect()
try:
psycopg2.extras.register_composite("type_ii", conn1, globally=True)
curs1 = conn1.cursor()
curs2 = conn2.cursor()
with conn1.cursor() as curs1, conn2.cursor() as curs2:
curs1.execute("select (1,2)::type_ii")
self.assertEqual(curs1.fetchone()[0], (1, 2))
curs2.execute("select (1,2)::type_ii")
@ -608,7 +603,7 @@ class AdaptTypeTestCase(ConnectingTestCase):
@skip_if_no_composite
def test_composite_namespace(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("""
select nspname from pg_namespace
where nspname = 'typens';
@ -633,7 +628,7 @@ class AdaptTypeTestCase(ConnectingTestCase):
t = psycopg2.extras.register_composite("type_isd", self.conn)
curs = self.conn.cursor()
with self.conn.cursor() as curs:
r1 = (10, 'hello', date(2011, 1, 2))
r2 = (20, 'world', date(2011, 1, 3))
curs.execute("select %s::type_isd[];", ([r1, r2],))
@ -652,7 +647,7 @@ class AdaptTypeTestCase(ConnectingTestCase):
def test_wrong_schema(self):
oid = self._create_type("type_ii", [("a", "integer"), ("b", "integer")])
c = CompositeCaster('type_ii', oid, [('a', 23), ('b', 23), ('c', 23)])
curs = self.conn.cursor()
with self.conn.cursor() as curs:
psycopg2.extensions.register_type(c.typecaster, curs)
curs.execute("select (1,2)::type_ii")
self.assertRaises(psycopg2.DataError, curs.fetchone)
@ -661,7 +656,7 @@ class AdaptTypeTestCase(ConnectingTestCase):
@skip_if_no_composite
@skip_before_postgres(8, 4)
def test_from_tables(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("""create table ctest1 (
id integer primary key,
temp int,
@ -710,7 +705,7 @@ class AdaptTypeTestCase(ConnectingTestCase):
conn = self.connect(connection_factory=RealDictConnection)
try:
register_composite('type_ii', conn)
curs = conn.cursor()
with conn.cursor() as curs:
curs.execute("select '(1,2)'::type_ii as x")
self.assertEqual(curs.fetchone()['x'], (1, 2))
finally:
@ -718,7 +713,7 @@ class AdaptTypeTestCase(ConnectingTestCase):
conn = self.connect(connection_factory=RealDictConnection)
try:
curs = conn.cursor()
with conn.cursor() as curs:
register_composite('type_ii', conn)
curs.execute("select '(1,2)'::type_ii as x")
self.assertEqual(curs.fetchone()['x'], (1, 2))
@ -739,7 +734,7 @@ class AdaptTypeTestCase(ConnectingTestCase):
self.assertEqual(t.name, 'type_isd')
self.assertEqual(t.oid, oid)
curs = self.conn.cursor()
with self.conn.cursor() as curs:
r = (10, 'hello', date(2011, 1, 2))
curs.execute("select %s::type_isd;", (r,))
v = curs.fetchone()[0]
@ -749,7 +744,7 @@ class AdaptTypeTestCase(ConnectingTestCase):
self.assertEqual(v['adate'], date(2011, 1, 2))
def _create_type(self, name, fields):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
try:
curs.execute("drop type %s cascade;" % name)
except psycopg2.ProgrammingError:
@ -776,7 +771,7 @@ def skip_if_no_json_type(f):
"""Skip a test if PostgreSQL json type is not available"""
@wraps(f)
def skip_if_no_json_type_(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("select oid from pg_type where typname = 'json'")
if not curs.fetchone():
return self.skipTest("json not available in test database")
@ -791,7 +786,7 @@ class JsonTestCase(ConnectingTestCase):
objs = [None, "te'xt", 123, 123.45,
u'\xe0\u20ac', ['a', 100], {'a': 100}]
curs = self.conn.cursor()
with self.conn.cursor() as curs:
for obj in enumerate(objs):
self.assertQuotedEqual(curs.mogrify("%s", (Json(obj),)),
psycopg2.extensions.QuotedString(json.dumps(obj)).getquoted())
@ -803,7 +798,7 @@ class JsonTestCase(ConnectingTestCase):
return float(obj)
return json.JSONEncoder.default(self, obj)
curs = self.conn.cursor()
with self.conn.cursor() as curs:
obj = Decimal('123.45')
def dumps(obj):
@ -822,7 +817,7 @@ class JsonTestCase(ConnectingTestCase):
def dumps(self, obj):
return json.dumps(obj, cls=DecimalEncoder)
curs = self.conn.cursor()
with self.conn.cursor() as curs:
obj = Decimal('123.45')
self.assertQuotedEqual(curs.mogrify("%s", (MyJson(obj),)), b"'123.45'")
@ -830,13 +825,13 @@ class JsonTestCase(ConnectingTestCase):
def test_register_on_dict(self):
psycopg2.extensions.register_adapter(dict, Json)
curs = self.conn.cursor()
with self.conn.cursor() as curs:
obj = {'a': 123}
self.assertQuotedEqual(
curs.mogrify("%s", (obj,)), b"""'{"a": 123}'""")
def test_type_not_available(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("select oid from pg_type where typname = 'json'")
if curs.fetchone():
return self.skipTest("json available in test database")
@ -846,8 +841,7 @@ class JsonTestCase(ConnectingTestCase):
@skip_before_postgres(9, 2)
def test_default_cast(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("""select '{"a": 100.0, "b": null}'::json""")
self.assertEqual(curs.fetchone()[0], {'a': 100.0, 'b': None})
@ -857,13 +851,13 @@ class JsonTestCase(ConnectingTestCase):
@skip_if_no_json_type
def test_register_on_connection(self):
psycopg2.extras.register_json(self.conn)
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("""select '{"a": 100.0, "b": null}'::json""")
self.assertEqual(curs.fetchone()[0], {'a': 100.0, 'b': None})
@skip_if_no_json_type
def test_register_on_cursor(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
psycopg2.extras.register_json(curs)
curs.execute("""select '{"a": 100.0, "b": null}'::json""")
self.assertEqual(curs.fetchone()[0], {'a': 100.0, 'b': None})
@ -872,7 +866,7 @@ class JsonTestCase(ConnectingTestCase):
@restore_types
def test_register_globally(self):
new, newa = psycopg2.extras.register_json(self.conn, globally=True)
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("""select '{"a": 100.0, "b": null}'::json""")
self.assertEqual(curs.fetchone()[0], {'a': 100.0, 'b': None})
@ -883,7 +877,7 @@ class JsonTestCase(ConnectingTestCase):
def loads(s):
return json.loads(s, parse_float=Decimal)
psycopg2.extras.register_json(self.conn, loads=loads)
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("""select '{"a": 100.0, "b": null}'::json""")
data = curs.fetchone()[0]
self.assert_(isinstance(data['a'], Decimal))
@ -899,7 +893,7 @@ class JsonTestCase(ConnectingTestCase):
new, newa = psycopg2.extras.register_json(
loads=loads, oid=oid, array_oid=array_oid)
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("""select '{"a": 100.0, "b": null}'::json""")
data = curs.fetchone()[0]
self.assert_(isinstance(data['a'], Decimal))
@ -907,8 +901,7 @@ class JsonTestCase(ConnectingTestCase):
@skip_before_postgres(9, 2)
def test_register_default(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
def loads(s):
return psycopg2.extras.json.loads(s, parse_float=Decimal)
psycopg2.extras.register_default_json(curs, loads=loads)
@ -926,14 +919,14 @@ class JsonTestCase(ConnectingTestCase):
@skip_if_no_json_type
def test_null(self):
psycopg2.extras.register_json(self.conn)
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("""select NULL::json""")
self.assertEqual(curs.fetchone()[0], None)
curs.execute("""select NULL::json[]""")
self.assertEqual(curs.fetchone()[0], None)
def test_no_array_oid(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
t1, t2 = psycopg2.extras.register_json(curs, oid=25)
self.assertEqual(t1.values[0], 25)
self.assertEqual(t2, None)
@ -956,13 +949,13 @@ class JsonTestCase(ConnectingTestCase):
@skip_before_postgres(8, 2)
def test_scs(self):
cnn_on = self.connect(options="-c standard_conforming_strings=on")
cur_on = cnn_on.cursor()
with cnn_on.cursor() as cur_on:
self.assertEqual(
cur_on.mogrify("%s", [psycopg2.extras.Json({'a': '"'})]),
b'\'{"a": "\\""}\'')
cnn_off = self.connect(options="-c standard_conforming_strings=off")
cur_off = cnn_off.cursor()
with cnn_off.cursor() as cur_off:
self.assertEqual(
cur_off.mogrify("%s", [psycopg2.extras.Json({'a': '"'})]),
b'E\'{"a": "\\\\""}\'')
@ -985,8 +978,7 @@ class JsonbTestCase(ConnectingTestCase):
return rv
def test_default_cast(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("""select '{"a": 100.0, "b": null}'::jsonb""")
self.assertEqual(curs.fetchone()[0], {'a': 100.0, 'b': None})
@ -995,12 +987,12 @@ class JsonbTestCase(ConnectingTestCase):
def test_register_on_connection(self):
psycopg2.extras.register_json(self.conn, loads=self.myloads, name='jsonb')
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("""select '{"a": 100.0, "b": null}'::jsonb""")
self.assertEqual(curs.fetchone()[0], {'a': 100.0, 'b': None, 'test': 1})
def test_register_on_cursor(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
psycopg2.extras.register_json(curs, loads=self.myloads, name='jsonb')
curs.execute("""select '{"a": 100.0, "b": null}'::jsonb""")
self.assertEqual(curs.fetchone()[0], {'a': 100.0, 'b': None, 'test': 1})
@ -1009,7 +1001,7 @@ class JsonbTestCase(ConnectingTestCase):
def test_register_globally(self):
new, newa = psycopg2.extras.register_json(self.conn,
loads=self.myloads, globally=True, name='jsonb')
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("""select '{"a": 100.0, "b": null}'::jsonb""")
self.assertEqual(curs.fetchone()[0], {'a': 100.0, 'b': None, 'test': 1})
@ -1020,7 +1012,7 @@ class JsonbTestCase(ConnectingTestCase):
return json.loads(s, parse_float=Decimal)
psycopg2.extras.register_json(self.conn, loads=loads, name='jsonb')
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("""select '{"a": 100.0, "b": null}'::jsonb""")
data = curs.fetchone()[0]
self.assert_(isinstance(data['a'], Decimal))
@ -1032,8 +1024,7 @@ class JsonbTestCase(ConnectingTestCase):
self.assertEqual(data['a'], 100.0)
def test_register_default(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
def loads(s):
return psycopg2.extras.json.loads(s, parse_float=Decimal)
@ -1050,7 +1041,7 @@ class JsonbTestCase(ConnectingTestCase):
self.assertEqual(data[0]['a'], Decimal('100.0'))
def test_null(self):
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("""select NULL::jsonb""")
self.assertEqual(curs.fetchone()[0], None)
curs.execute("""select NULL::jsonb[]""")
@ -1332,14 +1323,14 @@ class RangeCasterTestCase(ConnectingTestCase):
'daterange', 'tsrange', 'tstzrange')
def test_cast_null(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
for type in self.builtin_ranges:
cur.execute("select NULL::%s" % type)
r = cur.fetchone()[0]
self.assertEqual(r, None)
def test_cast_empty(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
for type in self.builtin_ranges:
cur.execute("select 'empty'::%s" % type)
r = cur.fetchone()[0]
@ -1347,7 +1338,7 @@ class RangeCasterTestCase(ConnectingTestCase):
self.assert_(r.isempty)
def test_cast_inf(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
for type in self.builtin_ranges:
cur.execute("select '(,)'::%s" % type)
r = cur.fetchone()[0]
@ -1357,7 +1348,7 @@ class RangeCasterTestCase(ConnectingTestCase):
self.assert_(r.upper_inf)
def test_cast_numbers(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
for type in ('int4range', 'int8range'):
cur.execute("select '(10,20)'::%s" % type)
r = cur.fetchone()[0]
@ -1382,7 +1373,7 @@ class RangeCasterTestCase(ConnectingTestCase):
self.assert_(not r.upper_inc)
def test_cast_date(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("select '(2000-01-01,2012-12-31)'::daterange")
r = cur.fetchone()[0]
self.assert_(isinstance(r, DateRange))
@ -1395,7 +1386,7 @@ class RangeCasterTestCase(ConnectingTestCase):
self.assert_(not r.upper_inc)
def test_cast_timestamp(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
ts1 = datetime(2000, 1, 1)
ts2 = datetime(2000, 12, 31, 23, 59, 59, 999)
cur.execute("select tsrange(%s, %s, '()')", (ts1, ts2))
@ -1410,7 +1401,7 @@ class RangeCasterTestCase(ConnectingTestCase):
self.assert_(not r.upper_inc)
def test_cast_timestamptz(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
ts1 = datetime(2000, 1, 1, tzinfo=FixedOffsetTimezone(600))
ts2 = datetime(2000, 12, 31, 23, 59, 59, 999,
tzinfo=FixedOffsetTimezone(600))
@ -1426,8 +1417,7 @@ class RangeCasterTestCase(ConnectingTestCase):
self.assert_(r.upper_inc)
def test_adapt_number_range(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
r = NumericRange(empty=True)
cur.execute("select %s::int4range", (r,))
r1 = cur.fetchone()[0]
@ -1453,8 +1443,7 @@ class RangeCasterTestCase(ConnectingTestCase):
self.assert_(r1.upper_inc)
def test_adapt_numeric_range(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
r = NumericRange(empty=True)
cur.execute("select %s::int4range", (r,))
r1 = cur.fetchone()[0]
@ -1480,8 +1469,7 @@ class RangeCasterTestCase(ConnectingTestCase):
self.assert_(r1.upper_inc)
def test_adapt_date_range(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
d1 = date(2012, 1, 1)
d2 = date(2012, 12, 31)
r = DateRange(d1, d2)
@ -1513,7 +1501,7 @@ class RangeCasterTestCase(ConnectingTestCase):
@restore_types
def test_register_range_adapter(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("create type textrange as range (subtype=text)")
rc = register_range('textrange', 'TextRange', cur)
@ -1539,7 +1527,7 @@ class RangeCasterTestCase(ConnectingTestCase):
self.assert_(r1.upper_inc)
def test_range_escaping(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("create type textrange as range (subtype=text)")
rc = register_range('textrange', 'TextRange', cur)
@ -1592,13 +1580,13 @@ class RangeCasterTestCase(ConnectingTestCase):
del ext.adapters[TextRange, ext.ISQLQuote]
def test_range_not_found(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
self.assertRaises(psycopg2.ProgrammingError,
register_range, 'nosuchrange', 'FailRange', cur)
@restore_types
def test_schema_range(self):
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("create schema rs")
cur.execute("create type r1 as range (subtype=text)")
cur.execute("create type r2 as range (subtype=text)")
@ -1624,6 +1612,32 @@ class RangeCasterTestCase(ConnectingTestCase):
cur.execute("rollback to savepoint x;")
class TestSolveConnCurs(ConnectingTestCase):
def test_pass_connection(self):
with _solve_conn_curs(self.conn) as (conn, curs):
self.assertIsInstance(conn, psycopg2.extensions.connection)
self.assertIsInstance(curs, psycopg2.extensions.cursor)
self.assertIs(conn, self.conn)
self.assertFalse(conn.closed)
self.assertFalse(curs.closed)
self.assertFalse(conn.closed)
self.assertTrue(curs.closed)
def test_pass_cursor(self):
with self.conn.cursor() as cursor:
with _solve_conn_curs(cursor) as (conn, curs):
self.assertIsInstance(conn, psycopg2.extensions.connection)
self.assertIsInstance(curs, psycopg2.extensions.cursor)
self.assertIs(conn, self.conn)
self.assertIsNot(curs, cursor)
self.assertFalse(conn.closed)
self.assertFalse(curs.closed)
self.assertFalse(conn.closed)
self.assertTrue(curs.closed)
self.assertFalse(conn.closed)
self.assertTrue(curs.closed)
def test_suite():
return unittest.TestLoader().loadTestsFromName(__name__)

116
tests/test_warnings.py Normal file
View File

@ -0,0 +1,116 @@
import re
import subprocess
import sys
import unittest
import warnings
import psycopg2
from .testconfig import dsn
from .testutils import skip_before_python
class WarningsTest(unittest.TestCase):
@skip_before_python(3)
def test_connection_not_closed(self):
def f():
psycopg2.connect(dsn)
msg = (
"^unclosed connection <connection object at 0x[0-9a-fA-F]+; dsn: '.*', "
"closed: 0>$"
)
with self.assertWarnsRegex(ResourceWarning, msg):
f()
@skip_before_python(3)
def test_cursor_not_closed(self):
def f():
conn = psycopg2.connect(dsn)
try:
conn.cursor()
finally:
conn.close()
msg = (
"^unclosed cursor <cursor object at 0x[0-9a-fA-F]+; closed: 0> for "
"connection <connection object at 0x[0-9a-fA-F]+; dsn: '.*', closed: 0>$"
)
with self.assertWarnsRegex(ResourceWarning, msg):
f()
def test_cursor_factory_returns_non_cursor(self):
def bad_factory(*args, **kwargs):
return object()
def f():
conn = psycopg2.connect(dsn)
try:
conn.cursor(cursor_factory=bad_factory)
finally:
conn.close()
with warnings.catch_warnings(record=True) as cm:
with self.assertRaises(TypeError):
f()
# No warning as no cursor was instantiated.
self.assertEquals(cm, [])
@skip_before_python(3)
def test_broken_close(self):
script = """
import psycopg2
class MyException(Exception):
pass
class MyCurs(psycopg2.extensions.cursor):
def close(self):
raise MyException
def f():
conn = psycopg2.connect(%(dsn)r)
try:
conn.cursor(cursor_factory=MyCurs, scrollable=True)
finally:
conn.close()
f()
""" % {"dsn": dsn}
p = subprocess.Popen(
[sys.executable, "-Walways", "-c", script],
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
)
output, _ = p.communicate()
output = output.decode()
# Normalize line endings.
output = "\n".join(output.splitlines())
self.assertRegex(
output,
re.compile(
r"^Exception ignored in: "
r"<cursor object at 0x[0-9a-fA-F]+; closed: 0>$",
re.M,
),
)
self.assertIn("\n__main__.MyException: \n", output)
self.assertRegex(
output,
re.compile(
r"ResourceWarning: unclosed cursor "
r"<cursor object at 0x[0-9a-fA-F]+; closed: 0> "
r"for connection "
r"<connection object at 0x[0-9a-fA-F]+; dsn: '.*', closed: 0>$",
re.M,
),
)
def test_suite():
return unittest.TestLoader().loadTestsFromName(__name__)
if __name__ == "__main__":
unittest.main()

30
tests/test_with.py
View File

@ -33,7 +33,7 @@ from .testutils import ConnectingTestCase, skip_before_postgres
class WithTestCase(ConnectingTestCase):
def setUp(self):
ConnectingTestCase.setUp(self)
curs = self.conn.cursor()
with self.conn.cursor() as curs:
try:
curs.execute("delete from test_with")
self.conn.commit()
@ -49,49 +49,49 @@ class WithConnectionTestCase(WithTestCase):
with self.conn as conn:
self.assert_(self.conn is conn)
self.assertEqual(conn.status, ext.STATUS_READY)
curs = conn.cursor()
with conn.cursor() as curs:
curs.execute("insert into test_with values (1)")
self.assertEqual(conn.status, ext.STATUS_BEGIN)
self.assertEqual(self.conn.status, ext.STATUS_READY)
self.assert_(not self.conn.closed)
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("select * from test_with")
self.assertEqual(curs.fetchall(), [(1,)])
def test_with_connect_idiom(self):
with self.connect() as conn:
self.assertEqual(conn.status, ext.STATUS_READY)
curs = conn.cursor()
with conn.cursor() as curs:
curs.execute("insert into test_with values (2)")
self.assertEqual(conn.status, ext.STATUS_BEGIN)
self.assertEqual(self.conn.status, ext.STATUS_READY)
self.assert_(not self.conn.closed)
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("select * from test_with")
self.assertEqual(curs.fetchall(), [(2,)])
def test_with_error_db(self):
def f():
with self.conn as conn:
curs = conn.cursor()
with conn.cursor() as curs:
curs.execute("insert into test_with values ('a')")
self.assertRaises(psycopg2.DataError, f)
self.assertEqual(self.conn.status, ext.STATUS_READY)
self.assert_(not self.conn.closed)
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("select * from test_with")
self.assertEqual(curs.fetchall(), [])
def test_with_error_python(self):
def f():
with self.conn as conn:
curs = conn.cursor()
with conn.cursor() as curs:
curs.execute("insert into test_with values (3)")
1 / 0
@ -99,7 +99,7 @@ class WithConnectionTestCase(WithTestCase):
self.assertEqual(self.conn.status, ext.STATUS_READY)
self.assert_(not self.conn.closed)
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("select * from test_with")
self.assertEqual(curs.fetchall(), [])
@ -120,13 +120,13 @@ class WithConnectionTestCase(WithTestCase):
super(MyConn, self).commit()
with self.connect(connection_factory=MyConn) as conn:
curs = conn.cursor()
with conn.cursor() as curs:
curs.execute("insert into test_with values (10)")
self.assertEqual(conn.status, ext.STATUS_READY)
self.assert_(commits)
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("select * from test_with")
self.assertEqual(curs.fetchall(), [(10,)])
@ -140,7 +140,7 @@ class WithConnectionTestCase(WithTestCase):
try:
with self.connect(connection_factory=MyConn) as conn:
curs = conn.cursor()
with conn.cursor() as curs:
curs.execute("insert into test_with values (11)")
1 / 0
except ZeroDivisionError:
@ -151,7 +151,7 @@ class WithConnectionTestCase(WithTestCase):
self.assertEqual(conn.status, ext.STATUS_READY)
self.assert_(rollbacks)
curs = conn.cursor()
with conn.cursor() as curs:
curs.execute("select * from test_with")
self.assertEqual(curs.fetchall(), [])
@ -168,7 +168,7 @@ class WithCursorTestCase(WithTestCase):
self.assertEqual(self.conn.status, ext.STATUS_READY)
self.assert_(not self.conn.closed)
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("select * from test_with")
self.assertEqual(curs.fetchall(), [(4,)])
@ -185,7 +185,7 @@ class WithCursorTestCase(WithTestCase):
self.assert_(not self.conn.closed)
self.assert_(curs.closed)
curs = self.conn.cursor()
with self.conn.cursor() as curs:
curs.execute("select * from test_with")
self.assertEqual(curs.fetchall(), [])

6
tests/testutils.py
View File

@ -228,7 +228,7 @@ def skip_if_no_uuid(f):
@wraps(f)
def skip_if_no_uuid_(self):
try:
cur = self.conn.cursor()
with self.conn.cursor() as cur:
cur.execute("select typname from pg_type where typname = 'uuid'")
has = cur.fetchone()
finally:
@ -248,7 +248,8 @@ def skip_if_tpc_disabled(f):
@wraps(f)
def skip_if_tpc_disabled_(self):
cnn = self.connect()
cur = cnn.cursor()
try:
with cnn.cursor() as cur:
try:
cur.execute("SHOW max_prepared_transactions;")
except psycopg2.ProgrammingError:
@ -256,6 +257,7 @@ def skip_if_tpc_disabled(f):
"server too old: two phase transactions not supported.")
else:
mtp = int(cur.fetchone()[0])
finally:
cnn.close()
if not mtp: