Merge pull request #467 from lepennec/cursor_description

Add description field to Python cursor to increase compatibility
duckdb · Mar 23, 2020 · a517936 · a517936
2 parents 04859b8 + a1ca808
commit a517936
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 0 deletions.
diff --git a/tools/pythonpkg/cursor.cpp b/tools/pythonpkg/cursor.cpp
@@ -27,6 +27,9 @@ static int duckdb_cursor_init(duckdb_Cursor *self, PyObject *args, PyObject *kwa
 	self->reset = 0;
 	self->rowcount = -1L;
 
+	Py_INCREF(Py_None);
+    Py_XSETREF(self->description, Py_None);
+
 	self->initialized = 1;
 
 	return 0;
@@ -37,6 +40,7 @@ static void duckdb_cursor_dealloc(duckdb_Cursor *self) {
 	duckdb_cursor_close(self, NULL);
 
 	Py_XDECREF(self->connection);
+	Py_XDECREF(self->description);
 	Py_TYPE(self)->tp_free((PyObject *)self);
 }
 
@@ -221,6 +225,9 @@ static PyObject *_duckdb_query_execute(duckdb_Cursor *self, int multiple, PyObje
 	PyObject *parameters = NULL;
 	PyObject *second_argument = NULL;
 
+	int numcols;
+	PyObject* descriptor;
+
 	bool need_transaction;
 
 	std::vector<duckdb::Value> params;
@@ -285,6 +292,10 @@ static PyObject *_duckdb_query_execute(duckdb_Cursor *self, int multiple, PyObje
 			goto error;
 		}
 	}
+
+	/* reset description and rowcount */
+    Py_INCREF(Py_None);
+    Py_SETREF(self->description, Py_None);
 	self->rowcount = 0L;
 	self->reset = 0;
 
@@ -324,6 +335,29 @@ static PyObject *_duckdb_query_execute(duckdb_Cursor *self, int multiple, PyObje
 		self->result = std::unique_ptr<duckdb::MaterializedQueryResult>(
 		    static_cast<duckdb::MaterializedQueryResult *>(res.release()));
 		Py_XDECREF(parameters);
+
+		numcols = self->result->collection.column_count();;
+		if (self->description == Py_None && numcols > 0) {
+            Py_SETREF(self->description, PyTuple_New(numcols));
+            if (!self->description) {
+                goto error;
+            }
+            for (int i = 0; i < numcols; i++) {
+                descriptor = PyTuple_New(7);
+                if (!descriptor) {
+                    goto error;
+                }
+                PyTuple_SetItem(descriptor, 0, PyUnicode_FromString(self->result->names[i].c_str()));
+                Py_INCREF(Py_None); PyTuple_SetItem(descriptor, 1, Py_None);
+                Py_INCREF(Py_None); PyTuple_SetItem(descriptor, 2, Py_None);
+                Py_INCREF(Py_None); PyTuple_SetItem(descriptor, 3, Py_None);
+                Py_INCREF(Py_None); PyTuple_SetItem(descriptor, 4, Py_None);
+                Py_INCREF(Py_None); PyTuple_SetItem(descriptor, 5, Py_None);
+                Py_INCREF(Py_None); PyTuple_SetItem(descriptor, 6, Py_None);
+                PyTuple_SetItem(self->description, i, descriptor);
+            }
+        }
+
 		self->rowcount = self->result->collection.count;
 		self->closed = 0;
 		self->offset = 0;
@@ -737,6 +771,7 @@ static struct PyMemberDef cursor_members[] = {
     {"connection", T_OBJECT, offsetof(duckdb_Cursor, connection), READONLY},
     //    {"lastrowid", T_OBJECT, offsetof(pysqlite_Cursor, lastrowid), READONLY},
     {"rowcount", T_LONG, offsetof(duckdb_Cursor, rowcount), READONLY},
+	{"description", T_OBJECT, offsetof(duckdb_Cursor, description), READONLY},
     {NULL}};
 
 static const char cursor_doc[] = PyDoc_STR("DuckDB database cursor class.");

diff --git a/tools/pythonpkg/cursor.h b/tools/pythonpkg/cursor.h
@@ -10,6 +10,8 @@ typedef struct {
 	uint64_t rowcount;
 	uint64_t offset;
 
+	PyObject* description;
+
 	int closed;
 	int reset;
 	int initialized;

diff --git a/tools/pythonpkg/tests/test_dbapi10.py b/tools/pythonpkg/tests/test_dbapi10.py
@@ -0,0 +1,8 @@
+# cursor description
+
+import numpy 
+
+class TestCursorDescription(object):
+    def test_description(self, duckdb_cursor):
+        description = duckdb_cursor.execute("SELECT * FROM integers").description
+        assert description == (('i', None, None, None, None, None, None),)