Add pg_query_is_utility_stmt() (#313)

Author: duckinator

Makefile

@@ -237,14 +237,15 @@ examples/normalize_error: examples/normalize_error.c $(ARLIB)
 examples/simple_plpgsql: examples/simple_plpgsql.c $(ARLIB)
 	$(CC) $(TEST_CFLAGS) -o $@ -g examples/simple_plpgsql.c $(ARLIB) $(TEST_LDFLAGS)
 
-TESTS = test/complex test/concurrency test/deparse test/fingerprint test/fingerprint_opts test/normalize test/normalize_utility test/parse test/parse_opts test/parse_protobuf test/parse_protobuf_opts test/parse_plpgsql test/scan test/split test/summary test/summary_truncate
+TESTS = test/complex test/concurrency test/deparse test/fingerprint test/fingerprint_opts test/is_utility_stmt test/normalize test/normalize_utility test/parse test/parse_opts test/parse_protobuf test/parse_protobuf_opts test/parse_plpgsql test/scan test/split test/summary test/summary_truncate
 test: $(TESTS)
 ifeq ($(VALGRIND),1)
 	$(VALGRIND_MEMCHECK) test/complex || (cat test/valgrind.log && false)
 	$(VALGRIND_MEMCHECK) test/concurrency || (cat test/valgrind.log && false)
 	$(VALGRIND_MEMCHECK) test/deparse || (cat test/valgrind.log && false)
 	$(VALGRIND_MEMCHECK) test/fingerprint || (cat test/valgrind.log && false)
 	$(VALGRIND_MEMCHECK) test/fingerprint_opts || (cat test/valgrind.log && false)
+	$(VALGRIND_MEMCHECK) test/is_utility_stmt || (cat test/valgrind.log && false)
 	$(VALGRIND_MEMCHECK) test/normalize || (cat test/valgrind.log && false)
 	$(VALGRIND_MEMCHECK) test/normalize_utility || (cat test/valgrind.log && false)
 	$(VALGRIND_MEMCHECK) test/parse || (cat test/valgrind.log && false)
@@ -264,6 +265,7 @@ else
 	test/deparse
 	test/fingerprint
 	test/fingerprint_opts
+	test/is_utility_stmt
 	test/normalize
 	test/normalize_utility
 	test/parse
@@ -297,6 +299,10 @@ test/fingerprint_opts: test/fingerprint_opts.c test/fingerprint_opts_tests.c $(A
 	# We have "-Isrc/" because this test uses pg_query_fingerprint_with_opts
 	$(CC) $(TEST_CFLAGS) -o $@ -Isrc/ test/fingerprint_opts.c $(ARLIB) $(TEST_LDFLAGS)
 
+test/is_utility_stmt: test/framework/main.c test/is_utility_stmt.c $(ARLIB)
+	# We have "-Isrc/postgres/include" because this test uses pg_query_summary_direct
+	$(CC) $(TEST_CFLAGS) -o $@ -Isrc/postgres/include test/framework/main.c test/is_utility_stmt.c $(ARLIB) $(TEST_LDFLAGS)
+
 test/normalize: test/normalize.c test/normalize_tests.c $(ARLIB)
 	$(CC) $(TEST_CFLAGS) -o $@ test/normalize.c $(ARLIB) $(TEST_LDFLAGS)
 

Makefile.msvc

@@ -50,11 +50,12 @@ examples/normalize_error: examples/normalize_error.c $(ARLIB)
 examples/simple_plpgsql: examples/simple_plpgsql.c $(ARLIB)
 	$(CC) $(CFLAGS) -o $@ examples/simple_plpgsql.c $(ARLIB)
 
-TESTS = test/deparse test/fingerprint test/fingerprint_opts test/normalize test/parse test/parse_opts test/parse_protobuf test/parse_protobuf_opts test/parse_plpgsql test/scan test/split test/summary test/summary_truncate
+TESTS = test/deparse test/fingerprint test/fingerprint_opts test/is_utility_stmt test/normalize test/parse test/parse_opts test/parse_protobuf test/parse_protobuf_opts test/parse_plpgsql test/scan test/split test/summary test/summary_truncate
 test: $(TESTS)
 	.\test\deparse
 	.\test\fingerprint
 	.\test\fingerprint_opts
+	.\test\is_utility_stmt
 	.\test\normalize
 	.\test\parse
 	.\test\parse_opts
@@ -85,6 +86,9 @@ test/fingerprint_opts: test/fingerprint_opts.c test/fingerprint_opts_tests.c $(A
 # We have "-Isrc/" because this test uses pg_query_fingerprint_with_opts
 	$(CC) $(CFLAGS) -o $@ -Isrc/ test/fingerprint_opts.c $(ARLIB)
 
+test/is_utility_stmt: test/framework/main.c test/is_utility_stmt.c $(ARLIB)
+	$(CC) $(CFLAGS) -o $@ test/framework/main.c test/is_utility_stmt.c $(ARLIB)
+
 test/normalize: test/normalize.c test/normalize_tests.c $(ARLIB)
 	$(CC) $(CFLAGS) -o $@ test/normalize.c $(ARLIB)
 

pg_query.h

@@ -1,6 +1,7 @@
 #ifndef PG_QUERY_H
 #define PG_QUERY_H
 
+#include <stdbool.h>
 #include <stdint.h>
 #include <sys/types.h>
 
@@ -15,6 +16,12 @@ typedef struct {
 	char* context; // additional context (optional, can be NULL)
 } PgQueryError;
 
+typedef struct {
+	int length;
+	bool *items;
+	PgQueryError* error;
+} PgQueryIsUtilityResult;
+
 typedef struct {
   size_t len;
   char* data;
@@ -135,6 +142,8 @@ PgQueryDeparseResult pg_query_deparse_protobuf(PgQueryProtobuf parse_tree);
 PgQueryDeparseResult pg_query_deparse_protobuf_opts(PgQueryProtobuf parse_tree, struct PostgresDeparseOpts opts);
 PgQueryDeparseCommentsResult pg_query_deparse_comments_for_query(const char *query);
 
+PgQueryIsUtilityResult pg_query_is_utility_stmt(const char *query);
+
 PgQuerySummaryParseResult pg_query_summary(const char* input, int parser_options, int truncate_limit);
 
 void pg_query_free_normalize_result(PgQueryNormalizeResult result);
@@ -146,6 +155,7 @@ void pg_query_free_deparse_comments_result(PgQueryDeparseCommentsResult result);
 void pg_query_free_protobuf_parse_result(PgQueryProtobufParseResult result);
 void pg_query_free_plpgsql_parse_result(PgQueryPlpgsqlParseResult result);
 void pg_query_free_fingerprint_result(PgQueryFingerprintResult result);
+void pg_query_free_is_utility_result(PgQueryIsUtilityResult result);
 void pg_query_free_summary_parse_result(PgQuerySummaryParseResult result);
 
 // Optional, cleans up the top-level memory context (automatically done for threads that exit)

src/pg_query_is_utility_stmt.c

@@ -0,0 +1,70 @@
+#include <stdbool.h>
+#include <stddef.h>
+#include <stdlib.h>
+
+#include "pg_query.h"
+#include "pg_query_internal.h"
+#include <nodes/nodes.h>
+#include <nodes/nodeFuncs.h>
+
+static
+bool
+is_utility_stmt_actual(RawStmt *raw_stmt)
+{
+	switch (nodeTag(raw_stmt->stmt))
+	{
+		case T_SelectStmt:
+		case T_InsertStmt:
+		case T_UpdateStmt:
+		case T_DeleteStmt:
+		case T_MergeStmt:
+			return false;
+
+		default:
+			return true;
+	}
+}
+
+PgQueryIsUtilityResult
+pg_query_is_utility_stmt(const char *query)
+{
+	PgQueryIsUtilityResult result = {0};
+	MemoryContext ctx = pg_query_enter_memory_context();
+
+	PgQueryInternalParsetreeAndError parsetree_and_error = pg_query_raw_parse(query, 0);
+
+	if (parsetree_and_error.error)
+	{
+		result.error = parsetree_and_error.error;
+	}
+	else
+	{
+		ListCell   *lc;
+
+		result.length = list_length(parsetree_and_error.tree);
+		result.items = malloc(sizeof(bool) * result.length);
+
+		foreach(lc, parsetree_and_error.tree)
+		{
+			RawStmt    *raw_stmt = lfirst_node(RawStmt, lc);
+
+			result.items[foreach_current_index(lc)] = is_utility_stmt_actual(raw_stmt);
+		}
+	}
+
+	if (parsetree_and_error.stderr_buffer)
+		free(parsetree_and_error.stderr_buffer);
+
+	pg_query_exit_memory_context(ctx);
+
+	return result;
+}
+
+void
+pg_query_free_is_utility_result(PgQueryIsUtilityResult result)
+{
+	if (result.error)
+		pg_query_free_error(result.error);
+
+	free(result.items);
+}

test/is_utility_stmt.c

@@ -0,0 +1,125 @@
+#include <pg_query.h>
+#include <stdbool.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+#include "test/framework/main.h"
+
+struct LastResult {
+	bool has_result;
+	PgQueryIsUtilityResult result;
+};
+static struct LastResult last_result = {false, {0}};
+
+static PgQueryIsUtilityResult
+is_utility_stmt(const char *msg)
+{
+	last_result.has_result = true;
+	last_result.result = pg_query_is_utility_stmt(msg);
+
+	return last_result.result;
+}
+
+void cleanup()
+{
+	if (!last_result.has_result)
+		return;
+
+	last_result.has_result = false;
+	pg_query_free_is_utility_result(last_result.result);
+}
+
+#define assert_result_equal(actual, expected_len, expected) assert_result_equal_impl(test_state, #actual, actual, expected_len, expected)
+static void
+assert_result_equal_impl(TestState * test_state, char *actual_str, PgQueryIsUtilityResult actual, size_t expected_len, bool expected[])
+{
+	bool all_match = (actual.length == expected_len);
+
+	for (size_t i = 0; (i < expected_len) && (i < actual.length) && all_match; i++) {
+		if (actual.items[i] != expected[i])
+			all_match = false;
+	}
+
+	if (all_match)
+		TEST_PASS();
+	else {
+		TEST_FAIL("Expected `actual` to equal `expected`, where\n");
+		printf("      actual: ");
+		for (size_t i = 0; i < actual.length; i++)
+			printf("%s ", actual.items[i] ? "true " : "false");
+
+		printf("\n    expected: ");
+		for (size_t i = 0; i < expected_len; i++)
+			printf("%s ", expected[i] ? "true " : "false");
+
+		puts("\n");
+	}
+}
+
+static void select_test(TestState* test_state)
+{
+	TEST_INIT();
+
+	assert_result_equal(is_utility_stmt("SELECT 1"), 1, (bool[]){false});
+}
+
+static void insert_test(TestState* test_state)
+{
+	TEST_INIT();
+
+	assert_result_equal(is_utility_stmt("INSERT INTO my_table VALUES(123)"), 1, (bool[]){false});
+}
+
+static void update_test(TestState* test_state)
+{
+	TEST_INIT();
+	assert_result_equal(is_utility_stmt("UPDATE my_table SET foo = 123"), 1, (bool[]){false});
+}
+
+static void delete_test(TestState* test_state)
+{
+	TEST_INIT();
+	assert_result_equal(is_utility_stmt("DELETE FROM my_table"), 1, (bool[]){false});
+}
+
+static void show_test(TestState* test_state)
+{
+	TEST_INIT();
+	assert_result_equal(is_utility_stmt("SHOW fsync"), 1, (bool[]){true});
+}
+
+static void set_test(TestState* test_state)
+{
+	TEST_INIT();
+	assert_result_equal(is_utility_stmt("SET fsync = off"), 1, (bool[]){true});
+}
+
+static void select2_test(TestState* test_state)
+{
+	TEST_INIT();
+	assert_result_equal(is_utility_stmt("SELECT 1; SELECT 2;"), 2, ((bool[]){false, false}));
+}
+
+static void select_show_test(TestState* test_state)
+{
+	TEST_INIT();
+	assert_result_equal(is_utility_stmt("SELECT 1; SHOW fsync;"), 2, ((bool[]){false, true}));
+}
+
+int
+main(int argc, char *argv[])
+{
+	TestFn	   *tests[] = {
+		&select_test,
+		&insert_test,
+		&update_test,
+		&delete_test,
+		&show_test,
+		&set_test,
+		&select2_test,
+		&select_show_test,
+		NULL
+	};
+
+	return test_run(argc, argv, tests, &cleanup);
+}