Merge pull request #622 from ged/add-array-dimensions

Add possibility to define the number of array dimensions to be encoded

ext/pg.h

@@ -206,6 +206,7 @@ typedef struct {
 	t_pg_coder comp;
 	t_pg_coder *elem;
 	int needs_quotation;
+	int dimensions;
 	char delimiter;
 } t_pg_composite_coder;
 

ext/pg_binary_encoder.c

@@ -320,6 +320,9 @@ pg_bin_enc_date(t_pg_coder *this, VALUE value, char *out, VALUE *intermediate, i
  * This encoder expects an Array of values or sub-arrays as input.
  * Other values are passed through as byte string without interpretation.
  *
+ * It is possible to enforce a number of dimensions to be encoded by #dimensions= .
+ * Deeper nested arrays are then passed to the elements encoder and less nested arrays raise an ArgumentError.
+ *
  * The accessors needs_quotation and delimiter are ignored for binary encoding.
  *
  */
@@ -346,7 +349,8 @@ pg_bin_enc_array(t_pg_coder *conv, VALUE value, char *out, VALUE *intermediate,
 				dim_sizes[ndim-1] = RARRAY_LENINT(el1);
 				nitems *= dim_sizes[ndim-1];
 				el2 = rb_ary_entry(el1, 0);
-				if (TYPE(el2) == T_ARRAY) {
+				if ( (this->dimensions < 0 || ndim < this->dimensions) &&
+						TYPE(el2) == T_ARRAY) {
 					ndim++;
 					if (ndim > MAXDIM)
 						rb_raise( rb_eArgError, "unsupported number of array dimensions: >%d", ndim );
@@ -356,6 +360,9 @@ pg_bin_enc_array(t_pg_coder *conv, VALUE value, char *out, VALUE *intermediate,
 				el1 = el2;
 			}
 		}
+		if( this->dimensions >= 0 && (ndim==0 ? 1 : ndim) != this->dimensions ){
+			rb_raise(rb_eArgError, "less array dimensions to encode (%d) than expected (%d)", ndim, this->dimensions);
+		}
 
 		if(out){
 			/* Second encoder pass -> write data to `out` */

ext/pg_coder.c

@@ -135,6 +135,7 @@ pg_composite_encoder_allocate( VALUE klass )
 	this->elem = NULL;
 	this->needs_quotation = 1;
 	this->delimiter = ',';
+	this->dimensions = -1;
 	rb_iv_set( self, "@elements_type", Qnil );
 	return self;
 }
@@ -157,6 +158,7 @@ pg_composite_decoder_allocate( VALUE klass )
 	this->elem = NULL;
 	this->needs_quotation = 1;
 	this->delimiter = ',';
+	this->dimensions = -1;
 	rb_iv_set( self, "@elements_type", Qnil );
 	return self;
 }
@@ -421,6 +423,49 @@ pg_coder_delimiter_get(VALUE self)
 	return rb_str_new(&this->delimiter, 1);
 }
 
+/*
+ * call-seq:
+ *    coder.dimensions = Integer
+ *    coder.dimensions = nil
+ *
+ * Set number of array dimensions to be encoded.
+ *
+ * This property ensures, that this number of dimensions is always encoded.
+ * If less dimensions than this number are in the given value, an ArgumentError is raised.
+ * If more dimensions than this number are in the value, the Array value is passed to the next encoder.
+ *
+ * Setting dimensions is especially useful, when a Record shall be encoded into an Array, since the Array encoder can not distinguish if the array shall be encoded as a higher dimension or as a record otherwise.
+ *
+ * The default is +nil+.
+ *
+ * See #dimensions
+ */
+static VALUE
+pg_coder_dimensions_set(VALUE self, VALUE dimensions)
+{
+	t_pg_composite_coder *this = RTYPEDDATA_DATA(self);
+	rb_check_frozen(self);
+	if(!NIL_P(dimensions) && NUM2INT(dimensions) < 0)
+		rb_raise( rb_eArgError, "dimensions must be nil or >= 0");
+	this->dimensions = NIL_P(dimensions) ? -1 : NUM2INT(dimensions);
+	return dimensions;
+}
+
+/*
+ * call-seq:
+ *    coder.dimensions -> Integer | nil
+ *
+ * Get number of enforced array dimensions or +nil+ if not set.
+ *
+ * See #dimensions=
+ */
+static VALUE
+pg_coder_dimensions_get(VALUE self)
+{
+	t_pg_composite_coder *this = RTYPEDDATA_DATA(self);
+	return this->dimensions < 0 ? Qnil : INT2NUM(this->dimensions);
+}
+
 /*
  * call-seq:
  *    coder.elements_type = coder
@@ -602,6 +647,8 @@ init_pg_coder(void)
 	 *
 	 * This is the base class for all type cast classes of PostgreSQL types,
 	 * that are made up of some sub type.
+	 *
+	 * See PG::TextEncoder::Array, PG::TextDecoder::Array, PG::BinaryEncoder::Array, PG::BinaryDecoder::Array, etc.
 	 */
 	rb_cPG_CompositeCoder = rb_define_class_under( rb_mPG, "CompositeCoder", rb_cPG_Coder );
 	rb_define_method( rb_cPG_CompositeCoder, "elements_type=", pg_coder_elements_type_set, 1 );
@@ -610,6 +657,8 @@ init_pg_coder(void)
 	rb_define_method( rb_cPG_CompositeCoder, "needs_quotation?", pg_coder_needs_quotation_get, 0 );
 	rb_define_method( rb_cPG_CompositeCoder, "delimiter=", pg_coder_delimiter_set, 1 );
 	rb_define_method( rb_cPG_CompositeCoder, "delimiter", pg_coder_delimiter_get, 0 );
+	rb_define_method( rb_cPG_CompositeCoder, "dimensions=", pg_coder_dimensions_set, 1 );
+	rb_define_method( rb_cPG_CompositeCoder, "dimensions", pg_coder_dimensions_get, 0 );
 
 	/* Document-class: PG::CompositeEncoder < PG::CompositeCoder */
 	rb_cPG_CompositeEncoder = rb_define_class_under( rb_mPG, "CompositeEncoder", rb_cPG_CompositeCoder );

ext/pg_text_encoder.c

@@ -537,14 +537,18 @@ quote_string(t_pg_coder *this, VALUE value, VALUE string, char *current_out, int
 }
 
 static char *
-write_array(t_pg_composite_coder *this, VALUE value, char *current_out, VALUE string, int quote, int enc_idx)
+write_array(t_pg_composite_coder *this, VALUE value, char *current_out, VALUE string, int quote, int enc_idx, int dimension)
 {
 	int i;
 
 	/* size of "{}" */
 	current_out = pg_rb_str_ensure_capa( string, 2, current_out, NULL );
 	*current_out++ = '{';
 
+	if( RARRAY_LEN(value) == 0 && this->dimensions >= 0 && dimension != this->dimensions ){
+		rb_raise(rb_eArgError, "less array dimensions to encode (%d) than expected (%d)", dimension, this->dimensions);
+	}
+
 	for( i=0; i<RARRAY_LEN(value); i++){
 		VALUE entry = rb_ary_entry(value, i);
 
@@ -554,17 +558,26 @@ write_array(t_pg_composite_coder *this, VALUE value, char *current_out, VALUE st
 		}
 
 		switch(TYPE(entry)){
-			case T_ARRAY:
-				current_out = write_array(this, entry, current_out, string, quote, enc_idx);
-			break;
 			case T_NIL:
+				if( this->dimensions >= 0 && dimension != this->dimensions ){
+					rb_raise(rb_eArgError, "less array dimensions to encode (%d) than expected (%d)", dimension, this->dimensions);
+				}
 				current_out = pg_rb_str_ensure_capa( string, 4, current_out, NULL );
 				*current_out++ = 'N';
 				*current_out++ = 'U';
 				*current_out++ = 'L';
 				*current_out++ = 'L';
 				break;
+			case T_ARRAY:
+				if( this->dimensions < 0 || dimension < this->dimensions ){
+					current_out = write_array(this, entry, current_out, string, quote, enc_idx, dimension+1);
+					break;
+				}
+				/* Number of dimensions reached -> handle array as normal value */
 			default:
+				if( this->dimensions >= 0 && dimension != this->dimensions ){
+					rb_raise(rb_eArgError, "less array dimensions to encode (%d) than expected (%d)", dimension, this->dimensions);
+				}
 				current_out = quote_string( this->elem, entry, string, current_out, quote, quote_array_buffer, this, enc_idx );
 		}
 	}
@@ -596,7 +609,7 @@ pg_text_enc_array(t_pg_coder *conv, VALUE value, char *out, VALUE *intermediate,
 		VALUE out_str = rb_str_new(NULL, 0);
 		PG_ENCODING_SET_NOCHECK(out_str, enc_idx);
 
-		end_ptr = write_array(this, value, RSTRING_PTR(out_str), out_str, this->needs_quotation, enc_idx);
+		end_ptr = write_array(this, value, RSTRING_PTR(out_str), out_str, this->needs_quotation, enc_idx, 1);
 
 		rb_str_set_len( out_str, end_ptr - RSTRING_PTR(out_str) );
 		*intermediate = out_str;

lib/pg/coder.rb

@@ -76,12 +76,13 @@ def to_h
 				elements_type: elements_type,
 				needs_quotation: needs_quotation?,
 				delimiter: delimiter,
+				dimensions: dimensions,
 			}
 		end
 
 		def inspect
 			str = super
-			str[-1,0] = " elements_type=#{elements_type.inspect} #{needs_quotation? ? 'needs' : 'no'} quotation"
+			str[-1,0] = " elements_type=#{elements_type.inspect} #{needs_quotation? ? 'needs' : 'no'} quotation#{dimensions && " #{dimensions} dimensions"}"
 			str
 		end
 	end

spec/pg/type_spec.rb

@@ -956,13 +956,75 @@ def expect_deprecated_coder_init
 
 						expect( binaryenc_text_array.encode([[[5,6]],[["6\"",7]],[[nil,5]]]) ).to eq( exp )
 					end
+
+					let!(:binaryenc_array_array) { PG::BinaryEncoder::Array.new elements_type: PG::BinaryEncoder::Array.new(elements_type: PG::BinaryEncoder::Int4.new(oid: 0x17), dimensions: 1), dimensions: 2 }
+
+					it 'encodes an array in an array of int4' do
+						exp = ["00000002" + "00000001" + "00000000" +
+						      "00000003" + "00000001" + "00000001" + "00000001" +
+
+						      "00000024" +
+						      "00000001" + "00000001" + "00000017" +
+						      "00000002" + "00000001" +
+						      "00000004" + "00000005" +
+						      "00000004" + "00000006" +
+
+						      "00000024" +
+						      "00000001" + "00000001" + "00000017" +
+						      "00000002" + "00000001" +
+						      "00000004" + "00000006" +
+						      "00000004" + "00000007" +
+
+						      "00000020" +
+						      "00000001" + "00000001" + "00000017" +
+						      "00000002" + "00000001" +
+						      "ffffffff" +
+						      "00000004" + "00000005"
+						      ].pack("H*")
+
+						expect( binaryenc_array_array.encode([[[5,6]],[[6,7]],[[nil,5]]]) ).to eq( exp )
+					end
 				end
 
 				context 'two dimensional arrays' do
 					it 'encodes an array of timestamps with sub arrays' do
 						expect( textenc_timestamp_array.encode([Time.new(2014,12,31),[nil, Time.new(2016,01,02, 23, 23, 59.99)]]) ).
 								to eq( %[{2014-12-31 00:00:00.000000000,{NULL,2016-01-02 23:23:59.990000000}}] )
 					end
+
+					context 'with dimensions' do
+						let!(:textenc_array_2dim) { textenc_string_array.dup.tap{|a| a.dimensions = 2} }
+						let!(:binaryenc_array_2dim) { binaryenc_array.dup.tap{|a| a.dimensions = 2} }
+
+						it 'encodes binary int array' do
+							binaryenc_array_2dim.encode([[1]])
+						end
+						it 'encodes text int array' do
+							expect( textenc_array_2dim.encode([[1]]) ).to eq( "{{1}}" )
+						end
+						it 'encodes empty array' do
+							binaryenc_array_2dim.encode([[]])
+						end
+						it 'encodes text empty array' do
+							expect( textenc_array_2dim.encode([[]]) ).to eq( "{{}}" )
+						end
+						it 'raises an error on 1 dim binary array input to int4' do
+							expect{ binaryenc_array_2dim.encode([1]) }.to raise_error( ArgumentError, /less array dimensions.*1.*2/)
+						end
+						it 'raises an error on 1 dim text array input to int4' do
+							expect{ textenc_array_2dim.encode([1]) }.to raise_error( ArgumentError, /less array dimensions.*1.*2/)
+						end
+
+						it 'raises an error on 0 dim array input to int4' do
+							expect{ binaryenc_array_2dim.encode([]) }.to raise_error( ArgumentError, /less array dimensions.*0.*2/)
+						end
+						it 'raises an error on 0 dim text array input to int4' do
+							expect{ textenc_array_2dim.encode([]) }.to raise_error( ArgumentError, /less array dimensions.*1.*2/)
+						end
+						it 'raises an error on 1 dim text array nil input' do
+							expect{ textenc_array_2dim.encode([nil]) }.to raise_error( ArgumentError, /less array dimensions.*1.*2/)
+						end
+					end
 				end
 
 				context 'one dimensional array' do
@@ -986,6 +1048,37 @@ def expect_deprecated_coder_init
 
 						expect( binaryenc_array.encode([nil, "6\""]) ).to eq( exp )
 					end
+
+					context 'with dimensions' do
+						let!(:textenc_array_1dim) { textenc_int_array.dup.tap{|a| a.dimensions = 1} }
+						let!(:binaryenc_array_1dim) { binaryenc_array.dup.tap{|a| a.dimensions = 1} }
+
+						it 'encodes an array' do
+							exp =["00000001" + "00000001" + "00000000" +
+							      "00000002" + "00000001" +
+							     "ffffffff" +
+							     "00000002" + "3622"
+							     ].pack("H*")
+
+							expect( binaryenc_array_1dim.encode([nil, "6\""]) ).to eq( exp )
+						end
+						it 'encodes an empty binary array' do
+							exp =["00000000" + "00000001" + "00000000"
+							     ].pack("H*")
+							expect( binaryenc_array_1dim.encode([]) ).to eq( exp )
+						end
+						it 'encodes an empty text array' do
+							expect( textenc_array_1dim.encode([]) ).to eq( "{}" )
+						end
+
+						let!(:binaryenc_int4_array_1dim) { PG::BinaryEncoder::Array.new elements_type: PG::BinaryEncoder::Int4.new, dimensions: 1 }
+						it 'raises an error on binary array input to int4' do
+							expect{ binaryenc_int4_array_1dim.encode([[1]]) }.to raise_error( NoMethodError, /to_i/)
+						end
+						it 'raises an error on text array input to int4' do
+							expect{ textenc_array_1dim.encode([[1]]) }.to raise_error( NoMethodError, /to_i/)
+						end
+					end
 				end
 
 				context 'other dimensional array' do
@@ -1091,7 +1184,8 @@ def expect_deprecated_coder_init
 			it "should respond to to_h" do
 				expect( textenc_int_array.to_h ).to eq( {
 					name: nil, oid: 0, format: 0, flags: 0,
-					elements_type: textenc_int, needs_quotation: false, delimiter: ','
+					elements_type: textenc_int, needs_quotation: false, delimiter: ',',
+					dimensions: nil
 				} )
 			end
 
@@ -1107,6 +1201,7 @@ def expect_deprecated_coder_init
 				expect( t.needs_quotation? ).to eq( true )
 				expect( t.delimiter ).to eq( ',' )
 				expect( t.elements_type ).to be_nil
+				expect( t.dimensions ).to be_nil
 			end
 
 			it "should deny changes when frozen" do
@@ -1117,6 +1212,7 @@ def expect_deprecated_coder_init
 				expect{ t.needs_quotation = true }.to raise_error(FrozenError)
 				expect{ t.delimiter = ","  }.to raise_error(FrozenError)
 				expect{ t.elements_type = nil }.to raise_error(FrozenError)
+				expect{ t.dimensions = 1 }.to raise_error(FrozenError)
 			end
 
 			it "should be shareable for Ractor", :ractor do