Auto-detect function argument types when dropping or updating

CHANGELOG.md

@@ -12,6 +12,12 @@ changelog, see the [commits] for each version via the version links.
 - Add `Function#signature` for PostgreSQL function identity (#207)
 - Add PostgreSQL versioning policy, officially supporting PostgreSQL 14-18 (#194)
 - Refactor Statements module to use explicit keyword arguments instead of `**options` hash (#186)
+- Auto-detect function argument types from `pg_proc` when dropping or updating
+  functions, fixing support for functions with parameters (#7)
+- Add `arguments:` option to `drop_function` and `update_function` for
+  targeting specific overloads of functions that share a name
+- Custom adapters that override `drop_function` or `update_function` must
+  now accept an `arguments:` keyword argument
 - Internal refactorings / improvements
   - Add scheduled EOL check for Ruby, Rails, and PostgreSQL (#205)
   - Add GitHub release creation to release task (#209)

lib/fx.rb

@@ -14,6 +14,12 @@
 # F(x) adds methods `ActiveRecord::Migration` to create and manage database
 # triggers and functions in Rails applications.
 module Fx
+  class Error < StandardError; end
+
+  # Raised when dropping an overloaded function without specifying which
+  # overload to target. Pass the `arguments:` option to disambiguate.
+  class AmbiguousFunctionError < Error; end
+
   # Hooks Fx into Rails.
   #
   # Enables fx migration methods, migration reversability, and `schema.rb`

lib/fx/adapters/postgres.rb

@@ -91,10 +91,13 @@ def create_trigger(sql_definition)
       #
       # @param name [String, Symbol] The name of the function.
       # @param sql_definition [String] The SQL schema for the function.
+      # @param arguments [String] Optional function argument types for
+      #   identifying overloaded functions (e.g. "integer, text"). This
+      #   option is specific to the Postgres adapter.
       #
       # @return [void]
-      def update_function(name, sql_definition)
-        drop_function(name)
+      def update_function(name, sql_definition, arguments: nil)
+        drop_function(name, arguments: arguments)
         create_function(sql_definition)
       end
 
@@ -121,11 +124,28 @@ def update_trigger(name, on:, sql_definition:)
       # This is typically called in a migration via
       # {Fx::Statements::Function#drop_function}.
       #
-      # @param name [String, Symbol] The name of the function to drop
+      # @param name [String, Symbol] The name of the function to drop.
+      # @param arguments [String] Optional function argument types for
+      #   identifying overloaded functions (e.g. "integer, text"). When not
+      #   provided, the argument types are looked up automatically from
+      #   pg_proc. If multiple overloads exist, an {Fx::AmbiguousFunctionError}
+      #   is raised. This option is specific to the Postgres adapter; custom
+      #   adapters that do not accept it will raise an ArgumentError.
       #
       # @return [void]
-      def drop_function(name)
-        execute("DROP FUNCTION #{name};")
+      def drop_function(name, arguments: nil)
+        function =
+          if arguments
+            Fx::Function.new(
+              "name" => name.to_s,
+              "definition" => "",
+              "arguments" => arguments
+            )
+          else
+            find_function(name)
+          end
+
+        execute("DROP FUNCTION #{function.signature};")
       end
 
       # Drops the trigger from the database
@@ -147,6 +167,24 @@ def drop_trigger(name, on:)
 
       delegate :execute, to: :connection
 
+      def find_function(name)
+        name_str = name.to_s
+        matches = functions.select { |function| function.name == name_str }
+
+        case matches.size
+        when 0
+          Fx::Function.new("name" => name_str, "definition" => "")
+        when 1
+          matches.first
+        else
+          signatures = matches.map(&:signature)
+          raise Fx::AmbiguousFunctionError, <<~MSG.chomp
+            Multiple definitions for function "#{name_str}": #{signatures.join(", ")}.
+            Specify which to drop: drop_function :#{name_str}, arguments: "<argument types>"
+          MSG
+        end
+      end
+
       def connection
         Fx::Adapters::Postgres::Connection.new(connectable.connection)
       end

lib/fx/statements.rb

@@ -10,6 +10,10 @@ module Statements
     # @param sql_definition [String] The SQL query for the function schema.
     #   If both `sql_definition` and `version` are provided,
     #   `sql_definition` takes precedence.
+    # @param arguments [String] Function argument types (e.g. "integer, text").
+    #   Not used during creation itself, but preserved by the command recorder
+    #   so that a rollback of this migration can pass the correct signature to
+    #   {#drop_function}. Only needed for overloaded functions.
     # @return [void] The database response from executing the create statement.
     #
     # @example Create from `db/functions/uppercase_users_name_v02.sql`
@@ -26,7 +30,7 @@ module Statements
     #     $$ LANGUAGE plpgsql;
     #   SQL
     #
-    def create_function(name, version: 1, sql_definition: nil, revert_to_version: nil)
+    def create_function(name, version: 1, sql_definition: nil, revert_to_version: nil, arguments: nil)
       validate_version_or_sql_definition_present!(version, sql_definition)
       sql_definition = resolve_sql_definition(sql_definition, name, version, :function)
 
@@ -39,13 +43,17 @@ def create_function(name, version: 1, sql_definition: nil, revert_to_version: ni
     # @param revert_to_version [Integer] Used to reverse the `drop_function`
     #   command on `rake db:rollback`. The provided version will be passed as
     #   the `version` argument to {#create_function}.
+    # @param arguments [String] Function argument types for identifying
+    #   overloaded functions (e.g. "integer, text"). When omitted, the
+    #   Postgres adapter auto-detects the signature from the database.
+    #   Custom adapters must accept this keyword to use it.
     # @return [void] The database response from executing the drop statement.
     #
     # @example Drop a function, rolling back to version 2 on rollback
     #   drop_function(:uppercase_users_name, revert_to_version: 2)
     #
-    def drop_function(name, revert_to_version: nil)
-      Fx.database.drop_function(name)
+    def drop_function(name, revert_to_version: nil, arguments: nil)
+      Fx.database.drop_function(name, **{arguments: arguments}.compact)
     end
 
     # Update a database function.
@@ -57,6 +65,10 @@ def drop_function(name, revert_to_version: nil)
     # @param sql_definition [String] The SQL query for the function schema.
     #   If both `sql_definition` and `version` are provided,
     #   `sql_definition` takes precedence.
+    # @param arguments [String] Function argument types for identifying
+    #   overloaded functions (e.g. "integer, text"). When omitted, the
+    #   Postgres adapter auto-detects the signature from the database.
+    #   Custom adapters must accept this keyword to use it.
     # @return [void] The database response from executing the create statement.
     #
     # @example Update function to a given version
@@ -77,12 +89,12 @@ def drop_function(name, revert_to_version: nil)
     #     $$ LANGUAGE plpgsql;
     #   SQL
     #
-    def update_function(name, version: nil, sql_definition: nil, revert_to_version: nil)
+    def update_function(name, version: nil, sql_definition: nil, revert_to_version: nil, arguments: nil)
       validate_version_or_sql_definition_present!(version, sql_definition)
 
       sql_definition = resolve_sql_definition(sql_definition, name, version, :function)
 
-      Fx.database.update_function(name, sql_definition)
+      Fx.database.update_function(name, sql_definition, **{arguments: arguments}.compact)
     end
 
     # Create a new database trigger.

spec/acceptance/user_manages_functions_spec.rb

@@ -54,4 +54,93 @@
     successfully "rails destroy fx:function adder"
     successfully "rake db:migrate"
   end
+
+  it "handles updating functions with arguments" do
+    successfully "rails generate fx:function multiply"
+    write_function_definition "multiply_v01", <<~SQL
+      CREATE FUNCTION multiply(x int, y int)
+      RETURNS int AS $$
+      BEGIN
+          RETURN x * y;
+      END;
+      $$ LANGUAGE plpgsql;
+    SQL
+    successfully "rake db:migrate"
+
+    result = execute("SELECT * FROM multiply(3, 4) AS result")
+    result["result"] = result["result"].to_i
+    expect(result).to eq("result" => 12)
+
+    successfully "rails generate fx:function multiply"
+    write_function_definition "multiply_v02", <<~SQL
+      CREATE FUNCTION multiply(x int, y int)
+      RETURNS int AS $$
+      BEGIN
+          RETURN x * y * 2;
+      END;
+      $$ LANGUAGE plpgsql;
+    SQL
+    successfully "rake db:migrate"
+
+    result = execute("SELECT * FROM multiply(3, 4) AS result")
+    result["result"] = result["result"].to_i
+    expect(result).to eq("result" => 24)
+
+    successfully "rake db:rollback"
+
+    result = execute("SELECT * FROM multiply(3, 4) AS result")
+    result["result"] = result["result"].to_i
+    expect(result).to eq("result" => 12)
+
+    successfully "rake db:rollback"
+
+    expect { execute("SELECT * FROM multiply(3, 4) AS result") }
+      .to raise_error(ActiveRecord::StatementInvalid)
+  end
+
+  it "handles dropping overloaded functions with explicit arguments" do
+    successfully "rails generate fx:function inc"
+    write_function_definition "inc_v01", <<~SQL
+      CREATE FUNCTION inc(x int)
+      RETURNS int AS $$
+      BEGIN RETURN x + 1; END;
+      $$ LANGUAGE plpgsql;
+    SQL
+    successfully "rake db:migrate"
+
+    execute <<~SQL
+      CREATE FUNCTION inc(x int, step int)
+      RETURNS int AS $$ BEGIN RETURN x + step; END; $$ LANGUAGE plpgsql;
+    SQL
+
+    result = execute("SELECT inc(5) AS result")
+    result["result"] = result["result"].to_i
+    expect(result).to eq("result" => 6)
+
+    result = execute("SELECT inc(5, 10) AS result")
+    result["result"] = result["result"].to_i
+    expect(result).to eq("result" => 15)
+
+    write_migration "drop_inc_one_arg", <<~RUBY
+      class DropIncOneArg < ActiveRecord::Migration[#{ActiveRecord::Migration.current_version}]
+        def change
+          drop_function :inc, arguments: "integer", revert_to_version: 1
+        end
+      end
+    RUBY
+    successfully "rake db:migrate"
+
+    expect { execute("SELECT inc(5) AS result") }
+      .to raise_error(ActiveRecord::StatementInvalid)
+
+    result = execute("SELECT inc(5, 10) AS result")
+    result["result"] = result["result"].to_i
+    expect(result).to eq("result" => 15)
+
+    successfully "rake db:rollback"
+
+    result = execute("SELECT inc(5) AS result")
+    result["result"] = result["result"].to_i
+    expect(result).to eq("result" => 6)
+  end
 end

spec/acceptance_helper.rb

@@ -61,6 +61,12 @@ def verify_identical_definitions(def_a, def_b)
     successfully "cmp #{def_a} #{def_b}"
   end
 
+  def write_migration(name, contents)
+    Dir.mkdir("db/migrate") unless Dir.exist?("db/migrate")
+    timestamp = Time.now.utc.strftime("%Y%m%d%H%M%S")
+    File.write("db/migrate/#{timestamp}_#{name}.rb", contents)
+  end
+
   def execute(command)
     ActiveRecord::Base.connection.execute(command).first
   end

spec/fx/adapters/postgres_spec.rb

@@ -53,7 +53,7 @@
 
   describe "#drop_function" do
     context "when the function has arguments" do
-      it "successfully drops a function with the entire function signature" do
+      it "successfully drops a function by looking up its signature" do
         adapter = Fx::Adapters::Postgres.new
         adapter.create_function(
           <<~SQL
@@ -91,6 +91,56 @@
         expect(adapter.functions.map(&:name)).not_to include("test")
       end
     end
+
+    context "when the function is overloaded" do
+      it "raises AmbiguousFunctionError" do
+        adapter = Fx::Adapters::Postgres.new
+        adapter.create_function(
+          <<~SQL
+            CREATE FUNCTION foo(x int)
+            RETURNS int AS $$
+            BEGIN RETURN x; END;
+            $$ LANGUAGE plpgsql;
+          SQL
+        )
+        adapter.create_function(
+          <<~SQL
+            CREATE FUNCTION foo(x int, y int)
+            RETURNS int AS $$
+            BEGIN RETURN x + y; END;
+            $$ LANGUAGE plpgsql;
+          SQL
+        )
+
+        expect { adapter.drop_function(:foo) }
+          .to raise_error(Fx::AmbiguousFunctionError, /Multiple definitions/)
+      end
+
+      it "drops the correct overload when arguments are specified" do
+        adapter = Fx::Adapters::Postgres.new
+        adapter.create_function(
+          <<~SQL
+            CREATE FUNCTION foo(x int)
+            RETURNS int AS $$
+            BEGIN RETURN x; END;
+            $$ LANGUAGE plpgsql;
+          SQL
+        )
+        adapter.create_function(
+          <<~SQL
+            CREATE FUNCTION foo(x int, y int)
+            RETURNS int AS $$
+            BEGIN RETURN x + y; END;
+            $$ LANGUAGE plpgsql;
+          SQL
+        )
+
+        adapter.drop_function(:foo, arguments: "int, int")
+
+        remaining = adapter.functions.select { |f| f.name == "foo" }
+        expect(remaining.length).to eq(1)
+      end
+    end
   end
 
   describe "#functions" do

spec/fx/command_recorder_spec.rb

@@ -17,6 +17,18 @@
 
       expect(recorder.commands).to eq([[:drop_function, [:test]]])
     end
+
+    it "reverts to drop_function preserving arguments" do
+      recorder = ActiveRecord::Migration::CommandRecorder.new
+
+      recorder.revert do
+        recorder.create_function :test, arguments: "integer, text"
+      end
+
+      expect(recorder.commands).to eq(
+        [[:drop_function, [:test, {arguments: "integer, text"}]]]
+      )
+    end
   end
 
   describe "#drop_function" do
@@ -38,6 +50,16 @@
       expect(recorder.commands).to eq([[:create_function, revert_args]])
     end
 
+    it "reverts to create_function preserving arguments" do
+      recorder = ActiveRecord::Migration::CommandRecorder.new
+      args = [:test, {revert_to_version: 3, arguments: "integer"}]
+      revert_args = [:test, {arguments: "integer", version: 3}]
+
+      recorder.revert { recorder.drop_function(*args) }
+
+      expect(recorder.commands).to eq([[:create_function, revert_args]])
+    end
+
     it "raises when reverting without revert_to_version set" do
       recorder = ActiveRecord::Migration::CommandRecorder.new
       args = [:test, {another_argument: 1}]
@@ -68,6 +90,16 @@
       expect(recorder.commands).to eq([[:update_function, revert_args]])
     end
 
+    it "reverts to update_function preserving arguments" do
+      recorder = ActiveRecord::Migration::CommandRecorder.new
+      args = [:test, {version: 2, revert_to_version: 1, arguments: "integer"}]
+      revert_args = [:test, {arguments: "integer", version: 1}]
+
+      recorder.revert { recorder.update_function(*args) }
+
+      expect(recorder.commands).to eq([[:update_function, revert_args]])
+    end
+
     it "raises when reverting without revert_to_version set" do
       recorder = ActiveRecord::Migration::CommandRecorder.new
       args = [:test, {version: 42, another_argument: 1}]