gitlab-pg_query 1.3.1 → 2.0.4

data/ext/pg_query/pg_query_readfuncs.h

@@ -0,0 +1,11 @@
+#ifndef PG_QUERY_READFUNCS_H
+#define PG_QUERY_READFUNCS_H
+
+#include "pg_query.h"
+
+#include "postgres.h"
+#include "nodes/pg_list.h"
+
+List * pg_query_protobuf_to_nodes(PgQueryProtobuf protobuf);
+
+#endif

data/ext/pg_query/pg_query_readfuncs_protobuf.c

@@ -0,0 +1,142 @@
+#include "pg_query_readfuncs.h"
+
+#include "nodes/nodes.h"
+#include "nodes/parsenodes.h"
+#include "nodes/pg_list.h"
+
+#include "protobuf/pg_query.pb-c.h"
+
+#define OUT_TYPE(typename, typename_c) PgQuery__##typename_c*
+
+#define READ_COND(typename, typename_c, typename_underscore, typename_underscore_upcase, typename_cast, outname) \
+	case PG_QUERY__NODE__NODE_##typename_underscore_upcase: \
+		return (Node *) _read##typename_c(msg->outname);
+
+#define READ_INT_FIELD(outname, outname_json, fldname) node->fldname = msg->outname;
+#define READ_UINT_FIELD(outname, outname_json, fldname) node->fldname = msg->outname;
+#define READ_LONG_FIELD(outname, outname_json, fldname) node->fldname = msg->outname;
+#define READ_FLOAT_FIELD(outname, outname_json, fldname) node->fldname = msg->outname;
+#define READ_BOOL_FIELD(outname, outname_json, fldname) node->fldname = msg->outname;
+
+#define READ_CHAR_FIELD(outname, outname_json, fldname) \
+	if (msg->outname != NULL && strlen(msg->outname) > 0) { \
+		node->fldname = msg->outname[0]; \
+	}
+
+#define READ_STRING_FIELD(outname, outname_json, fldname) \
+	if (msg->outname != NULL && strlen(msg->outname) > 0) { \
+		node->fldname = pstrdup(msg->outname); \
+	}
+
+#define READ_ENUM_FIELD(typename, outname, outname_json, fldname) \
+	node->fldname = _intToEnum##typename(msg->outname);
+
+#define READ_LIST_FIELD(outname, outname_json, fldname) \
+	{ \
+		if (msg->n_##outname > 0) \
+			node->fldname = list_make1(_readNode(msg->outname[0])); \
+	    for (int i = 1; i < msg->n_##outname; i++) \
+			node->fldname = lappend(node->fldname, _readNode(msg->outname[i])); \
+	}
+
+#define READ_BITMAPSET_FIELD(outname, outname_json, fldname) // FIXME
+
+#define READ_NODE_FIELD(outname, outname_json, fldname) \
+	node->fldname = *_readNode(msg->outname);
+
+#define READ_NODE_PTR_FIELD(outname, outname_json, fldname) \
+	if (msg->outname != NULL) { \
+		node->fldname = _readNode(msg->outname); \
+	}
+
+#define READ_EXPR_PTR_FIELD(outname, outname_json, fldname) \
+	if (msg->outname != NULL) { \
+		node->fldname = (Expr *) _readNode(msg->outname); \
+	}
+
+#define READ_VALUE_FIELD(outname, outname_json, fldname) \
+	if (msg->outname != NULL) { \
+		node->fldname = *((Value *) _readNode(msg->outname)); \
+	}
+
+#define READ_VALUE_PTR_FIELD(outname, outname_json, fldname) \
+	if (msg->outname != NULL) { \
+		node->fldname = (Value *) _readNode(msg->outname); \
+	}
+
+#define READ_SPECIFIC_NODE_FIELD(typename, typename_underscore, outname, outname_json, fldname) \
+	node->fldname = *_read##typename(msg->outname);
+
+#define READ_SPECIFIC_NODE_PTR_FIELD(typename, typename_underscore, outname, outname_json, fldname) \
+	if (msg->outname != NULL) { \
+		node->fldname = _read##typename(msg->outname); \
+	}
+
+static Node * _readNode(PgQuery__Node *msg);
+
+#include "pg_query_enum_defs.c"
+#include "pg_query_readfuncs_defs.c"
+
+static List * _readList(PgQuery__List *msg)
+{
+	List *node = NULL;
+	if (msg->n_items > 0)
+		node = list_make1(_readNode(msg->items[0]));
+	for (int i = 1; i < msg->n_items; i++)
+		node = lappend(node, _readNode(msg->items[i]));
+	return node;
+}
+
+static Node * _readNode(PgQuery__Node *msg)
+{
+	switch (msg->node_case)
+	{
+		#include "pg_query_readfuncs_conds.c"
+
+		case PG_QUERY__NODE__NODE_INTEGER:
+			return (Node *) makeInteger(msg->integer->ival);
+		case PG_QUERY__NODE__NODE_FLOAT:
+			return (Node *) makeFloat(pstrdup(msg->float_->str));
+		case PG_QUERY__NODE__NODE_STRING:
+			return (Node *) makeString(pstrdup(msg->string->str));
+		case PG_QUERY__NODE__NODE_BIT_STRING:
+			return (Node *) makeBitString(pstrdup(msg->bit_string->str));
+		case PG_QUERY__NODE__NODE_NULL:
+			{
+				Value *v = makeNode(Value);
+				v->type = T_Null;
+				return (Node *) v;
+			}
+		case PG_QUERY__NODE__NODE_LIST:
+			return (Node *) _readList(msg->list);
+		case PG_QUERY__NODE__NODE__NOT_SET:
+			return NULL;
+		default:
+			elog(ERROR, "unsupported protobuf node type: %d",
+				 (int) msg->node_case);
+	}
+}
+
+List * pg_query_protobuf_to_nodes(PgQueryProtobuf protobuf)
+{
+	PgQuery__ParseResult *result = NULL;
+	List * list = NULL;
+	size_t i = 0;
+
+	result = pg_query__parse_result__unpack(NULL, protobuf.len, (const uint8_t *) protobuf.data);
+
+	// TODO: Handle this by returning an error instead
+	Assert(result != NULL);
+
+	// TODO: Handle this by returning an error instead
+	Assert(result->version == PG_VERSION_NUM);
+
+	if (result->n_stmts > 0)
+		list = list_make1(_readRawStmt(result->stmts[0]));
+    for (i = 1; i < result->n_stmts; i++)
+		list = lappend(list, _readRawStmt(result->stmts[i]));
+
+	pg_query__parse_result__free_unpacked(result, NULL);
+
+	return list;
+}

data/ext/pg_query/pg_query_ruby.c

@@ -1,12 +1,18 @@
-#include "pg_query_ruby.h"
+#include "pg_query.h"
+#include "xxhash/xxhash.h"
+#include <ruby.h>
 
-void raise_ruby_parse_error(PgQueryParseResult result);
+void raise_ruby_parse_error(PgQueryProtobufParseResult result);
 void raise_ruby_normalize_error(PgQueryNormalizeResult result);
 void raise_ruby_fingerprint_error(PgQueryFingerprintResult result);
+void raise_ruby_scan_error(PgQueryScanResult result);
 
-VALUE pg_query_ruby_parse(VALUE self, VALUE input);
+VALUE pg_query_ruby_parse_protobuf(VALUE self, VALUE input);
+VALUE pg_query_ruby_deparse_protobuf(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);
+VALUE pg_query_ruby_hash_xxh3_64(VALUE self, VALUE input, VALUE seed);
 
 void Init_pg_query(void)
 {
@@ -14,12 +20,18 @@ void Init_pg_query(void)
 
 	cPgQuery = rb_const_get(rb_cObject, rb_intern("PgQuery"));
 
-	rb_define_singleton_method(cPgQuery, "_raw_parse", pg_query_ruby_parse, 1);
+	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, "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);
+	rb_define_singleton_method(cPgQuery, "hash_xxh3_64", pg_query_ruby_hash_xxh3_64, 2);
+	rb_define_const(cPgQuery, "PG_VERSION", rb_str_new2(PG_VERSION));
+	rb_define_const(cPgQuery, "PG_MAJORVERSION", rb_str_new2(PG_MAJORVERSION));
+	rb_define_const(cPgQuery, "PG_VERSION_NUM", INT2NUM(PG_VERSION_NUM));
 }
 
-void raise_ruby_parse_error(PgQueryParseResult result)
+void raise_ruby_parse_error(PgQueryProtobufParseResult result)
 {
 	VALUE cPgQuery, cParseError;
 	VALUE args[4];
@@ -32,7 +44,25 @@ void raise_ruby_parse_error(PgQueryParseResult result)
 	args[2] = INT2NUM(result.error->lineno);
 	args[3] = INT2NUM(result.error->cursorpos);
 
-	pg_query_free_parse_result(result);
+	pg_query_free_protobuf_parse_result(result);
+
+	rb_exc_raise(rb_class_new_instance(4, args, cParseError));
+}
+
+void raise_ruby_deparse_error(PgQueryDeparseResult 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_result(result);
 
 	rb_exc_raise(rb_class_new_instance(4, args, cParseError));
 }
@@ -73,21 +103,60 @@ void raise_ruby_fingerprint_error(PgQueryFingerprintResult result)
 	rb_exc_raise(rb_class_new_instance(4, args, cParseError));
 }
 
-VALUE pg_query_ruby_parse(VALUE self, VALUE input)
+void raise_ruby_scan_error(PgQueryScanResult result)
+{
+	VALUE cPgQuery, cScanError;
+	VALUE args[4];
+
+	cPgQuery   = rb_const_get(rb_cObject, rb_intern("PgQuery"));
+	cScanError = rb_const_get_at(cPgQuery, rb_intern("ScanError"));
+
+	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_scan_result(result);
+
+	rb_exc_raise(rb_class_new_instance(4, args, cScanError));
+}
+
+VALUE pg_query_ruby_parse_protobuf(VALUE self, VALUE input)
 {
 	Check_Type(input, T_STRING);
 
 	VALUE output;
-	PgQueryParseResult result = pg_query_parse(StringValueCStr(input));
+	PgQueryProtobufParseResult result = pg_query_parse_protobuf(StringValueCStr(input));
 
 	if (result.error) raise_ruby_parse_error(result);
 
 	output = rb_ary_new();
 
-	rb_ary_push(output, rb_str_new2(result.parse_tree));
+	rb_ary_push(output, rb_str_new(result.parse_tree.data, result.parse_tree.len));
 	rb_ary_push(output, rb_str_new2(result.stderr_buffer));
 
-	pg_query_free_parse_result(result);
+	pg_query_free_protobuf_parse_result(result);
+
+	return output;
+}
+
+VALUE pg_query_ruby_deparse_protobuf(VALUE self, VALUE input)
+{
+	Check_Type(input, T_STRING);
+
+	VALUE output;
+	PgQueryProtobuf pbuf = {0};
+	PgQueryDeparseResult result = {0};
+
+	pbuf.data = StringValuePtr(input);
+	pbuf.len = RSTRING_LEN(input);
+	result = pg_query_deparse_protobuf(pbuf);
+
+	if (result.error) raise_ruby_deparse_error(result);
+
+	output = rb_str_new2(result.query);
+
+	pg_query_free_deparse_result(result);
 
 	return output;
 }
@@ -117,8 +186,8 @@ VALUE pg_query_ruby_fingerprint(VALUE self, VALUE input)
 
 	if (result.error) raise_ruby_fingerprint_error(result);
 
-	if (result.hexdigest) {
-		output = rb_str_new2(result.hexdigest);
+	if (result.fingerprint_str) {
+		output = rb_str_new2(result.fingerprint_str);
 	} else {
 		output = Qnil;
 	}
@@ -127,3 +196,30 @@ VALUE pg_query_ruby_fingerprint(VALUE self, VALUE input)
 
 	return output;
 }
+
+VALUE pg_query_ruby_scan(VALUE self, VALUE input)
+{
+	Check_Type(input, T_STRING);
+
+	VALUE output;
+	PgQueryScanResult result = pg_query_scan(StringValueCStr(input));
+
+	if (result.error) raise_ruby_scan_error(result);
+
+	output = rb_ary_new();
+
+	rb_ary_push(output, rb_str_new(result.pbuf.data, result.pbuf.len));
+	rb_ary_push(output, rb_str_new2(result.stderr_buffer));
+
+	pg_query_free_scan_result(result);
+
+	return output;
+}
+
+VALUE pg_query_ruby_hash_xxh3_64(VALUE self, VALUE input, VALUE seed)
+{
+	Check_Type(input, T_STRING);
+	Check_Type(seed, T_FIXNUM);
+
+	return ULONG2NUM(XXH3_64bits_withSeed(StringValuePtr(input), RSTRING_LEN(input), NUM2ULONG(seed)));
+}

data/ext/pg_query/pg_query_scan.c

@@ -0,0 +1,173 @@
+#include "pg_query.h"
+#include "pg_query_internal.h"
+
+#include "parser/gramparse.h"
+#include "lib/stringinfo.h"
+
+#include "protobuf/pg_query.pb-c.h"
+
+#include <unistd.h>
+#include <fcntl.h>
+
+/* This is ugly. We need to access yyleng outside of scan.l, and casting yyscanner
+   to this internal struct seemed like one way to do it... */
+struct yyguts_t
+{
+  void *yyextra_r;
+  FILE *yyin_r, *yyout_r;
+  size_t yy_buffer_stack_top; /**< index of top of stack. */
+  size_t yy_buffer_stack_max; /**< capacity of stack. */
+  struct yy_buffer_state *yy_buffer_stack; /**< Stack as an array. */
+  char yy_hold_char;
+  size_t yy_n_chars;
+  size_t yyleng_r;
+};
+
+PgQueryScanResult pg_query_scan(const char* input)
+{
+  MemoryContext ctx = NULL;
+  PgQueryScanResult result = {0};
+  core_yyscan_t yyscanner;
+  core_yy_extra_type yyextra;
+  core_YYSTYPE yylval;
+  YYLTYPE    yylloc;
+  PgQuery__ScanResult scan_result = PG_QUERY__SCAN_RESULT__INIT;
+  PgQuery__ScanToken **output_tokens;
+  size_t token_count = 0;
+  size_t i;
+
+  ctx = pg_query_enter_memory_context();
+
+  MemoryContext parse_context = CurrentMemoryContext;
+
+  char stderr_buffer[STDERR_BUFFER_LEN + 1] = {0};
+#ifndef DEBUG
+  int stderr_global;
+  int stderr_pipe[2];
+#endif
+
+#ifndef DEBUG
+  // Setup pipe for stderr redirection
+  if (pipe(stderr_pipe) != 0) {
+    PgQueryError* error = malloc(sizeof(PgQueryError));
+
+    error->message = strdup("Failed to open pipe, too many open file descriptors")
+
+    result.error = error;
+
+    return result;
+  }
+
+  fcntl(stderr_pipe[0], F_SETFL, fcntl(stderr_pipe[0], F_GETFL) | O_NONBLOCK);
+
+  // Redirect stderr to the pipe
+  stderr_global = dup(STDERR_FILENO);
+  dup2(stderr_pipe[1], STDERR_FILENO);
+  close(stderr_pipe[1]);
+#endif
+
+  PG_TRY();
+  {
+    // Really this is stupid, we only run twice so we can pre-allocate the output array correctly
+    yyscanner = scanner_init(input, &yyextra, &ScanKeywords, ScanKeywordTokens);
+    for (;; token_count++)
+    {
+      if (core_yylex(&yylval, &yylloc, yyscanner) == 0) break;
+    }
+    scanner_finish(yyscanner);
+
+    output_tokens = malloc(sizeof(PgQuery__ScanToken *) * token_count);
+
+    /* initialize the flex scanner --- should match raw_parser() */
+    yyscanner = scanner_init(input, &yyextra, &ScanKeywords, ScanKeywordTokens);
+
+    /* Lex tokens  */
+    for (i = 0; ; i++)
+    {
+      int tok;
+      int keyword;
+
+      tok = core_yylex(&yylval, &yylloc, yyscanner);
+      if (tok == 0) break;
+
+      output_tokens[i] = malloc(sizeof(PgQuery__ScanToken));
+      pg_query__scan_token__init(output_tokens[i]);
+      output_tokens[i]->start = yylloc;
+      if (tok == SCONST || tok == BCONST || tok == XCONST || tok == IDENT || tok == C_COMMENT) {
+        output_tokens[i]->end = yyextra.yyllocend;
+      } else {
+        output_tokens[i]->end = yylloc + ((struct yyguts_t*) yyscanner)->yyleng_r;
+      }
+      output_tokens[i]->token = tok;
+
+      switch (tok) {
+      #define PG_KEYWORD(a,b,c) case b: output_tokens[i]->keyword_kind = c + 1; break;
+      #include "parser/kwlist.h"
+      default: output_tokens[i]->keyword_kind = 0;
+      }
+    }
+
+    scanner_finish(yyscanner);
+
+    scan_result.version = PG_VERSION_NUM;
+    scan_result.n_tokens = token_count;
+    scan_result.tokens = output_tokens;
+    result.pbuf.len = pg_query__scan_result__get_packed_size(&scan_result);
+    result.pbuf.data = malloc(result.pbuf.len);
+    pg_query__scan_result__pack(&scan_result, (void*) result.pbuf.data);
+
+    for (i = 0; i < token_count; i++) {
+      free(output_tokens[i]);
+    }
+    free(output_tokens);
+
+#ifndef DEBUG
+    // Save stderr for result
+    read(stderr_pipe[0], stderr_buffer, STDERR_BUFFER_LEN);
+#endif
+
+    result.stderr_buffer = strdup(stderr_buffer);
+  }
+  PG_CATCH();
+  {
+    ErrorData* error_data;
+    PgQueryError* error;
+
+    MemoryContextSwitchTo(parse_context);
+    error_data = CopyErrorData();
+
+    // Note: This is intentionally malloc so exiting the memory context doesn't free this
+    error = malloc(sizeof(PgQueryError));
+    error->message   = strdup(error_data->message);
+    error->filename  = strdup(error_data->filename);
+    error->funcname  = strdup(error_data->funcname);
+    error->context   = NULL;
+    error->lineno    = error_data->lineno;
+    error->cursorpos = error_data->cursorpos;
+
+    result.error = error;
+    FlushErrorState();
+  }
+  PG_END_TRY();
+
+#ifndef DEBUG
+  // Restore stderr, close pipe
+  dup2(stderr_global, STDERR_FILENO);
+  close(stderr_pipe[0]);
+  close(stderr_global);
+#endif
+
+  pg_query_exit_memory_context(ctx);
+
+  return result;
+}
+
+void pg_query_free_scan_result(PgQueryScanResult result)
+{
+  if (result.error) {
+    pg_query_free_error(result.error);
+  }
+
+  free(result.pbuf.data);
+  free(result.stderr_buffer);
+}