Make sure the statement is reset after a in step. Fixes #10

c_src/esqlite3_nif.c

@@ -57,6 +57,7 @@ typedef enum {
     cmd_prepare,
     cmd_bind,
     cmd_step,
+    cmd_reset,
     cmd_column_names,
     cmd_close,
     cmd_stop,
@@ -506,15 +507,35 @@ do_step(ErlNifEnv *env, sqlite3 *db, sqlite3_stmt *stmt)
 
     if(rc == SQLITE_ROW) 
         return make_row(env, stmt);
-    if(rc == SQLITE_DONE) 
-	    return make_atom(env, "$done");
     if(rc == SQLITE_BUSY)
         return make_atom(env, "$busy");
 
+    if(rc == SQLITE_DONE) { 
+        /* 
+         * Automatically reset the statement after a done so 
+         * column_names will work after the statement is done.
+         *
+         * Not resetting the statement can lead to vm crashes.
+         */
+        sqlite3_reset(stmt);
+        return make_atom(env, "$done");
+    }
+
     /* We use prepare_v2, so any error code can be returned. */
     return make_sqlite3_error_tuple(env, rc, db);
 }
 
+static ERL_NIF_TERM
+do_reset(ErlNifEnv *env, sqlite3 *db, sqlite3_stmt *stmt)
+{
+    int rc = sqlite3_reset(stmt);
+
+    if(rc == SQLITE_OK) 
+        return make_atom(env, "ok");
+
+    return make_sqlite3_error_tuple(env, rc, db);
+}
+
 static ERL_NIF_TERM
 do_column_names(ErlNifEnv *env, sqlite3_stmt *stmt)
 {
@@ -524,13 +545,20 @@ do_column_names(ErlNifEnv *env, sqlite3_stmt *stmt)
     ERL_NIF_TERM column_names;
      
     size = sqlite3_column_count(stmt);
-    array = (ERL_NIF_TERM *) malloc(sizeof(ERL_NIF_TERM) * size);
+    if(size <= 0)
+        return make_error_tuple(env, "no_columns");
 
+    array = (ERL_NIF_TERM *) malloc(sizeof(ERL_NIF_TERM) * size);
     if(!array)
         return make_error_tuple(env, "no_memory");
 
     for(i = 0; i < size; i++) {
         name = sqlite3_column_name(stmt, i);
+        if(name == NULL) {
+            free(array);
+            return make_error_tuple(env, "sqlite3_malloc_failure");
+        }
+
         array[i] = make_atom(env, name);
     }
 
@@ -566,6 +594,8 @@ evaluate_command(esqlite_command *cmd, esqlite_connection *conn)
 	    return do_prepare(cmd->env, conn, cmd->arg);
     case cmd_step:
 	    return do_step(cmd->env, conn->db, cmd->stmt);
+    case cmd_reset:
+	    return do_reset(cmd->env, conn->db, cmd->stmt);
     case cmd_bind:
 	    return do_bind(cmd->env, conn->db, cmd->stmt, cmd->arg);
     case cmd_column_names:
@@ -885,7 +915,45 @@ esqlite_step(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
 }
 
 /*
- * Step to a prepared statement
+ * Reset a prepared statement to its initial state
+ */
+static ERL_NIF_TERM 
+esqlite_reset(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
+{
+    esqlite_statement *stmt;
+    esqlite_command *cmd = NULL;
+    ErlNifPid pid;
+
+    if(argc != 3) 
+	    return enif_make_badarg(env);
+    if(!enif_get_resource(env, argv[0], esqlite_statement_type, (void **) &stmt))
+	    return enif_make_badarg(env);
+    if(!enif_is_ref(env, argv[1])) 
+	    return make_error_tuple(env, "invalid_ref");
+    if(!enif_get_local_pid(env, argv[2], &pid)) 
+	    return make_error_tuple(env, "invalid_pid"); 
+    if(!stmt->statement) 
+	    return make_error_tuple(env, "no_prepared_statement");
+
+    cmd = command_create();
+    if(!cmd) 
+	   return make_error_tuple(env, "command_create_failed");
+
+    cmd->type = cmd_reset;
+    cmd->ref = enif_make_copy(cmd->env, argv[1]);
+    cmd->pid = pid;
+    cmd->stmt = stmt->statement;
+
+    if(!stmt->connection) 
+	    return make_error_tuple(env, "no_connection");
+    if(!stmt->connection->commands)
+	    return make_error_tuple(env, "no_command_queue");
+
+    return push_command(env, stmt->connection, cmd);
+}
+
+/*
+ * Get the column names of the prepared statement.
  */
 static ERL_NIF_TERM 
 esqlite_column_names(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
@@ -993,6 +1061,7 @@ static ErlNifFunc nif_funcs[] = {
     {"prepare", 4, esqlite_prepare},
     {"insert", 4, esqlite_insert},
     {"step", 3, esqlite_step},
+    {"reset", 3, esqlite_reset},
     // TODO: {"esqlite_bind", 3, esqlite_bind_named},
     {"bind", 4, esqlite_bind},
     {"column_names", 3, esqlite_column_names},

src/esqlite3.erl

@@ -27,6 +27,7 @@
          insert/2,
          prepare/2, prepare/3,
          step/1, step/2, 
+         reset/1, 
          bind/2, bind/3, 
          fetchone/1,
          fetchall/1,
@@ -270,6 +271,14 @@ step(Stmt, Timeout) ->
     ok = esqlite3_nif:step(Stmt, Ref, self()),
     receive_answer(Ref, Timeout).
 
+%% @doc Reset the prepared statement back to its initial state.
+%%
+%% @spec reset(prepared_statement()) -> ok | {error, error_message()}
+reset(Stmt) ->
+    Ref = make_ref(),
+    ok = esqlite3_nif:reset(Stmt, Ref, self()),
+    receive_answer(Ref, ?DEFAULT_TIMEOUT).
+
 %% @doc Bind values to prepared statements
 %%
 %% @spec bind(prepared_statement(), value_list()) -> ok | {error, error_message()}

src/esqlite3_nif.erl

@@ -3,7 +3,7 @@
 
 %% @doc Low level erlang API for sqlite3 databases
 
-%% Copyright 2011, 2012, 2013 Maas-Maarten Zeeman
+%% Copyright 2011, 2012, 2013, 2014 Maas-Maarten Zeeman
 %%
 %% Licensed under the Apache License, Version 2.0 (the "License");
 %% you may not use this file except in compliance with the License.
@@ -28,6 +28,7 @@
     insert/4,
     prepare/4,
     step/3,
+    reset/3,
     finalize/3,
     bind/4,
     column_names/3,
@@ -91,6 +92,12 @@ prepare(_Db, _Ref, _Dest, _Sql) ->
 step(_Stmt, _Ref, _Dest) ->
     exit(nif_library_not_loaded).
 
+%% @doc
+%%
+%% @spec reset(statement(), reference(), pid()) -> ok | {error, message()}
+reset(_Stmt, _Ref, _Dest) ->
+    exit(nif_library_not_loaded).
+
 %% @doc
 %%
 %%

test/esqlite_test.erl

@@ -142,14 +142,69 @@ column_names_test() ->
     ok = esqlite3:exec("begin;", Db),
     ok = esqlite3:exec("create table test_table(one varchar(10), two int);", Db),
     ok = esqlite3:exec(["insert into test_table values(", "\"hello1\"", ",", "10" ");"], Db),
+    ok = esqlite3:exec(["insert into test_table values(", "\"hello2\"", ",", "20" ");"], Db),
     ok = esqlite3:exec("commit;", Db),
 
+    %% All columns
     {ok, Stmt} = esqlite3:prepare("select * from test_table", Db),
-    
+    {one, two} =  esqlite3:column_names(Stmt),
+    {row, {<<"hello1">>, 10}} = esqlite3:step(Stmt),
+    {one, two} =  esqlite3:column_names(Stmt),
+    {row, {<<"hello2">>, 20}} = esqlite3:step(Stmt),
+    {one, two} =  esqlite3:column_names(Stmt),
+    '$done' = esqlite3:step(Stmt),
     {one, two} =  esqlite3:column_names(Stmt),
 
+    %% One column
+    {ok, Stmt2} = esqlite3:prepare("select two from test_table", Db),
+    {two} =  esqlite3:column_names(Stmt2),
+    {row, {10}} = esqlite3:step(Stmt2),
+    {two} =  esqlite3:column_names(Stmt2),
+    {row, {20}} = esqlite3:step(Stmt2),
+    {two} =  esqlite3:column_names(Stmt2),
+    '$done' = esqlite3:step(Stmt2),
+    {two} =  esqlite3:column_names(Stmt2),
+
+    %% No columns
+    {ok, Stmt3} = esqlite3:prepare("values(1);", Db),
+    {column1} =  esqlite3:column_names(Stmt3),
+    {row, {1}} = esqlite3:step(Stmt3),
+    {column1} =  esqlite3:column_names(Stmt3),
+
+    %% Things get a bit weird when you retrieve the column name
+    %% when calling an aggragage function.
+    {ok, Stmt4} = esqlite3:prepare("select date('now');", Db),
+    {'date(\'now\')'} =  esqlite3:column_names(Stmt4),
+    {row, {Date}} = esqlite3:step(Stmt4),
+    true = is_binary(Date),
+
     ok.
 
+reset_test() ->
+    {ok, Db} = esqlite3:open(":memory:"),
+
+    {ok, Stmt} = esqlite3:prepare("select * from (values (1), (2));", Db),
+    {row, {1}} = esqlite3:step(Stmt),
+
+    ok = esqlite3:reset(Stmt),
+    {row, {1}} = esqlite3:step(Stmt),
+    {row, {2}} = esqlite3:step(Stmt),
+    '$done' = esqlite3:step(Stmt),
+
+    % After a done the statement is automatically reset.
+    {row, {1}} = esqlite3:step(Stmt),
+
+    % Calling reset multiple times...
+    ok = esqlite3:reset(Stmt),
+    ok = esqlite3:reset(Stmt),
+    ok = esqlite3:reset(Stmt),
+    ok = esqlite3:reset(Stmt),
+
+    {row, {1}} = esqlite3:step(Stmt),
+
+    ok.
+
+
 foreach_test() ->
     {ok, Db} = esqlite3:open(":memory:"),
     ok = esqlite3:exec("begin;", Db),