plpython.c

PostgreSQL 8.1.4的源码 适用于Linux下的开源数据库系统

static PySequenceMethods PLy_result_as_sequence = {
	(inquiry) PLy_result_length,	/* sq_length */
	(binaryfunc) 0,				/* sq_concat */
	(intargfunc) 0,				/* sq_repeat */
	(intargfunc) PLy_result_item,		/* sq_item */
	(intintargfunc) PLy_result_slice,	/* sq_slice */
	(intobjargproc) PLy_result_ass_item,		/* sq_ass_item */
	(intintobjargproc) PLy_result_ass_slice,	/* sq_ass_slice */
};

static PyTypeObject PLy_ResultType = {
	PyObject_HEAD_INIT(NULL)
	0,							/* ob_size */
	"PLyResult",				/* tp_name */
	sizeof(PLyResultObject),	/* tp_size */
	0,							/* tp_itemsize */

	/*
	 * methods
	 */
	(destructor) PLy_result_dealloc,	/* tp_dealloc */
	0,							/* tp_print */
	(getattrfunc) PLy_result_getattr,	/* tp_getattr */
	0,							/* tp_setattr */
	0,							/* tp_compare */
	0,							/* tp_repr */
	0,							/* tp_as_number */
	&PLy_result_as_sequence,	/* tp_as_sequence */
	0,							/* tp_as_mapping */
	0,							/* tp_hash */
	0,							/* tp_call */
	0,							/* tp_str */
	0,							/* tp_getattro */
	0,							/* tp_setattro */
	0,							/* tp_as_buffer */
	0,							/* tp_xxx4 */
	PLy_result_doc,				/* tp_doc */
};

static PyMethodDef PLy_result_methods[] = {
	{"nrows", PLy_result_nrows, METH_VARARGS, NULL},
	{"status", PLy_result_status, METH_VARARGS, NULL},
	{NULL, NULL, 0, NULL}
};

static PyMethodDef PLy_methods[] = {
	/*
	 * logging methods
	 */
	{"debug", PLy_debug, METH_VARARGS, NULL},
	{"log", PLy_log, METH_VARARGS, NULL},
	{"info", PLy_info, METH_VARARGS, NULL},
	{"notice", PLy_notice, METH_VARARGS, NULL},
	{"warning", PLy_warning, METH_VARARGS, NULL},
	{"error", PLy_error, METH_VARARGS, NULL},
	{"fatal", PLy_fatal, METH_VARARGS, NULL},

	/*
	 * create a stored plan
	 */
	{"prepare", PLy_spi_prepare, METH_VARARGS, NULL},

	/*
	 * execute a plan or query
	 */
	{"execute", PLy_spi_execute, METH_VARARGS, NULL},

	{NULL, NULL, 0, NULL}
};


/* plan object methods
 */
static PyObject *
PLy_plan_new(void)
{
	PLyPlanObject *ob;

	if ((ob = PyObject_NEW(PLyPlanObject, &PLy_PlanType)) == NULL)
		return NULL;

	ob->plan = NULL;
	ob->nargs = 0;
	ob->types = NULL;
	ob->args = NULL;

	return (PyObject *) ob;
}


static void
PLy_plan_dealloc(PyObject * arg)
{
	PLyPlanObject *ob = (PLyPlanObject *) arg;

	if (ob->plan)
		SPI_freeplan(ob->plan);
	if (ob->types)
		PLy_free(ob->types);
	if (ob->args)
	{
		int			i;

		for (i = 0; i < ob->nargs; i++)
			PLy_typeinfo_dealloc(&ob->args[i]);
		PLy_free(ob->args);
	}

	PyMem_DEL(arg);
}


static PyObject *
PLy_plan_getattr(PyObject * self, char *name)
{
	return Py_FindMethod(PLy_plan_methods, self, name);
}

static PyObject *
PLy_plan_status(PyObject * self, PyObject * args)
{
	if (PyArg_ParseTuple(args, ""))
	{
		Py_INCREF(Py_True);
		return Py_True;
		/* return PyInt_FromLong(self->status); */
	}
	PyErr_SetString(PLy_exc_error, "plan.status() takes no arguments");
	return NULL;
}



/* result object methods
 */

static PyObject *
PLy_result_new(void)
{
	PLyResultObject *ob;

	if ((ob = PyObject_NEW(PLyResultObject, &PLy_ResultType)) == NULL)
		return NULL;

	/* ob->tuples = NULL; */

	Py_INCREF(Py_None);
	ob->status = Py_None;
	ob->nrows = PyInt_FromLong(-1);
	ob->rows = PyList_New(0);

	return (PyObject *) ob;
}

static void
PLy_result_dealloc(PyObject * arg)
{
	PLyResultObject *ob = (PLyResultObject *) arg;

	Py_XDECREF(ob->nrows);
	Py_XDECREF(ob->rows);
	Py_XDECREF(ob->status);

	PyMem_DEL(ob);
}

static PyObject *
PLy_result_getattr(PyObject * self, char *name)
{
	return Py_FindMethod(PLy_result_methods, self, name);
}

static PyObject *
PLy_result_nrows(PyObject * self, PyObject * args)
{
	PLyResultObject *ob = (PLyResultObject *) self;

	Py_INCREF(ob->nrows);
	return ob->nrows;
}

static PyObject *
PLy_result_status(PyObject * self, PyObject * args)
{
	PLyResultObject *ob = (PLyResultObject *) self;

	Py_INCREF(ob->status);
	return ob->status;
}

static int
PLy_result_length(PyObject * arg)
{
	PLyResultObject *ob = (PLyResultObject *) arg;

	return PyList_Size(ob->rows);
}

static PyObject *
PLy_result_item(PyObject * arg, int idx)
{
	PyObject   *rv;
	PLyResultObject *ob = (PLyResultObject *) arg;

	rv = PyList_GetItem(ob->rows, idx);
	if (rv != NULL)
		Py_INCREF(rv);
	return rv;
}

static int
PLy_result_ass_item(PyObject * arg, int idx, PyObject * item)
{
	int			rv;
	PLyResultObject *ob = (PLyResultObject *) arg;

	Py_INCREF(item);
	rv = PyList_SetItem(ob->rows, idx, item);
	return rv;
}

static PyObject *
PLy_result_slice(PyObject * arg, int lidx, int hidx)
{
	PyObject   *rv;
	PLyResultObject *ob = (PLyResultObject *) arg;

	rv = PyList_GetSlice(ob->rows, lidx, hidx);
	if (rv == NULL)
		return NULL;
	Py_INCREF(rv);
	return rv;
}

static int
PLy_result_ass_slice(PyObject * arg, int lidx, int hidx, PyObject * slice)
{
	int			rv;
	PLyResultObject *ob = (PLyResultObject *) arg;

	rv = PyList_SetSlice(ob->rows, lidx, hidx, slice);
	return rv;
}

/* SPI interface
 */
static PyObject *
PLy_spi_prepare(PyObject * self, PyObject * args)
{
	PLyPlanObject *plan;
	PyObject   *list = NULL;
	PyObject   *volatile optr = NULL;
	char	   *query;
	void	   *tmpplan;
	MemoryContext oldcontext;

	/* Can't execute more if we have an unhandled error */
	if (PLy_error_in_progress)
	{
		PyErr_SetString(PLy_exc_error, "Transaction aborted.");
		return NULL;
	}

	if (!PyArg_ParseTuple(args, "s|O", &query, &list))
	{
		PyErr_SetString(PLy_exc_spi_error,
						"Invalid arguments for plpy.prepare()");
		return NULL;
	}

	if ((list) && (!PySequence_Check(list)))
	{
		PyErr_SetString(PLy_exc_spi_error,
					 "Second argument in plpy.prepare() must be a sequence");
		return NULL;
	}

	if ((plan = (PLyPlanObject *) PLy_plan_new()) == NULL)
		return NULL;

	oldcontext = CurrentMemoryContext;
	PG_TRY();
	{
		if (list != NULL)
		{
			int			nargs,
						i;

			nargs = PySequence_Length(list);
			if (nargs > 0)
			{
				plan->nargs = nargs;
				plan->types = PLy_malloc(sizeof(Oid) * nargs);
				plan->values = PLy_malloc(sizeof(Datum) * nargs);
				plan->args = PLy_malloc(sizeof(PLyTypeInfo) * nargs);

				/*
				 * the other loop might throw an exception, if PLyTypeInfo
				 * member isn't properly initialized the Py_DECREF(plan) will
				 * go boom
				 */
				for (i = 0; i < nargs; i++)
				{
					PLy_typeinfo_init(&plan->args[i]);
					plan->values[i] = (Datum) NULL;
				}

				for (i = 0; i < nargs; i++)
				{
					char	   *sptr;
					HeapTuple	typeTup;
					Form_pg_type typeStruct;

					optr = PySequence_GetItem(list, i);
					if (!PyString_Check(optr))
						elog(ERROR, "Type names must be strings.");
					sptr = PyString_AsString(optr);

					/*
					 * XXX should extend this to allow qualified type names
					 */
					typeTup = typenameType(makeTypeName(sptr));
					Py_DECREF(optr);
					optr = NULL;	/* this is important */

					plan->types[i] = HeapTupleGetOid(typeTup);
					typeStruct = (Form_pg_type) GETSTRUCT(typeTup);
					if (typeStruct->typtype != 'c')
						PLy_output_datum_func(&plan->args[i], typeTup);
					else
						elog(ERROR, "tuples not handled in plpy.prepare, yet.");
					ReleaseSysCache(typeTup);
				}
			}
		}

		plan->plan = SPI_prepare(query, plan->nargs, plan->types);
		if (plan->plan == NULL)
			elog(ERROR, "SPI_prepare failed: %s",
				 SPI_result_code_string(SPI_result));

		/* transfer plan from procCxt to topCxt */
		tmpplan = plan->plan;
		plan->plan = SPI_saveplan(tmpplan);
		SPI_freeplan(tmpplan);
		if (plan->plan == NULL)
			elog(ERROR, "SPI_saveplan failed: %s",
				 SPI_result_code_string(SPI_result));
	}
	PG_CATCH();
	{
		MemoryContextSwitchTo(oldcontext);
		PLy_error_in_progress = CopyErrorData();
		FlushErrorState();
		Py_DECREF(plan);
		Py_XDECREF(optr);
		if (!PyErr_Occurred())
			PyErr_SetString(PLy_exc_spi_error,
							"Unknown error in PLy_spi_prepare");
		/* XXX this oughta be replaced with errcontext mechanism */
		PLy_elog(WARNING, "in function %s:",
				 PLy_procedure_name(PLy_curr_procedure));
		return NULL;
	}
	PG_END_TRY();

	return (PyObject *) plan;
}

/* execute(query="select * from foo", limit=5)
 * execute(plan=plan, values=(foo, bar), limit=5)
 */
static PyObject *
PLy_spi_execute(PyObject * self, PyObject * args)
{
	char	   *query;
	PyObject   *plan;
	PyObject   *list = NULL;
	long		limit = 0;

	/* Can't execute more if we have an unhandled error */
	if (PLy_error_in_progress)
	{
		PyErr_SetString(PLy_exc_error, "Transaction aborted.");
		return NULL;
	}

	if (PyArg_ParseTuple(args, "s|l", &query, &limit))
		return PLy_spi_execute_query(query, limit);

	PyErr_Clear();

	if ((PyArg_ParseTuple(args, "O|Ol", &plan, &list, &limit)) &&
		(is_PLyPlanObject(plan)))
		return PLy_spi_execute_plan(plan, list, limit);

	PyErr_SetString(PLy_exc_error, "Expected a query or plan.");
	return NULL;
}

static PyObject *
PLy_spi_execute_plan(PyObject * ob, PyObject * list, long limit)
{
	volatile int nargs;
	int			i,
				rv;
	PLyPlanObject *plan;
	MemoryContext oldcontext;

	if (list != NULL)
	{
		if ((!PySequence_Check(list)) || (PyString_Check(list)))
		{
			char	   *msg = "plpy.execute() takes a sequence as its second argument";

			PyErr_SetString(PLy_exc_spi_error, msg);
			return NULL;
		}
		nargs = PySequence_Length(list);
	}
	else
		nargs = 0;

	plan = (PLyPlanObject *) ob;

	if (nargs != plan->nargs)
	{
		char	   *sv;
		PyObject   *so = PyObject_Str(list);

		if (!so)
			PLy_elog(ERROR, "function \"%s\" could not execute plan",
					 PLy_procedure_name(PLy_curr_procedure));
		sv = PyString_AsString(so);
		PLy_exception_set(PLy_exc_spi_error,
						  "Expected sequence of %d arguments, got %d. %s",
						  plan->nargs, nargs, sv);
		Py_DECREF(so);

		return NULL;
	}

	oldcontext = CurrentMemoryContext;
	PG_TRY();
	{
		char	   *nulls = palloc(nargs * sizeof(char));

		for (i = 0; i < nargs; i++)
		{
			PyObject   *elem,
					   *so;

			elem = PySequence_GetItem(list, i);
			if (elem != Py_None)
			{
				so = PyObject_Str(elem);
				if (!so)
					PLy_elog(ERROR, "function \"%s\" could not execute plan",
							 PLy_procedure_name(PLy_curr_procedure));
				Py_DECREF(elem);

				PG_TRY();
				{
					char *sv = PyString_AsString(so);

					plan->values[i] =
						FunctionCall3(&(plan->args[i].out.d.typfunc),
									  CStringGetDatum(sv),
								ObjectIdGetDatum(plan->args[i].out.d.typioparam),
									  Int32GetDatum(-1));
				}
				PG_CATCH();
				{
					Py_DECREF(so);
					PG_RE_THROW();
				}
				PG_END_TRY();

				Py_DECREF(so);
				nulls[i] = ' ';
			}
			else
			{
				Py_DECREF(elem);
				plan->values[i] = (Datum) 0;
				nulls[i] = 'n';
			}
		}

		rv = SPI_execute_plan(plan->plan, plan->values, nulls,
							  PLy_curr_procedure->fn_readonly, limit);

		pfree(nulls);
	}
	PG_CATCH();
	{
		MemoryContextSwitchTo(oldcontext);
		PLy_error_in_progress = CopyErrorData();
		FlushErrorState();

		/*
		 * cleanup plan->values array
		 */
		for (i = 0; i < nargs; i++)
		{
			if (!plan->args[i].out.d.typbyval &&
				(plan->values[i] != (Datum) NULL))
			{
				pfree(DatumGetPointer(plan->values[i]));
				plan->values[i] = (Datum) NULL;
			}
		}

		if (!PyErr_Occurred())
			PyErr_SetString(PLy_exc_error,
							"Unknown error in PLy_spi_execute_plan");
		/* XXX this oughta be replaced with errcontext mechanism */
		PLy_elog(WARNING, "in function %s:",
				 PLy_procedure_name(PLy_curr_procedure));
		return NULL;
	}
	PG_END_TRY();

	for (i = 0; i < nargs; i++)
	{
		if (!plan->args[i].out.d.typbyval &&
			(plan->values[i] != (Datum) NULL))
		{
			pfree(DatumGetPointer(plan->values[i]));
			plan->values[i] = (Datum) NULL;