Adding define functions

spec/migrator/create_function_statement_spec.cr

@@ -0,0 +1,28 @@
+require "../spec_helper"
+
+describe Avram::Migrator::CreateFunctionStatement do
+  it "builds the proper SQL for creating a function" do
+    sql = <<-SQL
+    IF NEW.updated_at IS NULL OR NEW.updated_at = OLD.updated_at THEN
+      NEW.updated_at := now();
+    END IF;
+    RETURN NEW;
+    SQL
+    statement = Avram::Migrator::CreateFunctionStatement.new("set_updated_at", sql)
+
+    full_statement = statement.build
+    full_statement.should contain "CREATE OR REPLACE FUNCTION set_updated_at()"
+    full_statement.should contain "RETURNS trigger"
+    full_statement.should contain "IF NEW.updated_at IS NULL OR NEW.updated_at = OLD.updated_at THEN"
+  end
+
+  it "builds more complex functions" do
+    sql = "RETURN i + 1;"
+    statement = Avram::Migrator::CreateFunctionStatement.new("increment(i integer)", sql, returns: "integer")
+
+    full_statement = statement.build
+    full_statement.should contain "CREATE OR REPLACE FUNCTION increment(i integer)"
+    full_statement.should contain "RETURNS integer"
+    full_statement.should contain "RETURN i + 1;"
+  end
+end

spec/migrator/drop_function_statement_spec.cr

@@ -0,0 +1,8 @@
+require "../spec_helper"
+
+describe Avram::Migrator::DropFunctionStatement do
+  it "builds the proper SQL for dropping a function" do
+    statement = Avram::Migrator::DropFunctionStatement.new("set_updated_at")
+    statement.build.should eq %{DROP FUNCTION IF EXISTS "set_updated_at" CASCADE;}
+  end
+end

src/avram/migrator/create_function_statement.cr

@@ -0,0 +1,23 @@
+class Avram::Migrator::CreateFunctionStatement
+  def initialize(@name : String, @body : String, @returns : String = "trigger")
+  end
+
+  def function_name
+    if @name.ends_with?(')')
+      @name
+    else
+      "#{@name}()"
+    end
+  end
+
+  def build
+    <<-SQL
+    CREATE OR REPLACE FUNCTION #{function_name}
+      RETURNS #{@returns} AS $$
+    BEGIN
+      #{@body}
+    END
+    $$ LANGUAGE 'plpgsql';
+    SQL
+  end
+end

src/avram/migrator/drop_function_statement.cr

@@ -0,0 +1,10 @@
+class Avram::Migrator::DropFunctionStatement
+  def initialize(@name : String)
+  end
+
+  def build
+    <<-SQL
+    DROP FUNCTION IF EXISTS "#{@name}" CASCADE;
+    SQL
+  end
+end

src/avram/migrator/statement_helpers.cr

@@ -59,4 +59,12 @@ module Avram::Migrator::StatementHelpers
   def update_extension(name : String, to : String? = nil)
     prepared_statements << Avram::Migrator::AlterExtensionStatement.new(name, to: to).build
   end
+
+  def create_function(name : String, body : String, returns : String = "trigger")
+    prepared_statements << Avram::Migrator::CreateFunctionStatement.new(name, body: body, returns: returns).build
+  end
+
+  def drop_function(name : String)
+    prepared_statements << Avram::Migrator::DropFunctionStatement.new(name).build
+  end
 end