--- /dev/null
+/*
+** 2019-04-26
+**
+** The author disclaims copyright to this source code. In place of
+** a legal notice, here is a blessing:
+**
+** May you do good and not evil.
+** May you find forgiveness for yourself and forgive others.
+** May you share freely, never taking more than you give.
+**
+*************************************************************************
+**
+** This file implements a virtual-table that can be used to access a
+** sharded table implemented as the UNION ALL of various separate tables.
+*/
+#if !defined(SQLITEINT_H)
+#include "sqlite3ext.h"
+#endif
+SQLITE_EXTENSION_INIT1
+#include <string.h>
+#include <assert.h>
+#include <math.h>
+
+/* shardvtab_vtab is a subclass of sqlite3_vtab which is
+** underlying representation of the virtual table
+*/
+typedef struct shardvtab_vtab shardvtab_vtab;
+struct shardvtab_vtab {
+ sqlite3_vtab base; /* Base class - must be first */
+ sqlite3 *db; /* The database connection */
+ char *zView; /* Name of view that implements the shard */
+ int nCol; /* Number of columns in the view */
+ char **azCol; /* Names of the columns, individually malloced */
+};
+
+/* shardvtab_cursor is a subclass of sqlite3_vtab_cursor which will
+** serve as the underlying representation of a cursor that scans
+** over rows of the result
+*/
+typedef struct shardvtab_cursor shardvtab_cursor;
+struct shardvtab_cursor {
+ sqlite3_vtab_cursor base; /* Base class - must be first */
+ sqlite3_stmt *pStmt; /* Prepared statement to access the shard */
+ int rcLastStep; /* Last return from sqlite3_step() */
+};
+
+/*
+** The shardvtabConnect() method is invoked to create a new
+** shard virtual table.
+**
+** Think of this routine as the constructor for shardvtab_vtab objects.
+**
+** All this routine needs to do is:
+**
+** (1) Allocate the shardvtab_vtab object and initialize all fields.
+**
+** (2) Tell SQLite (via the sqlite3_declare_vtab() interface) what the
+** result set of queries against the virtual table will look like.
+*/
+static int shardvtabConnect(
+ sqlite3 *db,
+ void *pAux,
+ int argc, const char *const*argv,
+ sqlite3_vtab **ppVtab,
+ char **pzErr
+){
+ shardvtab_vtab *pNew;
+ int rc;
+ char *zSql;
+ sqlite3_str *pSchema;
+ sqlite3_stmt *pStmt = 0;
+ const char *zView = 0;
+ char **azCol = 0;
+ int nCol = 0;
+ char cSep;
+ int i;
+
+ if( argc!=4 || argv[0]==0 ){
+ *pzErr = sqlite3_mprintf("one argument requires: the name of a view");
+ return SQLITE_ERROR;
+ }
+ zView = argv[3];
+ zSql = sqlite3_mprintf("SELECT * FROM \"%w\"", zView);
+ if( zSql==0 ){
+ return SQLITE_NOMEM;
+ }
+ rc = sqlite3_prepare_v2(db, zSql, -1, &pStmt, 0);
+ sqlite3_free(zSql);
+ if( rc ){
+ *pzErr = sqlite3_mprintf("not a valid view: \"%w\"", zView);
+ return SQLITE_NOMEM;
+ }
+ pSchema = sqlite3_str_new(db);
+ if( pSchema==0 ){
+ sqlite3_finalize(pStmt);
+ return SQLITE_NOMEM;
+ }
+ sqlite3_str_appendall(pSchema, "CREATE TABLE x");
+ cSep = '(';
+ for(i=0; i<sqlite3_column_count(pStmt); i++){
+ const char *zName = sqlite3_column_name(pStmt,i);
+ char **azNew = sqlite3_realloc64(azCol, sizeof(azCol[0])*(i+1));
+ if( azNew==0 ){
+ rc = SQLITE_NOMEM;
+ goto shardvtab_connect_error;
+ }
+ sqlite3_str_appendf(pSchema, "%c\"%w\"", cSep, zName);
+ cSep = ',';
+ azCol = azNew;
+ azCol[nCol] = sqlite3_mprintf("%s", zName);
+ if( azCol[nCol]==0 ){
+ rc = SQLITE_NOMEM;
+ goto shardvtab_connect_error;
+ }
+ nCol++;
+ }
+ sqlite3_str_appendall(pSchema, ")");
+ sqlite3_finalize(pStmt);
+ pStmt = 0;
+ zSql = sqlite3_str_finish(pSchema);
+ pSchema = 0;
+ if( zSql==0 ){
+ rc = SQLITE_NOMEM;
+ goto shardvtab_connect_error;
+ }
+ rc = sqlite3_declare_vtab(db, zSql);
+ sqlite3_free(zSql);
+ if( rc!=SQLITE_OK ){
+ goto shardvtab_connect_error;
+ }else{
+ size_t n = strlen(zView) + 1;
+ pNew = sqlite3_malloc64( sizeof(*pNew) + n );
+ *ppVtab = (sqlite3_vtab*)pNew;
+ if( pNew==0 ){
+ rc = SQLITE_NOMEM;
+ goto shardvtab_connect_error;
+ }
+ memset(pNew, 0, sizeof(*pNew));
+ pNew->db = db;
+ pNew->zView = (char*)&pNew[1];
+ memcpy(pNew->zView, zView, n);
+ pNew->nCol = nCol;
+ pNew->azCol = azCol;
+ }
+ return SQLITE_OK;
+
+shardvtab_connect_error:
+ sqlite3_finalize(pStmt);
+ for(i=0; i<nCol; i++) sqlite3_free(azCol[i]);
+ sqlite3_free(azCol);
+ sqlite3_free(sqlite3_str_finish(pSchema));
+ return rc;
+}
+
+/*
+** This method is the destructor for shardvtab_vtab objects.
+*/
+static int shardvtabDisconnect(sqlite3_vtab *pVtab){
+ int i;
+ shardvtab_vtab *p = (shardvtab_vtab*)pVtab;
+ for(i=0; i<p->nCol; i++) sqlite3_free(p->azCol[i]);
+ sqlite3_free(p->azCol);
+ sqlite3_free(p);
+ return SQLITE_OK;
+}
+
+/*
+** Constructor for a new shardvtab_cursor object.
+*/
+static int shardvtabOpen(sqlite3_vtab *p, sqlite3_vtab_cursor **ppCursor){
+ shardvtab_cursor *pCur;
+ pCur = sqlite3_malloc( sizeof(*pCur) );
+ if( pCur==0 ) return SQLITE_NOMEM;
+ memset(pCur, 0, sizeof(*pCur));
+ *ppCursor = &pCur->base;
+ return SQLITE_OK;
+}
+
+/*
+** Destructor for a shardvtab_cursor.
+*/
+static int shardvtabClose(sqlite3_vtab_cursor *cur){
+ shardvtab_cursor *pCur = (shardvtab_cursor*)cur;
+ sqlite3_finalize(pCur->pStmt);
+ sqlite3_free(pCur);
+ return SQLITE_OK;
+}
+
+
+/*
+** Advance a shardvtab_cursor to its next row of output.
+*/
+static int shardvtabNext(sqlite3_vtab_cursor *cur){
+ shardvtab_cursor *pCur = (shardvtab_cursor*)cur;
+ int rc;
+ rc = pCur->rcLastStep = sqlite3_step(pCur->pStmt);
+ if( rc==SQLITE_ROW || rc==SQLITE_DONE ) return SQLITE_OK;
+ return rc;
+}
+
+/*
+** Return values of columns for the row at which the shardvtab_cursor
+** is currently pointing.
+*/
+static int shardvtabColumn(
+ sqlite3_vtab_cursor *cur, /* The cursor */
+ sqlite3_context *ctx, /* First argument to sqlite3_result_...() */
+ int i /* Which column to return */
+){
+ shardvtab_cursor *pCur = (shardvtab_cursor*)cur;
+ sqlite3_result_value(ctx, sqlite3_column_value(pCur->pStmt, i));
+ return SQLITE_OK;
+}
+
+/*
+** Return the rowid for the current row. In this implementation, the
+** rowid is the same as the output value.
+*/
+static int shardvtabRowid(sqlite3_vtab_cursor *cur, sqlite_int64 *pRowid){
+ *pRowid = 0;
+ return SQLITE_OK;
+}
+
+/*
+** Return TRUE if the cursor has been moved off of the last
+** row of output.
+*/
+static int shardvtabEof(sqlite3_vtab_cursor *cur){
+ shardvtab_cursor *pCur = (shardvtab_cursor*)cur;
+ return pCur->rcLastStep!=SQLITE_ROW;
+}
+
+/*
+** This method is called to "rewind" the shardvtab_cursor object back
+** to the first row of output. This method is always called at least
+** once prior to any call to shardvtabColumn() or shardvtabRowid() or
+** shardvtabEof().
+*/
+static int shardvtabFilter(
+ sqlite3_vtab_cursor *pVtabCursor,
+ int idxNum, const char *idxStr,
+ int argc, sqlite3_value **argv
+){
+ shardvtab_cursor *pCur = (shardvtab_cursor *)pVtabCursor;
+ shardvtab_vtab *pTab = (shardvtab_vtab *)pVtabCursor->pVtab;
+ int rc;
+ sqlite3_finalize(pCur->pStmt);
+ pCur->pStmt = 0;
+ rc = sqlite3_prepare_v2(pTab->db, idxStr, -1, &pCur->pStmt, 0);
+ if( rc==SQLITE_OK ){
+ int i;
+ for(i=0; i<argc; i++){
+ sqlite3_bind_value(pCur->pStmt, i+1, argv[i]);
+ }
+ }else{
+ sqlite3_finalize(pCur->pStmt);
+ pCur->pStmt = 0;
+ }
+ pCur->rcLastStep = rc;
+ return rc;
+}
+
+/*
+** SQLite will invoke this method one or more times while planning a query
+** that uses the virtual table. This routine needs to create
+** a query plan for each invocation and compute an estimated cost for that
+** plan.
+*/
+static int shardvtabBestIndex(
+ sqlite3_vtab *tab,
+ sqlite3_index_info *p
+){
+ shardvtab_vtab *pTab = (shardvtab_vtab*)tab;
+ int i;
+ int n;
+ sqlite3_stmt *pStmt;
+ int rc;
+ sqlite3_str *pSql;
+ char *zSep = "WHERE";
+ char *zSql;
+ pSql = sqlite3_str_new(pTab->db);
+ if( pSql==0 ) return SQLITE_NOMEM;
+ sqlite3_str_appendf(pSql, "SELECT * FROM \"%w\"", pTab->zView);
+ for(i=n=0; i<p->nConstraint; i++){
+ const char *zOp;
+ int iCol;
+ if( p->aConstraint[i].usable==0 ) continue;
+ iCol = p->aConstraint[i].iColumn;
+ if( iCol<0 ) continue;
+ zOp = 0;
+ switch( p->aConstraint[i].op ){
+ case SQLITE_INDEX_CONSTRAINT_EQ: zOp = "=="; break;
+ case SQLITE_INDEX_CONSTRAINT_GT: zOp = ">"; break;
+ case SQLITE_INDEX_CONSTRAINT_LE: zOp = "<="; break;
+ case SQLITE_INDEX_CONSTRAINT_LT: zOp = "<"; break;
+ case SQLITE_INDEX_CONSTRAINT_GE: zOp = ">="; break;
+ case SQLITE_INDEX_CONSTRAINT_MATCH: zOp = "MATCH"; break;
+ case SQLITE_INDEX_CONSTRAINT_LIKE: zOp = "LIKE"; break;
+ case SQLITE_INDEX_CONSTRAINT_GLOB: zOp = "GLOB"; break;
+ case SQLITE_INDEX_CONSTRAINT_REGEXP: zOp = "REGEXP"; break;
+ case SQLITE_INDEX_CONSTRAINT_NE: zOp = "<>"; break;
+ case SQLITE_INDEX_CONSTRAINT_IS: zOp = "IS"; break;
+ }
+ if( zOp ){
+ n++;
+ p->aConstraintUsage[i].argvIndex = n;
+ sqlite3_str_appendf(pSql, " %s (\"%w\" %s ?%d)",
+ zSep, pTab->azCol[iCol], zOp, n);
+ zSep = "AND";
+ }
+ }
+ zSql = sqlite3_str_finish(pSql);
+ if( zSql==0 ){
+ return SQLITE_NOMEM;
+ }
+ rc = sqlite3_prepare_v2(pTab->db, zSql, -1, &pStmt, 0);
+ if( rc==SQLITE_OK ){
+ int x = sqlite3_stmt_status(pStmt, SQLITE_STMTSTATUS_EST_COST, 0);
+ p->estimatedCost = pow(2.0, 0.1*x);
+ p->estimatedRows =
+ sqlite3_stmt_status(pStmt, SQLITE_STMTSTATUS_EST_ROWS, 0);
+ p->idxStr = zSql;
+ p->needToFreeIdxStr = 1;
+ }else{
+ sqlite3_free(zSql);
+ }
+ sqlite3_finalize(pStmt);
+ return rc;
+}
+
+/*
+** This following structure defines all the methods for the
+** virtual table.
+*/
+static sqlite3_module shardvtabModule = {
+ /* iVersion */ 0,
+ /* xCreate */ shardvtabConnect,
+ /* xConnect */ shardvtabConnect,
+ /* xBestIndex */ shardvtabBestIndex,
+ /* xDisconnect */ shardvtabDisconnect,
+ /* xDestroy */ shardvtabDisconnect,
+ /* xOpen */ shardvtabOpen,
+ /* xClose */ shardvtabClose,
+ /* xFilter */ shardvtabFilter,
+ /* xNext */ shardvtabNext,
+ /* xEof */ shardvtabEof,
+ /* xColumn */ shardvtabColumn,
+ /* xRowid */ shardvtabRowid,
+ /* xUpdate */ 0,
+ /* xBegin */ 0,
+ /* xSync */ 0,
+ /* xCommit */ 0,
+ /* xRollback */ 0,
+ /* xFindMethod */ 0,
+ /* xRename */ 0,
+ /* xSavepoint */ 0,
+ /* xRelease */ 0,
+ /* xRollbackTo */ 0,
+ /* xShadowName */ 0
+};
+
+
+#ifdef _WIN32
+__declspec(dllexport)
+#endif
+int sqlite3_shardvtab_init(
+ sqlite3 *db,
+ char **pzErrMsg,
+ const sqlite3_api_routines *pApi
+){
+ int rc = SQLITE_OK;
+ SQLITE_EXTENSION_INIT2(pApi);
+ rc = sqlite3_create_module(db, "shardvtab", &shardvtabModule, 0);
+ return rc;
+}