diff --git a/c_src/esqlite3_nif.c b/c_src/esqlite3_nif.c index 10397b9..b5734ca 100644 --- a/c_src/esqlite3_nif.c +++ b/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(int i = 0; i < size; i++) + array[i] = make_cell(env, statement, i); - for(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) { - /* - * 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"); + if (chunk_size > 0) + rc = sqlite3_step(stmt); } - /* We use prepare_v2, so any error code can be returned. */ - return make_sqlite3_error_tuple(env, rc, db); + 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. + */ + sqlite3_reset(stmt); + 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); + } + + 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}, diff --git a/src/esqlite3.erl b/src/esqlite3.erl index 9a9a91e..5b69c83 100644 --- a/src/esqlite3.erl +++ b/src/esqlite3.erl @@ -31,6 +31,7 @@ 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]). @@ -38,6 +39,7 @@ -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,79 @@ 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. -%% Try the step, when the database is busy, --spec try_step(statement(), non_neg_integer()) -> - '$done' | - term(). -try_step(_Statement, Tries) when Tries > 5 -> +%% 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 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); + {Status, Rows} -> + {Status, Rows ++ Rest} end. %% @doc Execute Sql statement, returns the number of affected rows. @@ -280,7 +302,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. diff --git a/src/esqlite3_nif.erl b/src/esqlite3_nif.erl index d5e480f..2f56464 100644 --- a/src/esqlite3_nif.erl +++ b/src/esqlite3_nif.erl @@ -27,7 +27,7 @@ changes/3, insert/4, prepare/4, - step/4, + multi_step/5, reset/4, finalize/4, bind/5, @@ -90,8 +90,8 @@ prepare(_Db, _Ref, _Dest, _Sql) -> %% @doc %% -%% @spec step(statement(), reference(), pid()) -> ok | {error, message()} -step(_Db, _Stmt, _Ref, _Dest) -> +%% @spec multi_step(statement(), pos_integer(), reference(), pid()) -> {term(), list(tuple)} | {error, message()} +multi_step(_Db, _Stmt, _Chunk_Size, _Ref, _Dest) -> erlang:nif_error(nif_library_not_loaded). %% @doc