From 830cd347686b23ab6efe7e80bf98f12a106a89d9 Mon Sep 17 00:00:00 2001 From: Aleksandr Golubov Date: Tue, 28 Nov 2023 23:35:08 +0500 Subject: [PATCH] retrieving columns name from the result --- c_src/nif.cpp | 37 ++++++++++++++++++++++++ lib/duckdbex.ex | 18 ++++++++++-- lib/nif.ex | 7 +++-- mix.exs | 2 +- test/duckdbex_test.exs | 7 +++++ test/nif/columns_test.exs | 60 +++++++++++++++++++++++++++++++++++++++ 6 files changed, 126 insertions(+), 5 deletions(-) create mode 100644 test/nif/columns_test.exs diff --git a/c_src/nif.cpp b/c_src/nif.cpp index 704bfcb..2efec5b 100644 --- a/c_src/nif.cpp +++ b/c_src/nif.cpp @@ -239,6 +239,32 @@ execute_statement(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]) { return nif::make_ok_tuple(env, resource_builder.make_and_release_resource(env)); } +static ERL_NIF_TERM +columns(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]) { + if (argc != 1) + return enif_make_badarg(env); + + erlang_resource* result = nullptr; + if(!enif_get_resource(env, argv[0], query_result_nif_type, (void**)&result)) + return enif_make_badarg(env); + + if (result->data->HasError()) { + auto error = result->data->GetError(); + return nif::make_error_tuple(env, error); + } + + if (duckdb::idx_t columns_count = result->data->ColumnCount()) { + std::vector columns(columns_count); + for (duckdb::idx_t col = 0; col < columns_count; col++) { + duckdb::string column_name = result->data->ColumnName(col); + columns[col] = nif::make_binary_term(env, column_name); + } + return enif_make_list_from_array(env, &columns[0], columns.size()); + } else { + return enif_make_list(env, 0); + } +} + static ERL_NIF_TERM fetch_chunk(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]) { if (argc != 1) @@ -248,6 +274,11 @@ fetch_chunk(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]) { if(!enif_get_resource(env, argv[0], query_result_nif_type, (void**)&result)) return enif_make_badarg(env); + if (result->data->HasError()) { + auto error = result->data->GetError(); + return nif::make_error_tuple(env, error); + } + std::vector rows; duckdb::unique_ptr chunk; @@ -284,6 +315,11 @@ fetch_all(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]) { if(!enif_get_resource(env, argv[0], query_result_nif_type, (void**)&result)) return enif_make_badarg(env); + if (result->data->HasError()) { + auto error = result->data->GetError(); + return nif::make_error_tuple(env, error); + } + std::vector rows; duckdb::unique_ptr chunk; @@ -562,6 +598,7 @@ static ErlNifFunc nif_funcs[] = { {"prepare_statement", 2, prepare_statement, ERL_NIF_DIRTY_JOB_IO_BOUND}, {"execute_statement", 1, execute_statement, ERL_NIF_DIRTY_JOB_IO_BOUND}, {"execute_statement", 2, execute_statement, ERL_NIF_DIRTY_JOB_IO_BOUND}, + {"columns", 1, columns, ERL_NIF_DIRTY_JOB_IO_BOUND}, {"fetch_chunk", 1, fetch_chunk, ERL_NIF_DIRTY_JOB_IO_BOUND}, {"fetch_all", 1, fetch_all, ERL_NIF_DIRTY_JOB_IO_BOUND}, {"appender", 2, appender, ERL_NIF_DIRTY_JOB_IO_BOUND}, diff --git a/lib/duckdbex.ex b/lib/duckdbex.ex index 4defcf0..1e77233 100644 --- a/lib/duckdbex.ex +++ b/lib/duckdbex.ex @@ -139,6 +139,20 @@ defmodule Duckdbex do def execute_statement(statement, args) when is_reference(statement) and is_list(args), do: Duckdbex.NIF.execute_statement(statement, args) + @doc """ + Returns columns names from the query result. + + ## Examples + + iex> {:ok, db} = Duckdbex.open() + iex> {:ok, conn} = Duckdbex.connection(db) + iex> {:ok, res} = Duckdbex.query(conn, "SELECT 1 as 'my_name';") + iex> ["my_name"] = Duckdbex.columns(res) + """ + @spec columns(query_result()) :: list() | {:error, reason()} + def columns(query_result) when is_reference(query_result), + do: Duckdbex.NIF.columns(query_result) + @doc """ Fetches a data chunk from the query result. @@ -151,7 +165,7 @@ defmodule Duckdbex do iex> {:ok, res} = Duckdbex.query(conn, "SELECT 1;") iex> [[1]] = Duckdbex.fetch_chunk(res) """ - @spec fetch_chunk(query_result()) :: :ok | {:error, reason()} + @spec fetch_chunk(query_result()) :: list() | {:error, reason()} def fetch_chunk(query_result) when is_reference(query_result), do: Duckdbex.NIF.fetch_chunk(query_result) @@ -167,7 +181,7 @@ defmodule Duckdbex do iex> {:ok, res} = Duckdbex.query(conn, "SELECT 1;") iex> [[1]] = Duckdbex.fetch_all(res) """ - @spec fetch_all(query_result()) :: :ok | {:error, reason()} + @spec fetch_all(query_result()) :: list() | {:error, reason()} def fetch_all(query_result) when is_reference(query_result), do: Duckdbex.NIF.fetch_all(query_result) diff --git a/lib/nif.ex b/lib/nif.ex index 9e091bb..f765145 100644 --- a/lib/nif.ex +++ b/lib/nif.ex @@ -34,10 +34,13 @@ defmodule Duckdbex.NIF do @spec execute_statement(statement(), list()) :: {:ok, query_result()} | {:error, reason()} def execute_statement(_statement, _args), do: :erlang.nif_error(:not_loaded) - @spec fetch_chunk(query_result()) :: :ok | {:error, reason()} + @spec columns(query_result()) :: list(binary()) | {:error, reason()} + def columns(_query_result), do: :erlang.nif_error(:not_loaded) + + @spec fetch_chunk(query_result()) :: list() | {:error, reason()} def fetch_chunk(_query_result), do: :erlang.nif_error(:not_loaded) - @spec fetch_all(query_result()) :: :ok | {:error, reason()} + @spec fetch_all(query_result()) :: list() | {:error, reason()} def fetch_all(_query_result), do: :erlang.nif_error(:not_loaded) @spec appender(connection(), binary()) :: {:ok, appender()} | {:error, reason()} diff --git a/mix.exs b/mix.exs index 7a00074..052a4cd 100644 --- a/mix.exs +++ b/mix.exs @@ -1,7 +1,7 @@ defmodule Duckdbex.MixProject do use Mix.Project - @version "0.2.6" + @version "0.2.7" @duckdb_version "0.9.2" def project do diff --git a/test/duckdbex_test.exs b/test/duckdbex_test.exs index 94bd922..732069d 100644 --- a/test/duckdbex_test.exs +++ b/test/duckdbex_test.exs @@ -32,6 +32,13 @@ defmodule DuckdbexTest do assert {:ok, _res} = Duckdbex.query(conn, "SELECT 1 WHERE 1 = $1;", [1]) end + test "columns/1" do + assert {:ok, db} = Duckdbex.open() + assert {:ok, conn} = Duckdbex.connection(db) + assert {:ok, result} = Duckdbex.query(conn, "SELECT 1 as 'one_column_name', 2 WHERE 1 = $1;", [1]) + assert ["one_column_name", "2"] = Duckdbex.columns(result) + end + test "fetch_chunk/1" do assert {:ok, db} = Duckdbex.open() assert {:ok, conn} = Duckdbex.connection(db) diff --git a/test/nif/columns_test.exs b/test/nif/columns_test.exs new file mode 100644 index 0000000..f8c899a --- /dev/null +++ b/test/nif/columns_test.exs @@ -0,0 +1,60 @@ +defmodule Duckdbex.Nif.ColumnsTest do + use ExUnit.Case + + alias Duckdbex.NIF + + setup ctx do + {:ok, db} = NIF.open(":memory:", nil) + {:ok, conn} = NIF.connection(db) + Map.put(ctx, :conn, conn) + end + + test "when table is empty, returns columns names", %{conn: conn} do + {:ok, _} = + Duckdbex.NIF.query(conn, """ + CREATE TABLE columns_test(count BIGINT, is_ready BOOLEAN, name VARCHAR); + """) + + {:ok, result_ref} = Duckdbex.NIF.query(conn, "SELECT * FROM columns_test") + + assert ["count", "is_ready", "name"] = Duckdbex.NIF.columns(result_ref) + end + + test "when table has rows, returns columns names", %{conn: conn} do + {:ok, _} = + Duckdbex.NIF.query(conn, """ + CREATE TABLE columns_test(count BIGINT, is_ready BOOLEAN, name VARCHAR); + """) + + {:ok, _} = + Duckdbex.NIF.query(conn, """ + INSERT INTO columns_test VALUES (1, true, 'one'), (2, true, 'two'); + """) + + {:ok, result_ref} = Duckdbex.NIF.query(conn, "SELECT name, count FROM columns_test") + + assert ["name", "count"] = Duckdbex.NIF.columns(result_ref) + end + + test "when select at different column name, returns specified column name", %{conn: conn} do + {:ok, _} = + Duckdbex.NIF.query(conn, """ + CREATE TABLE columns_test(count BIGINT, is_ready BOOLEAN, name VARCHAR); + """) + + {:ok, _} = + Duckdbex.NIF.query(conn, """ + INSERT INTO columns_test VALUES (1, true, 'one'), (2, true, 'two'); + """) + + {:ok, result_ref} = Duckdbex.NIF.query(conn, "SELECT name as my_name FROM columns_test") + + assert ["my_name"] = Duckdbex.NIF.columns(result_ref) + end + + test "when select constants, returns constanst itself as columns names", %{conn: conn} do + {:ok, result_ref} = Duckdbex.NIF.query(conn, "SELECT 1, 'two', 3.14") + + assert ["1", "'two'", "3.14"] = Duckdbex.NIF.columns(result_ref) + end +end