Adds support for query timeouts

Author: MichaelBaker

lib/sqlitex/query.ex

@@ -32,7 +32,7 @@ defmodule Sqlitex.Query do
   def query(db, sql, opts \\ []) do
     with {:ok, stmt} <- Statement.prepare(db, sql, opts),
          {:ok, stmt} <- Statement.bind_values(stmt, Keyword.get(opts, :bind, []), opts),
-         {:ok, res} <- Statement.fetch_all(stmt, Keyword.get(opts, :into, [])),
+         {:ok, res} <- Statement.fetch_all(stmt, Keyword.get(opts, :db_timeout, 5_000), Keyword.get(opts, :into, [])),
     do: {:ok, res}
   end
 
@@ -77,7 +77,7 @@ defmodule Sqlitex.Query do
   def query_rows(db, sql, opts \\ []) do
     with {:ok, stmt} <- Statement.prepare(db, sql, opts),
          {:ok, stmt} <- Statement.bind_values(stmt, Keyword.get(opts, :bind, []), opts),
-         {:ok, rows} <- Statement.fetch_all(stmt, :raw_list),
+         {:ok, rows} <- Statement.fetch_all(stmt, Keyword.get(opts, :db_timeout, 5_000), :raw_list),
     do: {:ok, %{rows: rows, columns: stmt.column_names, types: stmt.column_types}}
   end
 

lib/sqlitex/server.ex

@@ -178,7 +178,7 @@ defmodule Sqlitex.Server do
 
     with {%Cache{} = new_cache, stmt} <- Cache.prepare(stmt_cache, sql, db_opts),
          {:ok, stmt} <- Statement.bind_values(stmt, Keyword.get(opts, :bind, []), db_opts),
-         {:ok, rows} <- Statement.fetch_all(stmt, Keyword.get(opts, :into, [])),
+         {:ok, rows} <- Statement.fetch_all(stmt, Keyword.get(opts, :db_timeout, 5_000), Keyword.get(opts, :into, [])),
     do: {:ok, rows, new_cache}
   end
 
@@ -187,7 +187,7 @@ defmodule Sqlitex.Server do
 
     with {%Cache{} = new_cache, stmt} <- Cache.prepare(stmt_cache, sql, db_opts),
          {:ok, stmt} <- Statement.bind_values(stmt, Keyword.get(opts, :bind, []), db_opts),
-         {:ok, rows} <- Statement.fetch_all(stmt, :raw_list),
+         {:ok, rows} <- Statement.fetch_all(stmt, Keyword.get(opts, :db_timeout, 5_000), :raw_list),
     do: {:ok,
          %{rows: rows, columns: stmt.column_names, types: stmt.column_types},
          new_cache}

lib/sqlitex/statement.ex

@@ -21,7 +21,7 @@ defmodule Sqlitex.Statement do
   iex(6)> Sqlitex.Statement.exec(statement)
   :ok
   iex(7)> {:ok, statement} = Sqlitex.Statement.prepare(db, "SELECT * FROM data;")
-  iex(8)> Sqlitex.Statement.fetch_all(statement)
+  iex(8)> Sqlitex.Statement.fetch_all(statement, 1_000)
   {:ok, [[id: 1, name: "hello"]]}
   iex(9)> Sqlitex.close(db)
   :ok
@@ -46,7 +46,7 @@ defmodule Sqlitex.Statement do
   and store it separately in the `Statement` struct. Only the portion of the query
   preceding the returning clause is passed to SQLite's prepare function.
 
-  Later, when such a statement struct is passed to `fetch_all/2` or `fetch_all!/2`
+  Later, when such a statement struct is passed to `fetch_all/3` or `fetch_all!/3`
   the returning clause is parsed and the query is performed with the following
   additional logic:
 
@@ -73,6 +73,8 @@ defmodule Sqlitex.Statement do
             column_names: [],
             column_types: []
 
+  @chunk_size 5000
+
   alias Sqlitex.Config
 
   @doc """
@@ -170,35 +172,36 @@ defmodule Sqlitex.Statement do
   ## Parameters
 
   * `statement` - The statement to run.
+  * `timeout` - The query timeout to be passed to esqlite.
   * `into` - The collection to put the results into. Defaults to an empty list.
 
   ## Returns
 
   * `{:ok, results}`
   * `{:error, error}`
   """
-  def fetch_all(statement, into \\ []) do
-    case raw_fetch_all(statement) do
+  def fetch_all(statement, timeout, into \\ []) do
+    case raw_fetch_all(statement, timeout) do
       {:error, _} = other -> other
       raw_data ->
         {:ok, Row.from(statement.column_types, statement.column_names, raw_data, into)}
     end
   end
 
-  defp raw_fetch_all(%__MODULE__{returning: nil, statement: statement}) do
-    :esqlite3.fetchall(statement)
+  defp raw_fetch_all(%__MODULE__{returning: nil, statement: statement}, timeout) do
+    :esqlite3.fetchall(statement, @chunk_size, timeout)
   end
-  defp raw_fetch_all(statement) do
-    returning_query(statement)
+  defp raw_fetch_all(statement, timeout) do
+    returning_query(statement, timeout)
   end
 
   @doc """
-  Same as `fetch_all/2` but raises a Sqlitex.Statement.FetchAllError on error.
+  Same as `fetch_all/3` but raises a Sqlitex.Statement.FetchAllError on error.
 
   Returns the results otherwise.
   """
-  def fetch_all!(statement, into \\ []) do
-    case fetch_all(statement, into) do
+  def fetch_all!(statement, timeout, into \\ []) do
+    case fetch_all(statement, timeout, into) do
       {:ok, results} -> results
       {:error, reason} -> raise Sqlitex.Statement.FetchAllError, reason: reason
     end
@@ -338,11 +341,11 @@ defmodule Sqlitex.Statement do
     {:error, :invalid_returning_clause}
   end
 
-  defp returning_query(%__MODULE__{database: db} = stmt) do
+  defp returning_query(%__MODULE__{database: db} = stmt, timeout) do
     sp = "sp_#{random_id()}"
     {:ok, _} = db_exec(db, "SAVEPOINT #{sp}")
 
-    case returning_query_in_savepoint(sp, stmt) do
+    case returning_query_in_savepoint(sp, stmt, timeout) do
       {:error, _} = error ->
         rollback(db, sp)
         error
@@ -354,7 +357,7 @@ defmodule Sqlitex.Statement do
 
   defp returning_query_in_savepoint(sp, %__MODULE__{database: db,
                                                     statement: statement,
-                                                    returning: {table, cols, cmd, ref}})
+                                                    returning: {table, cols, cmd, ref}}, timeout)
   do
     temp_table = "t_#{random_id()}"
     temp_fields = Enum.join(cols, ", ")
@@ -371,7 +374,7 @@ defmodule Sqlitex.Statement do
 
     with {:ok, _} = db_exec(db, "CREATE TEMP TABLE #{temp_table} (#{temp_fields})"),
          {:ok, _} = db_exec(db, trigger),
-         result = :esqlite3.fetchall(statement),
+         result = :esqlite3.fetchall(statement, @chunk_size, timeout),
          {:ok, rows} = db_exec(db, "SELECT #{column_names} FROM #{temp_table}"),
          {:ok, _} = db_exec(db, "DROP TRIGGER IF EXISTS #{trigger_name}"),
          {:ok, _} = db_exec(db, "DROP TABLE IF EXISTS #{temp_table}")

test/statement_test.exs

@@ -7,7 +7,7 @@ defmodule StatementTest do
 
     result = db
               |> Sqlitex.Statement.prepare!("PRAGMA user_version;")
-              |> Sqlitex.Statement.fetch_all!
+              |> Sqlitex.Statement.fetch_all!(1_000)
 
     assert result == [[user_version: 0]]
   end
@@ -20,7 +20,7 @@ defmodule StatementTest do
     stmt = Sqlitex.Statement.prepare!(db, "INSERT INTO x(str) VALUES (?1) "
                                           <> ";--RETURNING ON INSERT x,id")
 
-    rows = Sqlitex.Statement.fetch_all!(stmt)
+    rows = Sqlitex.Statement.fetch_all!(stmt, 1_000)
     assert rows == [[id: 1]]
   end
 
@@ -32,7 +32,7 @@ defmodule StatementTest do
     stmt = Sqlitex.Statement.prepare!(db, "INSERT INTO x(str) VALUES (?1) "
                                           <> ";--RETURNING ON INSERT x,id")
 
-    rows = Sqlitex.Statement.fetch_all!(stmt, :raw_list)
+    rows = Sqlitex.Statement.fetch_all!(stmt, 1_000, :raw_list)
     assert rows == [[1]]
   end
 
@@ -44,7 +44,7 @@ defmodule StatementTest do
     stmt = Sqlitex.Statement.prepare!(db, "INSERT INTO x(str) VALUES ('x'),('y'),('z') "
                                           <> ";--RETURNING ON INSERT x,id")
 
-    rows = Sqlitex.Statement.fetch_all!(stmt)
+    rows = Sqlitex.Statement.fetch_all!(stmt, 1_000)
     assert rows == [[id: 1], [id: 2], [id: 3]]
   end
 
@@ -56,7 +56,7 @@ defmodule StatementTest do
     stmt = Sqlitex.Statement.prepare!(db, "INSERT INTO x(str) VALUES ('x'),('y'),('z') "
                                           <> ";--RETURNING ON INSERT x,id")
 
-    rows = Sqlitex.Statement.fetch_all!(stmt, :raw_list)
+    rows = Sqlitex.Statement.fetch_all!(stmt, 1_000, :raw_list)
     assert rows == [[1], [2], [3]]
   end
 
@@ -71,7 +71,7 @@ defmodule StatementTest do
     stmt = Sqlitex.Statement.prepare!(db, "INSERT INTO x(str) VALUES ('x') "
                                           <> ";--RETURNING ON INSERT x,id")
 
-    result = Sqlitex.Statement.fetch_all(stmt, :raw_list)
+    result = Sqlitex.Statement.fetch_all(stmt, 1_000, :raw_list)
     assert result == {:error, {:constraint, 'UNIQUE constraint failed: x.str'}}
   end
 end