Deparser: Support pretty-printing and other deparser options

Author: lfittl

ext/pg_query/pg_query_ruby.c

@@ -1,4 +1,5 @@
 #include "pg_query.h"
+#include "postgres_deparse.h"
 #include "xxhash/xxhash.h"
 #include <ruby.h>
 
@@ -10,6 +11,8 @@ void raise_ruby_split_error(PgQuerySplitResult result);
 
 VALUE pg_query_ruby_parse_protobuf(VALUE self, VALUE input);
 VALUE pg_query_ruby_deparse_protobuf(VALUE self, VALUE input);
+VALUE pg_query_ruby_deparse_protobuf_opts(VALUE self, VALUE input, VALUE pretty_print, VALUE comments, VALUE indent_size, VALUE max_line_length, VALUE trailing_newline, VALUE commas_start_of_line);
+VALUE pg_query_ruby_deparse_comments_for_query(VALUE self, VALUE input);
 VALUE pg_query_ruby_normalize(VALUE self, VALUE input);
 VALUE pg_query_ruby_fingerprint(VALUE self, VALUE input);
 VALUE pg_query_ruby_scan(VALUE self, VALUE input);
@@ -24,6 +27,8 @@ __attribute__((visibility ("default"))) void Init_pg_query(void)
 
 	rb_define_singleton_method(cPgQuery, "parse_protobuf", pg_query_ruby_parse_protobuf, 1);
 	rb_define_singleton_method(cPgQuery, "deparse_protobuf", pg_query_ruby_deparse_protobuf, 1);
+	rb_define_singleton_method(cPgQuery, "deparse_protobuf_opts", pg_query_ruby_deparse_protobuf_opts, 7);
+	rb_define_singleton_method(cPgQuery, "deparse_comments_for_query", pg_query_ruby_deparse_comments_for_query, 1);
 	rb_define_singleton_method(cPgQuery, "normalize", pg_query_ruby_normalize, 1);
 	rb_define_singleton_method(cPgQuery, "fingerprint", pg_query_ruby_fingerprint, 1);
 	rb_define_singleton_method(cPgQuery, "_raw_scan", pg_query_ruby_scan, 1);
@@ -70,6 +75,24 @@ void raise_ruby_deparse_error(PgQueryDeparseResult result)
 	rb_exc_raise(rb_class_new_instance(4, args, cParseError));
 }
 
+void raise_ruby_deparse_comments_error(PgQueryDeparseCommentsResult result)
+{
+	VALUE cPgQuery, cParseError;
+	VALUE args[4];
+
+	cPgQuery    = rb_const_get(rb_cObject, rb_intern("PgQuery"));
+	cParseError = rb_const_get_at(cPgQuery, rb_intern("ParseError"));
+
+	args[0] = rb_str_new2(result.error->message);
+	args[1] = rb_str_new2(result.error->filename);
+	args[2] = INT2NUM(result.error->lineno);
+	args[3] = INT2NUM(result.error->cursorpos);
+
+	pg_query_free_deparse_comments_result(result);
+
+	rb_exc_raise(rb_class_new_instance(4, args, cParseError));
+}
+
 void raise_ruby_normalize_error(PgQueryNormalizeResult result)
 {
 	VALUE cPgQuery, cParseError;
@@ -182,6 +205,83 @@ VALUE pg_query_ruby_deparse_protobuf(VALUE self, VALUE input)
 	return output;
 }
 
+VALUE pg_query_ruby_deparse_protobuf_opts(VALUE self, VALUE input, VALUE pretty_print, VALUE comments, VALUE indent_size, VALUE max_line_length, VALUE trailing_newline, VALUE commas_start_of_line)
+{
+	Check_Type(input, T_STRING);
+	Check_Type(comments, T_ARRAY);
+	Check_Type(indent_size, T_FIXNUM);
+	Check_Type(max_line_length, T_FIXNUM);
+
+	VALUE output;
+	PgQueryProtobuf pbuf = {0};
+	PgQueryDeparseResult result = {0};
+	PostgresDeparseOpts deparse_opts = {0};
+	deparse_opts.pretty_print = RTEST(pretty_print);
+	deparse_opts.indent_size = NUM2INT(indent_size);
+	deparse_opts.max_line_length = NUM2INT(max_line_length);
+	deparse_opts.trailing_newline = RTEST(trailing_newline);
+	deparse_opts.commas_start_of_line = RTEST(commas_start_of_line);
+	deparse_opts.comments = malloc(RARRAY_LEN(comments) * sizeof(PostgresDeparseComment *));
+	deparse_opts.comment_count = RARRAY_LEN(comments);
+	VALUE *comments_arr = RARRAY_PTR(comments);
+
+	for (int i = 0; i < RARRAY_LEN(comments); i++)
+	{
+		PostgresDeparseComment* comment  = malloc(sizeof(PostgresDeparseComment));
+		VALUE str_ref                    = rb_ivar_get(comments_arr[i], rb_intern("@str"));
+		comment->match_location          = NUM2INT(rb_ivar_get(comments_arr[i], rb_intern("@match_location")));
+		comment->newlines_before_comment = NUM2INT(rb_ivar_get(comments_arr[i], rb_intern("@newlines_before_comment")));
+		comment->newlines_after_comment  = NUM2INT(rb_ivar_get(comments_arr[i], rb_intern("@newlines_after_comment")));
+		comment->str                     = StringValueCStr(str_ref);
+
+		deparse_opts.comments[i] = comment;
+	}
+
+	pbuf.data = StringValuePtr(input);
+	pbuf.len = RSTRING_LEN(input);
+	result = pg_query_deparse_protobuf_opts(pbuf, deparse_opts);
+
+	if (result.error) raise_ruby_deparse_error(result);
+
+	output = rb_str_new2(result.query);
+
+	pg_query_free_deparse_result(result);
+	for (int i = 0; i < deparse_opts.comment_count; i++)
+		free(deparse_opts.comments[i]);
+	free(deparse_opts.comments);
+
+	return output;
+}
+
+VALUE pg_query_ruby_deparse_comments_for_query(VALUE self, VALUE input)
+{
+	Check_Type(input, T_STRING);
+
+	VALUE cPgQuery, cDeparseComment;
+
+	cPgQuery        = rb_const_get(rb_cObject, rb_intern("PgQuery"));
+	cDeparseComment = rb_const_get_at(cPgQuery, rb_intern("DeparseComment"));
+
+	VALUE output = rb_ary_new();
+	PgQueryDeparseCommentsResult result = pg_query_deparse_comments_for_query(StringValueCStr(input));
+	if (result.error) raise_ruby_deparse_comments_error(result);
+
+	for (int i = 0; i < result.comment_count; i++)
+	{
+		PostgresDeparseComment* comment = result.comments[i];
+		VALUE c = rb_class_new_instance(0, NULL, cDeparseComment);
+		rb_ivar_set(c, rb_intern("@str"), rb_str_new2(comment->str));
+		rb_ivar_set(c, rb_intern("@match_location"), INT2NUM(comment->match_location));
+		rb_ivar_set(c, rb_intern("@newlines_before_comment"), INT2NUM(comment->newlines_before_comment));
+		rb_ivar_set(c, rb_intern("@newlines_after_comment"), INT2NUM(comment->newlines_after_comment));
+		rb_ary_push(output, c);
+	}
+
+	pg_query_free_deparse_comments_result(result);
+
+	return output;
+}
+
 VALUE pg_query_ruby_normalize(VALUE self, VALUE input)
 {
 	Check_Type(input, T_STRING);

lib/pg_query/deparse.rb

@@ -1,18 +1,39 @@
 module PgQuery
   class ParserResult
-    def deparse
-      PgQuery.deparse(@tree)
+    def deparse(opts: nil)
+      PgQuery.deparse(@tree, opts: opts)
     end
   end
 
+  class DeparseComment
+    attr_accessor :match_location, :newlines_before_comment, :newlines_after_comment, :str
+  end
+
+  DeparseOpts = Struct.new(:pretty_print, :comments, :indent_size, :max_line_length,
+                           :trailing_newline, :commas_start_of_line, keyword_init: true)
+
   # Reconstruct all of the parsed queries into their original form
-  def self.deparse(tree)
-    if PgQuery::ParseResult.method(:encode).arity == 1
-      PgQuery.deparse_protobuf(PgQuery::ParseResult.encode(tree)).force_encoding('UTF-8')
-    elsif PgQuery::ParseResult.method(:encode).arity == -1
-      PgQuery.deparse_protobuf(PgQuery::ParseResult.encode(tree, recursion_limit: 1_000)).force_encoding('UTF-8')
+  def self.deparse(tree, opts: nil)
+    protobuf_encoded = if PgQuery::ParseResult.method(:encode).arity == 1
+                         PgQuery::ParseResult.encode(tree)
+                       elsif PgQuery::ParseResult.method(:encode).arity == -1
+                         PgQuery::ParseResult.encode(tree, recursion_limit: 1_000)
+                       else
+                         raise ArgumentError, 'Unsupported protobuf Ruby API'
+                       end
+
+    if opts
+      PgQuery.deparse_protobuf_opts(
+        protobuf_encoded,
+        opts.pretty_print,
+        opts.comments || [],
+        opts.indent_size || 0,
+        opts.max_line_length || 0,
+        opts.trailing_newline,
+        opts.commas_start_of_line
+      ).force_encoding('UTF-8')
     else
-      raise ArgumentError, 'Unsupported protobuf Ruby API'
+      PgQuery.deparse_protobuf(protobuf_encoded).force_encoding('UTF-8')
     end
   end
 

spec/lib/pg_query/deparse_pretty_print_spec.rb

@@ -0,0 +1,39 @@
+require 'spec_helper'
+
+describe PgQuery do
+  describe '.deparse' do
+    subject { PgQuery.parse(query).deparse(opts: PgQuery::DeparseOpts.new(pretty_print: true, comments: PgQuery.deparse_comments_for_query(query))) }
+
+    context 'SELECT' do
+      context 'basic statement' do
+        let(:query) do
+          <<~Q
+            SELECT a AS b
+            FROM x
+            WHERE
+                y = 5
+                AND z = y
+          Q
+        end
+
+        it { is_expected.to eq query.strip }
+      end
+    end
+
+    context 'comments' do
+      context 'multi-line' do
+        let(:query) do
+          <<~Q
+            --
+            -- Name: EXTENSION btree_gist; Type: COMMENT; Schema: -; Owner: -
+            --
+
+            COMMENT ON EXTENSION btree_gist IS 'support for indexing common datatypes in GiST'
+          Q
+        end
+
+        it { is_expected.to eq query.strip }
+      end
+    end
+  end
+end