feat: expose scan function

Author: icehaunter

c_src/libpg_query_ex.c

@@ -113,9 +113,62 @@ static ERL_NIF_TERM deparse_query(ErlNifEnv *env, int argc,
   }
 }
 
+static ERL_NIF_TERM scan_query(ErlNifEnv *env, int argc,
+                             const ERL_NIF_TERM argv[]) {
+  ErlNifBinary query;
+  ERL_NIF_TERM term;
+
+  if (argc == 1 && enif_inspect_binary(env, argv[0], &query)) {
+    // add one more byte for the null termination
+    char statement[query.size + 1];
+
+    strncpy(statement, (char *)query.data, query.size);
+
+    // terminate the string
+    statement[query.size] = 0;
+
+    PgQueryScanResult result = pg_query_scan(statement);
+
+    if (result.error) {
+      ERL_NIF_TERM error_map = enif_make_new_map(env);
+
+      if (!enif_make_map_put(
+        env,
+        error_map,
+        enif_make_atom(env, "message"), 
+        make_binary(env, result.error->message),
+        &error_map
+      )) {
+        return enif_raise_exception(env, make_binary(env, "failed to update map"));
+      }
+
+      if (result.error->cursorpos > 0 && !enif_make_map_put(
+        env,
+        error_map,
+        enif_make_atom(env, "cursorpos"), 
+        // drop the cursorpos by one, so it's zero-indexed
+        enif_make_int(env, result.error->cursorpos - 1),
+        &error_map
+      )) {
+        return enif_raise_exception(env, make_binary(env, "failed to update map"));
+      }
+
+      term = enif_make_tuple2(env, enif_make_atom(env, "error"), error_map);
+    } else {
+      term = result_tuple(env, "ok", result.pbuf.data, result.pbuf.len);
+    }
+    pg_query_free_scan_result(result);
+
+    return term;
+  } else {
+    return enif_make_badarg(env);
+  }
+}
+
 static ErlNifFunc funcs[] = {
   {"parse_query", 1, parse_query},
-  {"deparse_query", 1, deparse_query}
+  {"deparse_query", 1, deparse_query},
+  {"scan_query", 1, scan_query}
 };
 
 ERL_NIF_INIT(Elixir.PgQuery.Parser, funcs, NULL, NULL, NULL, NULL)

lib/pg_query.ex

@@ -25,4 +25,18 @@ defmodule PgQuery do
   """
   @spec protobuf_to_query!(%PgQuery.ParseResult{}) :: String.t() | no_return()
   defdelegate protobuf_to_query!(parse_result), to: PgQuery.Parser
+
+  @doc """
+  Scans the binary statement `stmt` into tokens.
+  """
+  @spec scan(String.t()) :: {:ok, %PgQuery.ScanResult{}} | {:error, error()}
+  defdelegate scan(stmt), to: PgQuery.Parser
+
+  @doc """
+  Scans the binary statement `stmt` into tokens.
+
+  Raises if the statement is invalid.
+  """
+  @spec scan!(String.t()) :: %PgQuery.ScanResult{} | no_return()
+  defdelegate scan!(stmt), to: PgQuery.Parser
 end

lib/pg_query/parser.ex

@@ -48,6 +48,23 @@ defmodule PgQuery.Parser do
     end
   end
 
+  def scan(query) when is_binary(query) do
+    with {:ok, proto} <- scan_query(query) do
+      Protox.decode(proto, PgQuery.ScanResult)
+    end
+  end
+
+  def scan!(query) when is_binary(query) do
+    case scan_query(query) do
+      {:ok, proto} ->
+        Protox.decode!(proto, PgQuery.ScanResult)
+
+      {:error, %{message: message}} ->
+        raise Error, message: message
+    end
+  end
+
   def parse_query(_query), do: :erlang.nif_error(:nif_not_loaded)
   def deparse_query(_encoded_proto), do: :erlang.nif_error(:nif_not_loaded)
+  def scan_query(_query), do: :erlang.nif_error(:nif_not_loaded)
 end

mix.exs

@@ -5,7 +5,7 @@ defmodule PgQuery.MixProject do
     [
       app: :pg_query_ex,
       elixir: "~> 1.13",
-      version: "0.6.0",
+      version: "0.6.1",
       start_permanent: Mix.env() == :prod,
       compilers: [:elixir_make] ++ Mix.compilers(),
       deps: deps(),

test/pg_query_test.exs

@@ -39,4 +39,30 @@ defmodule PgQueryTest do
     assert {:ok, query2} = PgQuery.protobuf_to_query(ast)
     assert query == query2
   end
+
+  test "parses and deparses a query with bang function" do
+    query = "CREATE TABLE a (id int8 PRIMARY KEY)"
+    ast = PgQuery.parse!(query)
+    query2 = PgQuery.protobuf_to_query!(ast)
+    assert query == query2
+  end
+
+  test "scans a query" do
+    query = "SELECT * FROM users WHERE id = 1"
+    assert {:ok, scan_result} = PgQuery.scan(query)
+
+    # Basic structure validation
+    assert %PgQuery.ScanResult{tokens: tokens} = scan_result
+    assert is_list(tokens)
+    assert length(tokens) == 8
+
+    # Check some expected tokens
+    assert Enum.any?(tokens, fn token ->
+             token.token == :SELECT
+           end)
+
+    assert Enum.any?(tokens, fn token ->
+             token.token == :FROM
+           end)
+  end
 end