pg 1.1.4 → 1.2.0

data/ext/pg_result.c

@@ -6,15 +6,22 @@
 
 #include "pg.h"
 
-
 VALUE rb_cPGresult;
+static VALUE sym_symbol, sym_string, sym_static_symbol;
 
-static void pgresult_gc_free( t_pg_result * );
 static VALUE pgresult_type_map_set( VALUE, VALUE );
-static VALUE pgresult_s_allocate( VALUE );
 static t_pg_result *pgresult_get_this( VALUE );
 static t_pg_result *pgresult_get_this_safe( VALUE );
 
+#if defined(HAVE_PQRESULTMEMORYSIZE)
+
+static ssize_t
+pgresult_approx_size(const PGresult *result)
+{
+	return PQresultMemorySize(result);
+}
+
+#else
 
 #define PGRESULT_DATA_BLOCKSIZE 2048
 typedef struct pgresAttValue
@@ -44,7 +51,7 @@ count_leading_zero_bits(unsigned int x)
 }
 
 static ssize_t
-pgresult_approx_size(PGresult *result)
+pgresult_approx_size(const PGresult *result)
 {
 	int num_fields = PQnfields(result);
 	ssize_t size = 0;
@@ -94,7 +101,29 @@ pgresult_approx_size(PGresult *result)
 
 	return size;
 }
+#endif
 
+/*
+ * GC Mark function
+ */
+static void
+pgresult_gc_mark( t_pg_result *this )
+{
+	int i;
+
+	rb_gc_mark( this->connection );
+	rb_gc_mark( this->typemap );
+	rb_gc_mark( this->tuple_hash );
+	rb_gc_mark( this->field_map );
+
+	for( i=0; i < this->nfields; i++ ){
+		rb_gc_mark( this->fnames[i] );
+	}
+}
+
+/*
+ * GC Free function
+ */
 static void
 pgresult_clear( t_pg_result *this )
 {
@@ -104,15 +133,41 @@ pgresult_clear( t_pg_result *this )
 		rb_gc_adjust_memory_usage(-this->result_size);
 #endif
 	}
+	this->result_size = 0;
+	this->nfields = -1;
 	this->pgresult = NULL;
 }
 
+static void
+pgresult_gc_free( t_pg_result *this )
+{
+	pgresult_clear( this );
+	xfree(this);
+}
+
 static size_t
 pgresult_memsize( t_pg_result *this )
 {
+	/* Ideally the memory 'this' is pointing to should be taken into account as well.
+	 * However we don't want to store two memory sizes in t_pg_result just for reporting by ObjectSpace.memsize_of.
+	 */
 	return this->result_size;
 }
 
+static const rb_data_type_t pgresult_type = {
+	"pg",
+	{
+		(void (*)(void*))pgresult_gc_mark,
+		(void (*)(void*))pgresult_gc_free,
+		(size_t (*)(const void *))pgresult_memsize,
+	},
+	0, 0,
+#ifdef RUBY_TYPED_FREE_IMMEDIATELY
+	RUBY_TYPED_FREE_IMMEDIATELY,
+#endif
+};
+
+
 /*
  * Global functions
  */
@@ -124,12 +179,10 @@ static VALUE
 pg_new_result2(PGresult *result, VALUE rb_pgconn)
 {
 	int nfields = result ? PQnfields(result) : 0;
-	VALUE self = pgresult_s_allocate( rb_cPGresult );
+	VALUE self;
 	t_pg_result *this;
 
 	this = (t_pg_result *)xmalloc(sizeof(*this) +  sizeof(*this->fnames) * nfields);
-	RTYPEDDATA_DATA(self) = this;
-
 	this->pgresult = result;
 	this->connection = rb_pgconn;
 	this->typemap = pg_typemap_all_strings;
@@ -137,18 +190,21 @@ pg_new_result2(PGresult *result, VALUE rb_pgconn)
 	this->nfields = -1;
 	this->tuple_hash = Qnil;
 	this->field_map = Qnil;
-
-	PG_ENCODING_SET_NOCHECK(self, ENCODING_GET(rb_pgconn));
+	this->flags = 0;
+	self = TypedData_Wrap_Struct(rb_cPGresult, &pgresult_type, this);
 
 	if( result ){
 		t_pg_connection *p_conn = pg_get_connection(rb_pgconn);
 		VALUE typemap = p_conn->type_map_for_results;
-
 		/* Type check is done when assigned to PG::Connection. */
 		t_typemap *p_typemap = DATA_PTR(typemap);
 
+		this->enc_idx = p_conn->enc_idx;
 		this->typemap = p_typemap->funcs.fit_to_result( typemap, self );
 		this->p_typemap = DATA_PTR( this->typemap );
+		this->flags = p_conn->flags;
+	} else {
+		this->enc_idx = rb_locale_encindex();
 	}
 
 	return self;
@@ -159,22 +215,38 @@ pg_new_result(PGresult *result, VALUE rb_pgconn)
 {
 	VALUE self = pg_new_result2(result, rb_pgconn);
 	t_pg_result *this = pgresult_get_this(self);
-	t_pg_connection *p_conn = pg_get_connection(rb_pgconn);
 
 	this->autoclear = 0;
 
-	if( p_conn->guess_result_memsize ){
-		/* Approximate size of underlying pgresult memory storage and account to ruby GC */
-		this->result_size = pgresult_approx_size(result);
+	/* Estimate size of underlying pgresult memory storage and account to ruby GC.
+	 * There's no need to adjust the GC for xmalloc'ed memory, but libpq is using libc malloc() ruby doesn't know about.
+	 */
+	/* TODO: If someday most systems provide PQresultMemorySize(), it's questionable to store result_size in t_pg_result in addition to the value already stored in PGresult.
+	 * For now the memory savings don't justify the ifdefs necessary to support both cases.
+	 */
+	this->result_size = pgresult_approx_size(result);
 
 #ifdef HAVE_RB_GC_ADJUST_MEMORY_USAGE
-		rb_gc_adjust_memory_usage(this->result_size);
+	rb_gc_adjust_memory_usage(this->result_size);
 #endif
-	}
 
 	return self;
 }
 
+static VALUE
+pg_copy_result(t_pg_result *this)
+{
+	int nfields = this->nfields == -1 ? (this->pgresult ? PQnfields(this->pgresult) : 0) : this->nfields;
+	size_t len = sizeof(*this) +  sizeof(*this->fnames) * nfields;
+	t_pg_result *copy;
+
+	copy = (t_pg_result *)xmalloc(len);
+	memcpy(copy, this, len);
+	this->result_size = 0;
+
+	return TypedData_Wrap_Struct(rb_cPGresult, &pgresult_type, copy);
+}
+
 VALUE
 pg_new_result_autoclear(PGresult *result, VALUE rb_pgconn)
 {
@@ -229,7 +301,7 @@ pg_result_check( VALUE self )
 		}
 	}
 
-	PG_ENCODING_SET_NOCHECK( error, ENCODING_GET(self) );
+	PG_ENCODING_SET_NOCHECK( error, this->enc_idx );
 
 	sqlstate = PQresultErrorField( this->pgresult, PG_DIAG_SQLSTATE );
 	klass = lookup_error_class( sqlstate );
@@ -305,37 +377,6 @@ pgresult_autoclear_p( VALUE self )
  * DATA pointer functions
  */
 
-/*
- * GC Mark function
- */
-static void
-pgresult_gc_mark( t_pg_result *this )
-{
-	int i;
-
-	if( !this ) return;
-	rb_gc_mark( this->connection );
-	rb_gc_mark( this->typemap );
-	rb_gc_mark( this->tuple_hash );
-	rb_gc_mark( this->field_map );
-
-	for( i=0; i < this->nfields; i++ ){
-		rb_gc_mark( this->fnames[i] );
-	}
-}
-
-/*
- * GC Free function
- */
-static void
-pgresult_gc_free( t_pg_result *this )
-{
-	if( !this ) return;
-	pgresult_clear( this );
-
-	xfree(this);
-}
-
 /*
  * Fetch the PG::Result object data pointer and check it's
  * PGresult data pointer for sanity.
@@ -365,32 +406,30 @@ pgresult_get(VALUE self)
 	return this->pgresult;
 }
 
-
-static const rb_data_type_t pgresult_type = {
-	"pg",
-	{
-		(void (*)(void*))pgresult_gc_mark,
-		(void (*)(void*))pgresult_gc_free,
-		(size_t (*)(const void *))pgresult_memsize,
-	},
-	0, 0,
-#ifdef RUBY_TYPED_FREE_IMMEDIATELY
-	RUBY_TYPED_FREE_IMMEDIATELY,
-#endif
-};
-
-/*
- * Document-method: allocate
- *
- * call-seq:
- *   PG::Result.allocate -> result
- */
-static VALUE
-pgresult_s_allocate( VALUE klass )
+static VALUE pg_cstr_to_sym(char *cstr, unsigned int flags, int enc_idx)
 {
-	VALUE self = TypedData_Wrap_Struct( klass, &pgresult_type, NULL );
-
-	return self;
+	VALUE fname;
+#ifdef TRUFFLERUBY
+	if( flags & (PG_RESULT_FIELD_NAMES_SYMBOL | PG_RESULT_FIELD_NAMES_STATIC_SYMBOL) ){
+#else
+	if( flags & PG_RESULT_FIELD_NAMES_SYMBOL ){
+		rb_encoding *enc = rb_enc_from_index(enc_idx);
+		fname = rb_check_symbol_cstr(cstr, strlen(cstr), enc);
+		if( fname == Qnil ){
+			fname = rb_str_new2(cstr);
+			PG_ENCODING_SET_NOCHECK(fname, enc_idx);
+			fname = rb_str_intern(fname);
+		}
+	} else if( flags & PG_RESULT_FIELD_NAMES_STATIC_SYMBOL ){
+#endif
+		rb_encoding *enc = rb_enc_from_index(enc_idx);
+		fname = ID2SYM(rb_intern3(cstr, strlen(cstr), enc));
+	} else {
+		fname = rb_str_new2(cstr);
+		PG_ENCODING_SET_NOCHECK(fname, enc_idx);
+		fname = rb_obj_freeze(fname);
+	}
+	return fname;
 }
 
 static void pgresult_init_fnames(VALUE self)
@@ -402,12 +441,9 @@ static void pgresult_init_fnames(VALUE self)
 		int nfields = PQnfields(this->pgresult);
 
 		for( i=0; i<nfields; i++ ){
-			VALUE fname = rb_tainted_str_new2(PQfname(this->pgresult, i));
-			PG_ENCODING_SET_NOCHECK(fname, ENCODING_GET(self));
-			this->fnames[i] = rb_obj_freeze(fname);
+			char *cfname = PQfname(this->pgresult, i);
+			this->fnames[i] = pg_cstr_to_sym(cfname, this->flags, this->enc_idx);
 			this->nfields = i + 1;
-
-			RB_GC_GUARD(fname);
 		}
 		this->nfields = nfields;
 	}
@@ -470,8 +506,9 @@ pgresult_result_status(VALUE self)
 static VALUE
 pgresult_res_status(VALUE self, VALUE status)
 {
-	VALUE ret = rb_tainted_str_new2(PQresStatus(NUM2INT(status)));
-	PG_ENCODING_SET_NOCHECK(ret, ENCODING_GET(self));
+	t_pg_result *this = pgresult_get_this_safe(self);
+	VALUE ret = rb_str_new2(PQresStatus(NUM2INT(status)));
+	PG_ENCODING_SET_NOCHECK(ret, this->enc_idx);
 	return ret;
 }
 
@@ -484,10 +521,39 @@ pgresult_res_status(VALUE self, VALUE status)
 static VALUE
 pgresult_error_message(VALUE self)
 {
-	VALUE ret = rb_tainted_str_new2(PQresultErrorMessage(pgresult_get(self)));
-	PG_ENCODING_SET_NOCHECK(ret, ENCODING_GET(self));
+	t_pg_result *this = pgresult_get_this_safe(self);
+	VALUE ret = rb_str_new2(PQresultErrorMessage(this->pgresult));
+	PG_ENCODING_SET_NOCHECK(ret, this->enc_idx);
+	return ret;
+}
+
+#ifdef HAVE_PQRESULTVERBOSEERRORMESSAGE
+/*
+ * call-seq:
+ *    res.verbose_error_message( verbosity, show_context ) -> String
+ *
+ * Returns a reformatted version of the error message associated with a PGresult object.
+ *
+ * Available since PostgreSQL-9.6
+ */
+static VALUE
+pgresult_verbose_error_message(VALUE self, VALUE verbosity, VALUE show_context)
+{
+	t_pg_result *this = pgresult_get_this_safe(self);
+	VALUE ret;
+	char *c_str;
+
+	c_str = PQresultVerboseErrorMessage(this->pgresult, NUM2INT(verbosity), NUM2INT(show_context));
+	if(!c_str)
+		rb_raise(rb_eNoMemError, "insufficient memory to format error message");
+
+	ret = rb_str_new2(c_str);
+	PQfreemem(c_str);
+	PG_ENCODING_SET_NOCHECK(ret, this->enc_idx);
+
 	return ret;
 }
+#endif
 
 /*
  * call-seq:
@@ -538,14 +604,14 @@ pgresult_error_message(VALUE self)
 static VALUE
 pgresult_error_field(VALUE self, VALUE field)
 {
-	PGresult *result = pgresult_get( self );
+	t_pg_result *this = pgresult_get_this_safe(self);
 	int fieldcode = NUM2INT( field );
-	char * fieldstr = PQresultErrorField( result, fieldcode );
+	char * fieldstr = PQresultErrorField( this->pgresult, fieldcode );
 	VALUE ret = Qnil;
 
 	if ( fieldstr ) {
-		ret = rb_tainted_str_new2( fieldstr );
-		PG_ENCODING_SET_NOCHECK( ret, ENCODING_GET(self ));
+		ret = rb_str_new2( fieldstr );
+		PG_ENCODING_SET_NOCHECK( ret, this->enc_idx );
 	}
 
 	return ret;
@@ -583,24 +649,25 @@ pgresult_nfields(VALUE self)
 
 /*
  * call-seq:
- *    res.fname( index ) -> String
+ *    res.fname( index ) -> String or Symbol
  *
  * Returns the name of the column corresponding to _index_.
+ * Depending on #field_name_type= it's a String or Symbol.
+ *
  */
 static VALUE
 pgresult_fname(VALUE self, VALUE index)
 {
-	VALUE fname;
-	PGresult *result = pgresult_get(self);
+	t_pg_result *this = pgresult_get_this_safe(self);
 	int i = NUM2INT(index);
+	char *cfname;
 
-	if (i < 0 || i >= PQnfields(result)) {
+	if (i < 0 || i >= PQnfields(this->pgresult)) {
 		rb_raise(rb_eArgError,"invalid field number %d", i);
 	}
 
-	fname = rb_tainted_str_new2(PQfname(result, i));
-	PG_ENCODING_SET_NOCHECK(fname, ENCODING_GET(self));
-	return rb_obj_freeze(fname);
+	cfname = PQfname(this->pgresult, i);
+	return pg_cstr_to_sym(cfname, this->flags, this->enc_idx);
 }
 
 /*
@@ -899,8 +966,9 @@ pgresult_paramtype(VALUE self, VALUE param_number)
 static VALUE
 pgresult_cmd_status(VALUE self)
 {
-	VALUE ret = rb_tainted_str_new2(PQcmdStatus(pgresult_get(self)));
-	PG_ENCODING_SET_NOCHECK(ret, ENCODING_GET(self));
+	t_pg_result *this = pgresult_get_this_safe(self);
+	VALUE ret = rb_str_new2(PQcmdStatus(this->pgresult));
+	PG_ENCODING_SET_NOCHECK(ret, this->enc_idx);
 	return ret;
 }
 
@@ -1100,8 +1168,12 @@ static VALUE
 pgresult_field_values( VALUE self, VALUE field )
 {
 	PGresult *result = pgresult_get( self );
-	const char *fieldname = StringValueCStr( field );
-	int fnum = PQfnumber( result, fieldname );
+	const char *fieldname;
+	int fnum;
+
+	if( RB_TYPE_P(field, T_SYMBOL) ) field = rb_sym_to_s( field );
+	fieldname = StringValueCStr( field );
+	fnum = PQfnumber( result, fieldname );
 
 	if ( fnum < 0 )
 		rb_raise( rb_eIndexError, "no such field '%s' in result", fieldname );
@@ -1144,6 +1216,25 @@ pgresult_tuple_values(VALUE self, VALUE index)
 	}
 }
 
+static void ensure_init_for_tuple(VALUE self)
+{
+	t_pg_result *this = pgresult_get_this_safe(self);
+
+	if( this->field_map == Qnil ){
+		int i;
+		VALUE field_map = rb_hash_new();
+
+		if( this->nfields == -1 )
+			pgresult_init_fnames( self );
+
+		for( i = 0; i < this->nfields; i++ ){
+			rb_hash_aset(field_map, this->fnames[i], INT2FIX(i));
+		}
+		rb_obj_freeze(field_map);
+		this->field_map = field_map;
+	}
+}
+
 /*
  *  call-seq:
  *     res.tuple( n )   -> PG::Tuple
@@ -1164,19 +1255,7 @@ pgresult_tuple(VALUE self, VALUE index)
 	if ( tuple_num < 0 || tuple_num >= num_tuples )
 		rb_raise( rb_eIndexError, "Index %d is out of range", tuple_num );
 
-	if( this->field_map == Qnil ){
-		int i;
-		VALUE field_map = rb_hash_new();
-
-		if( this->nfields == -1 )
-			pgresult_init_fnames( self );
-
-		for( i = 0; i < this->nfields; i++ ){
-			rb_hash_aset(field_map, this->fnames[i], INT2FIX(i));
-		}
-		rb_obj_freeze(field_map);
-		this->field_map = field_map;
-	}
+	ensure_init_for_tuple(self);
 
   return pg_tuple_new(self, tuple_num);
 }
@@ -1208,7 +1287,7 @@ pgresult_each(VALUE self)
  * call-seq:
  *    res.fields() -> Array
  *
- * Returns an array of Strings representing the names of the fields in the result.
+ * Depending on #field_name_type= returns an array of strings or symbols representing the names of the fields in the result.
  */
 static VALUE
 pgresult_fields(VALUE self)
@@ -1307,11 +1386,17 @@ yield_tuple(VALUE self, int ntuples, int nfields)
 {
 	int tuple_num;
 	t_pg_result *this = pgresult_get_this(self);
-	VALUE result = pg_new_result(this->pgresult, this->connection);
+	VALUE copy;
 	UNUSED(nfields);
 
+	/* make a copy of the base result, that is bound to the PG::Tuple */
+	copy = pg_copy_result(this);
+	/* The copy is now owner of the PGresult and is responsible to PQclear it.
+	 * We clear the pgresult here, so that it's not double freed on error within yield. */
+	this->pgresult = NULL;
+
 	for(tuple_num = 0; tuple_num < ntuples; tuple_num++) {
-		VALUE tuple = pgresult_tuple(result, INT2FIX(tuple_num));
+		VALUE tuple = pgresult_tuple(copy, INT2FIX(tuple_num));
 		rb_yield( tuple );
 	}
 }
@@ -1392,8 +1477,6 @@ pgresult_stream_any(VALUE self, void (*yielder)(VALUE, int, int))
  *     # do something with each received row of the second query
  *   end
  *   conn.get_result  # => nil   (no more results)
- *
- * Available since PostgreSQL-9.2
  */
 static VALUE
 pgresult_stream_each(VALUE self)
@@ -1410,8 +1493,6 @@ pgresult_stream_each(VALUE self)
  *
  * This method works equally to #stream_each , but yields an Array of
  * values.
- *
- * Available since PostgreSQL-9.2
  */
 static VALUE
 pgresult_stream_each_row(VALUE self)
@@ -1426,21 +1507,81 @@ pgresult_stream_each_row(VALUE self)
  * Yields each row of the result set in single row mode.
  *
  * This method works equally to #stream_each , but yields a PG::Tuple object.
- *
- * Available since PostgreSQL-9.2
  */
 static VALUE
 pgresult_stream_each_tuple(VALUE self)
 {
+	/* allocate VALUEs that are shared between all streamed tuples */
+	ensure_init_for_tuple(self);
+
 	return pgresult_stream_any(self, yield_tuple);
 }
 
+/*
+ * call-seq:
+ *    res.field_name_type = Symbol
+ *
+ * Set type of field names specific to this result.
+ * It can be set to one of:
+ * * +:string+ to use String based field names
+ * * +:symbol+ to use Symbol based field names
+ * * +:static_symbol+ to use pinned Symbol (can not be garbage collected) - Don't use this, it will probably removed in future.
+ *
+ * The default is retrieved from PG::Connection#field_name_type , which defaults to +:string+ .
+ *
+ * This setting affects several result methods:
+ * * keys of Hash returned by #[] , #each and #stream_each
+ * * #fields
+ * * #fname
+ * * field names used by #tuple and #stream_each_tuple
+ *
+ * The type of field names can only be changed before any of the affected methods have been called.
+ *
+ */
+static VALUE
+pgresult_field_name_type_set(VALUE self, VALUE sym)
+{
+	t_pg_result *this = pgresult_get_this(self);
+	if( this->nfields != -1 ) rb_raise(rb_eArgError, "field names are already materialized");
+
+	this->flags &= ~PG_RESULT_FIELD_NAMES_MASK;
+	if( sym == sym_symbol ) this->flags |= PG_RESULT_FIELD_NAMES_SYMBOL;
+	else if ( sym == sym_static_symbol ) this->flags |= PG_RESULT_FIELD_NAMES_STATIC_SYMBOL;
+	else if ( sym == sym_string );
+	else rb_raise(rb_eArgError, "invalid argument %+"PRIsVALUE, sym);
+
+	return sym;
+}
+
+/*
+ * call-seq:
+ *    res.field_name_type -> Symbol
+ *
+ * Get type of field names.
+ *
+ * See description at #field_name_type=
+ */
+static VALUE
+pgresult_field_name_type_get(VALUE self)
+{
+	t_pg_result *this = pgresult_get_this(self);
+	if( this->flags & PG_RESULT_FIELD_NAMES_SYMBOL ){
+		return sym_symbol;
+	} else if( this->flags & PG_RESULT_FIELD_NAMES_STATIC_SYMBOL ){
+		return sym_static_symbol;
+	} else {
+		return sym_string;
+	}
+}
 
 void
 init_pg_result()
 {
-	rb_cPGresult = rb_define_class_under( rb_mPG, "Result", rb_cObject );
-	rb_define_alloc_func( rb_cPGresult, pgresult_s_allocate );
+	sym_string = ID2SYM(rb_intern("string"));
+	sym_symbol = ID2SYM(rb_intern("symbol"));
+	sym_static_symbol = ID2SYM(rb_intern("static_symbol"));
+
+	rb_cPGresult = rb_define_class_under( rb_mPG, "Result", rb_cData );
 	rb_include_module(rb_cPGresult, rb_mEnumerable);
 	rb_include_module(rb_cPGresult, rb_mPGconstants);
 
@@ -1449,6 +1590,10 @@ init_pg_result()
 	rb_define_method(rb_cPGresult, "res_status", pgresult_res_status, 1);
 	rb_define_method(rb_cPGresult, "error_message", pgresult_error_message, 0);
 	rb_define_alias( rb_cPGresult, "result_error_message", "error_message");
+#ifdef HAVE_PQRESULTVERBOSEERRORMESSAGE
+	rb_define_method(rb_cPGresult, "verbose_error_message", pgresult_verbose_error_message, 2);
+	rb_define_alias( rb_cPGresult, "result_verbose_error_message", "verbose_error_message");
+#endif
 	rb_define_method(rb_cPGresult, "error_field", pgresult_error_field, 1);
 	rb_define_alias( rb_cPGresult, "result_error_field", "error_field" );
 	rb_define_method(rb_cPGresult, "clear", pg_result_clear, 0);
@@ -1496,6 +1641,7 @@ init_pg_result()
 	rb_define_method(rb_cPGresult, "stream_each", pgresult_stream_each, 0);
 	rb_define_method(rb_cPGresult, "stream_each_row", pgresult_stream_each_row, 0);
 	rb_define_method(rb_cPGresult, "stream_each_tuple", pgresult_stream_each_tuple, 0);
-}
-
 
+	rb_define_method(rb_cPGresult, "field_name_type=", pgresult_field_name_type_set, 1 );
+	rb_define_method(rb_cPGresult, "field_name_type", pgresult_field_name_type_get, 0 );
+}