retrieving columns name from the result

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<duckdb::QueryResult>* 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<ERL_NIF_TERM> 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<ERL_NIF_TERM> rows;
 
   duckdb::unique_ptr<duckdb::DataChunk> 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<ERL_NIF_TERM> rows;
 
   duckdb::unique_ptr<duckdb::DataChunk> 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},

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)
 

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()}

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

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)

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