Added function protobuf_query_array

mpartel · mpartel · commit daa68914ea28 · 2020-02-26T16:16:47.000+02:00
diff --git a/Makefile b/Makefile
@@ -4,16 +4,16 @@
 # To release a new version, change these numbers and
 # run ./build-and-test.sh to update the test cases.
 EXT_VERSION_MAJOR = 0
-EXT_VERSION_MINOR = 1
-EXT_VERSION_PATCHLEVEL = 2
+EXT_VERSION_MINOR = 2
+EXT_VERSION_PATCHLEVEL = 0
 EXT_VERSION = $(EXT_VERSION_MAJOR).$(EXT_VERSION_MINOR).$(EXT_VERSION_PATCHLEVEL)
 
 PROTOBUF_ROOT=third_party/protobuf
 PROTOC=$(PROTOBUF_ROOT)/src/protoc
 
 MODULE_big = postgres_protobuf
 EXTENSION = postgres_protobuf
-DATA = postgres_protobuf--0.1.sql
+DATA = postgres_protobuf--0.1.sql postgres_protobuf--0.1--0.2.sql
 DOCS = README.md
 REGRESS = postgres_protobuf
 OBJS=$(patsubst %.cpp, %.o, $(wildcard *.cpp))
diff --git a/README.md b/README.md
@@ -13,7 +13,9 @@ Examples:
 ```sql
 SELECT protobuf_query('MyProto:some_submessage.some_map[some_key]', my_proto_column) FROM ...;
 
--- or for multiple results
+-- or for multiple results as an array
+SELECT protobuf_query_array('MyProto:some_repeated_field[*].some_map[*].some_field', my_proto_column) FROM ...;
+-- or as rows
 SELECT protobuf_query_multi('MyProto:some_repeated_field[*].some_map[*].some_field', my_proto_column) FROM ...;
 ```
 
@@ -96,7 +98,8 @@ SELECT protobuf_query('path.to.Message:path.to.field', protobuf_as_byte_array) A
 The following functions are defined:
 
 - `protobuf_query(query, protobuf)` returns the first matching field in the protobuf, or NULL if missing or proto3 default.
-- `protobuf_query_multi(query, protobuf)` returns all matching fields in the protobuf. Missing or proto3 default values are not returned.
+- `protobuf_query_array(query, protobuf)` returns all matching fields in the protobuf as a text array. Missing or proto3 default values are not returned.
+- `protobuf_query_multi(query, protobuf)` returns all matching fields in the protobuf as a set of rows. Missing or proto3 default values are not returned.
 - `protobuf_to_json_text(protobuf_type, protobuf)` converts the protobuf to a JSON string, assuming it's of the given type.
 - `protobuf_from_json_text(protobuf_type, json_str)` parses a protobuf from a JSON string, assuming it's of the given type.
 - `protobuf_extension_version()` returns the extension version `X.Y.Z` as a number `X*10000+Y*100+Z`.
diff --git a/expected/postgres_protobuf.out b/expected/postgres_protobuf.out
@@ -5,7 +5,7 @@
 CREATE EXTENSION postgres_protobuf;
 SELECT protobuf_extension_version() AS result;
 result
-102
+200
 (1 row)
 INSERT INTO protobuf_file_descriptor_sets (name, file_descriptor_set) VALUES ('default', '\x0afb0c0a196d61696e5f64657363726970746f725f7365742e70726f746f1209706770622e7465737422f7070a0e4578616d706c654d657373616765122c0a077363616c61727318012001280b32122e706770622e746573742e5363616c61727352077363616c61727312250a0e72657065617465645f696e743332180220032805520d7265706561746564496e74333212270a0f72657065617465645f737472696e67180320032809520e7265706561746564537472696e67123c0a05696e6e657218042001280b32262e706770622e746573742e4578616d706c654d6573736167652e496e6e65724d6573736167655205696e6e6572124d0a0e72657065617465645f696e6e657218052003280b32262e706770622e746573742e4578616d706c654d6573736167652e496e6e65724d657373616765520d7265706561746564496e6e6572122a0a07616e5f656e756d18062001280e32112e706770622e746573742e416e456e756d5206616e456e756d124a0a0b6d61705f7374723273747218072003280b32292e706770622e746573742e4578616d706c654d6573736167652e4d617053747232737472456e747279520a6d617053747232737472124a0a0b6d61705f696e743273747218082003280b32292e706770622e746573742e4578616d706c654d6573736167652e4d6170496e7432737472456e747279520a6d6170496e7432737472124a0a0b6d61705f696e7432696e7418092003280b32292e706770622e746573742e4578616d706c654d6573736167652e4d6170496e7432696e74456e747279520a6d6170496e7432696e7412500a0d6d61705f73747232696e6e6572180a2003280b322b2e706770622e746573742e4578616d706c654d6573736167652e4d617053747232696e6e6572456e747279520c6d617053747232696e6e65721a3d0a0f4d617053747232737472456e74727912100a036b657918012001280952036b657912140a0576616c7565180220012809520576616c75653a0238011a3d0a0f4d6170496e7432737472456e74727912100a036b657918012001280552036b657912140a0576616c7565180220012809520576616c75653a0238011a3d0a0f4d6170496e7432696e74456e74727912100a036b657918012001280552036b657912140a0576616c7565180220012805520576616c75653a0238011a670a114d617053747232696e6e6572456e74727912100a036b657918012001280952036b6579123c0a0576616c756518022001280b32262e706770622e746573742e4578616d706c654d6573736167652e496e6e65724d657373616765520576616c75653a0238011a520a0c496e6e65724d657373616765121b0a09696e6e65725f7374721801200128095208696e6e657253747212250a0e696e6e65725f7265706561746564180220032809520d696e6e657252657065617465642296040a075363616c61727312210a0c646f75626c655f6669656c64180120012801520b646f75626c654669656c64121f0a0b666c6f61745f6669656c64180220012802520a666c6f61744669656c64121f0a0b696e7433325f6669656c64180320012805520a696e7433324669656c64121f0a0b696e7436345f6669656c64180420012803520a696e7436344669656c6412210a0c75696e7433325f6669656c6418052001280d520b75696e7433324669656c6412210a0c75696e7436345f6669656c64180620012804520b75696e7436344669656c6412210a0c73696e7433325f6669656c64180720012811520b73696e7433324669656c6412210a0c73696e7436345f6669656c64180820012812520b73696e7436344669656c6412230a0d666978656433325f6669656c64180920012807520c666978656433324669656c6412230a0d666978656436345f6669656c64180a20012806520c666978656436344669656c6412250a0e73666978656433325f6669656c64180b2001280f520d73666978656433324669656c6412250a0e73666978656436345f6669656c64180c20012810520d73666978656436344669656c64121d0a0a626f6f6c5f6669656c64180d200128085209626f6f6c4669656c6412210a0c737472696e675f6669656c64180e20012809520b737472696e674669656c64121f0a0b62797465735f6669656c64180f2001280c520a62797465734669656c642a380a06416e456e756d120e0a0a456e756d56616c7565301000120e0a0a456e756d56616c7565311001120e0a0a456e756d56616c7565321002620670726f746f33'::BYTEA);
 INSERT INTO protobuf_file_descriptor_sets (name, file_descriptor_set) VALUES ('other', '\x0a6f0a1a6f746865725f64657363726970746f725f7365742e70726f746f120f706770622e746573742e6f7468657222380a154d657373616765496e4f7468657244657363536574121f0a0b696e7433325f6669656c64180120012805520a696e7433324669656c64620670726f746f33'::BYTEA);
@@ -424,6 +424,37 @@ result
 {"scalars":{"int32Field":123}}
 (1 row)
 --
+-- Array queries
+--
+SELECT protobuf_query_array('pgpb.test.ExampleMessage:repeated_int32[*]', '\x12037bc803'::BYTEA) AS result;
+result
+{123,456}
+(1 row)
+SELECT protobuf_query_array('pgpb.test.ExampleMessage:map_int2str[*]', '\x4207087b1203414141420808c8031203424242'::BYTEA) AS result;
+result
+{AAA,BBB}
+(1 row)
+SELECT protobuf_query_array('pgpb.test.ExampleMessage:map_int2str|keys', '\x4207087b1203414141420808c8031203424242'::BYTEA) AS result;
+result
+{123,456}
+(1 row)
+SELECT protobuf_query_array('pgpb.test.ExampleMessage:', '\x0a02187b'::BYTEA) AS result;
+result
+{"{\"scalars\":{\"int32Field\":123}}"}
+(1 row)
+SELECT protobuf_query_array('pgpb.test.ExampleMessage:scalars', '\x0a02187b'::BYTEA) AS result;
+result
+{"{\"int32Field\":123}"}
+(1 row)
+SELECT protobuf_query_array('pgpb.test.ExampleMessage:scalars.int32_field', '\x0a02187b'::BYTEA) AS result;
+result
+{123}
+(1 row)
+SELECT protobuf_query_array('pgpb.test.ExampleMessage:scalars.string_field', '\x0a02187b'::BYTEA) AS result;
+result
+{}
+(1 row)
+--
 -- Converting to JSON
 --
 SELECT protobuf_to_json_text('pgpb.test.ExampleMessage', '\x0a02187b'::BYTEA) AS result;
diff --git a/generate_test_cases.rb b/generate_test_cases.rb
@@ -318,6 +318,22 @@ def pg_quote(s)
     end
   end
 
+  section "Array queries" do
+    with_proto('repeated_int32: 123, repeated_int32: 456') do
+      test_sql("SELECT protobuf_query_array('pgpb.test.ExampleMessage:repeated_int32[*]', #{pg_proto}) AS result;", ['{123,456}'])
+    end
+    with_proto('map_int2str: { key: 123, value: "AAA" }, map_int2str { key: 456, value: "BBB" }') do
+      test_sql("SELECT protobuf_query_array('pgpb.test.ExampleMessage:map_int2str[*]', #{pg_proto}) AS result;", ['{AAA,BBB}'])
+      test_sql("SELECT protobuf_query_array('pgpb.test.ExampleMessage:map_int2str|keys', #{pg_proto}) AS result;", ['{123,456}'])
+    end
+    with_proto('scalars { int32_field: 123 }') do
+      test_sql("SELECT protobuf_query_array('pgpb.test.ExampleMessage:', #{pg_proto}) AS result;", ['{"{\"scalars\":{\"int32Field\":123}}"}'])
+      test_sql("SELECT protobuf_query_array('pgpb.test.ExampleMessage:scalars', #{pg_proto}) AS result;", ['{"{\"int32Field\":123}"}'])
+      test_sql("SELECT protobuf_query_array('pgpb.test.ExampleMessage:scalars.int32_field', #{pg_proto}) AS result;", ['{123}'])
+      test_sql("SELECT protobuf_query_array('pgpb.test.ExampleMessage:scalars.string_field', #{pg_proto}) AS result;", ['{}'])
+    end
+  end
+
   section "Converting to JSON" do
     with_proto('scalars { int32_field: 123 }') do
       test_sql("SELECT protobuf_to_json_text('pgpb.test.ExampleMessage', #{pg_proto}) AS result;", ['{"scalars":{"int32Field":123}}'])
diff --git a/postgres_protobuf--0.1--0.2.sql b/postgres_protobuf--0.1--0.2.sql
@@ -0,0 +1,10 @@
+-- complain if script is sourced in psql, rather than via CREATE EXTENSION
+\echo Use "CREATE EXTENSION postgres_protobuf" to load this file. \quit
+
+CREATE FUNCTION protobuf_query_array(
+    IN TEXT,  -- Query
+    IN BYTEA  -- Binary protobuf
+)
+    RETURNS TEXT[]
+    AS 'MODULE_PATHNAME'
+    LANGUAGE C STRICT STABLE;
diff --git a/postgres_protobuf.control b/postgres_protobuf.control
@@ -1,4 +1,4 @@
 comment = 'Protocol buffers for PostgreSQL'
-default_version = '0.1'
+default_version = '0.2'
 module_pathname = '$libdir/postgres_protobuf'
 relocatable = true
diff --git a/postgres_protobuf.cpp b/postgres_protobuf.cpp
@@ -13,8 +13,10 @@ extern "C" {
 #include <postgres.h>
 
 #include <access/htup_details.h>
+#include <catalog/pg_type.h>
 #include <fmgr.h>
 #include <funcapi.h>
+#include <utils/lsyscache.h>
 }  // extern "C"
 
 namespace postgres_protobuf {
@@ -88,6 +90,7 @@ extern "C" {
 PG_FUNCTION_INFO_V1(protobuf_extension_version);
 PG_FUNCTION_INFO_V1(protobuf_query);
 PG_FUNCTION_INFO_V1(protobuf_query_multi);
+PG_FUNCTION_INFO_V1(protobuf_query_array);
 PG_FUNCTION_INFO_V1(protobuf_to_json_text);
 PG_FUNCTION_INFO_V1(protobuf_from_json_text);
 
@@ -145,6 +148,61 @@ Datum protobuf_query(PG_FUNCTION_ARGS) {
   }
 }
 
+Datum protobuf_query_array(PG_FUNCTION_ARGS) {
+  using namespace querying;
+
+  assert(PG_NARGS() == 2);
+
+  try {
+    text* query_text = PG_GETARG_TEXT_P(0);
+    std::string query_str(VARDATA_ANY(query_text),
+                          VARSIZE_ANY_EXHDR(query_text));
+    querying::Query query(query_str, std::nullopt);
+    PGPROTO_DEBUG("Query parsed");
+
+    bytea* proto_bytea = PG_GETARG_BYTEA_P(1);
+    const uint8* proto_data =
+        reinterpret_cast<const uint8*>(VARDATA_ANY(proto_bytea));
+    size_t proto_len = VARSIZE_ANY_EXHDR(proto_bytea);
+    const auto rows = query.Run(proto_data, proto_len);
+    PGPROTO_DEBUG("Query ran. Results: %lu", rows.size());
+    Datum* elements = static_cast<Datum*>(palloc0_or_throw_bad_alloc(sizeof(Datum) * rows.size()));
+    for (size_t i = 0; i < rows.size(); ++i) {
+      const std::string& row = rows[i];
+      size_t size = VARHDRSZ + row.size();
+      bytea* p = static_cast<bytea*>(palloc0_or_throw_bad_alloc(size));
+      SET_VARSIZE(p, size);
+      memcpy(VARDATA(p), row.data(), row.size());
+      elements[i] = reinterpret_cast<Datum>(p);
+    }
+    int16 typlen;
+    bool typbyval;
+    char typalign;
+    get_typlenbyvalalign(TEXTOID, &typlen, &typbyval, &typalign);
+    ArrayType* result = construct_array(elements, rows.size(), TEXTOID, typlen, typbyval, typalign);
+    PG_RETURN_ARRAYTYPE_P(result);
+  } catch (const std::bad_alloc& e) {
+    ereport(ERROR, (errcode(ERRCODE_OUT_OF_MEMORY), errmsg("out of memory")));
+  } catch (const BadProto& e) {
+    // TODO: is this a good error code?
+    ereport(ERROR, (errcode(ERRCODE_INVALID_BINARY_REPRESENTATION),
+                    errmsg("invalid protobuf: %s", e.msg.c_str())));
+  } catch (const BadQuery& e) {
+    // TODO: is this a good error code?
+    ereport(ERROR, (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+                    errmsg("invalid query: %s", e.msg.c_str())));
+  } catch (const RecursionDepthExceeded& e) {
+    // TODO: is this a good error code?
+    // TODO: make the limit configurable
+    ereport(ERROR, (errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
+                    errmsg("protobuf recursion depth exceeded")));
+  } catch (...) {
+    ereport(ERROR,
+            (errcode(ERRCODE_INTERNAL_ERROR),
+             errmsg("unknown C++ exception in postgres_protobuf extension")));
+  }
+}
+
 Datum protobuf_query_multi(PG_FUNCTION_ARGS) {
   using namespace querying;
 
diff --git a/sql/postgres_protobuf.sql b/sql/postgres_protobuf.sql
@@ -249,6 +249,16 @@ SELECT protobuf_query('pgpb.test.ExampleMessage:an_enum', '\x3002'::BYTEA) AS re
 --
 SELECT protobuf_query('pgpb.test.ExampleMessage:', '\x0a02187b'::BYTEA) AS result;
 --
+-- Array queries
+--
+SELECT protobuf_query_array('pgpb.test.ExampleMessage:repeated_int32[*]', '\x12037bc803'::BYTEA) AS result;
+SELECT protobuf_query_array('pgpb.test.ExampleMessage:map_int2str[*]', '\x4207087b1203414141420808c8031203424242'::BYTEA) AS result;
+SELECT protobuf_query_array('pgpb.test.ExampleMessage:map_int2str|keys', '\x4207087b1203414141420808c8031203424242'::BYTEA) AS result;
+SELECT protobuf_query_array('pgpb.test.ExampleMessage:', '\x0a02187b'::BYTEA) AS result;
+SELECT protobuf_query_array('pgpb.test.ExampleMessage:scalars', '\x0a02187b'::BYTEA) AS result;
+SELECT protobuf_query_array('pgpb.test.ExampleMessage:scalars.int32_field', '\x0a02187b'::BYTEA) AS result;
+SELECT protobuf_query_array('pgpb.test.ExampleMessage:scalars.string_field', '\x0a02187b'::BYTEA) AS result;
+--
 -- Converting to JSON
 --
 SELECT protobuf_to_json_text('pgpb.test.ExampleMessage', '\x0a02187b'::BYTEA) AS result;
