Add field_type information to result

Towards trilogy-libraries#99, this exposes additional field data on the query result via a `Result#columns` method. We were already using this data in the C extension to determine how to cast the column values. We now add a pointer to the `column_info` struct onto the result and then lazily add the values to the `Result` instance when called. This method could eventually replace the `#fields` method which at the moment only includes the column name. Co-authored-by: Daniel Colson <[email protected]> Co-authored-by: Charlotte Wen <[email protected]>
ipc103 · Sep 28, 2023 · 2e9e86b · 2e9e86b
1 parent bed8602
commit 2e9e86b
Show file tree

Hide file tree

Showing 4 changed files with 97 additions and 12 deletions.
diff --git a/contrib/ruby/ext/trilogy-ruby/cext.c b/contrib/ruby/ext/trilogy-ruby/cext.c
@@ -18,7 +18,7 @@
 VALUE Trilogy_CastError;
 static VALUE Trilogy_BaseConnectionError, Trilogy_ProtocolError, Trilogy_SSLError, Trilogy_QueryError,
     Trilogy_ConnectionClosedError, Trilogy_ConnectionRefusedError, Trilogy_ConnectionResetError,
-    Trilogy_TimeoutError, Trilogy_SyscallError, Trilogy_Result, Trilogy_EOFError;
+    Trilogy_TimeoutError, Trilogy_SyscallError, Trilogy_Result, Trilogy_Result_Column, Trilogy_EOFError;
 
 static ID id_socket, id_host, id_port, id_username, id_password, id_found_rows, id_connect_timeout, id_read_timeout,
     id_write_timeout, id_keepalive_enabled, id_keepalive_idle, id_keepalive_interval, id_keepalive_count,
@@ -34,6 +34,11 @@ struct trilogy_ctx {
     VALUE encoding;
 };
 
+struct trilogy_result_ctx {
+    struct column_info *column_info;
+    uint64_t column_count;
+};
+
 static void mark_trilogy(void *ptr)
 {
     struct trilogy_ctx *ctx = ptr;
@@ -49,6 +54,25 @@ static void free_trilogy(void *ptr)
     xfree(ptr);
 }
 
+static void free_trilogy_result(void *ptr)
+{
+    struct trilogy_result_ctx *ctx = ptr;
+    if (ctx->column_info != NULL) {
+        xfree(ctx->column_info);
+    }
+    xfree(ptr);
+}
+
+static size_t trilogy_result_memsize(const void *ptr)
+{
+    const struct trilogy_result_ctx *ctx = ptr;
+    size_t memsize = sizeof(struct trilogy_result_ctx);
+    if (ctx->column_info != NULL) {
+        memsize += sizeof(struct column_info) * ctx->column_count;
+    }
+    return memsize;
+}
+
 static size_t trilogy_memsize(const void *ptr) {
     const struct trilogy_ctx *ctx = ptr;
     size_t memsize = sizeof(struct trilogy_ctx);
@@ -69,6 +93,15 @@ static const rb_data_type_t trilogy_data_type = {
     .flags = RUBY_TYPED_FREE_IMMEDIATELY | RUBY_TYPED_WB_PROTECTED
 };
 
+static const rb_data_type_t trilogy_result_data_type = {
+    .wrap_struct_name = "trilogy_result",
+    .function = {
+        .dfree = free_trilogy_result,
+        .dsize = trilogy_result_memsize,
+    },
+    .flags = RUBY_TYPED_FREE_IMMEDIATELY | RUBY_TYPED_WB_PROTECTED
+};
+
 static struct trilogy_ctx *get_ctx(VALUE obj)
 {
     struct trilogy_ctx *ctx;
@@ -183,6 +216,22 @@ static VALUE allocate_trilogy(VALUE klass)
     return obj;
 }
 
+static VALUE allocate_trilogy_result(VALUE klass)
+{
+    struct trilogy_result_ctx *ctx;
+
+    VALUE obj = TypedData_Make_Struct(klass, struct trilogy_result_ctx, &trilogy_result_data_type, ctx);
+
+    return obj;
+}
+
+static struct trilogy_result_ctx *get_trilogy_result_ctx(VALUE obj)
+{
+    struct trilogy_result_ctx *ctx;
+    TypedData_Get_Struct(obj, struct trilogy_result_ctx, &trilogy_result_data_type, ctx);
+    return ctx;
+}
+
 static int flush_writes(struct trilogy_ctx *ctx)
 {
     while (1) {
@@ -765,10 +814,9 @@ static VALUE read_query_response(VALUE vargs)
         rb_ivar_set(result, id_ivar_last_insert_id, Qnil);
         rb_ivar_set(result, id_ivar_affected_rows, Qnil);
     }
-
-    VALUE rb_column_info;
-    struct column_info *column_info = ALLOCV_N(struct column_info, rb_column_info, column_count);
-
+    struct trilogy_result_ctx *trilogy_result_ctx = get_trilogy_result_ctx(result);
+    trilogy_result_ctx->column_info = ALLOC_N(struct column_info, column_count);
+    trilogy_result_ctx->column_count = column_count;
     for (uint64_t i = 0; i < column_count; i++) {
         trilogy_column_t column;
 
@@ -798,11 +846,12 @@ static VALUE read_query_response(VALUE vargs)
 
         rb_ary_push(column_names, column_name);
 
-        column_info[i].type = column.type;
-        column_info[i].flags = column.flags;
-        column_info[i].len = column.len;
-        column_info[i].charset = column.charset;
-        column_info[i].decimals = column.decimals;
+        trilogy_result_ctx->column_info[i].name = column_name;
+        trilogy_result_ctx->column_info[i].type = column.type;
+        trilogy_result_ctx->column_info[i].flags = column.flags;
+        trilogy_result_ctx->column_info[i].len = column.len;
+        trilogy_result_ctx->column_info[i].charset = column.charset;
+        trilogy_result_ctx->column_info[i].decimals = column.decimals;
     }
 
     VALUE rb_row_values;
@@ -829,12 +878,12 @@ static VALUE read_query_response(VALUE vargs)
 
         if (args->cast_options->flatten_rows) {
             for (uint64_t i = 0; i < column_count; i++) {
-                rb_ary_push(rows, rb_trilogy_cast_value(row_values + i, column_info + i, args->cast_options));
+                rb_ary_push(rows, rb_trilogy_cast_value(row_values + i, trilogy_result_ctx->column_info + i, args->cast_options));
             }
         } else {
             VALUE row = rb_ary_new2(column_count);
             for (uint64_t i = 0; i < column_count; i++) {
-                rb_ary_push(row, rb_trilogy_cast_value(row_values + i, column_info + i, args->cast_options));
+                rb_ary_push(row, rb_trilogy_cast_value(row_values + i, trilogy_result_ctx->column_info + i, args->cast_options));
             }
             rb_ary_push(rows, row);
         }
@@ -1095,6 +1144,21 @@ static VALUE rb_trilogy_write_timeout_set(VALUE self, VALUE write_timeout)
     return write_timeout;
 }
 
+static VALUE rb_trilogy_result_columns(VALUE self)
+{
+    struct trilogy_result_ctx *trilogy_result_ctx = get_trilogy_result_ctx(self);
+    VALUE cols = rb_ary_new();
+    for (uint64_t i = 0; i < trilogy_result_ctx->column_count; i++) {
+        VALUE obj = rb_funcall(
+            Trilogy_Result_Column, rb_intern("new"), 6, trilogy_result_ctx->column_info[i].name,
+            rb_int_new(trilogy_result_ctx->column_info[i].type), rb_int_new(trilogy_result_ctx->column_info[i].len),
+            rb_int_new(trilogy_result_ctx->column_info[i].flags),
+            rb_int_new(trilogy_result_ctx->column_info[i].charset), rb_int_new(trilogy_result_ctx->column_info[i].decimals));
+        rb_ary_push(cols, obj);
+    }
+    return cols;    
+}
+
 static VALUE rb_trilogy_server_status(VALUE self) { return LONG2FIX(get_open_ctx(self)->conn.server_status); }
 
 static VALUE rb_trilogy_server_version(VALUE self) { return rb_str_new_cstr(get_open_ctx(self)->server_version); }
@@ -1172,7 +1236,12 @@ RUBY_FUNC_EXPORTED void Init_cext()
     rb_global_variable(&Trilogy_ConnectionClosedError);
 
     Trilogy_Result = rb_const_get(Trilogy, rb_intern("Result"));
+    rb_define_alloc_func(Trilogy_Result, allocate_trilogy_result);
     rb_global_variable(&Trilogy_Result);
+    rb_define_private_method(Trilogy_Result, "_columns", rb_trilogy_result_columns, 0);
+
+    Trilogy_Result_Column = rb_const_get(Trilogy_Result, rb_intern("Column"));
+    rb_global_variable(&Trilogy_Result_Column);
 
     Trilogy_SyscallError = rb_const_get(Trilogy, rb_intern("SyscallError"));
     rb_global_variable(&Trilogy_SyscallError);

diff --git a/contrib/ruby/ext/trilogy-ruby/trilogy-ruby.h b/contrib/ruby/ext/trilogy-ruby/trilogy-ruby.h
@@ -21,6 +21,7 @@ struct rb_trilogy_cast_options {
 };
 
 struct column_info {
+    VALUE name;
     TRILOGY_TYPE_t type;
     TRILOGY_CHARSET_t charset;
     uint32_t len;

diff --git a/contrib/ruby/lib/trilogy/result.rb b/contrib/ruby/lib/trilogy/result.rb
@@ -28,6 +28,18 @@ def each(&bk)
       rows.each(&bk)
     end
 
+    def columns
+      @columns ||= _columns
+    end
+
     include Enumerable
+
+    class Column
+      attr_reader :name, :type, :length, :flags, :charset, :decimals
+
+      def initialize(name, type, length, flags, charset, decimals)
+        @name, @type, @length, @flags, @charset, @decimals = name, type, length, flags, charset, decimals
+      end
+    end
   end
 end
diff --git a/contrib/ruby/test/client_test.rb b/contrib/ruby/test/client_test.rb
@@ -254,6 +254,8 @@ def test_trilogy_query_values
     result = client.query_with_flags("SELECT id, int_test FROM trilogy_test", client.query_flags | Trilogy::QUERY_FLAGS_FLATTEN_ROWS)
 
     assert_equal ["id", "int_test"], result.fields
+    assert_equal 2, result.columns.count
+    assert_equal ["id", "int_test"], result.columns.map(&:name)
     assert_equal [1, 4, 2, 3, 3, 1], result.rows
   end
 
@@ -314,6 +316,7 @@ def test_trilogy_query_result_object
     result = client.query "SELECT 1 AS a, 2 AS b"
 
     assert_equal ["a", "b"], result.fields
+    assert_equal ["a", "b"], result.columns.map(&:name)
     assert_equal [[1, 2]], result.rows
     assert_equal [{ "a" => 1, "b" => 2 }], result.each_hash.to_a
     assert_equal [[1, 2]], result.to_a