samba-mirror/source3/python/py_tdb.c

/* 
   Python wrappers for TDB module

   Copyright (C) Tim Potter, 2002-2003
   
     ** NOTE! The following LGPL license applies to the tdb python
     ** scripting library. This does NOT imply that all of Samba is 
     ** released under the LGPL
   
   This library 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 2 of the License, or (at your option) any later version.

   This library 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.
   
   You should have received a copy of the GNU Lesser General Public
   License along with this library; if not, write to the Free Software
   Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
*/

#include "includes.h"

/* This symbol is used in both includes.h and Python.h which causes an
   annoying compiler warning. */

#ifdef HAVE_FSTAT
#undef HAVE_FSTAT
#endif

#include "Python.h"

/* Tdb exception */

PyObject *py_tdb_error;

/* tdb handle object */

typedef struct {
	PyObject_HEAD
	TDB_CONTEXT *tdb;
} tdb_hnd_object;

PyTypeObject tdb_hnd_type;
     
PyObject *new_tdb_hnd_object(TDB_CONTEXT *tdb)
{
	tdb_hnd_object *obj;

	obj = PyObject_New(tdb_hnd_object, &tdb_hnd_type);
	obj->tdb = tdb;

	return (PyObject *)obj;
}

PyObject *py_tdb_close(PyObject *self, PyObject *args)
{
	tdb_hnd_object *obj;

	if (!PyArg_ParseTuple(args, "O!", &tdb_hnd_type, &obj))
		return NULL;

	if (tdb_close(obj->tdb) == -1) {
		obj->tdb = NULL;
		PyErr_SetString(py_tdb_error, strerror(errno));
		return NULL;
	}

	obj->tdb = NULL;

	Py_INCREF(Py_None);
	return Py_None;
}

PyObject *py_tdb_open(PyObject *self, PyObject *args, PyObject *kw)
{
	static char *kwlist[] = { "name", "hash_size", "tdb_flags",
				  "open_flags", "mode", NULL };
	char *name;
	int hash_size = 0, flags = TDB_DEFAULT, open_flags = -1, open_mode = 0600;	
	TDB_CONTEXT *tdb;

	if (!PyArg_ParseTupleAndKeywords(
		    args, kw, "s|iiii", kwlist, &name, &hash_size, &flags,
		    &open_flags, &open_mode))
		return NULL;

	/* Default open_flags to read/write */

	if (open_flags == -1) {
		if (access(name, W_OK) == -1)
			open_flags = O_RDONLY;
		else
			open_flags = O_RDWR;
	}

	if (!(tdb = tdb_open(name, hash_size, flags, open_flags, open_mode))) {
		PyErr_SetString(py_tdb_error, strerror(errno));
		return NULL;
	}

	return new_tdb_hnd_object(tdb);
}

/*
 * Allow a tdb to act as a python mapping (dictionary)
 */

static int tdb_traverse_count(TDB_CONTEXT *tdb, TDB_DATA key, TDB_DATA value,
			      void *state)
{
	/* Do nothing - tdb_traverse will return the number of records
           traversed. */

	return 0;
}

static int tdb_hnd_length(tdb_hnd_object *obj)
{
	int result;

	result = tdb_traverse(obj->tdb, tdb_traverse_count, NULL);

	return result;
}

static PyObject *tdb_hnd_subscript(tdb_hnd_object *obj, PyObject *key)
{
	TDB_DATA drec, krec;
	PyObject *result;

	if (!PyArg_Parse(key, "s#", &krec.dptr, &krec.dsize))
		return NULL;

	drec = tdb_fetch(obj->tdb, krec);

	if (!drec.dptr) {
		PyErr_SetString(PyExc_KeyError,
				PyString_AsString(key));
		return NULL;
	}

	result = PyString_FromStringAndSize(drec.dptr, drec.dsize);
	free(drec.dptr);

	return result;
}
	
static int tdb_ass_subscript(tdb_hnd_object *obj, PyObject *key, PyObject *value)
{
	TDB_DATA krec, drec;

        if (!PyArg_Parse(key, "s#", &krec.dptr, &krec.dsize)) {
		PyErr_SetString(PyExc_TypeError,
				"tdb mappings have string indices only");
		return -1;
	}

        if (!obj->tdb) {
		PyErr_SetString(
			py_tdb_error, "tdb object has been closed"); 
		return -1; 
        }

	if (!value) {

		/* Delete value */

		if (tdb_delete(obj->tdb, krec) == -1) {
			PyErr_SetString(PyExc_KeyError,
					PyString_AsString(value));
			return -1;
		}

	} else {

		/* Set value */

		if (!PyArg_Parse(value, "s#", &drec.dptr, &drec.dsize)) {
			PyErr_SetString(PyExc_TypeError,
				    "tdb mappings have string elements only");
			return -1;
		}

		errno = 0;

		if (tdb_store(obj->tdb, krec, drec, 0) < 0 ) {
			if (errno != 0)
				PyErr_SetFromErrno(py_tdb_error);
			else
				PyErr_SetString(
					py_tdb_error, 
					(char *)tdb_errorstr(obj->tdb));

			return -1;
		}
	}

	return 0;
} 

static PyMappingMethods tdb_mapping = {
	(inquiry) tdb_hnd_length,
	(binaryfunc) tdb_hnd_subscript,
	(objobjargproc) tdb_ass_subscript
};

/*
 * Utility methods
 */

/* Return non-zero if a given key exists in the tdb */

PyObject *py_tdb_hnd_has_key(PyObject *self, PyObject *args)
{
	tdb_hnd_object *obj = (tdb_hnd_object *)self;
	TDB_DATA key;

	if (!PyArg_ParseTuple(args, "s#", &key.dptr, &key.dsize))
		return NULL;

        if (!obj->tdb) {
		PyErr_SetString(
			py_tdb_error, "tdb object has been closed"); 
		return NULL;
        }	

	return PyInt_FromLong(tdb_exists(obj->tdb, key));
}

/* Return a list of keys in the tdb */

static int tdb_traverse_keys(TDB_CONTEXT *tdb, TDB_DATA key, TDB_DATA value,
			     void *state)
{
	PyObject *key_list = (PyObject *)state;

	PyList_Append(key_list, 
		      PyString_FromStringAndSize(key.dptr, key.dsize));

	return 0;
}

PyObject *py_tdb_hnd_keys(PyObject *self, PyObject *args)
{
	tdb_hnd_object *obj = (tdb_hnd_object *)self;
	PyObject *key_list = PyList_New(0);

        if (!obj->tdb) {
		PyErr_SetString(py_tdb_error, "tdb object has been closed"); 
		return NULL;
        }	

	if (tdb_traverse(obj->tdb, tdb_traverse_keys, key_list) == -1) {
		PyErr_SetString(py_tdb_error, "error traversing tdb");
		Py_DECREF(key_list);
		return NULL;
	}

	return key_list;	
}

PyObject *py_tdb_hnd_first_key(PyObject *self, PyObject *args)
{
	tdb_hnd_object *obj = (tdb_hnd_object *)self;
	TDB_DATA key;

        if (!obj->tdb) {
		PyErr_SetString(py_tdb_error, "tdb object has been closed"); 
		return NULL;
        }	

	key = tdb_firstkey(obj->tdb);

	return Py_BuildValue("s#", key.dptr, key.dsize);
}

PyObject *py_tdb_hnd_next_key(PyObject *self, PyObject *py_oldkey)
{
	tdb_hnd_object *obj = (tdb_hnd_object *)self;
	TDB_DATA key, oldkey;

        if (!obj->tdb) {
		PyErr_SetString(py_tdb_error, "tdb object has been closed"); 
		return NULL;
        }	

	if (!PyArg_Parse(py_oldkey, "s#", &oldkey.dptr, &oldkey.dsize))
		return NULL;

	key = tdb_nextkey(obj->tdb, oldkey);

	return Py_BuildValue("s#", key.dptr, key.dsize);
}

/*
 * Locking routines
 */

PyObject *py_tdb_hnd_lock_all(PyObject *self, PyObject *args)
{
	tdb_hnd_object *obj = (tdb_hnd_object *)self;
	int result;

        if (!obj->tdb) {
		PyErr_SetString(py_tdb_error, "tdb object has been closed"); 
		return NULL;
        }	

	result = tdb_lockall(obj->tdb);

	return PyInt_FromLong(result != -1);
}

PyObject *py_tdb_hnd_unlock_all(PyObject *self, PyObject *args)
{
	tdb_hnd_object *obj = (tdb_hnd_object *)self;

        if (!obj->tdb) {
		PyErr_SetString(py_tdb_error, "tdb object has been closed"); 
		return NULL;
        }	

	tdb_unlockall(obj->tdb);

	Py_INCREF(Py_None);
	return Py_None;
}

/* Return an array of keys from a python object which must be a string or a
   list of strings. */

static BOOL make_lock_list(PyObject *py_keys, TDB_DATA **keys, int *num_keys)
{
	/* Are we a list or a string? */

	if (!PyList_Check(py_keys) && !PyString_Check(py_keys)) {
		PyErr_SetString(PyExc_TypeError, "arg must be list of string");
		return False;
	}

	if (PyList_Check(py_keys)) {
		int i;

		/* Turn python list into array of keys */
		
		*num_keys = PyList_Size(py_keys);
		*keys = (TDB_DATA *)SMB_XMALLOC_ARRAY(TDB_DATA, (*num_keys));
		
		for (i = 0; i < *num_keys; i++) {
			PyObject *key = PyList_GetItem(py_keys, i);
			
			if (!PyString_Check(key)) {
				PyErr_SetString(
					PyExc_TypeError,
					"list elements must be strings");
				return False;
			}

			PyArg_Parse(key, "s#", &(*keys)[i].dptr, 
				    &(*keys)[i].dsize);
		}

	} else {

		/* Turn python string into a single key */

		*keys = (TDB_DATA *)SMB_XMALLOC_P(TDB_DATA);
		*num_keys = 1;
		PyArg_Parse(py_keys, "s#", &(*keys)->dptr, &(*keys)->dsize);
	}

	return True;
}

/*
 * tdb traversal
 */

struct traverse_info {
	PyObject *callback;
	PyObject *state;
};

static int tdb_traverse_traverse(TDB_CONTEXT *tdb, TDB_DATA key, TDB_DATA value,
				 void *state)
{
	struct traverse_info *info = state;
	PyObject *arglist, *py_result;
	int result;

	arglist = Py_BuildValue("(s#s#O)", key.dptr, key.dsize, value.dptr,
				value.dsize, info->state);

	py_result = PyEval_CallObject(info->callback, arglist);

	Py_DECREF(arglist);
	
	if (!PyInt_Check(py_result)) {
		result = 1;	/* Hmm - non-integer object returned by callback */
		goto done;
	}

	result = PyInt_AsLong(py_result);

done:
	Py_DECREF(py_result);
	return result;
}

PyObject *py_tdb_hnd_traverse(PyObject *self, PyObject *args, PyObject *kw)
{
	tdb_hnd_object *obj = (tdb_hnd_object *)self;
	static char *kwlist[] = { "traverse_fn", "state", NULL };
	PyObject *state = Py_None, *callback;
	struct traverse_info info;
	int result;

	if (!PyArg_ParseTupleAndKeywords(
		    args, kw, "O|O", kwlist, &callback, &state))
		return NULL;

	if (!PyCallable_Check(callback)) {
		PyErr_SetString(PyExc_TypeError, "parameter must be callable");
		return NULL;
        }

	Py_INCREF(callback);
	Py_INCREF(state);

	info.callback = callback;
	info.state = state;

	result = tdb_traverse(obj->tdb, tdb_traverse_traverse, &info);

	Py_DECREF(callback);
	Py_DECREF(state);

	return PyInt_FromLong(result);
}

PyObject *py_tdb_hnd_chainlock(PyObject *self, PyObject *args)
{
	tdb_hnd_object *obj = (tdb_hnd_object *)self;
	TDB_DATA key;
	int result;

        if (!obj->tdb) {
		PyErr_SetString(py_tdb_error, "tdb object has been closed"); 
		return NULL;
        }	

	if (!PyArg_ParseTuple(args, "s#", &key.dptr, &key.dsize))
		return NULL;

	result = tdb_chainlock(obj->tdb, key);

	return PyInt_FromLong(result != -1);
}

PyObject *py_tdb_hnd_chainunlock(PyObject *self, PyObject *args)
{
	tdb_hnd_object *obj = (tdb_hnd_object *)self;
	TDB_DATA key;
	int result;

        if (!obj->tdb) {
		PyErr_SetString(py_tdb_error, "tdb object has been closed"); 
		return NULL;
        }	

	if (!PyArg_ParseTuple(args, "s#", &key.dptr, &key.dsize))
		return NULL;

	result = tdb_chainunlock(obj->tdb, key);

	return PyInt_FromLong(result != -1);
}

PyObject *py_tdb_hnd_lock_bystring(PyObject *self, PyObject *args)
{
	tdb_hnd_object *obj = (tdb_hnd_object *)self;
	int result, timeout = 30;
	char *s;

        if (!obj->tdb) {
		PyErr_SetString(py_tdb_error, "tdb object has been closed"); 
		return NULL;
        }	

	if (!PyArg_ParseTuple(args, "s|i", &s, &timeout))
		return NULL;

	result = tdb_lock_bystring(obj->tdb, s, timeout);

	return PyInt_FromLong(result != -1);
}

PyObject *py_tdb_hnd_unlock_bystring(PyObject *self, PyObject *args)
{
	tdb_hnd_object *obj = (tdb_hnd_object *)self;
	char *s;

        if (!obj->tdb) {
		PyErr_SetString(py_tdb_error, "tdb object has been closed"); 
		return NULL;
        }	

	if (!PyArg_ParseTuple(args, "s", &s))
		return NULL;

	tdb_unlock_bystring(obj->tdb, s);

	Py_INCREF(Py_None);
	return Py_None;
}

/* 
 * Method dispatch table for this module
 */

static PyMethodDef tdb_methods[] = {
	{ "open", (PyCFunction)py_tdb_open, METH_VARARGS | METH_KEYWORDS },
	{ "close", (PyCFunction)py_tdb_close, METH_VARARGS },
	{ NULL }
};

/* 
 * Methods on a tdb object
 */

static PyMethodDef tdb_hnd_methods[] = {
	{ "keys", (PyCFunction)py_tdb_hnd_keys, METH_VARARGS },
	{ "has_key", (PyCFunction)py_tdb_hnd_has_key, METH_VARARGS },
	{ "first_key", (PyCFunction)py_tdb_hnd_first_key, METH_VARARGS },
	{ "next_key", (PyCFunction)py_tdb_hnd_next_key, METH_VARARGS },
	{ "lock_all", (PyCFunction)py_tdb_hnd_lock_all, METH_VARARGS },
	{ "unlock_all", (PyCFunction)py_tdb_hnd_unlock_all, METH_VARARGS },
	{ "traverse", (PyCFunction)py_tdb_hnd_traverse, METH_VARARGS | METH_KEYWORDS },
	{ "chainlock", (PyCFunction)py_tdb_hnd_chainlock, METH_VARARGS | METH_KEYWORDS },
	{ "chainunlock", (PyCFunction)py_tdb_hnd_chainunlock, METH_VARARGS | METH_KEYWORDS },
	{ "lock_bystring", (PyCFunction)py_tdb_hnd_lock_bystring, METH_VARARGS | METH_KEYWORDS },
	{ "unlock_bystring", (PyCFunction)py_tdb_hnd_unlock_bystring, METH_VARARGS | METH_KEYWORDS },
	{ NULL }
};

/* Deallocate a tdb handle object */

static void tdb_hnd_dealloc(PyObject* self)
{
        tdb_hnd_object *hnd = (tdb_hnd_object *)self;

	if (hnd->tdb) {
		tdb_close(hnd->tdb);
		hnd->tdb = NULL;
	}
}

/* Return tdb handle attributes */

static PyObject *tdb_hnd_getattr(PyObject *self, char *attrname)
{
	return Py_FindMethod(tdb_hnd_methods, self, attrname);
}

static char tdb_hnd_type_doc[] = 
"Python wrapper for tdb.";

PyTypeObject tdb_hnd_type = {
	PyObject_HEAD_INIT(NULL)
	0,
	"tdb",
	sizeof(tdb_hnd_object),
	0,
	tdb_hnd_dealloc,	/* tp_dealloc*/
	0,			/* tp_print*/
	tdb_hnd_getattr,	/* tp_getattr*/
	0,			/* tp_setattr*/
	0,			/* tp_compare*/
	0,			/* tp_repr*/
	0,			/* tp_as_number*/
	0,			/* tp_as_sequence*/
	&tdb_mapping,		/* tp_as_mapping*/
	0,			/* tp_hash */
	0,			/* tp_call */
	0,			/* tp_str */
	0,			/* tp_getattro */
	0,			/* tp_setattro */
	0,			/* tp_as_buffer*/
	Py_TPFLAGS_DEFAULT,	/* tp_flags */
	tdb_hnd_type_doc,	/* tp_doc */
};

/* Constants */

static struct const_vals {
	char *name;
	uint32 value;
} module_const_vals[] = {

        /* Flags for tdb_open() */

	{ "TDB_DEFAULT", TDB_DEFAULT },
	{ "TDB_CLEAR_IF_FIRST", TDB_CLEAR_IF_FIRST },
	{ "TDB_INTERNAL", TDB_INTERNAL },
	{ "TDB_NOLOCK", TDB_NOLOCK },
	{ "TDB_NOMMAP", TDB_NOMMAP },
	{ "TDB_CONVERT", TDB_CONVERT },
	{ "TDB_BIGENDIAN", TDB_BIGENDIAN },
	
	{ NULL },
};

static void const_init(PyObject *dict)
{
	struct const_vals *tmp;
	PyObject *obj;

	for (tmp = module_const_vals; tmp->name; tmp++) {
		obj = PyInt_FromLong(tmp->value);
		PyDict_SetItemString(dict, tmp->name, obj);
		Py_DECREF(obj);
	}
}

/* Module initialisation */

void inittdb(void)
{
	PyObject *module, *dict;

	/* Initialise module */

	module = Py_InitModule("tdb", tdb_methods);
	dict = PyModule_GetDict(module);

	py_tdb_error = PyErr_NewException("tdb.error", NULL, NULL);
	PyDict_SetItemString(dict, "error", py_tdb_error);

	/* Initialise policy handle object */

	tdb_hnd_type.ob_type = &PyType_Type;

	PyDict_SetItemString(dict, "tdb.hnd", 
			     (PyObject *)&tdb_hnd_type);

	/* Initialise constants */

	const_init(dict);
}