diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index e96e9e3..f8f10a3 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -43,12 +43,12 @@ jobs: experimental: true continue-on-error: ${{matrix.experimental}} steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v4 - uses: ruby/setup-ruby@v1 with: ruby-version: ${{ matrix.ruby }} - name: Bundler cache - uses: actions/cache@v2 + uses: actions/cache@v4 with: path: vendor/bundle key: ${{ runner.os }}-${{ matrix.ruby }}-gems-${{ hashFiles('**/Gemfile.lock') }} @@ -73,7 +73,7 @@ jobs: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v4 - name: Release Gem uses: discourse/publish-rubygems-action@v2 diff --git a/CHANGELOG b/CHANGELOG index 74ebc48..403cc8e 100644 --- a/CHANGELOG +++ b/CHANGELOG @@ -1,3 +1,8 @@ +Unreleased + +- FEATURE: add optional `mini_sql-pg_native` companion gem for faster PostgreSQL row materialization +- FIX: initialize PostgreSQL type map before streaming queries enter single-row mode + 2024-08-21 - 1.6.0 - FEATURE: optionaly allow encoding pg arrays efficiently diff --git a/Gemfile b/Gemfile index 554e8a1..9fb97e7 100644 --- a/Gemfile +++ b/Gemfile @@ -5,4 +5,4 @@ source "https://rubygems.org" git_source(:github) { |repo_name| "https://github.com/#{repo_name}" } # Specify your gem's dependencies in mini_sql.gemspec -gemspec +gemspec name: "mini_sql" diff --git a/README.md b/README.md index 956133f..9dbaa20 100644 --- a/README.md +++ b/README.md @@ -8,6 +8,16 @@ Add this line to your application's Gemfile: gem 'mini_sql' ``` +For faster PostgreSQL row materialization, add the optional native companion gem: + +```ruby +gem 'mini_sql' +gem 'mini_sql-pg_native' +``` + +`mini_sql-pg_native` only affects the PostgreSQL adapter and can be disabled with +`MINI_SQL_PG_NATIVE=0`. + And then execute: $ bundle diff --git a/Rakefile b/Rakefile index dbefe50..d4b707f 100644 --- a/Rakefile +++ b/Rakefile @@ -1,6 +1,7 @@ # frozen_string_literal: true -require "bundler/gem_tasks" +require "bundler/gem_helper" +Bundler::GemHelper.install_tasks(name: "mini_sql") require "rake/testtask" if RUBY_ENGINE == 'jruby' # Excluding sqlite3 tests diff --git a/bench/pg_native_materialize_only_perf.rb b/bench/pg_native_materialize_only_perf.rb new file mode 100644 index 0000000..d23f634 --- /dev/null +++ b/bench/pg_native_materialize_only_perf.rb @@ -0,0 +1,52 @@ +# frozen_string_literal: true + +require 'pg' +require 'benchmark/ips' +require_relative '../lib/mini_sql' + +conn = PG.connect(dbname: ENV.fetch('PGDATABASE', 'discourse_sql_ft'), user: ENV.fetch('PGUSER', 'agent')) +conn.async_exec <<~SQL + DROP TABLE IF EXISTS mini_sql_pg_native_bench; + CREATE UNLOGGED TABLE mini_sql_pg_native_bench AS + SELECT + g AS id, + 'title ' || g AS title, + 'slug-' || g AS slug, + (g % 1000) AS user_id, + (g % 100) AS category_id, + (g % 2 = 0) AS visible, + now() - (g || ' seconds')::interval AS created_at, + now() AS updated_at, + repeat('x', 32) AS payload, + g::numeric / 10 AS score + FROM generate_series(1, 50000) g; +SQL + +mini = MiniSql::Connection.get(conn) +mode = ENV['MINI_SQL_PG_NATIVE'] == '0' ? 'ruby' : 'pg_native' +puts "mode=#{mode} pg_native_defined=#{defined?(MiniSql::Postgres::Native::RowMaterializer).inspect} ruby=#{RUBY_VERSION} pg=#{PG.library_version}" + +RESULTS = { + 'mat_1k_2col' => conn.async_exec('select id, title from mini_sql_pg_native_bench order by id limit 1000'), + 'mat_1k_10col' => conn.async_exec('select id, title, slug, user_id, category_id, visible, created_at, updated_at, payload, score from mini_sql_pg_native_bench order by id limit 1000'), + 'mat_10k_10col' => conn.async_exec('select id, title, slug, user_id, category_id, visible, created_at, updated_at, payload, score from mini_sql_pg_native_bench order by id limit 10000'), +}.freeze + +RESULTS.each_value { |r| r.type_map = mini.type_map } +cache = MiniSql::Postgres::DeserializerCache.new +RESULTS.each_value { |r| 3.times { cache.materialize(r) } } + +Benchmark.ips do |x| + x.config(warmup: Float(ENV.fetch('WARMUP', '2')), time: Float(ENV.fetch('TIME', '5'))) + + RESULTS.each do |name, result| + x.report(name) do |n| + while n > 0 + rows = cache.materialize(result) + first = rows[0] + raise 'bad row' unless first && first.id + n -= 1 + end + end + end +end diff --git a/bench/pg_native_materialize_perf.rb b/bench/pg_native_materialize_perf.rb new file mode 100644 index 0000000..d1f96f7 --- /dev/null +++ b/bench/pg_native_materialize_perf.rb @@ -0,0 +1,53 @@ +# frozen_string_literal: true + +require 'pg' +require 'benchmark/ips' +require_relative '../lib/mini_sql' + +conn = PG.connect(dbname: ENV.fetch('PGDATABASE', 'discourse_sql_ft'), user: ENV.fetch('PGUSER', 'agent')) +conn.async_exec <<~SQL + DROP TABLE IF EXISTS mini_sql_pg_native_bench; + CREATE UNLOGGED TABLE mini_sql_pg_native_bench AS + SELECT + g AS id, + 'title ' || g AS title, + 'slug-' || g AS slug, + (g % 1000) AS user_id, + (g % 100) AS category_id, + (g % 2 = 0) AS visible, + now() - (g || ' seconds')::interval AS created_at, + now() AS updated_at, + repeat('x', 32) AS payload, + g::numeric / 10 AS score + FROM generate_series(1, 50000) g; + CREATE INDEX ON mini_sql_pg_native_bench(id); +SQL + +mini = MiniSql::Connection.get(conn) +mode = ENV['MINI_SQL_PG_NATIVE'] == '0' ? 'ruby' : 'pg_native' +puts "mode=#{mode} pg_native_defined=#{defined?(MiniSql::Postgres::Native::RowMaterializer).inspect} ruby=#{RUBY_VERSION} pg=#{PG.library_version}" + +QUERIES = { + 'query_1k_2col' => 'select id, title from mini_sql_pg_native_bench order by id limit 1000', + 'query_1k_10col' => 'select id, title, slug, user_id, category_id, visible, created_at, updated_at, payload, score from mini_sql_pg_native_bench order by id limit 1000', + 'query_10k_10col' => 'select id, title, slug, user_id, category_id, visible, created_at, updated_at, payload, score from mini_sql_pg_native_bench order by id limit 10000', +}.freeze + +# Warm caches: PostgreSQL buffer cache, pg type map, mini_sql materializer cache. +QUERIES.each_value { |sql| 3.times { mini.query(sql) } } + +Benchmark.ips do |x| + x.config(warmup: Float(ENV.fetch('WARMUP', '2')), time: Float(ENV.fetch('TIME', '5'))) + + QUERIES.each do |name, sql| + x.report(name) do |n| + while n > 0 + rows = mini.query(sql) + # Touch values so dead-code elimination has nowhere cute to hide. + first = rows[0] + raise 'bad row' unless first && first.id + n -= 1 + end + end + end +end diff --git a/bench/pg_native_simple_materialize_only_perf.rb b/bench/pg_native_simple_materialize_only_perf.rb new file mode 100644 index 0000000..1a88a4d --- /dev/null +++ b/bench/pg_native_simple_materialize_only_perf.rb @@ -0,0 +1,46 @@ +# frozen_string_literal: true + +require 'pg' +require 'benchmark/ips' +require_relative '../lib/mini_sql' + +conn = PG.connect(dbname: ENV.fetch('PGDATABASE', 'discourse_sql_ft'), user: ENV.fetch('PGUSER', 'agent')) +conn.async_exec <<~SQL + DROP TABLE IF EXISTS mini_sql_pg_native_simple_bench; + CREATE UNLOGGED TABLE mini_sql_pg_native_simple_bench AS + SELECT + g AS id, + 'title ' || g AS title, + 'slug-' || g AS slug, + (g % 1000) AS user_id, + (g % 100) AS category_id, + repeat('x', 32) AS payload + FROM generate_series(1, 50000) g; +SQL + +mini = MiniSql::Connection.get(conn) +mode = ENV['MINI_SQL_PG_NATIVE'] == '0' ? 'ruby' : 'pg_native' +puts "mode=#{mode} pg_native_defined=#{defined?(MiniSql::Postgres::Native::RowMaterializer).inspect} ruby=#{RUBY_VERSION} pg=#{PG.library_version}" + +RESULTS = { + 'simple_mat_1k_2col' => conn.async_exec('select id, title from mini_sql_pg_native_simple_bench order by id limit 1000'), + 'simple_mat_1k_6col' => conn.async_exec('select id, title, slug, user_id, category_id, payload from mini_sql_pg_native_simple_bench order by id limit 1000'), + 'simple_mat_10k_6col' => conn.async_exec('select id, title, slug, user_id, category_id, payload from mini_sql_pg_native_simple_bench order by id limit 10000'), +}.freeze + +RESULTS.each_value { |r| r.type_map = mini.type_map } +cache = MiniSql::Postgres::DeserializerCache.new +RESULTS.each_value { |r| 3.times { cache.materialize(r) } } + +Benchmark.ips do |x| + x.config(warmup: Float(ENV.fetch('WARMUP', '2')), time: Float(ENV.fetch('TIME', '5'))) + RESULTS.each do |name, result| + x.report(name) do |n| + while n > 0 + rows = cache.materialize(result) + raise 'bad row' unless rows[0]&.id + n -= 1 + end + end + end +end diff --git a/ext/mini_sql/pg_native/extconf.rb b/ext/mini_sql/pg_native/extconf.rb new file mode 100644 index 0000000..8f520cb --- /dev/null +++ b/ext/mini_sql/pg_native/extconf.rb @@ -0,0 +1,11 @@ +# frozen_string_literal: true + +require "mkmf" +require "rubygems" + +pg_spec = Gem::Specification.find_by_name("pg") +# rubocop:disable Style/GlobalVars +$INCFLAGS << " -I#{File.join(pg_spec.full_gem_path, "ext")}" +# rubocop:enable Style/GlobalVars + +create_makefile("mini_sql/pg_native/pg_native") diff --git a/ext/mini_sql/pg_native/pg_native.c b/ext/mini_sql/pg_native/pg_native.c new file mode 100644 index 0000000..51179d8 --- /dev/null +++ b/ext/mini_sql/pg_native/pg_native.c @@ -0,0 +1,143 @@ +#include "ruby.h" +#include "pg.h" + +typedef struct { + VALUE klass; + long nfields; + ID *ivars; +} row_materializer_t; + +static VALUE cRowMaterializer; + +static void rm_mark(void *ptr) { + row_materializer_t *rm = (row_materializer_t *)ptr; + if (rm) rb_gc_mark_movable(rm->klass); +} + +static void rm_compact(void *ptr) { + row_materializer_t *rm = (row_materializer_t *)ptr; + if (rm) rm->klass = rb_gc_location(rm->klass); +} + +static void rm_free(void *ptr) { + row_materializer_t *rm = (row_materializer_t *)ptr; + if (rm) { + xfree(rm->ivars); + xfree(rm); + } +} + +static size_t rm_memsize(const void *ptr) { + const row_materializer_t *rm = (const row_materializer_t *)ptr; + return rm ? sizeof(row_materializer_t) + (sizeof(ID) * rm->nfields) : 0; +} + +static const rb_data_type_t rm_type = { + "MiniSql::Postgres::Native::RowMaterializer", + { rm_mark, rm_free, rm_memsize, rm_compact, }, + 0, 0, RUBY_TYPED_FREE_IMMEDIATELY +}; + +static VALUE rm_alloc(VALUE klass) { + row_materializer_t *rm; + VALUE obj = TypedData_Make_Struct(klass, row_materializer_t, &rm_type, rm); + rm->klass = Qnil; + rm->nfields = 0; + rm->ivars = NULL; + return obj; +} + +static VALUE rm_initialize(VALUE self, VALUE row_class, VALUE fields) { + row_materializer_t *rm; + TypedData_Get_Struct(self, row_materializer_t, &rm_type, rm); + + Check_Type(fields, T_ARRAY); + rm->klass = row_class; + rm->nfields = RARRAY_LEN(fields); + xfree(rm->ivars); + rm->ivars = ALLOC_N(ID, rm->nfields); + + for (long i = 0; i < rm->nfields; i++) { + VALUE field = rb_obj_as_string(rb_ary_entry(fields, i)); + VALUE ivar_name = rb_str_plus(rb_str_new_cstr("@"), field); + rm->ivars[i] = rb_intern_str(ivar_name); + } + + return self; +} + +static void check_pgresult(t_pg_result *pgresult) { + if (!pgresult->pgresult) { + rb_raise(rb_eNoResultError, "PG::Result has been cleared"); + } +} + +static VALUE rm_materialize(VALUE self, VALUE result, VALUE index) { + row_materializer_t *rm; + TypedData_Get_Struct(self, row_materializer_t, &rm_type, rm); + + int idx = NUM2INT(index); + t_pg_result *pgresult = pgresult_get_this(result); + check_pgresult(pgresult); + + int rows = PQntuples(pgresult->pgresult); + if (idx < 0 || idx >= rows) { + rb_raise(rb_eArgError, "invalid tuple number %d", idx); + } + + VALUE row = rb_obj_alloc(rm->klass); + long nfields = rm->nfields; + ID *ivars = rm->ivars; + t_typemap *typemap = pgresult->p_typemap; + + for (long col = 0; col < nfields; col++) { + VALUE val = typemap->funcs.typecast_result_value(typemap, result, idx, col); + rb_ivar_set(row, ivars[col], val); + } + + return row; +} + +static VALUE rm_materialize_all(VALUE self, VALUE result) { + row_materializer_t *rm; + TypedData_Get_Struct(self, row_materializer_t, &rm_type, rm); + + t_pg_result *pgresult = pgresult_get_this(result); + check_pgresult(pgresult); + + int rows = PQntuples(pgresult->pgresult); + long nfields = rm->nfields; + ID *ivars = rm->ivars; + t_typemap *typemap = pgresult->p_typemap; + VALUE ary = rb_ary_new_capa(rows); + + for (int idx = 0; idx < rows; idx++) { + VALUE row = rb_obj_alloc(rm->klass); + for (long col = 0; col < nfields; col++) { + VALUE val = typemap->funcs.typecast_result_value(typemap, result, idx, col); + rb_ivar_set(row, ivars[col], val); + } + rb_ary_store(ary, idx, row); + } + + return ary; +} + +static VALUE rm_row_class(VALUE self) { + row_materializer_t *rm; + TypedData_Get_Struct(self, row_materializer_t, &rm_type, rm); + return rm->klass; +} + +void Init_pg_native(void) { + VALUE mMiniSql = rb_define_module("MiniSql"); + VALUE mPostgres = rb_define_module_under(mMiniSql, "Postgres"); + VALUE mNative = rb_define_module_under(mPostgres, "Native"); + + cRowMaterializer = rb_define_class_under(mNative, "RowMaterializer", rb_cObject); + rb_define_alloc_func(cRowMaterializer, rm_alloc); + rb_define_method(cRowMaterializer, "initialize", rm_initialize, 2); + rb_define_method(cRowMaterializer, "materialize", rm_materialize, 2); + rb_define_method(cRowMaterializer, "materialize_all", rm_materialize_all, 1); + rb_define_method(cRowMaterializer, "row_class", rm_row_class, 0); +} diff --git a/lib/mini_sql/mysql/connection.rb b/lib/mini_sql/mysql/connection.rb index c08b027..6e0e948 100644 --- a/lib/mini_sql/mysql/connection.rb +++ b/lib/mini_sql/mysql/connection.rb @@ -29,8 +29,8 @@ def query_array(sql, *params) end def exec(sql, *params) - run(sql, :array, params) - raw_connection.affected_rows + result = run(sql, :array, params) + result ? result.count : raw_connection.affected_rows end def query(sql, *params) diff --git a/lib/mini_sql/pg_native.rb b/lib/mini_sql/pg_native.rb new file mode 100644 index 0000000..513af24 --- /dev/null +++ b/lib/mini_sql/pg_native.rb @@ -0,0 +1,4 @@ +# frozen_string_literal: true + +require "pg" +require "mini_sql/pg_native/pg_native" diff --git a/lib/mini_sql/postgres/connection.rb b/lib/mini_sql/postgres/connection.rb index f801f7d..1037164 100644 --- a/lib/mini_sql/postgres/connection.rb +++ b/lib/mini_sql/postgres/connection.rb @@ -120,6 +120,7 @@ def query_each(sql, *params) sql = param_encoder.encode(sql, *params) end + tm = type_map raw_connection.send_query(sql) raw_connection.set_single_row_mode @@ -127,23 +128,25 @@ def query_each(sql, *params) result = raw_connection.get_result break if !result - result.check - - if result.ntuples == 0 - # skip, this happens at the end when we get totals - else - materializer ||= deserializer_cache.materializer(result) - result.type_map = type_map - i = 0 - # technically we should only get 1 row here - # but protect against future batching changes - while i < result.ntuples - yield materializer.materialize(result, i) - i += 1 + begin + result.check + + if result.ntuples == 0 + # skip, this happens at the end when we get totals + else + materializer ||= deserializer_cache.materializer(result) + result.type_map = tm + i = 0 + # technically we should only get 1 row here + # but protect against future batching changes + while i < result.ntuples + yield materializer.materialize(result, i) + i += 1 + end end + ensure + result.clear end - - result.clear end end @@ -153,6 +156,7 @@ def query_each_hash(sql, *params) sql = param_encoder.encode(sql, *params) end + tm = type_map raw_connection.send_query(sql) raw_connection.set_single_row_mode @@ -160,23 +164,25 @@ def query_each_hash(sql, *params) result = raw_connection.get_result break if !result - result.check - - if result.ntuples == 0 - # skip, this happens at the end when we get totals - else - result.type_map = type_map - i = 0 + begin + result.check - # technically we should only get 1 row here - # but protect against future batching changes - while i < result.ntuples - yield result[i] - i += 1 + if result.ntuples == 0 + # skip, this happens at the end when we get totals + else + result.type_map = tm + i = 0 + + # technically we should only get 1 row here + # but protect against future batching changes + while i < result.ntuples + yield result[i] + i += 1 + end end + ensure + result.clear end - - result.clear end end diff --git a/lib/mini_sql/postgres/deserializer_cache.rb b/lib/mini_sql/postgres/deserializer_cache.rb index 99db7fe..bda19ea 100644 --- a/lib/mini_sql/postgres/deserializer_cache.rb +++ b/lib/mini_sql/postgres/deserializer_cache.rb @@ -1,5 +1,11 @@ # frozen_string_literal: true +begin + require "mini_sql/pg_native" unless ENV["MINI_SQL_PG_NATIVE"] == "0" +rescue LoadError + # mini_sql-pg_native is an optional companion gem. +end + module MiniSql module Postgres class DeserializerCache @@ -28,29 +34,28 @@ def materializer(result) def materialize(result, decorator_module = nil) return [] if result.ntuples == 0 - key = result.fields.join(',') - - # trivial fast LRU implementation - materializer = @cache.delete(key) - if materializer - @cache[key] = materializer - else - materializer = @cache[key] = new_row_materializer(result) - @cache.shift if @cache.length > @max_size - end + materializer = materializer(result) if decorator_module - materializer = materializer.decorated(decorator_module) + if materializer.respond_to?(:row_class) + materializer = materializer.row_class.decorated(decorator_module) + else + materializer = materializer.decorated(decorator_module) + end end - i = 0 - r = [] - # quicker loop - while i < result.ntuples - r << materializer.materialize(result, i) - i += 1 + if materializer.respond_to?(:materialize_all) + materializer.materialize_all(result) + else + i = 0 + r = [] + # quicker loop + while i < result.ntuples + r << materializer.materialize(result, i) + i += 1 + end + r end - r end private @@ -67,19 +72,25 @@ def new_row_materializer(result) i += 1 end - Class.new do + row_class = Class.new do extend MiniSql::Decoratable include MiniSql::Result attr_accessor(*fields) + end - instance_eval <<~RUBY - def materialize(pg_result, index) - r = self.new - #{col = -1; fields.map { |f| "r.#{f} = pg_result.getvalue(index, #{col += 1})" }.join("; ")} - r - end - RUBY + row_class.instance_eval <<~RUBY + def materialize(pg_result, index) + r = self.new + #{col = -1; fields.map { |f| "r.#{f} = pg_result.getvalue(index, #{col += 1})" }.join("; ")} + r + end + RUBY + + if defined?(MiniSql::Postgres::Native::RowMaterializer) && ENV["MINI_SQL_PG_NATIVE"] != "0" + MiniSql::Postgres::Native::RowMaterializer.new(row_class, fields) + else + row_class end end end diff --git a/mini_sql-pg_native.gemspec b/mini_sql-pg_native.gemspec new file mode 100644 index 0000000..b9e7ba9 --- /dev/null +++ b/mini_sql-pg_native.gemspec @@ -0,0 +1,35 @@ +# frozen_string_literal: true + +lib = File.expand_path("lib", __dir__) +$LOAD_PATH.unshift(lib) unless $LOAD_PATH.include?(lib) +require "mini_sql/version" + +Gem::Specification.new do |spec| + spec.name = "mini_sql-pg_native" + spec.version = MiniSql::VERSION + spec.authors = ["Sam Saffron"] + spec.email = ["sam.saffron@gmail.com"] + + spec.summary = "Optional native PostgreSQL row materializer for mini_sql" + spec.description = "Optional native PostgreSQL row materializer for mini_sql" + spec.homepage = "https://github.com/discourse/mini_sql" + spec.license = "MIT" + + spec.metadata = { + "bug_tracker_uri" => "https://github.com/discourse/mini_sql/issues", + "source_code_uri" => "https://github.com/discourse/mini_sql", + "changelog_uri" => "https://github.com/discourse/mini_sql/blob/main/CHANGELOG" + } + + # rubocop:disable Discourse/NoChdir + spec.files = Dir.chdir(__dir__) do + `git ls-files -z ext/mini_sql/pg_native lib/mini_sql/pg_native.rb lib/mini_sql/version.rb LICENSE.txt`.split("\x0") + end + # rubocop:enable Discourse/NoChdir + spec.require_paths = ["lib"] + spec.extensions = ["ext/mini_sql/pg_native/extconf.rb"] + + spec.required_ruby_version = ">= 2.7" + spec.add_dependency "mini_sql", "= #{MiniSql::VERSION}" + spec.add_dependency "pg", ">= 1.6", "< 1.7" +end diff --git a/mini_sql.gemspec b/mini_sql.gemspec index e3444e0..e467219 100644 --- a/mini_sql.gemspec +++ b/mini_sql.gemspec @@ -27,7 +27,12 @@ Gem::Specification.new do |spec| # The `git ls-files -z` loads the files in the RubyGem that have been added into git. # rubocop:disable Discourse/NoChdir spec.files = Dir.chdir(File.expand_path('..', __FILE__)) do - `git ls-files -z`.split("\x0").reject { |f| f.match(%r{^(test|spec|features)/}) } + `git ls-files -z`.split("\x0").reject do |f| + f.match(%r{^(test|spec|features)/}) || + f.match?(%r{^ext/mini_sql/pg_native/}) || + f == "lib/mini_sql/pg_native.rb" || + f == "mini_sql-pg_native.gemspec" + end end # rubocop:enable Discourse/NoChdir spec.require_paths = ["lib"] diff --git a/test/mini_sql/postgres/native_materializer_test.rb b/test/mini_sql/postgres/native_materializer_test.rb new file mode 100644 index 0000000..2fa8551 --- /dev/null +++ b/test/mini_sql/postgres/native_materializer_test.rb @@ -0,0 +1,154 @@ +# frozen_string_literal: true + +$LOAD_PATH.unshift File.expand_path("../../../lib", __dir__) + +require "bigdecimal" +require "minitest/autorun" +require "pg" +require "mini_sql" +require "mini_sql/postgres/deserializer_cache" + +class MiniSqlPostgresNativeMaterializerTest < Minitest::Test + DECORATOR = Module.new do + def amount_price + amount + price + end + end + + def setup + skip "mini_sql native materializer is disabled" if ENV["MINI_SQL_PG_NATIVE"] == "0" + skip "mini_sql-pg_native is not loaded" unless defined?(MiniSql::Postgres::Native::RowMaterializer) + + @connection = pg_connection + rescue PG::Error => e + skip "PostgreSQL test connection unavailable: #{e.message}" + end + + def teardown + @connection&.raw_connection&.close + end + + def test_query_materializes_typed_values_with_native_materializer + row = @connection.query(<<~SQL).first + select + 1::int4 as one, + 9223372036854775807::int8 as big, + 12.34::numeric as amount, + '2020-02-03 04:05:06'::timestamp as happened_at, + null::text as missing + SQL + + assert_equal 1, row.one + assert_equal 9_223_372_036_854_775_807, row.big + assert_equal BigDecimal("12.34"), row.amount + assert_instance_of Time, row.happened_at + assert_nil row.missing + assert_native_materializer_for("select 1::int4 as one") + end + + def test_query_each_uses_native_materializer + rows = [] + @connection.query_each("select generate_series(1, 3)::int4 as n") { |row| rows << row } + + assert_equal [1, 2, 3], rows.map(&:n) + assert_native_materializer_for("select 1::int4 as n") + end + + def test_query_decorator_falls_back_to_decorated_row_class + decorated = @connection.query_decorator(DECORATOR, "select 2::int4 amount, 3::int4 price").first + + assert_equal 5, decorated.amount_price + assert_equal DECORATOR, decorated.class.decorator + + plain = @connection.query("select 2::int4 amount, 3::int4 price").first + assert_nil plain.class.decorator + refute_respond_to plain, :amount_price + end + + def test_unnamed_columns_and_duplicate_aliases_match_ruby_path + unnamed = @connection.query("select 1, 2 two, 3").first + assert_equal 1, unnamed.column0 + assert_equal 2, unnamed.two + assert_equal 3, unnamed.column2 + + duplicate = @connection.query("select 1::int4 a, 2::int4 a").first + assert_equal 2, duplicate.a + assert_equal({ a: 2 }, duplicate.to_h) + end + + def test_invalid_field_names_raise_like_ruby_materializer + assert_parity_with_ruby_path("select 1::int4 \"a-b\"") do + @connection.query("select 1::int4 \"a-b\"") + end + end + + def test_custom_type_map_is_used_by_native_materializer + map = PG::TypeMapByOid.new + cnn = pg_connection(type_map: map) + + assert_equal "1", cnn.query("select 1::int4 a").first.a + ensure + cnn&.raw_connection&.close + end + + def test_materializer_rejects_cleared_results + result = @connection.raw_connection.async_exec("select 1::int4 a") + result.type_map = @connection.type_map + materializer = @connection.deserializer_cache.materializer(result) + assert_instance_of MiniSql::Postgres::Native::RowMaterializer, materializer + + result.clear + + assert_raises(PG::Error) { materializer.materialize(result, 0) } + assert_raises(PG::Error) { materializer.materialize_all(result) } + end + + def test_gc_compact_keeps_cached_materializer_valid + @connection.query("select 1::int4 a") + + if GC.respond_to?(:verify_compaction_references) + GC.verify_compaction_references(double_heap: true, toward: :empty) + elsif GC.respond_to?(:compact) + GC.compact + else + skip "GC compaction is unavailable" + end + + assert_equal 2, @connection.query("select 2::int4 a").first.a + end + + private + + def pg_connection(options = {}) + args = { dbname: "test_mini_sql" } + %i[port host password user].each do |name| + if (val = ENV["MINI_SQL_PG_#{name.upcase}"]) + args[name] = val + end + end + + MiniSql::Connection.get(PG.connect(**args), options) + end + + def assert_native_materializer_for(sql) + result = @connection.raw_connection.async_exec(sql) + materializer = @connection.deserializer_cache.materializer(result) + assert_instance_of MiniSql::Postgres::Native::RowMaterializer, materializer + ensure + result&.clear + end + + def assert_parity_with_ruby_path(sql) + native_error = assert_raises(NameError, SyntaxError) { yield } + + previous = ENV["MINI_SQL_PG_NATIVE"] + ENV["MINI_SQL_PG_NATIVE"] = "0" + ruby_connection = pg_connection(deserializer_cache: MiniSql::Postgres::DeserializerCache.new) + ruby_error = assert_raises(NameError, SyntaxError) { ruby_connection.query(sql) } + + assert_equal ruby_error.class, native_error.class + ensure + ruby_connection&.raw_connection&.close + ENV["MINI_SQL_PG_NATIVE"] = previous + end +end