Writing a Custom SQLite Function (in C) - Part 2

In part 1 we built a trivial SQLite function in C. Our goal was to cover the most critical parts of the SQLite API. One consequence of using a simple example is that we sidestepped any conversation about memory management. This second part looks specifically at some memory management that you'll likely need to be aware of.

Memory

To execute and read the result of our increment function, we used the sqlite3_bind_int64 and sqlite3_column_int functions respectively. As the name implies, these functions let us bind an integer to an SQLite statement and read an integer from a result. Integers have two important properties: their size is known upfront (a 64 bit integer always requires 64 bits of memory) and they are small. Types that have these two properties, including composite types (say a Point struct with two integer fields, X and Y) require straightforward memory management. When a function begins execution, the memory that we know it will require (essentially the sum of bytes needed for all the ints, floats, bytes, ...) is allocated upfront. This is called the "stack". Parameters are copied onto the stack and local variables (think a counter in a for loop) are initialized onit. Importantly, when the function returns, that entire stack is deallocated. Hence the memory is automatically allocated when the function starts and deallocate when the function ends.

In most programs though, not all of our memory needs are going to be known upfront. The classic example is processing a user's input; for example their favorite color where "red", "blue" and "yellow" all require different amounts of storage space. The most common strategy for dealing with this is to allocate memory on the "heap" which is tied to the entire program's lifetime, as opposed to the execution of individual functions. For small and shortlived programs, we can allocate memory on the heap, as needed, never freeing it, and let the OS reclaim the space when the program exit. But for long-running programs, space on the heap has to be managed. Specifically, because memory is a finite resource, we want memory allocated on the heap to be deallocated once it is no longer needed so it can be re-used . This is largely what a garbage collector does: it figures out when heap memory can be deallocated.

In C, we have no garbage collector. We are responsible for deallocating any memory we allocate on the heap. One common way to mess this up is to use the memory after we've deallocated it (a.k.a use-after-free). Similarly common is to not free the memory at all. In fact, it's easy to mess things up so badly that we lose all references to allocated memory and thus can't free it even if we wanted to (until the OS reclaims it on program exit). This is known as a memory leak.

To allocate memory on the heap, we use the malloc function, specifying how much memory we want. On success, it returns a pointer to the newly allocated memory. We use free on this pointer to deallocate the memory. At first glance, this sounds straightforward; in practice, it can be opaque. Maybe it's just because I don't touch C very often, but when I do, I spend a lot of time worrying about whether a function allocated memory and, if it did, whether I'm responsible for freeing it.

SQLite Strings

The real life example that inspired these posts was the development of a function, current_user_id, which wrapped a select user_id from tmp_session query. I wanted to encapsulate this query to make it usable within CTEs, triggers, default values and so on, while hiding the implementation details. user_id, being a string, required understanding of how to interact with heap-allocated memory from SQLite.

How the tmp_session table is populated isn't important. We're just using it as an example to get a heap-allocated value (in this case, a string) from SQLite. In my real use-case, it's a temporary table (thus scoped to an individual connection) which always has 0 or 1 rows depending on whether the request is being made from an anonymous or a logged-in user.

The first thing we'll do is register our function and create the skeleton:

static void current_user_id(sqlite3_context*, int, sqlite3_value**);

//...

static int run(sqlite3 *db) {
  // current_user_id is called with 0 arguments, thus the 3rd parameter is 0
  int rc = sqlite3_create_function(db, "current_user_id", 0, SQLITE_UTF8, NULL, &current_user_id, NULL, NULL);
  if (rc != SQLITE_OK) {
    print_error("create_current_user_id: %s", db);
    return 2;
  }
  return 0;
}

static void current_user_id(sqlite3_context *context, int argc, sqlite3_value **argv) {

}

The body of current_user_id will borrow from the code we wrote in Part 1 to execute the increment function. We need to create a statement, step through it (again, just once) and clean up:

static void current_user_id(sqlite3_context *context, int argc, sqlite3_value **argv) {
  // We get our db connection from the context
  sqlite3 *db = sqlite3_context_db_handle(context);

  sqlite3_stmt *stmt;
  int rc = sqlite3_prepare_v2(db, "select user_id from tmp_session", -1, &stmt, 0);
  if (rc != SQLITE_OK) {
    // Something bad happened, we'll come back to look at this later
    return;
  }

  rc = sqlite3_step(stmt);
  if (rc == SQLITE_ROW) {
    char *user_id = sqlite3_column_text(stmt, 0);
    // !!HERE!! What do we do with user_id
  } else if (rc == SQLITE_DONE) {
    // There's no user, we'll return null
    sqlite3_result_null(context);
  } else {
    // something bad happened, we'll come back to look at this later.
  }

  sqlite3_finalize(stmt);
}

The sqlite3_context_db_handle function is new and is used to get a reference to our connection from a sqlite3_context*. We need the connection to create our statement. The query that we're executing doesn't have any parameters/placeholders, so we don't need to bind anything. If our query returned a row, we called sqlite3_column_text to get the user_id value. If our query didn't have any rows, we called sqlite3_result_null to have our function return null.

Besides adding error handling (which we'll do later), we need to figure out how to return the user_id from our function. In Part 1 we called sqlite3_result_int to return an integer. Just above, we called sqlite3_result_null to return null. Maybe there's a sqlite3_result_text? In fact there is, but it's slightly more complicated to use. Here's the working code:

char *user_id = sqlite3_column_text(stmt, 0);
sqlite3_result_text(context, user_id, -1, SQLITE_TRANSIENT);

Compared to sqlite3_result_int, sqlite3_result_text takes 2 additional parameters. The 3rd parameter is the length of the string. In the above we're being lazy and passing -1, which means SQLite will scan user_id for a null terminator to determine the length. The 4th parameter can either be SQLITE_TRANSIENT or SQLITE_STATIC and specifically exists to inform SQLite about how it should manage the memory associated with the 2nd parameter, the char*. SQLITE_STATIC means the memory is allocated in constant space, that it won't suddenly disappear or change. SQLITE_STATIC is a powerful promise that allows sqlite3_result_text to simply reference the supplied char *. On the other hand, SQLITE_TRANSIENT makes no such guarantees. It basically informs sqlite3_result_text that after it returns, the memory referenced by the 2nd parameter could be changed or freed. Thus, when SQLITE_TRANSIENT is specified, as we've done above, sqlite3_result_text will allocate its own memory to hold the contents of the 2nd parameter and copy the value. Because it creates a copy, it no longer cares what happens to the original.

The above raises an important question: should we use SQLITE_TRANSIENT or SQLITE_STATIC? I mean SQLITE_STATIC seems preferable because it avoids having to allocate len(user_id) space and then having to create a copy. But we can only specify SQLITE_STATIC if we're sure that the memory referenced by user_id is, in fact, static. In our code above, the reference in question is returned by sqlite3_column_text. We could take an educated guess about how this behaves: it seems unlikely that the memory associated with an sqlite3_stmt* would outlive the call to sqlite3_finalize(stmt). To be sure, we need to refer to the documentation:

The pointers returned [from the sqlite3_column_* functions] are valid until (...) sqlite3_step() or sqlite3_reset() or sqlite3_finalize() is called. The memory space used to hold strings and BLOBs is freed automatically.

Our educated guess was only partially correct. In cases where we expect multiple rows and thus call sqlite3_step multiple times, the memory would be valid only until the next iteration. In our case, since sqlite3_step is only called once, the memory is valid only until sqlite3_finalize is called. In any case, we need the data to be returned out by our function, well beyond the point where sqlite3_finalize is called. Therefore, we must use SQLITE_TRANSIENT.

Thankfully, the SQLite API handles our cases easily. The memory ownership from stepping through a result is straightforward (returned pointers are valid only until the next step and/or finalize) and functions like sqlite3_result_text can be told to take ownership, by creating a copy, of the data we pass in.

What if we made things more complicated and did more than just take data from one SQLite function and pass it into another? For example, let's say that we wanted to create a function that returned the user_id for use in other code (essentially, refactoring the above logic)? In that case, we'd need to do what sqlite3_result_text does: allocate space for a copy of the value:

static char *get_user_id(sqlite3_context *context) {
  char *user_id = NULL;
  sqlite3 *db = sqlite3_context_db_handle(context);

  sqlite3_stmt *stmt;
  int rc = sqlite3_prepare_v2(db, "select user_id from tmp_session", -1, &stmt, 0);
  if (rc != SQLITE_OK) {
    // Something bad happened, we'll come back to look at this later
    return NULL;
  }

  rc = sqlite3_step(stmt);
  if (rc == SQLITE_ROW) {
    size_t len = sqlite3_column_bytes(stmt, 0);
    if (len != 0) {
      user_id = (char*)sqlite3_malloc64(len);
      if (user_id) {
        // if we successfully allocated the necessary memory
        // we can copy the value from the result into it
        memcpy(user_id, sqlite3_column_text(stmt, 0), len);
      } else {
        // sqlite3 has a helper method specifically for returning
        // a memory allocation failure
        sqlite3_result_error_nomem(context);
      }
    }
  }
  sqlite3_finalize(stmt);
  return user_id;
}

There's a lot to unpack, but it's essentially the same code as before, except for the part where we allocate memory and copy the contents of sqlite3_column_text to the new memory. Do note the use of the sqlite3_column_bytes function which returns the length of bytes for a text or blob column. We need this to know how much memory to allocate. (We could have used this function instead of passing -1 to the sqlite3_result_text function.)

Our function calls sqlite3_malloc64, which is like C's malloc but allocate the memory via SQLite (i.e. SQLite might call malloc itself, or might be able to return memory from a cache/buffer). However, we don't call sqlite3_free. The caller of this function will be responsible for freeing the returned pointer (or for passing the responsibility to its caller).

Error Handling

In the above example, we use the sqlite3_result_error_nomem function to return an error. This is a special function specifically designed for returning memory failures (like sqlite3_malloc64 returning NULL). But there are a lot of other places in our code where we'd like to return an error, for example, if the call to sqlite3_prepare_v2 fails. This could happen if the tmp_session table didn't exist. Just like we have a sqlite3_result_int and sqlite3_result_text, we also have a sqlite3_result_error.

An initial way to do error handling would be to return a static error message:

static char *get_user_id(sqlite3_context *context) {
  // ...
  int rc = sqlite3_prepare_v2(db, "select user_id from tmp_session", -1, &stmt, 0);
  if (rc != SQLITE_OK) {
    // ADDED THIS
    sqlite3_result_error(context, "failed to select user_id from tmp_session", -1);
    return NULL;
  }
  // ...

In some cases, this might be good enough, but there's a good chance we'll want the underlying error as well. We need to craft a better error message - it needs to be dynamic:

static char *get_user_id(sqlite3_context *context) {
  // ...
  int rc = sqlite3_prepare_v2(db, "select user_id from tmp_session", -1, &stmt, 0);
    const char* err = sqlite3_mprintf("get_user_id.prepare - %s", sqlite3_errmsg(db));
    sqlite3_result_error(context, err, -1);
    sqlite3_free((void *)err);
    return NULL;
  }
  // ...

The sqlite3_mprintf allocates memory and returns a pointer. It is our responsibility to free it when it is no longer needed. We do this (free the memory) after calling sqlite3_result_error and passing the error message to it. But is this ok? How do we know that sqlite3_result_error is OK with us freeing memory we just passed to it? Unfortunately, we need to either examine the implementation or hope the documentation can help. The SQLite documentation is, thankfully, clear on the subject:

The sqlite3_result_error() and sqlite3_result_error16() routines make a private copy of the error message text before they return. Hence, the calling function can deallocate or modify the text after they return without harm

Therefore, our above code is correct. We need to allocate memory to create a dynamic error message (sqlite3_mprintf takes care of that) and we can free it after passing it to sqlite3_result_error because that function creates a copy (which it internally owns and is responsible for).

Conclusion

The SQLite API is intuitive and friendly. Memory management and data ownership work in a logical and predictable manner. Still, there's a significant burden to understand where memory is allocated and who is responsible for freeing it. The most important thing to remember is that a pointer returned from an sqlite3_column_* function is only valid until the next call to sqlite3_step, sqlite3_reset or sqlite3_finalize. If you want data to exist beyond this, you'll need to copy it.

Furthermore, I believe the error handling pattern, while simple, is also common. And it showcases the importance of understanding who owns what memory (and the importance of good documentation). In this case, we had to know that sqlite3_result_error creates a copy of the supplied error to know that our error string could be freed.