New 'withhold' parameter for connection.cursor()

onlinemediagroup · Aug 10, 2011 · 479bdf7 · 479bdf7
1 parent 30a046c
commit 479bdf7
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 6 deletions.
diff --git a/psycopg/connection_type.c b/psycopg/connection_type.c
@@ -53,15 +53,28 @@ static PyObject *
 psyco_conn_cursor(connectionObject *self, PyObject *args, PyObject *keywds)
 {
     const char *name = NULL;
-    PyObject *obj, *factory = NULL;
+    PyObject *obj, *factory = NULL, *withhold = NULL;
 
-    static char *kwlist[] = {"name", "cursor_factory", NULL};
+    static char *kwlist[] = {"name", "cursor_factory", "withhold", NULL};
 
-    if (!PyArg_ParseTupleAndKeywords(args, keywds, "|sO", kwlist,
-                                     &name, &factory)) {
+    if (!PyArg_ParseTupleAndKeywords(args, keywds, "|sOO", kwlist,
+                                     &name, &factory, &withhold)) {
         return NULL;
     }
 
+    if (withhold != NULL) {
+        if (withhold == Py_True && name == NULL) {
+            PyErr_SetString(ProgrammingError,
+                "'withhold=True can be specified only for named cursors");
+            return NULL;
+        }
+        if (withhold != NULL && withhold != Py_True && withhold != Py_False) {
+            PyErr_SetString(ProgrammingError,
+                "'withhold should be True or False");
+            return NULL;
+        }
+    }
+
     EXC_IF_CONN_CLOSED(self);
 
     if (self->status != CONN_STATUS_READY &&
@@ -95,6 +108,9 @@ psyco_conn_cursor(connectionObject *self, PyObject *args, PyObject *keywds)
         Py_DECREF(obj);
         return NULL;
     }
+
+    if (withhold == Py_True)
+        ((cursorObject*)obj)->withhold = 1;
 
     Dprintf("psyco_conn_cursor: new cursor at %p: refcnt = "
         FORMAT_CODE_PY_SSIZE_T,
@@ -525,7 +541,7 @@ psyco_conn_set_session(connectionObject *self, PyObject *args, PyObject *kwargs)
 
 
 #define psyco_conn_autocommit_doc \
-"set or return the autocommit status."
+"Set or return the autocommit status."
 
 static PyObject *
 psyco_conn_autocommit_get(connectionObject *self)

diff --git a/psycopg2.cproj b/psycopg2.cproj
@@ -237,6 +237,7 @@
     <Compile Include="psycopg\green.c" />
     <Compile Include="psycopg\notify_type.c" />
     <Compile Include="psycopg\xid_type.c" />
+    <Compile Include="psycopg\bytes_format.c" />
   </ItemGroup>
   <ProjectExtensions>
     <MonoDevelop>

diff --git a/psycopg2.sln b/psycopg2.sln
@@ -35,6 +35,10 @@ Global
 		$4.inheritsSet = VisualStudio
 		$4.inheritsScope = text/plain
 		$4.scope = text/x-readme
+		$0.TextStylePolicy = $5
+		$5.inheritsSet = VisualStudio
+		$5.inheritsScope = text/plain
+		$5.scope = text/plain
 		name = psycopg2
 	EndGlobalSection
 EndGlobal
diff --git a/tests/test_cursor.py b/tests/test_cursor.py
@@ -159,6 +159,11 @@ def test_invalid_name(self):
         self.assertEqual(curs.fetchall(), [(10,), (20,), (30,)])
 
     def test_withhold(self):
+        self.assertRaises(psycopg2.ProgrammingError, self.conn.cursor, 
+                          withhold=True)    
+        self.assertRaises(psycopg2.ProgrammingError, self.conn.cursor, 
+                          "NAME", withhold="")    
+
         curs = self.conn.cursor()
         curs.execute("drop table if exists withhold")
         curs.execute("create table withhold (data int)")
@@ -173,7 +178,14 @@ def test_withhold(self):
         curs.execute("select data from withhold order by data")
         self.conn.commit()
         self.assertEqual(curs.fetchall(), [(10,), (20,), (30,)])
-
+        curs.close()
+
+        curs = self.conn.cursor("W", withhold=True)
+        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()
         curs.execute("drop table withhold")
         self.conn.commit()