Skip to content

Instantly share code, notes, and snippets.

@maxim maxim/stream_ext.ex
Created Jun 19, 2018

Embed
What would you like to do?
defmodule StreamExt do
@doc """
Stream records from Ecto in batches.
Supported options are
* batch_size - how many rows to fetch at once
* strategy - either :offset or :id
* :offset - uses SQL offset to fetch pages of results
This is slower, but works for any query.
* :id - uses "where id > last_id" to fetch pages of results
This is faster, but requires the select to be a tuple with
incremental id as its first element.
"""
def ecto(repo, query, options) do
batch_size = options |> Keyword.get(:batch_size, 1000)
strategy = options |> Keyword.get(:strategy, :offset)
case strategy do
:offset -> ecto_in_batches(repo, query, batch_size) |> Stream.concat()
:id -> ecto_in_batches_by_id(repo, query, batch_size) |> Stream.concat()
end
end
@doc """
Stream data in batches from an Ecto repo using a query.
"""
def ecto_in_batches(repo, query, batch_size \\ 1000) do
import Ecto.Query, only: [from: 1, from: 2]
Stream.unfold(0, fn
:done ->
nil
offset ->
results =
repo.all(from(_ in query, offset: ^offset, limit: ^batch_size))
if length(results) < batch_size,
do: {results, :done},
else: {results, offset + batch_size}
end)
end
def ecto_in_batches_by_id(repo, query, batch_size \\ 1000) do
import Ecto.Query, only: [from: 2]
Stream.unfold(-1, fn
:done ->
nil
last_id ->
results =
repo.all(
from(
r in query,
where: r.id > ^last_id,
limit: ^batch_size
)
)
{new_last_id, count} =
results
|> Enum.reduce({-1, 0}, fn r, {_, i} -> {elem(r, 0), i + 1} end)
if count < batch_size,
do: {results, :done},
else: {results, new_last_id}
end)
end
end
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.