pg 0.18.0.pre20140820094244 → 0.18.0.pre20141017155815

data/ext/pg_result.c

@@ -1,6 +1,6 @@
 /*
  * pg_result.c - PG::Result class extension
- * $Id: pg_result.c,v 738d4c4125a1 2014/07/23 18:10:52 mina $
+ * $Id: pg_result.c,v f0b7f99b1dd5 2014/10/06 08:31:02 kanis $
  *
  */
 
@@ -9,8 +9,12 @@
 
 VALUE rb_cPGresult;
 
-static void pgresult_gc_free( PGresult * );
-static PGresult* pgresult_get( VALUE );
+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 );
+
 
 
 /*
@@ -23,16 +27,40 @@ static PGresult* pgresult_get( VALUE );
 VALUE
 pg_new_result(PGresult *result, VALUE rb_pgconn)
 {
-	PGconn *conn = pg_get_pgconn( rb_pgconn );
-	VALUE val = Data_Wrap_Struct(rb_cPGresult, NULL, pgresult_gc_free, result);
-#ifdef M17N_SUPPORTED
-	rb_encoding *enc = pg_conn_enc_get( conn );
-	ENCODING_SET( val, rb_enc_to_index(enc) );
-#endif
+	VALUE self = pgresult_s_allocate( rb_cPGresult );
+	t_pg_result *this = pgresult_get_this(self);
+	t_pg_connection *p_conn = pg_get_connection(rb_pgconn);
+
+	PG_ENCODING_SET_NOCHECK(self, ENCODING_GET(rb_pgconn));
+
+	this->pgresult = result;
+	this->connection = rb_pgconn;
+	if( result ){
+		VALUE typemap = p_conn->type_map_for_results;
+
+		if( !NIL_P(typemap) ){
+			t_typemap *p_typemap;
+
+			/* Type check is done when assigned to PG::Connection. */
+			p_typemap = DATA_PTR(typemap);
+
+			typemap = p_typemap->fit_to_result( typemap, self );
+			this->p_typemap = DATA_PTR( typemap );
+		}
 
-	rb_iv_set( val, "@connection", rb_pgconn );
+		this->typemap = typemap;
+	}
+
+	return self;
+}
 
-	return val;
+VALUE
+pg_new_result_autoclear(PGresult *result, VALUE rb_pgconn)
+{
+	VALUE self = pg_new_result(result, rb_pgconn);
+	t_pg_result *this = pgresult_get_this(self);
+	this->autoclear = 1;
+	return self;
 }
 
 /*
@@ -44,24 +72,18 @@ pg_new_result(PGresult *result, VALUE rb_pgconn)
 VALUE
 pg_result_check( VALUE self )
 {
+	t_pg_result *this = pgresult_get_this(self);
 	VALUE error, exception, klass;
-	VALUE rb_pgconn = rb_iv_get( self, "@connection" );
-	PGconn *conn = pg_get_pgconn(rb_pgconn);
-	PGresult *result;
-#ifdef M17N_SUPPORTED
-	rb_encoding *enc = pg_conn_enc_get( conn );
-#endif
+	PGconn *conn = pg_get_pgconn(this->connection);
 	char * sqlstate;
 
-	Data_Get_Struct(self, PGresult, result);
-
-	if(result == NULL)
+	if(this->pgresult == NULL)
 	{
 		error = rb_str_new2( PQerrorMessage(conn) );
 	}
 	else
 	{
-		switch (PQresultStatus(result))
+		switch (PQresultStatus(this->pgresult))
 		{
 		case PGRES_TUPLES_OK:
 		case PGRES_COPY_OUT:
@@ -78,22 +100,20 @@ pg_result_check( VALUE self )
 		case PGRES_BAD_RESPONSE:
 		case PGRES_FATAL_ERROR:
 		case PGRES_NONFATAL_ERROR:
-			error = rb_str_new2( PQresultErrorMessage(result) );
+			error = rb_str_new2( PQresultErrorMessage(this->pgresult) );
 			break;
 		default:
 			error = rb_str_new2( "internal error : unknown result status." );
 		}
 	}
 
-#ifdef M17N_SUPPORTED
-	rb_enc_set_index( error, rb_enc_to_index(enc) );
-#endif
+	PG_ENCODING_SET_NOCHECK( error, ENCODING_GET(self) );
 
-	sqlstate = PQresultErrorField( result, PG_DIAG_SQLSTATE );
+	sqlstate = PQresultErrorField( this->pgresult, PG_DIAG_SQLSTATE );
 	klass = lookup_error_class( sqlstate );
 	exception = rb_exc_new3( klass, error );
-	rb_iv_set( exception, "@connection", rb_pgconn );
-	rb_iv_set( exception, "@result", result ? self : Qnil );
+	rb_iv_set( exception, "@connection", this->connection );
+	rb_iv_set( exception, "@result", this->pgresult ? self : Qnil );
 	rb_exc_raise( exception );
 
 	/* Not reached */
@@ -112,41 +132,119 @@ pg_result_check( VALUE self )
  *    res.clear() -> nil
  *
  * Clears the PG::Result object as the result of the query.
+ *
+ * If PG::Result#autoclear? is true then the result is marked as cleared
+ * and the underlying C struct will be cleared automatically by libpq.
+ *
  */
 VALUE
 pg_result_clear(VALUE self)
 {
-	PQclear(pgresult_get(self));
-	DATA_PTR(self) = NULL;
+	t_pg_result *this = pgresult_get_this(self);
+	if( !this->autoclear )
+		PQclear(pgresult_get(self));
+	this->pgresult = NULL;
 	return Qnil;
 }
 
+/*
+ * call-seq:
+ *    res.cleared?      -> boolean
+ *
+ * Returns +true+ if the backend result memory has been free'd.
+ */
+VALUE
+pgresult_cleared_p( VALUE self )
+{
+	t_pg_result *this = pgresult_get_this(self);
+	return this->pgresult ? Qfalse : Qtrue;
+}
 
+/*
+ * call-seq:
+ *    res.autoclose?      -> boolean
+ *
+ * Returns +true+ if the underlying C struct will be cleared automatically by libpq.
+ * Elsewise the result is cleared by PG::Result#clear or by the GC when it's no longer in use.
+ *
+ */
+VALUE
+pgresult_autoclear_p( VALUE self )
+{
+	t_pg_result *this = pgresult_get_this(self);
+	return this->autoclear ? Qtrue : Qfalse;
+}
 
 /*
  * DATA pointer functions
  */
 
+/*
+ * GC Mark function
+ */
+static void
+pgresult_gc_mark( t_pg_result *this )
+{
+	rb_gc_mark( this->connection );
+	rb_gc_mark( this->typemap );
+}
+
 /*
  * GC Free function
  */
 static void
-pgresult_gc_free( PGresult *result )
+pgresult_gc_free( t_pg_result *this )
 {
-	if(result != NULL)
-		PQclear(result);
+	if(this->pgresult != NULL && !this->autoclear)
+		PQclear(this->pgresult);
+
+	xfree(this);
 }
 
 /*
- * Fetch the data pointer for the result object
+ * Fetch the PG::Result object data pointer and check it's
+ * PGresult data pointer for sanity.
  */
-static PGresult*
+static t_pg_result *
+pgresult_get_this_safe( VALUE self )
+{
+	t_pg_result *this = pgresult_get_this(self);
+
+	if (this->pgresult == NULL) rb_raise(rb_ePGerror, "result has been cleared");
+	return this;
+}
+
+/*
+ * Fetch the PGresult pointer for the result object and check validity
+ */
+PGresult*
 pgresult_get(VALUE self)
 {
-	PGresult *result;
-	Data_Get_Struct(self, PGresult, result);
-	if (result == NULL) rb_raise(rb_ePGerror, "result has been cleared");
-	return result;
+	t_pg_result *this = pgresult_get_this(self);
+
+	if (this->pgresult == NULL) rb_raise(rb_ePGerror, "result has been cleared");
+	return this->pgresult;
+}
+
+/*
+ * Document-method: allocate
+ *
+ * call-seq:
+ *   PG::Result.allocate -> result
+ */
+static VALUE
+pgresult_s_allocate( VALUE klass )
+{
+	t_pg_result *this;
+	VALUE self = Data_Make_Struct( klass, t_pg_result, pgresult_gc_mark, pgresult_gc_free, this );
+
+	this->pgresult = NULL;
+	this->connection = Qnil;
+	this->typemap = Qnil;
+	this->p_typemap = DATA_PTR( pg_default_typemap );
+	this->autoclear = 0;
+
+	return self;
 }
 
 
@@ -205,7 +303,7 @@ static VALUE
 pgresult_res_status(VALUE self, VALUE status)
 {
 	VALUE ret = rb_tainted_str_new2(PQresStatus(NUM2INT(status)));
-	ASSOCIATE_INDEX(ret, self);
+	PG_ENCODING_SET_NOCHECK(ret, ENCODING_GET(self));
 	return ret;
 }
 
@@ -219,7 +317,7 @@ static VALUE
 pgresult_error_message(VALUE self)
 {
 	VALUE ret = rb_tainted_str_new2(PQresultErrorMessage(pgresult_get(self)));
-	ASSOCIATE_INDEX(ret, self);
+	PG_ENCODING_SET_NOCHECK(ret, ENCODING_GET(self));
 	return ret;
 }
 
@@ -279,7 +377,7 @@ pgresult_error_field(VALUE self, VALUE field)
 
 	if ( fieldstr ) {
 		ret = rb_tainted_str_new2( fieldstr );
-		ASSOCIATE_INDEX( ret, self );
+		PG_ENCODING_SET_NOCHECK( ret, ENCODING_GET(self ));
 	}
 
 	return ret;
@@ -327,7 +425,7 @@ pgresult_fname(VALUE self, VALUE index)
 		rb_raise(rb_eArgError,"invalid field number %d", i);
 	}
 	fname = rb_tainted_str_new2(PQfname(result, i));
-	ASSOCIATE_INDEX(fname, self);
+	PG_ENCODING_SET_NOCHECK(fname, ENCODING_GET(self));
 	return fname;
 }
 
@@ -361,9 +459,9 @@ pgresult_fnumber(VALUE self, VALUE name)
 
 	Check_Type(name, T_STRING);
 
-	n = PQfnumber(pgresult_get(self), StringValuePtr(name));
+	n = PQfnumber(pgresult_get(self), StringValueCStr(name));
 	if (n == -1) {
-		rb_raise(rb_eArgError,"Unknown field: %s", StringValuePtr(name));
+		rb_raise(rb_eArgError,"Unknown field: %s", StringValueCStr(name));
 	}
 	return INT2FIX(n);
 }
@@ -515,30 +613,6 @@ pgresult_fsize(VALUE self, VALUE index)
 }
 
 
-static VALUE
-pgresult_value(VALUE self, PGresult *result, int tuple_num, int field_num)
-{
-    VALUE val;
-    if ( PQgetisnull(result, tuple_num, field_num) ) {
-        return Qnil;
-    }
-    else {
-        val = rb_tainted_str_new( PQgetvalue(result, tuple_num, field_num ),
-                                  PQgetlength(result, tuple_num, field_num) );
-
-#ifdef M17N_SUPPORTED
-        /* associate client encoding for text format only */
-        if ( 0 == PQfformat(result, field_num) ) {
-            ASSOCIATE_INDEX( val, self );
-        } else {
-            rb_enc_associate( val, rb_ascii8bit_encoding() );
-        }
-#endif
-
-        return val;
-    }
-}
-
 /*
  * call-seq:
  *    res.getvalue( tup_num, field_num )
@@ -549,18 +623,17 @@ pgresult_value(VALUE self, PGresult *result, int tuple_num, int field_num)
 static VALUE
 pgresult_getvalue(VALUE self, VALUE tup_num, VALUE field_num)
 {
-	PGresult *result;
+	t_pg_result *this = pgresult_get_this_safe(self);
 	int i = NUM2INT(tup_num);
 	int j = NUM2INT(field_num);
 
-	result = pgresult_get(self);
-	if(i < 0 || i >= PQntuples(result)) {
+	if(i < 0 || i >= PQntuples(this->pgresult)) {
 		rb_raise(rb_eArgError,"invalid tuple number %d", i);
 	}
-	if(j < 0 || j >= PQnfields(result)) {
+	if(j < 0 || j >= PQnfields(this->pgresult)) {
 		rb_raise(rb_eArgError,"invalid field number %d", j);
 	}
-	return pgresult_value(self, result, i, j);
+	return this->p_typemap->typecast_result_value(self, i, j);
 }
 
 /*
@@ -653,7 +726,7 @@ static VALUE
 pgresult_cmd_status(VALUE self)
 {
 	VALUE ret = rb_tainted_str_new2(PQcmdStatus(pgresult_get(self)));
-	ASSOCIATE_INDEX(ret, self);
+	PG_ENCODING_SET_NOCHECK(ret, ENCODING_GET(self));
 	return ret;
 }
 
@@ -707,20 +780,20 @@ pgresult_oid_value(VALUE self)
 static VALUE
 pgresult_aref(VALUE self, VALUE index)
 {
-	PGresult *result = pgresult_get(self);
+	t_pg_result *this = pgresult_get_this_safe(self);
 	int tuple_num = NUM2INT(index);
 	int field_num;
 	VALUE fname;
 	VALUE tuple;
 
-	if ( tuple_num < 0 || tuple_num >= PQntuples(result) )
+	if ( tuple_num < 0 || tuple_num >= PQntuples(this->pgresult) )
 		rb_raise( rb_eIndexError, "Index %d is out of range", tuple_num );
 
 	tuple = rb_hash_new();
-	for ( field_num = 0; field_num < PQnfields(result); field_num++ ) {
-		fname = rb_tainted_str_new2( PQfname(result,field_num) );
-		ASSOCIATE_INDEX(fname, self);
-		rb_hash_aset( tuple, fname, pgresult_value(self, result, tuple_num, field_num) );
+	for ( field_num = 0; field_num < PQnfields(this->pgresult); field_num++ ) {
+		fname = rb_tainted_str_new2( PQfname(this->pgresult,field_num) );
+		PG_ENCODING_SET_NOCHECK(fname, ENCODING_GET(self));
+		rb_hash_aset( tuple, fname, this->p_typemap->typecast_result_value(self, tuple_num, field_num) );
 	}
 	return tuple;
 }
@@ -734,18 +807,18 @@ pgresult_aref(VALUE self, VALUE index)
 static VALUE
 pgresult_each_row(VALUE self)
 {
-	PGresult* result = (PGresult*) pgresult_get(self);
+	t_pg_result *this = pgresult_get_this_safe(self);
 	int row;
 	int field;
-	int num_rows = PQntuples(result);
-	int num_fields = PQnfields(result);
+	int num_rows = PQntuples(this->pgresult);
+	int num_fields = PQnfields(this->pgresult);
 
 	for ( row = 0; row < num_rows; row++ ) {
 		VALUE new_row = rb_ary_new2(num_fields);
 
 		/* populate the row */
 		for ( field = 0; field < num_fields; field++ ) {
-		    rb_ary_store( new_row, field, pgresult_value(self, result, row, field) );
+			rb_ary_store( new_row, field, this->p_typemap->typecast_result_value(self, row, field) );
 		}
 		rb_yield( new_row );
 	}
@@ -753,6 +826,34 @@ pgresult_each_row(VALUE self)
 	return Qnil;
 }
 
+/*
+ * call-seq:
+ *    res.values -> Array
+ *
+ * Returns all tuples as an array of arrays.
+ */
+static VALUE
+pgresult_values(VALUE self)
+{
+	t_pg_result *this = pgresult_get_this_safe(self);
+	int row;
+	int field;
+	int num_rows = PQntuples(this->pgresult);
+	int num_fields = PQnfields(this->pgresult);
+	VALUE results = rb_ary_new2( num_rows );
+
+	for ( row = 0; row < num_rows; row++ ) {
+		VALUE new_row = rb_ary_new2(num_fields);
+
+		/* populate the row */
+		for ( field = 0; field < num_fields; field++ ) {
+			rb_ary_store( new_row, field, this->p_typemap->typecast_result_value(self, row, field) );
+		}
+		rb_ary_store( results, row, new_row );
+	}
+
+	return results;
+}
 
 /*
  * Make a Ruby array out of the encoded values from the specified
@@ -761,28 +862,16 @@ pgresult_each_row(VALUE self)
 static VALUE
 make_column_result_array( VALUE self, int col )
 {
-	PGresult *result = pgresult_get( self );
-	int rows = PQntuples( result );
+	t_pg_result *this = pgresult_get_this_safe(self);
+	int rows = PQntuples( this->pgresult );
 	int i;
-	VALUE val = Qnil;
 	VALUE results = rb_ary_new2( rows );
 
-	if ( col >= PQnfields(result) )
+	if ( col >= PQnfields(this->pgresult) )
 		rb_raise( rb_eIndexError, "no column %d in result", col );
 
 	for ( i=0; i < rows; i++ ) {
-		val = rb_tainted_str_new( PQgetvalue(result, i, col),
-		                          PQgetlength(result, i, col) );
-
-#ifdef M17N_SUPPORTED
-		/* associate client encoding for text format only */
-		if ( 0 == PQfformat(result, col) ) {
-			ASSOCIATE_INDEX( val, self );
-		} else {
-			rb_enc_associate( val, rb_ascii8bit_encoding() );
-		}
-#endif
-
+		VALUE val = this->p_typemap->typecast_result_value(self, i, col);
 		rb_ary_store( results, i, val );
 	}
 
@@ -817,7 +906,7 @@ static VALUE
 pgresult_field_values( VALUE self, VALUE field )
 {
 	PGresult *result = pgresult_get( self );
-	const char *fieldname = StringValuePtr( field );
+	const char *fieldname = StringValueCStr( field );
 	int fnum = PQfnumber( result, fieldname );
 
 	if ( fnum < 0 )
@@ -861,18 +950,78 @@ pgresult_fields(VALUE self)
 
 	for ( i = 0; i < n; i++ ) {
 		VALUE val = rb_tainted_str_new2(PQfname(result, i));
-		ASSOCIATE_INDEX(val, self);
+		PG_ENCODING_SET_NOCHECK(val, ENCODING_GET(self));
 		rb_ary_store( fields, i, val );
 	}
 
 	return fields;
 }
 
+/*
+ * call-seq:
+ *    res.type_map = typemap
+ *
+ * Set the TypeMap that is used for type casts of result values to ruby objects.
+ *
+ * All value retrieval methods will respect the type map and will do the
+ * type casts from PostgreSQL's wire format to Ruby objects on the fly,
+ * according to the rules and decoders defined in the given typemap.
+ *
+ * +typemap+ can be:
+ * * a kind of PG::TypeMap
+ * * +nil+ - to type cast all result values to String.
+ *
+ */
+static VALUE
+pgresult_type_map_set(VALUE self, VALUE typemap)
+{
+	t_pg_result *this = pgresult_get_this(self);
+
+	if( NIL_P(typemap) ){
+		this->p_typemap = DATA_PTR( pg_default_typemap );
+	} else {
+		t_typemap *p_typemap;
+
+		if ( !rb_obj_is_kind_of(typemap, rb_cTypeMap) ) {
+			rb_raise( rb_eTypeError, "wrong argument type %s (expected kind of PG::TypeMap)",
+					rb_obj_classname( typemap ) );
+		}
+		Data_Get_Struct(typemap, t_typemap, p_typemap);
+
+		typemap = p_typemap->fit_to_result( typemap, self );
+		this->p_typemap = DATA_PTR( typemap );
+	}
+
+	this->typemap = typemap;
+
+	return typemap;
+}
+
+/*
+ * call-seq:
+ *    res.type_map -> value
+ *
+ * Returns the TypeMap that is currently set for type casts of result values to ruby objects.
+ *
+ * Returns either:
+ * * a kind of PG::TypeMap or
+ * * +nil+ - when no type map is set.
+ *
+ */
+static VALUE
+pgresult_type_map_get(VALUE self)
+{
+	t_pg_result *this = pgresult_get_this(self);
+
+	return this->typemap;
+}
+
 
 void
 init_pg_result()
 {
 	rb_cPGresult = rb_define_class_under( rb_mPG, "Result", rb_cObject );
+	rb_define_alloc_func( rb_cPGresult, pgresult_s_allocate );
 	rb_include_module(rb_cPGresult, rb_mEnumerable);
 	rb_include_module(rb_cPGresult, rb_mPGconstants);
 
@@ -913,8 +1062,14 @@ init_pg_result()
 	rb_define_method(rb_cPGresult, "each", pgresult_each, 0);
 	rb_define_method(rb_cPGresult, "fields", pgresult_fields, 0);
 	rb_define_method(rb_cPGresult, "each_row", pgresult_each_row, 0);
+	rb_define_method(rb_cPGresult, "values", pgresult_values, 0);
 	rb_define_method(rb_cPGresult, "column_values", pgresult_column_values, 1);
 	rb_define_method(rb_cPGresult, "field_values", pgresult_field_values, 1);
+	rb_define_method(rb_cPGresult, "cleared?", pgresult_cleared_p, 0);
+	rb_define_method(rb_cPGresult, "autoclear?", pgresult_autoclear_p, 0);
+
+	rb_define_method(rb_cPGresult, "type_map=", pgresult_type_map_set, 1);
+	rb_define_method(rb_cPGresult, "type_map", pgresult_type_map_get, 0);
 }