3 files changed, 158 insertions, 0 deletions
diff --git a/psycopg/green.c b/psycopg/green.c
new file mode 100644
index 0000000..874c264
--- /dev/null
+++ b/psycopg/green.c
@@ -0,0 +1,94 @@
+/* green.c - cooperation with coroutine libraries.
+ *
+ * Copyright (C) 2010 Daniele Varrazzo <daniele.varrazzo@gmail.com>
+ *
+ * This file is part of psycopg.
+ *
+ * psycopg2 is free software: you can redistribute it and/or modify it
+ * under the terms of the GNU Lesser General Public License as published
+ * by the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * In addition, as a special exception, the copyright holders give
+ * permission to link this program with the OpenSSL library (or with
+ * modified versions of OpenSSL that use the same license as OpenSSL),
+ * and distribute linked combinations including the two.
+ *
+ * You must obey the GNU Lesser General Public License in all respects for
+ * all of the code used other than OpenSSL.
+ *
+ * psycopg2 is distributed in the hope that it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public
+ * License for more details.
+ */
+
+#define PSYCOPG_MODULE
+#include "psycopg/config.h"
+#include "psycopg/python.h"
+#include "psycopg/psycopg.h"
+#include "psycopg/green.h"
+
+HIDDEN PyObject *wait_callback = NULL;
+
+/* Register a callback function to block waiting for data.
+ *
+ * The function is exported by the _psycopg module.
+ */
+PyObject *
+psyco_set_wait_callback(PyObject *self, PyObject *obj)
+{
+    Py_XDECREF(wait_callback);
+
+    if (obj != Py_None) {
+        wait_callback = obj;
+        Py_INCREF(obj);
+    }
+    else {
+        wait_callback = NULL;
+    }
+
+    Py_INCREF(Py_None);
+    return Py_None;
+}
+
+
+/* Return nonzero if a wait callback should be called. */
+int
+psyco_green()
+{
+#ifdef PSYCOPG_EXTENSIONS
+    return (NULL != wait_callback);
+#else
+    return 0;
+#endif
+}
+
+/* Block waiting for data available in an async connection.
+ *
+ * This function assumes `wait_callback` to be available:
+ * raise `InterfaceError` if it is not. Use `psyco_green()` to check if
+ * the function is to be called.
+ *
+ * The function returns the return value of the called function.
+ */
+PyObject *
+psyco_wait(PyObject *conn, PyObject *curs)
+{
+    PyObject *rv;
+    PyObject *cb;
+
+    Dprintf("psyco_wait");
+    cb = wait_callback;
+    if (!cb) {
+        PyErr_SetString(OperationalError, "wait callback not available");
+        return NULL;
+    }
+
+    Py_INCREF(cb);
+    rv = PyObject_CallFunctionObjArgs(cb, conn, curs, NULL);
+    Py_DECREF(cb);
+
+    return rv;
+}
+
diff --git a/psycopg/green.h b/psycopg/green.h
new file mode 100644
index 0000000..2899927
--- /dev/null
+++ b/psycopg/green.h
@@ -0,0 +1,58 @@
+/* green.c - cooperation with coroutine libraries.
+ *
+ * Copyright (C) 2010 Daniele Varrazzo <daniele.varrazzo@gmail.com>
+ *
+ * This file is part of psycopg.
+ *
+ * psycopg2 is free software: you can redistribute it and/or modify it
+ * under the terms of the GNU Lesser General Public License as published
+ * by the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * In addition, as a special exception, the copyright holders give
+ * permission to link this program with the OpenSSL library (or with
+ * modified versions of OpenSSL that use the same license as OpenSSL),
+ * and distribute linked combinations including the two.
+ *
+ * You must obey the GNU Lesser General Public License in all respects for
+ * all of the code used other than OpenSSL.
+ *
+ * psycopg2 is distributed in the hope that it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public
+ * License for more details.
+ */
+
+#ifndef PSYCOPG_GREEN_H
+#define PSYCOPG_GREEN_H 1
+
+struct PyObject;
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#define psyco_set_wait_callback_doc \
+"set_wait_callback(f) -- Register a callback function to block waiting for data.\n" \
+"\n" \
+"The callback must should signature :samp:`fun({conn}, {cur}=None)` and\n" \
+"is called to wait for data available whenever a blocking function from the\n" \
+"libpq is called.  Use `!register_wait_function(None)` to revert to the\n" \
+"original behaviour (using blocking libpq functions).\n" \
+"\n" \
+"The function is an hook to allow coroutine-based libraries (such as\n" \
+"eventlet_ or gevent_) to switch when Psycopg is blocked, allowing\n" \
+"other coroutines to run concurrently.\n" \
+"\n" \
+"See `~psycopg2.extras.wait_select()` for an example of a wait callback\n" \
+"implementation.\n"
+
+HIDDEN PyObject *psyco_set_wait_callback(PyObject *self, PyObject *obj);
+HIDDEN int psyco_green(void);
+HIDDEN PyObject *psyco_wait(PyObject *conn, PyObject *curs);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif /* !defined(PSYCOPG_GREEN_H) */
diff --git a/psycopg/psycopgmodule.c b/psycopg/psycopgmodule.c
index 528c592..60c9c04 100644
--- a/psycopg/psycopgmodule.c
+++ b/psycopg/psycopgmodule.c
@@ -32,6 +32,7 @@
 #include "psycopg/psycopg.h"
 #include "psycopg/connection.h"
 #include "psycopg/cursor.h"
+#include "psycopg/green.h"
 #include "psycopg/lobject.h"
 #include "psycopg/typecast.h"
 #include "psycopg/microprotocols.h"
@@ -693,6 +694,11 @@ static PyMethodDef psycopgMethods[] = {
      METH_VARARGS, psyco_IntervalFromMx_doc},
 #endif
 
+#ifdef PSYCOPG_EXTENSIONS
+    {"set_wait_callback",  (PyCFunction)psyco_set_wait_callback,
+     METH_O, psyco_set_wait_callback_doc},
+#endif
+
     {NULL, NULL, 0, NULL}        /* Sentinel */
 };