python/cpython#93421: simulate DB API rowcount attribute

CMakeLists.txt

cmake_minimum_required(VERSION 3.14)

project(
  TestSQLite
  VERSION 0.1.0
  LANGUAGES C
)

find_package(SQLite3 REQUIRED)

add_executable(test main.c)
set_property(TARGET test PROPERTY C_STANDARD 99)
if (NOT ${SQLite3_FOUND})
  message(FATAL_ERROR "Unable to find SQLite3")
endif()
target_include_directories(test PRIVATE ${SQLite3_INCLUDE_DIRS})
target_link_libraries(test sqlite3)

main.c

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

#include <sqlite3.h>

typedef struct {
    sqlite3_stmt *stmt;
    const char *query;
} stmt;
static stmt *stmt_cache = NULL;

void print_error(sqlite3 *db) {
    fprintf(stderr, "rc: (%d) %s\n", sqlite3_errcode(db), sqlite3_errmsg(db));
}

sqlite3_stmt *fetch_stmt(sqlite3 *db, const char *query) {
    int i = 0;
    while (stmt_cache[i].query != NULL) {
        if (strcmp(stmt_cache[i].query, query) == 0) {
            printf("(cached)");
            return stmt_cache[i].stmt;
        }
        i++;
    }
    stmt_cache[i].query = query;
    int rc = sqlite3_prepare_v2(db, query, -1, &stmt_cache[i].stmt, NULL);
    assert(rc == SQLITE_OK);
    return stmt_cache[i].stmt;
}

sqlite3_stmt *new_stmt(sqlite3 *db, const char *query) {
    sqlite3_stmt *stmt;
    int rc = sqlite3_prepare_v2(db, query, -1, &stmt, NULL);
    assert(rc == SQLITE_OK);
    return stmt;
}

int step_thru(sqlite3 *db, const char *query) {
    int use_cache = 1;  // set to 0 to disable statement cache
    int rc = SQLITE_OK;
    printf("executing '%s' ", query);
    sqlite3_stmt *stmt = use_cache ? fetch_stmt(db, query) : new_stmt(db, query);
    printf("\n");
    assert(stmt != NULL);
    sqlite3_reset(stmt);
    int saved_changes = sqlite3_total_changes(db);
    do {
        rc = sqlite3_step(stmt);
        printf("- step, changes=%d, total_changes=%d\n",
               sqlite3_changes(db), sqlite3_total_changes(db));
    } while (rc == SQLITE_ROW);

    // simulate the Python sqlite3 cursor rowcount attribute
    if (!sqlite3_stmt_readonly(stmt)) {
        if (rc == SQLITE_DONE) {
            printf("- rowcount (changes): %d, (total changes): %d\n",
                   sqlite3_changes(db), sqlite3_total_changes(db) - saved_changes);
        }
    }

    rc = sqlite3_reset(stmt);
    if (!use_cache) {
        rc = sqlite3_finalize(stmt);
    }
    return rc;
}

int do_queries(sqlite3 *db, int n, const char **queries) {
    int rc = SQLITE_OK;
    for (int i = 0; i < n; i++) {
        rc = step_thru(db, queries[i]);
        if (rc != SQLITE_OK) {
            return rc;
        }
    }
    return rc;
}

void prepare_cache(const int n) {
    stmt_cache = (stmt *)calloc(n, sizeof(stmt));
    memset(stmt_cache, 0, n * sizeof(stmt));
}

void cleanup_cache(const int n) {
    for (int i = 0; i < n; i++) {
        sqlite3_stmt *stmt = stmt_cache[i].stmt;
        if (stmt) {
            (void)sqlite3_finalize(stmt);
        }
    }
    if (stmt_cache) {
        free(stmt_cache);
    }
}

int main(int argc, char **argv) {
    sqlite3 *db = NULL;

    printf("Compiled with SQLite " SQLITE_VERSION "\n");
    printf("Executing with SQLite %s\n", sqlite3_libversion());

    int rc = sqlite3_open(":memory:", &db);
    assert(rc == SQLITE_OK);

    const char *queries[] = {
        "CREATE TABLE some_table ("
        " id INTEGER NOT NULL,"
        " value VARCHAR(40) NOT NULL,"
        " PRIMARY KEY (id)"
        ")",
        "BEGIN",
        "INSERT INTO some_table (id, value) VALUES (1, 'v1')",
        "INSERT INTO some_table (id, value) VALUES (2, 'v2')",
        "INSERT INTO some_table (id, value) VALUES (3, 'v3')",
        "UPDATE some_table SET value='v2' WHERE id<3 RETURNING id",
        "DROP TABLE some_table",
        "COMMIT",
    };
    const size_t n = sizeof(queries) / sizeof(queries[0]);
    prepare_cache(n);

    rc = do_queries(db, n, queries);
    if (rc != SQLITE_OK) {
        goto exit;
    }

    rc = do_queries(db, n, queries);
    if (rc != SQLITE_OK) {
        goto exit;
    }

exit:
    cleanup_cache(n);

    rc = sqlite3_close(db);
    if (rc != SQLITE_OK) {
        print_error(db);
        return 1;
    }
    return 0;
}