feat: add an option to deparse a query from AST (#5)

Author: icehaunter

c_src/libpg_query_ex.c

@@ -76,6 +76,46 @@ static ERL_NIF_TERM parse_query(ErlNifEnv *env, int argc,
   }
 }
 
-static ErlNifFunc funcs[] = {{"parse_query", 1, parse_query}};
+static ERL_NIF_TERM deparse_query(ErlNifEnv *env, int argc,
+                               const ERL_NIF_TERM argv[]) {
+  ErlNifBinary proto;
+  ERL_NIF_TERM term;
+
+  if (argc == 1 && enif_inspect_binary(env, argv[0], &proto)) {
+    PgQueryProtobuf parse_tree;
+    parse_tree.data = (char *)proto.data;
+    parse_tree.len = proto.size;
+
+    PgQueryDeparseResult result = pg_query_deparse_protobuf(parse_tree);
+
+    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"));
+      }
+
+      term = enif_make_tuple2(env, enif_make_atom(env, "error"), error_map);
+    } else {
+      term = result_tuple(env, "ok", result.query, strlen(result.query));
+    }
+    pg_query_free_deparse_result(result);
+
+    return term;
+  } else {
+    return enif_make_badarg(env);
+  }
+}
+
+static ErlNifFunc funcs[] = {
+  {"parse_query", 1, parse_query},
+  {"deparse_query", 1, deparse_query}
+};
 
 ERL_NIF_INIT(Elixir.PgQuery.Parser, funcs, NULL, NULL, NULL, NULL)

lib/pg_query.ex

@@ -13,4 +13,16 @@ defmodule PgQuery do
   """
   @spec parse(String.t()) :: {:ok, %PgQuery.ParseResult{}} | {:error, error()}
   defdelegate parse(stmt), to: PgQuery.Parser
+
+  @doc """
+  Deparses the Protobuf AST `parse_result` into a query string.
+  """
+  @spec protobuf_to_query(%PgQuery.ParseResult{}) :: {:ok, String.t()} | {:error, error()}
+  defdelegate protobuf_to_query(parse_result), to: PgQuery.Parser
+
+  @doc """
+  Deparses the Protobuf AST `parse_result` into a query string.
+  """
+  @spec protobuf_to_query!(%PgQuery.ParseResult{}) :: String.t() | no_return()
+  defdelegate protobuf_to_query!(parse_result), to: PgQuery.Parser
 end

lib/pg_query/parser.ex

@@ -27,10 +27,27 @@ defmodule PgQuery.Parser do
       {:ok, proto} ->
         Protox.decode!(proto, PgQuery.ParseResult)
 
-      {:error, reason} ->
-        raise Error, message: reason
+      {:error, %{message: message}} ->
+        raise Error, message: message
+    end
+  end
+
+  def protobuf_to_query(%PgQuery.ParseResult{} = parse_result) do
+    with {:ok, encoded} <- Protox.encode(parse_result) do
+      deparse_query(IO.iodata_to_binary(encoded))
+    end
+  end
+
+  def protobuf_to_query!(%PgQuery.ParseResult{} = parse_result) do
+    case protobuf_to_query(parse_result) do
+      {:ok, query} ->
+        query
+
+      {: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)
 end

test/pg_query_test.exs

@@ -32,4 +32,11 @@ defmodule PgQueryTest do
 
     assert binary_part(query, cursorpos, 4) == "snot"
   end
+
+  test "parses and deparses a query" do
+    query = "CREATE TABLE a (id int8 PRIMARY KEY)"
+    assert {:ok, ast} = PgQuery.parse(query)
+    assert {:ok, query2} = PgQuery.protobuf_to_query(ast)
+    assert query == query2
+  end
 end