From cc4a0ff798416aaa4e639c12f01450c23ddbd15b Mon Sep 17 00:00:00 2001
From: Jeremy Evans <code@jeremyevans.net>
Date: Wed, 18 Dec 2024 10:44:47 -0800
Subject: [PATCH] Support column aliases with data types on PostgreSQL, useful
 for selecting from functions returning records

Previously, you had to use Sequel.lit to handle this type of
derived column list.  While this is useful when selecting from
any database functions that return records, in terms of working
with other parts of Sequel, this makes usage of to_recordset in
the pg_json extension simpler, as shown by the spec change.
---
 CHANGELOG                              |  2 ++
 lib/sequel/adapters/shared/postgres.rb | 19 +++++++++++++++++++
 lib/sequel/dataset/sql.rb              |  7 ++++++-
 spec/adapters/postgres_spec.rb         |  2 +-
 spec/core/mock_adapter_spec.rb         |  6 ++++++
 5 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG b/CHANGELOG
index 27e567cf1..9fa82008c 100644
--- a/CHANGELOG
+++ b/CHANGELOG
@@ -1,5 +1,7 @@
 === master
 
+* Support column aliases with data types on PostgreSQL, useful for selecting from functions returning records (jeremyevans)
+
 * Break ties in timestamp migrator version handling using lexicographic sort of rest of migration filename (jeremyevans)
 
 * Fix strict_unused_block warnings when running specs on Ruby 3.4 (jeremyevans)
diff --git a/lib/sequel/adapters/shared/postgres.rb b/lib/sequel/adapters/shared/postgres.rb
index 8ffac22c4..f7e6a20ac 100644
--- a/lib/sequel/adapters/shared/postgres.rb
+++ b/lib/sequel/adapters/shared/postgres.rb
@@ -2388,6 +2388,25 @@ def delete_using_sql(sql)
         join_from_sql(:USING, sql)
       end
 
+      # Handle column aliases containing data types, useful for selecting from functions
+      # that return the record data type.
+      def derived_column_list_sql_append(sql, column_aliases)
+        c = false
+        comma = ', '
+        column_aliases.each do |a|
+          sql << comma if c
+          if a.is_a?(Array)
+            raise Error, "column aliases specified as arrays must have only 2 elements, the first is alias name and the second is data type" unless a.length == 2
+            a, type = a
+            identifier_append(sql, a)
+            sql << " " << db.cast_type_literal(type).to_s
+          else
+            identifier_append(sql, a)
+          end
+          c ||= true
+        end
+      end
+    
       # Add ON CONFLICT clause if it should be used
       def insert_conflict_sql(sql)
         if opts = @opts[:insert_conflict]
diff --git a/lib/sequel/dataset/sql.rb b/lib/sequel/dataset/sql.rb
index 189e0eb88..b5b2e95a4 100644
--- a/lib/sequel/dataset/sql.rb
+++ b/lib/sequel/dataset/sql.rb
@@ -1032,7 +1032,7 @@ def as_sql_append(sql, aliaz, column_aliases=nil)
       if column_aliases
         raise Error, "#{db.database_type} does not support derived column lists" unless supports_derived_column_lists?
         sql << '('
-        identifier_list_append(sql, column_aliases)
+        derived_column_list_sql_append(sql, column_aliases)
         sql << ')'
       end
     end
@@ -1165,6 +1165,11 @@ def delete_from_sql(sql)
       end
     end
 
+    # Append the column aliases to the SQL.
+    def derived_column_list_sql_append(sql, column_aliases)
+      identifier_list_append(sql, column_aliases)
+    end
+    
     # Disable caching of SQL for the current dataset
     def disable_sql_caching!
       cache_set(:_no_cache_sql, true)
diff --git a/spec/adapters/postgres_spec.rb b/spec/adapters/postgres_spec.rb
index cd70bd6eb..d61620e35 100644
--- a/spec/adapters/postgres_spec.rb
+++ b/spec/adapters/postgres_spec.rb
@@ -4116,7 +4116,7 @@ def left_item_id
       @db.get(ja.typeof).must_equal 'array'
       @db.from(ja.array_elements_text.as(:v)).select_map(:v).map{|s| s.gsub(' ', '')}.must_equal ['2', '3', '["a","b"]']
       @db.from(jo.to_record.as(:v, [Sequel.lit('a integer'), Sequel.lit('b text')])).select_map(:a).must_equal [1]
-      @db.from(pg_json.call([{'a'=>1, 'b'=>1}]).op.to_recordset.as(:v, [Sequel.lit('a integer'), Sequel.lit('b integer')])).select_map(:a).must_equal [1]
+      @db.from(pg_json.call([{'a'=>1, 'b'=>1}]).op.to_recordset.as(:v, [[:a, Integer], [:b, Integer]])).select_map(:a).must_equal [1]
 
       if json_type == :jsonb
         @db.get(jo.has_key?('a')).must_equal true
diff --git a/spec/core/mock_adapter_spec.rb b/spec/core/mock_adapter_spec.rb
index 0cab029fb..4217acbfc 100644
--- a/spec/core/mock_adapter_spec.rb
+++ b/spec/core/mock_adapter_spec.rb
@@ -913,6 +913,12 @@ def @db.schema(x) [[:id, {:primary_key=>false, :auto_increment=>false}]] end
     db.immediate_constraints(:server=>:test)
     db.sqls.must_equal ['SET CONSTRAINTS ALL IMMEDIATE -- test']
   end
+
+  it "should raise if an invalid number of elements is provided when using derived column lists with an array" do
+    @db.from{c.function.as(:d, [[:a, :b], [:e, :f]])}.sql.must_equal 'SELECT * FROM c() AS "d"("a" b, "e" f)'
+    proc{@db.from{c.function.as(:d, [[:a], [:e, :f]])}.sql}.must_raise Sequel::Error
+    proc{@db.from{c.function.as(:d, [[:a, :b, :c], [:e, :f]])}.sql}.must_raise Sequel::Error
+  end
 end
 
 describe "MySQL support" do