Issue #21718: cursor.description is now available for queries using CTEs

According to PEP 249, cursor.description must be
available for any SELECT statements, such as those
that use CTEs.

Backported from f67fa9c898

Additional test cases added by me.

Lib/sqlite3/test/types.py

@@ -274,6 +274,45 @@ class ColNamesTests(unittest.TestCase):
         self.cur.execute("select * from test where 0 = 1")
         self.assertEqual(self.cur.description[0][0], "x")
 
+    def CheckCursorDescriptionInsert(self):
+        self.cur.execute("insert into test values (1)")
+        self.assertIsNone(self.cur.description)
+
+
+@unittest.skipIf(sqlite.sqlite_version_info < (3, 8, 3), "CTEs not supported")
+class CommonTableExpressionTests(unittest.TestCase):
+
+    def setUp(self):
+        self.con = sqlite.connect(":memory:")
+        self.cur = self.con.cursor()
+        self.cur.execute("create table test(x foo)")
+
+    def tearDown(self):
+        self.cur.close()
+        self.con.close()
+
+    def CheckCursorDescriptionCTESimple(self):
+        self.cur.execute("with one as (select 1) select * from one")
+        self.assertIsNotNone(self.cur.description)
+        self.assertEqual(self.cur.description[0][0], "1")
+
+    def CheckCursorDescriptionCTESMultipleColumns(self):
+        self.cur.execute("insert into test values(1)")
+        self.cur.execute("insert into test values(2)")
+        self.cur.execute("with testCTE as (select * from test) select * from testCTE")
+        self.assertIsNotNone(self.cur.description)
+        self.assertEqual(self.cur.description[0][0], "x")
+
+    def CheckCursorDescriptionCTE(self):
+        self.cur.execute("insert into test values (1)")
+        self.cur.execute("with bar as (select * from test) select * from test where x = 1")
+        self.assertIsNotNone(self.cur.description)
+        self.assertEqual(self.cur.description[0][0], "x")
+        self.cur.execute("with bar as (select * from test) select * from test where x = 2")
+        self.assertIsNotNone(self.cur.description)
+        self.assertEqual(self.cur.description[0][0], "x")
+
+
 class ObjectAdaptationTests(unittest.TestCase):
     def cast(obj):
         return float(obj)
@@ -372,7 +411,8 @@ def suite():
     adaptation_suite = unittest.makeSuite(ObjectAdaptationTests, "Check")
     bin_suite = unittest.makeSuite(BinaryConverterTests, "Check")
     date_suite = unittest.makeSuite(DateTimeTests, "Check")
-    return unittest.TestSuite((sqlite_type_suite, decltypes_type_suite, colnames_type_suite, adaptation_suite, bin_suite, date_suite))
+    cte_suite = unittest.makeSuite(CommonTableExpressionTests, "Check")
+    return unittest.TestSuite((sqlite_type_suite, decltypes_type_suite, colnames_type_suite, adaptation_suite, bin_suite, date_suite, cte_suite))
 
 def test():
     runner = unittest.TextTestRunner()

Misc/NEWS

@@ -46,6 +46,8 @@ Core and Builtins
 Library
 -------
 
+- Issue #21718: cursor.description is now available for queries using CTEs.
+
 - Issue #2466: posixpath.ismount now correctly recognizes mount points which
   the user does not have permission to access.
 

Modules/_sqlite/cursor.c

@@ -646,12 +646,11 @@ PyObject* _pysqlite_query_execute(pysqlite_Cursor* self, int multiple, PyObject*
             goto error;
         }
 
-        if (rc == SQLITE_ROW || (rc == SQLITE_DONE && statement_type == STATEMENT_SELECT)) {
+        if (rc == SQLITE_ROW || rc == SQLITE_DONE) {
-            if (self->description == Py_None) {
+            Py_BEGIN_ALLOW_THREADS
-                Py_BEGIN_ALLOW_THREADS
+            numcols = sqlite3_column_count(self->statement->st);
-                numcols = sqlite3_column_count(self->statement->st);
+            Py_END_ALLOW_THREADS
-                Py_END_ALLOW_THREADS
+            if (self->description == Py_None && numcols > 0) {
-
                 Py_SETREF(self->description, PyTuple_New(numcols));
                 if (!self->description) {
                     goto error;