Merge pull request mmzeeman#43 from TenderPro/BulkRowsReading

Bulk rows reading

Author: mmzeeman

c_src/esqlite3_nif.c

@@ -54,7 +54,7 @@ typedef enum {
     cmd_changes,
     cmd_prepare,
     cmd_bind,
-    cmd_step,
+    cmd_multi_step,
     cmd_reset,
     cmd_column_names,
     cmd_column_types,
@@ -476,49 +476,66 @@ make_cell(ErlNifEnv *env, sqlite3_stmt *statement, unsigned int i)
 }
 
 static ERL_NIF_TERM
-make_row(ErlNifEnv *env, sqlite3_stmt *statement) 
+make_row(ErlNifEnv *env, sqlite3_stmt *statement, ERL_NIF_TERM *array, int size)
 {
-    int i, size;
-    ERL_NIF_TERM *array;
-    ERL_NIF_TERM row;
-     
-    size = sqlite3_column_count(statement);
-    array = (ERL_NIF_TERM *) enif_alloc(sizeof(ERL_NIF_TERM)*size);
-
-    if(!array) 
-	    return make_error_tuple(env, "no_memory");
+    if(!array)
+        return make_error_tuple(env, "no_memory");
 
-    for(i = 0; i < size; i++) 
-	    array[i] = make_cell(env, statement, i);
+    for(int i = 0; i < size; i++) 
+        array[i] = make_cell(env, statement, i);
 
-    row = make_row_tuple(env, enif_make_tuple_from_array(env, array, size));
-    enif_free(array);
-    return row;
+    return enif_make_tuple_from_array(env, array, size);
 }
 
 static ERL_NIF_TERM
-do_step(ErlNifEnv *env, sqlite3 *db, sqlite3_stmt *stmt)
+do_multi_step(ErlNifEnv *env, sqlite3 *db, sqlite3_stmt *stmt, const ERL_NIF_TERM arg)
 {
+    ERL_NIF_TERM status;
+    ERL_NIF_TERM rows = enif_make_list_from_array(env, NULL, 0);
+    ERL_NIF_TERM *rowBuffer = NULL;
+    int rowBufferSize = 0;
+
+    int chunk_size = 0;
+    enif_get_int(env, arg, &chunk_size);
+
     int rc = sqlite3_step(stmt);
+    while (rc == SQLITE_ROW && chunk_size-- > 0)
+    {
+        if (!rowBufferSize)
+            rowBufferSize = sqlite3_column_count(stmt);
+        if (rowBuffer == NULL)
+            rowBuffer = (ERL_NIF_TERM *) enif_alloc(sizeof(ERL_NIF_TERM)*rowBufferSize);
 
-    if(rc == SQLITE_ROW) 
-        return make_row(env, stmt);
-    if(rc == SQLITE_BUSY)
-        return make_atom(env, "$busy");
+        rows = enif_make_list_cell(env, make_row(env, stmt, rowBuffer, rowBufferSize), rows);
 
-    if(rc == SQLITE_DONE) { 
+        if (chunk_size > 0)
+            rc = sqlite3_step(stmt);
+    }
+
+    switch(rc) {
+    case SQLITE_ROW:
+        status = make_atom(env, "rows");
+        break;
+    case SQLITE_BUSY:
+        status = make_atom(env, "$busy");
+        break;
+    case 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.
-         */
+        * 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");
+        status = make_atom(env, "$done");
+        break;
+    default:
+        /* We use prepare_v2, so any error code can be returned. */
+        return make_sqlite3_error_tuple(env, rc, db);
     }
 
-    /* We use prepare_v2, so any error code can be returned. */
-    return make_sqlite3_error_tuple(env, rc, db);
+    enif_free(rowBuffer);
+    return enif_make_tuple2(env, status, rows);
 }
 
 static ERL_NIF_TERM
@@ -630,8 +647,8 @@ evaluate_command(esqlite_command *cmd, esqlite_connection *conn)
 	    return do_changes(cmd->env, conn, cmd->arg);
     case cmd_prepare:
 	    return do_prepare(cmd->env, conn, cmd->arg);
-    case cmd_step:
-	    return do_step(cmd->env, conn->db, stmt->statement);
+    case cmd_multi_step:
+        return do_multi_step(cmd->env, conn->db, stmt->statement, cmd->arg);
     case cmd_reset:
 	    return do_reset(cmd->env, conn->db, stmt->statement);
     case cmd_bind:
@@ -679,7 +696,7 @@ esqlite_connection_run(void *arg)
 	        enif_send(NULL, &cmd->pid, cmd->env, make_answer(cmd, evaluate_command(cmd, db)));
         }
 
-	    command_destroy(cmd);    
+	    command_destroy(cmd);
     }
 
     return NULL;
@@ -916,39 +933,47 @@ esqlite_bind(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
 }
 
 /*
- * Step to a prepared statement
+ * Multi step to a prepared statement
  */
 static ERL_NIF_TERM 
-esqlite_step(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
+esqlite_multi_step(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
 {
     esqlite_connection *conn;
     esqlite_statement *stmt;
     esqlite_command *cmd = NULL;
     ErlNifPid pid;
+    int chunk_size = 0;
 
-    if(argc != 4) 
-	    return enif_make_badarg(env);
+    if(argc != 5) 
+        return enif_make_badarg(env);
 
     if(!enif_get_resource(env, argv[0], esqlite_connection_type, (void **) &conn))
-	    return enif_make_badarg(env);
+        return enif_make_badarg(env);
+    
     if(!enif_get_resource(env, argv[1], esqlite_statement_type, (void **) &stmt))
-	    return enif_make_badarg(env);
-    if(!enif_is_ref(env, argv[2])) 
-	    return make_error_tuple(env, "invalid_ref");
-    if(!enif_get_local_pid(env, argv[3], &pid)) 
-	    return make_error_tuple(env, "invalid_pid"); 
+        return enif_make_badarg(env);
+
+    if(!enif_get_int(env, argv[2], &chunk_size))
+        return make_error_tuple(env, "invalid_chunk_size");
+    
+    if(!enif_is_ref(env, argv[3])) 
+        return make_error_tuple(env, "invalid_ref");
+    
+    if(!enif_get_local_pid(env, argv[4], &pid)) 
+        return make_error_tuple(env, "invalid_pid");
 
     if(!stmt->statement) 
-	    return make_error_tuple(env, "no_prepared_statement");
+        return make_error_tuple(env, "no_prepared_statement");
 
     cmd = command_create();
     if(!cmd) 
-	   return make_error_tuple(env, "command_create_failed");
+        return make_error_tuple(env, "command_create_failed");
 
-    cmd->type = cmd_step;
-    cmd->ref = enif_make_copy(cmd->env, argv[2]);
+    cmd->type = cmd_multi_step;
+    cmd->ref = enif_make_copy(cmd->env, argv[3]);
     cmd->pid = pid;
     cmd->stmt = enif_make_copy(cmd->env, argv[1]);
+    cmd->arg = enif_make_copy(cmd->env, argv[2]);
 
     return push_command(env, conn, cmd);
 }
@@ -1133,7 +1158,7 @@ static ErlNifFunc nif_funcs[] = {
     {"changes", 3, esqlite_changes},
     {"prepare", 4, esqlite_prepare},
     {"insert", 4, esqlite_insert},
-    {"step", 4, esqlite_step},
+    {"multi_step", 5, esqlite_multi_step},
     {"reset", 4, esqlite_reset},
     // TODO: {"esqlite_bind", 3, esqlite_bind_named},
     {"bind", 5, esqlite_bind},

src/esqlite3.erl

@@ -31,13 +31,15 @@
          bind/2, bind/3,
          fetchone/1,
          fetchall/1,
+         fetchall/2,
          column_names/1, column_names/2,
          column_types/1, column_types/2,
          close/1, close/2]).
 
 -export([q/2, q/3, map/3, foreach/3]).
 
 -define(DEFAULT_TIMEOUT, 5000).
+-define(DEFAULT_CHUNK_SIZE, 5000).
 
 %%
 -type connection() :: {connection, reference(), term()}.
@@ -123,19 +125,19 @@ foreach(F, Sql, Connection) ->
       Row :: tuple(),
       ColumnNames :: tuple().
 foreach_s(F, Statement) when is_function(F, 1) ->
-    case try_step(Statement, 0) of
-        '$done' -> ok;
+    case try_multi_step(Statement, 1, [], 0) of
+        {'$done', []} -> ok;
         {error, _} = E -> F(E);
-        {row, Row} ->
+        {rows, [Row | []]} ->
             F(Row),
             foreach_s(F, Statement)
     end;
 foreach_s(F, Statement) when is_function(F, 2) ->
     ColumnNames = column_names(Statement),
-    case try_step(Statement, 0) of
-        '$done' -> ok;
+    case try_multi_step(Statement, 1, [], 0) of
+        {'$done', []} -> ok;
         {error, _} = E -> F([], E);
-        {row, Row} ->
+        {rows, [Row | []]} ->
             F(ColumnNames, Row),
             foreach_s(F, Statement)
     end.
@@ -147,59 +149,82 @@ foreach_s(F, Statement) when is_function(F, 2) ->
       ColumnNames :: tuple(),
       Type :: term().
 map_s(F, Statement) when is_function(F, 1) ->
-    case try_step(Statement, 0) of
-        '$done' -> [];
+    case try_multi_step(Statement, 1, [], 0) of
+        {'$done', []} -> [];
         {error, _} = E -> F(E);
-        {row, Row} ->
+        {rows, [Row | []]} ->
             [F(Row) | map_s(F, Statement)]
     end;
 map_s(F, Statement) when is_function(F, 2) ->
     ColumnNames = column_names(Statement),
-    case try_step(Statement, 0) of
-        '$done' -> [];
+    case try_multi_step(Statement, 1, [], 0) of
+        {'$done', []} -> [];
         {error, _} = E -> F([], E);
-        {row, Row} ->
+        {rows, [Row | []]} ->
             [F(ColumnNames, Row) | map_s(F, Statement)]
     end.
 
 %%
--spec fetchone(statement()) -> tuple().
+%%-spec fetchone(statement()) -> tuple().
 fetchone(Statement) ->
-    case try_step(Statement, 0) of
-        '$done' -> ok;
+    case try_multi_step(Statement, 1, [], 0) of
+        {'$done', []} -> ok;
         {error, _} = E -> E;
-        {row, Row} -> Row
+        {rows, [Row | []]} -> Row
     end.
 
-%%
+%% @doc Fetch all records
+%% @param Statement is prepared sql statement
+%% @spec fetchall(statement()) -> list(tuple()) | {error, term()}.
 -spec fetchall(statement()) ->
                       list(tuple()) |
                       {error, term()}.
 fetchall(Statement) ->
-    case try_step(Statement, 0) of
-        '$done' ->
-            [];
-        {error, _} = E -> E;
-        {row, Row} ->
-            case fetchall(Statement) of
-                {error, _} = E -> E;
-                Rest -> [Row | Rest]
-            end
+    fetchall(Statement, ?DEFAULT_CHUNK_SIZE).
+
+%% @doc Fetch all records
+%% @param Statement is prepared sql statement
+%% @param ChunkSize is a count of rows to read from sqlite and send to erlang process in one bulk.
+%%        Decrease this value if rows are heavy. Default value is 5000 (DEFAULT_CHUNK_SIZE).
+%% @spec fetchall(statement()) -> list(tuple()) | {error, term()}.
+-spec fetchall(statement(), pos_integer()) ->
+                      list(tuple()) |
+                      {error, term()}.
+fetchall(Statement, ChunkSize) ->
+    case fetchall_internal(Statement, ChunkSize, []) of
+        {'$done', Rows} -> lists:reverse(Rows);
+        {error, _} = E -> E
+    end.
+
+%% return rows in revers order
+-spec fetchall_internal(statement(), pos_integer(), list(tuple())) ->
+                {'$done', list(tuple())} |
+                {error, term()}.
+fetchall_internal(Statement, ChunkSize, Rest) ->
+    case try_multi_step(Statement, ChunkSize, Rest, 0) of
+        {rows, Rows} -> fetchall_internal(Statement, ChunkSize, Rows);
+        Else -> Else
     end.
 
-%% Try the step, when the database is busy,
--spec try_step(statement(), non_neg_integer()) -> 
-                      '$done' |
-                      term().
-try_step(_Statement, Tries) when Tries > 5 ->
+%% Try a number of steps, when the database is busy,
+%% return rows in revers order
+-spec try_multi_step(statement(), pos_integer(), list(tuple()), non_neg_integer()) -> 
+                {rows, list(tuple())} |
+                {'$done', list(tuple())} |
+                {error, term()}.
+try_multi_step(_Statement, _ChunkSize, _Rest, Tries) when Tries > 5 ->
     throw(too_many_tries);
-try_step(Statement, Tries) ->
-    case esqlite3:step(Statement) of
-        '$busy' ->
+try_multi_step(Statement, ChunkSize, Rest, Tries) ->
+    case multi_step(Statement, ChunkSize) of
+        {'$busy', Rows} -> %% core can fetch a number of rows (rows < ChunkSize) per 'multi_step' call and then get busy...
+            erlang:display({"busy", Tries}),
             timer:sleep(100 * Tries),
-            try_step(Statement, Tries + 1);
-        Something ->
-            Something
+            try_multi_step(Statement, ChunkSize, Rows ++ Rest, Tries + 1);
+        {rows, Rows} ->
+            {rows, Rows ++ Rest};
+        {'$done', Rows} ->
+            {'$done', Rows ++ Rest};
+        Else -> Else
     end.
 
 %% @doc Execute Sql statement, returns the number of affected rows.
@@ -280,7 +305,29 @@ step(Stmt) ->
 -spec step(term(), timeout()) -> tuple() | '$busy' | '$done'.
 step({statement, Stmt, {connection, _, Conn}}, Timeout) ->
     Ref = make_ref(),
-    ok = esqlite3_nif:step(Conn, Stmt, Ref, self()),
+    ok = esqlite3_nif:multi_step(Conn, Stmt, 1, Ref, self()),
+    case receive_answer(Ref, Timeout) of
+        {rows, [Row | []]} -> {row, Row};
+        {'$done', []} -> '$done';
+        {'$busy', []} -> '$busy';
+        Else -> Else
+    end.
+
+%% make multiple sqlite steps per call
+%% return rows in reverse order
+multi_step(Stmt, ChunkSize) ->
+    multi_step(Stmt, ChunkSize, ?DEFAULT_TIMEOUT).
+
+%% make multiple sqlite steps per call
+%% return rows in reverse order
+-spec multi_step(term(), pos_integer(), timeout()) ->
+                {rows, list(tuple())} |
+                {'$busy', list(tuple())} |
+                {'$done', list(tuple())} |
+                {error, term()}.
+multi_step({statement, Stmt, {connection, _, Conn}}, ChunkSize, Timeout) ->
+    Ref = make_ref(),
+    ok = esqlite3_nif:multi_step(Conn, Stmt, ChunkSize, Ref, self()),
     receive_answer(Ref, Timeout).
 
 %% @doc Reset the prepared statement back to its initial state.