datastore_sqlite.c

一个C语言写的快速贝叶斯垃圾邮件过滤工具

/**
 * \file datastore_sqlite.c SQLite 3 database driver back-end
 * \author Matthias Andree <[e-mail]>
 * \date 2004, 2005
 *
 * This file handles a static table named "bogofilter" in a SQLite3
 * database. The table has two "BLOB"-typed columns, key and value.
 *
 * GNU GENERAL PUBLIC LICENSE v2
 */

#include "common.h"

#include <errno.h>
#include <sqlite3.h>

#include "datastore_db.h"

#include "error.h"
#include "rand_sleep.h"
#include "xmalloc.h"
#include "xstrdup.h"

/** Structure to hold database handle and associated data. */
struct dbhsqlite_t {
    char *path;	   /**< directory to hold database */
    char *name;	   /**< database file name */
    sqlite3 *db;   /**< pointer to SQLite3 handle */
    sqlite3_stmt *select; /**< prepared SELECT statement for DB retrieval */
    sqlite3_stmt *insert; /**< prepared INSERT OR REPLACE for DB update */
    sqlite3_stmt *delete; /**< prepared DELETE statement */
    bool created;  /**< gets set by db_open if it created the database new */
    bool swapped;  /**< if endian swapped on disk vs. current host */
};

/** Convenience shortcut to avoid typing "struct dbh_t" */
typedef struct dbhsqlite_t dbh_t;

static const char *ENDIAN32 = ".ENDIAN32";

void db_flush(void *unused) { (void)unused; }

static int sql_txn_begin(void *vhandle);
static int sql_txn_abort(void *vhandle);
static int sql_txn_commit(void *vhandle);
static u_int32_t sql_pagesize(bfpath *bfp);

/** The layout of the bogofilter table, formatted as SQL statement.
 *
 * The additional index, although making writes a bit slower, speeds up
 * queries noticably as it improves locality of referenced data and
 * reduces complexity of the retrieval of the value column.
 */
#define LAYOUT \
	"CREATE TABLE bogofilter (" \
	"   key   BLOB PRIMARY KEY," \
	"   value BLOB);" \
	"CREATE INDEX bfidx ON bogofilter(key,value);"
/*
 * another experimental layout is as follows,
 * but does not appear to make a lot of difference
 * performance-wise (evaluation in other environments
 * is required though):
 *
#define LAYOUT \
    "CREATE TABLE bogofilter (key BLOB, value BLOB); " \
    "CREATE INDEX bfidx ON bogofilter(key,value);" \
    "CREATE TRIGGER bfuniquekey BEFORE INSERT ON bogofilter " \
    " FOR EACH ROW WHEN EXISTS(SELECT key FROM bogofilter WHERE (key=NEW.key) LIMIT 1) " \
    " BEGIN UPDATE bogofilter SET value=NEW.value WHERE (key=NEW.key); SELECT RAISE(IGNORE); END;"
#endif
 */

dsm_t dsm_sqlite = {
    /* public -- used in datastore.c */
    &sql_txn_begin,
    &sql_txn_abort,
    &sql_txn_commit,

    /* private -- used in datastore_db_*.c */
    NULL,	/* dsm_env_init         */
    NULL,	/* dsm_cleanup          */
    NULL,	/* dsm_cleanup_lite     */
    NULL,	/* dsm_get_env_dbe      */
    NULL,	/* dsm_database_name    */
    NULL,	/* dsm_recover_open     */
    NULL,	/* dsm_auto_commit_flags*/
    NULL,	/* dsm_get_rmw_flag     */
    NULL,	/* dsm_lock             */
    NULL,	/* dsm_common_close     */
    NULL,	/* dsm_sync             */
    NULL,	/* dsm_log_flush        */
    &sql_pagesize,/* dsm_pagesize       */
    NULL,	/* dsm_purgelogs        */
    NULL,	/* dsm_checkpoint       */
    NULL,	/* dsm_recover          */
    NULL,	/* dsm_remove           */
    NULL,	/* dsm_verify           */
    NULL,	/* dsm_list_logfiles    */
    NULL	/* dsm_leafpages        */
};

dsm_t *dsm = &dsm_sqlite;

/** The command to begin a regular transaction. */
#define BEGIN \
	"BEGIN TRANSACTION;"

/* real functions */
/** Initialize database handle and return it.
 * \returns non-NULL, as it exits with EX_ERROR in case of trouble. */
static dbh_t *dbh_init(bfpath *bfp)
{
    dbh_t *handle;

    dsm = &dsm_sqlite;

    handle = xmalloc(sizeof(dbh_t));
    memset(handle, 0, sizeof(dbh_t));

    handle->name = xstrdup(bfp->filepath);

    return handle;
}

/** Free internal database handle \a dbh. */
static void free_dbh(dbh_t *dbh) {
    if (!dbh)
	return;
    xfree(dbh->name);
    xfree(dbh->path);
    xfree(dbh);
}

/** Executes the SQL statement \a cmd on the database \a db and returns
 * the sqlite3_exec return code. If the return code is nonzero, this
 * routine will have printed an error message.
 */
static int sqlexec(sqlite3 *db, const char *cmd) {
    char *e = NULL;
    int rc;

    rc = sqlite3_exec(db, cmd, NULL, NULL, &e);
    if (rc) {
	print_error(__FILE__, __LINE__,
		"Error executing \"%s\": %s (#%d)\n",
		cmd, e ? e : "NULL", rc);
	if (e)
	    sqlite3_free(e);
    }
    return rc;
}

static sqlite3_stmt *sqlprep(dbh_t *dbh, const char *cmd, bool bailout /** exit on error? */) {
    const char *tail; /* dummy */
    sqlite3_stmt *ptr;
    if (sqlite3_prepare(dbh->db, cmd, strlen(cmd), &ptr, &tail) != SQLITE_OK) {
	print_error(__FILE__, __LINE__, "cannot compile %s: %s\n", cmd, sqlite3_errmsg(dbh->db));
	if (bailout)
	    exit(EX_ERROR);
	return NULL;
    }
    return ptr;
}

/** Short trace handler function, passed to SQLite if debugging is
 * enabled. */
static void db_trace(void *userdata /** unused */,
	const char *log /** log message */) {
    (void)userdata;
    fprintf(dbgout, "SQLite[%ld]: %s\n", (long)getpid(), log);
}

/** Foreach function, we call \a hook for
 * each (key, value) tuple in the database.
 */
static int db_loop(sqlite3 *db,	/**< SQLite3 database handle */
	const char *cmd,	/**< SQL command to obtain data */
	db_foreach_t hook,	/**< if non-NULL, called for each value */
	void *userdata		/**  this is passed to the \a hook */
	) {
    const char *tail;
    sqlite3_stmt *stmt;
    int rc;
    bool loop, found = false;
    dbv_t key, val;

    /* sqlite3_exec doesn't allow us to retrieve BLOBs */
    rc = sqlite3_prepare(db, cmd, strlen(cmd), &stmt, &tail);
    if (rc) {
	print_error(__FILE__, __LINE__,
		"Error preparing \"%s\": %s (#%d)\n",
		cmd, sqlite3_errmsg(db), rc);
	sqlite3_finalize(stmt);
	return rc;
    }
    loop = true;
    while (loop) {
	rc = sqlite3_step(stmt);
	switch (rc) {
	    case SQLITE_ROW:
		found = true;
		if (hook != NULL)
		{
		    key.leng = sqlite3_column_bytes(stmt, /* column */ 0);
		    key.data = xmalloc(key.leng);
		    memcpy(key.data, sqlite3_column_blob(stmt, 0), key.leng);

		    val.leng = sqlite3_column_bytes(stmt, /* column */ 1);
		    val.data = xmalloc(val.leng);
		    memcpy(val.data, sqlite3_column_blob(stmt, 1), val.leng);

		    /* skip ENDIAN32 token */
		    if (key.leng != strlen(ENDIAN32)
			    || memcmp(key.data, ENDIAN32, key.leng) != 0)
			rc = hook(&key, &val, userdata);
		    else
			rc = 0;

		    xfree(val.data);
		    xfree(key.data);
		    if (rc) {
			sqlite3_finalize(stmt);
			return rc;
		    }
		}
		break;
	    case SQLITE_DONE:
		loop = false;
		break;
	    default:
		print_error(__FILE__, __LINE__, "Error executing \"%s\": %s (#%d)\n",
		cmd, sqlite3_errmsg(db), rc);
		sqlite3_finalize(stmt);
		return rc;
	}
    }
    /* free resources */
    sqlite3_finalize(stmt);
    return found ? 0 : DS_NOTFOUND;
}

/** This busy handler just sleeps a while and retries */
static int busyhandler(void *dummy, int count)
{
    (void)dummy;
    (void)count;
    rand_sleep(1000, 1000000);
    return 1;
}

static void check_sqlite_version(void)
{
    unsigned int vmaj, vmin, vpl;
    static int complained;
    const char *v;

    if (complained)
	return;
    complained = 1;
    v = sqlite3_libversion();
    sscanf(v, "%u.%u.%u", &vmaj, &vmin, &vpl);
    if (vmaj > 3) return;
    if (vmaj == 3 && vmin > 2) return;
    if (vmaj == 3 && vmin == 2 && vpl >= 6) return;
    if (!getenv("BF_USE_OLD_SQLITE"))
	fprintf(stderr,
		"\n"
		"WARNING: please update sqlite to 3.2.6 or newer.\n"
		"\n");
}

void *db_open(void *dummyenv, bfpath *bfp, dbmode_t mode)
{
    int rc;
    dbh_t *dbh;
    dbv_t k, v;

    (void)dummyenv;

    check_sqlite_version();

    dbh = dbh_init(bfp);

    /* open database file */
    if (DEBUG_DATABASE(1) || getenv("BF_DEBUG_DB")) {
	fprintf(dbgout, "SQLite: db_open(%s)\n", dbh->name);
	fflush(dbgout);
    }