new_dbsync/pg_vtab.c

/*
 * Minimal SQLite virtual table module backed by PostgreSQL via libpq.
 *
 * Usage (after building as shared object):
 *   sqlite3> .load ./pg_vtab.so
 *   sqlite> CREATE VIRTUAL TABLE t1 USING pg_vtab('host=... dbname=... user=... password=...', 'schema.table');
 *   sqlite> SELECT * FROM t1 LIMIT 10;
 *
 * Notes:
 * - This is a minimal read-only implementation.
 * - It only supports full table scans (no constraint pushdown).
 * - All values are returned as TEXT.
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>

#include <libpq-fe.h>
#include <sqlite3.h>

/* vtab object */
typedef struct PgVTab {
    sqlite3_vtab base;    /* must be first */
    char *conninfo;
    char *schema_table;   /* qualified table name */
    int ncol;
    char **colnames;
} PgVTab;

/* cursor object */
typedef struct PgVTabCursor {
    sqlite3_vtab_cursor base; /* must be first */
    PGconn *pgconn;
    PGresult *res;
    int row;                /* current row index in res */
    sqlite3_int64 rowid;    /* generated rowid */
} PgVTabCursor;

/* forward declarations */
static int pgConnect(sqlite3 *db, void *pAux, int argc, const char *const *argv, sqlite3_vtab **ppVTab, char **pzErr);
static int pgDisconnect(sqlite3_vtab *pVTab);
static int pgBestIndex(sqlite3_vtab *pVTab, sqlite3_index_info *pIdx);
static int pgOpen(sqlite3_vtab *pVTab, sqlite3_vtab_cursor **ppCursor);
static int pgClose(sqlite3_vtab_cursor *cur);
static int pgFilter(sqlite3_vtab_cursor *pVtabCursor, int idxNum, const char *idxStr, int argc, sqlite3_value **argv);
static int pgNext(sqlite3_vtab_cursor *cur);
static int pgEof(sqlite3_vtab_cursor *cur);
static int pgColumn(sqlite3_vtab_cursor *cur, sqlite3_context *ctx, int i);
static int pgRowid(sqlite3_vtab_cursor *cur, sqlite3_int64 *pRowid);

static sqlite3_module PgModule = {
    0,                       /* iVersion */
    pgConnect,
    pgConnect,               /* xCreate == xConnect for this module */
    pgBestIndex,
    pgDisconnect,
    pgDisconnect,            /* xDestroy == xDisconnect */
    pgOpen,
    pgClose,
    pgFilter,
    pgNext,
    pgEof,
    pgColumn,
    pgRowid,
    NULL, NULL, NULL, NULL, NULL, NULL, NULL
};

/* helper: strdup-like */
static char *str_dup(const char *s){
    if(!s) return NULL;
    size_t n = strlen(s)+1;
    char *r = (char*)malloc(n);
    if(r) memcpy(r,s,n);
    return r;
}

/* read column names from POSTGRES by executing SELECT * FROM table LIMIT 0 */
static int fetch_pg_columns(const char *conninfo, const char *schema_table, int *pncol, char ***pnames, char **pErr){
    PGconn *conn = PQconnectdb(conninfo);
    if (PQstatus(conn) != CONNECTION_OK) {
        if(pErr) *pErr = str_dup(PQerrorMessage(conn));
        PQfinish(conn);
        return SQLITE_ERROR;
    }
    char q[1024];
    snprintf(q, sizeof(q), "SELECT * FROM %s LIMIT 0", schema_table);
    PGresult *res = PQexec(conn, q);
    if (PQresultStatus(res) != PGRES_TUPLES_OK && PQresultStatus(res) != PGRES_COMMAND_OK) {
        if(pErr) *pErr = str_dup(PQerrorMessage(conn));
        PQclear(res);
        PQfinish(conn);
        return SQLITE_ERROR;
    }
    int n = PQnfields(res);
    char **names = (char**)malloc(sizeof(char*)*n);
    for(int i=0;i<n;i++){
        names[i] = str_dup(PQfname(res,i));
    }
    PQclear(res);
    PQfinish(conn);
    *pncol = n;
    *pnames = names;
    return SQLITE_OK;
}

static int pgConnect(sqlite3 *db, void *pAux, int argc, const char *const *argv, sqlite3_vtab **ppVTab, char **pzErr){
    /* argv: module name, db name, conninfo, schema.table */
    if(argc < 4){
        if(pzErr) *pzErr = str_dup("pg_vtab requires connection string and qualified table name");
        return SQLITE_ERROR;
    }
    const char *conninfo = argv[2];
    const char *schema_table = argv[3];

    int ncol = 0;
    char **names = NULL;
    int rc = fetch_pg_columns(conninfo, schema_table, &ncol, &names, pzErr);
    if(rc != SQLITE_OK) return rc;

    /* build CREATE TABLE declaration with TEXT columns */
    size_t needed = 64;
    for(int i=0;i<ncol;i++) needed += strlen(names[i]) + 8;
    char *decl = (char*)malloc(needed);
    if(!decl){
        if(pzErr) *pzErr = str_dup("out of memory");
        return SQLITE_NOMEM;
    }
    strcpy(decl, "CREATE TABLE x(");
    for(int i=0;i<ncol;i++){
        if(i) strcat(decl, ",");
        strcat(decl, names[i]);
        strcat(decl, " TEXT");
    }
    strcat(decl, ")");

    rc = sqlite3_declare_vtab(db, decl);
    free(decl);
    if(rc != SQLITE_OK){
        if(pzErr) *pzErr = str_dup("sqlite3_declare_vtab failed");
        for(int i=0;i<ncol;i++) free(names[i]); free(names);
        return rc;
    }

    PgVTab *vtab = (PgVTab*)sqlite3_malloc(sizeof(PgVTab));
    if(!vtab) return SQLITE_NOMEM;
    memset(vtab,0,sizeof(PgVTab));
    vtab->conninfo = str_dup(conninfo);
    vtab->schema_table = str_dup(schema_table);
    vtab->ncol = ncol;
    vtab->colnames = names;

    *ppVTab = (sqlite3_vtab*)vtab;
    return SQLITE_OK;
}

static int pgDisconnect(sqlite3_vtab *pVTab){
    PgVTab *v = (PgVTab*)pVTab;
    if(!v) return SQLITE_OK;
    if(v->conninfo) free(v->conninfo);
    if(v->schema_table) free(v->schema_table);
    for(int i=0;i<v->ncol;i++) free(v->colnames[i]);
    if(v->colnames) free(v->colnames);
    sqlite3_free(v);
    return SQLITE_OK;
}

static int pgBestIndex(sqlite3_vtab *pVTab, sqlite3_index_info *pIdx){
    /* do not support indexes or constraints; full table scan */
    pIdx->estimatedCost = 1e6;
    return SQLITE_OK;
}

static int pgOpen(sqlite3_vtab *pVTab, sqlite3_vtab_cursor **ppCursor){
    PgVTabCursor *cur = (PgVTabCursor*)sqlite3_malloc(sizeof(PgVTabCursor));
    if(!cur) return SQLITE_NOMEM;
    memset(cur,0,sizeof(PgVTabCursor));
    *ppCursor = &cur->base;
    cur->pgconn = NULL;
    cur->res = NULL;
    cur->row = 0;
    cur->rowid = 0;
    return SQLITE_OK;
}

static int pgClose(sqlite3_vtab_cursor *curp){
    PgVTabCursor *cur = (PgVTabCursor*)curp;
    if(cur->res) PQclear(cur->res);
    if(cur->pgconn) PQfinish(cur->pgconn);
    sqlite3_free(cur);
    return SQLITE_OK;
}

static int pgFilter(sqlite3_vtab_cursor *pVtabCursor, int idxNum, const char *idxStr, int argc, sqlite3_value **argv){
    PgVTabCursor *cur = (PgVTabCursor*)pVtabCursor;
    PgVTab *vtab = (PgVTab*)(pVtabCursor->pVtab);
    /* connect to pg */
    if(cur->res){ PQclear(cur->res); cur->res = NULL; }
    if(cur->pgconn){ PQfinish(cur->pgconn); cur->pgconn = NULL; }
    cur->pgconn = PQconnectdb(vtab->conninfo);
    if(PQstatus(cur->pgconn) != CONNECTION_OK){
        /* can't report error via sqlite3_errmsg here, return error */
        return SQLITE_ERROR;
    }

    /* simple select all */
    char q[1024];
    snprintf(q, sizeof(q), "SELECT * FROM %s", vtab->schema_table);
    cur->res = PQexec(cur->pgconn, q);
    if(PQresultStatus(cur->res) != PGRES_TUPLES_OK && PQresultStatus(cur->res) != PGRES_COMMAND_OK){
        PQclear(cur->res); cur->res = NULL; return SQLITE_ERROR;
    }
    cur->row = 0;
    cur->rowid = 0;
    return SQLITE_OK;
}

static int pgNext(sqlite3_vtab_cursor *curp){
    PgVTabCursor *cur = (PgVTabCursor*)curp;
    cur->row++;
    cur->rowid++;
    return SQLITE_OK;
}

static int pgEof(sqlite3_vtab_cursor *curp){
    PgVTabCursor *cur = (PgVTabCursor*)curp;
    if(!cur->res) return 1;
    int nt = PQntuples(cur->res);
    return cur->row >= nt;
}

static int pgColumn(sqlite3_vtab_cursor *curp, sqlite3_context *ctx, int i){
    PgVTabCursor *cur = (PgVTabCursor*)curp;
    if(!cur->res) { sqlite3_result_null(ctx); return SQLITE_OK; }
    if(PQgetisnull(cur->res, cur->row, i)){
        sqlite3_result_null(ctx);
    } else {
        const char *val = PQgetvalue(cur->res, cur->row, i);
        sqlite3_result_text(ctx, val, -1, SQLITE_TRANSIENT);
    }
    return SQLITE_OK;
}

static int pgRowid(sqlite3_vtab_cursor *curp, sqlite3_int64 *pRowid){
    PgVTabCursor *cur = (PgVTabCursor*)curp;
    *pRowid = cur->rowid;
    return SQLITE_OK;
}

/* Entry point called by sqlite3 when loading extension */
int sqlite3_pg_vtab_init(sqlite3 *db, char **pzErrMsg, const sqlite3_api_routines *pApi){
    SQLITE_EXTENSION_INIT2(pApi)}
    int rc = sqlite3_create_module(db, "pg_vtab", &PgModule, NULL);
    if(rc != SQLITE_OK){
        if(pzErrMsg) *pzErrMsg = str_dup("failed to register pg_vtab module");
    }
    return rc;
}