Support partial indexes with a WHERE clause

spec/avram/migrator/create_index_statement_spec.cr

@@ -37,4 +37,21 @@ describe Avram::Migrator::CreateIndexStatement do
       statement.should eq %(CREATE INDEX custom_index_name ON users USING btree ("email");)
     end
   end
+
+  context "partial index with a WHERE clause" do
+    it "appends the predicate after the column list" do
+      statement = Avram::Migrator::CreateIndexStatement.new(:users, :email, where: "email IS NOT NULL").build
+      statement.should eq %(CREATE INDEX users_email_index ON users USING btree ("email") WHERE email IS NOT NULL;)
+    end
+
+    it "generates a conditional UNIQUE partial index" do
+      statement = Avram::Migrator::CreateIndexStatement.new(:servers, columns: [:col_a, :col_b, :col_c], unique: true, where: "col_d >= 5").build
+      statement.should eq %(CREATE UNIQUE INDEX servers_col_a_col_b_col_c_index ON servers USING btree ("col_a", "col_b", "col_c") WHERE col_d >= 5;)
+    end
+
+    it "composes with CONCURRENTLY" do
+      statement = Avram::Migrator::CreateIndexStatement.new(:users, :email, concurrently: true, where: "email IS NOT NULL").build
+      statement.should eq %(CREATE INDEX CONCURRENTLY users_email_index ON users USING btree ("email") WHERE email IS NOT NULL;)
+    end
+  end
 end

spec/avram/migrator/index_statement_helpers_spec.cr

@@ -0,0 +1,77 @@
+require "../../spec_helper"
+
+private class IndexHelperHarness
+  include Avram::Migrator::StatementHelpers
+
+  @table_name : Avram::TableName = :users
+  getter prepared_statements = [] of String
+
+  def added_indexes : Array(String)
+    index_statements
+  end
+end
+
+describe Avram::Migrator::StatementHelpers do
+  describe "#create_index with a partial-index where:" do
+    it "passes a raw String predicate straight through via where_raw" do
+      helper = IndexHelperHarness.new
+      helper.create_index(:users, :email, where_raw: "email IS NOT NULL")
+      helper.prepared_statements.last.should eq %(CREATE INDEX users_email_index ON users USING btree ("email") WHERE email IS NOT NULL;)
+    end
+
+    it "raises when given both where: and where_raw:" do
+      helper = IndexHelperHarness.new
+      expect_raises(ArgumentError) do
+        helper.create_index(:users, :name, where: UserQuery.new.age(18), where_raw: "age = 18")
+      end
+    end
+
+    it "inlines an Avram::Queryable's where conditions as literals" do
+      helper = IndexHelperHarness.new
+      helper.create_index(:users, :name, unique: true, where: UserQuery.new.age(18))
+      helper.prepared_statements.last.should eq %(CREATE UNIQUE INDEX users_name_index ON users USING btree ("name") WHERE "users"."age" = '18';)
+    end
+
+    it "single-quotes boolean values inlined from a query" do
+      helper = IndexHelperHarness.new
+      helper.create_index(:users, :name, where: UserQuery.new.available_for_hire(true))
+      helper.prepared_statements.last.should eq %(CREATE INDEX users_name_index ON users USING btree ("name") WHERE "users"."available_for_hire" = 'true';)
+    end
+
+    it "keeps OR conjunctions and drops order_by/limit from the predicate" do
+      query = UserQuery.new.age(18).or(&.name("Paul")).order_by(:name, :asc).limit(5)
+      helper = IndexHelperHarness.new
+      helper.create_index(:users, :name, where: query)
+      helper.prepared_statements.last.should eq %(CREATE INDEX users_name_index ON users USING btree ("name") WHERE "users"."age" = '18' OR "users"."name" = 'Paul';)
+    end
+
+    it "composes where: with a custom index name" do
+      helper = IndexHelperHarness.new
+      helper.create_index(:users, :email, name: "active_email_idx", where: UserQuery.new.available_for_hire(true))
+      helper.prepared_statements.last.should eq %(CREATE INDEX active_email_idx ON users USING btree ("email") WHERE "users"."available_for_hire" = 'true';)
+    end
+
+    it "raises when given a query with no where conditions" do
+      helper = IndexHelperHarness.new
+      expect_raises(Avram::InvalidQueryError) do
+        helper.create_index(:users, :name, where: UserQuery.new)
+      end
+    end
+  end
+end
+
+describe Avram::Migrator::IndexStatementHelpers do
+  describe "#add_index with a partial-index where:" do
+    it "passes a raw String predicate straight through via where_raw" do
+      helper = IndexHelperHarness.new
+      helper.add_index(:email, where_raw: "email IS NOT NULL")
+      helper.added_indexes.last.should eq %(CREATE INDEX users_email_index ON users USING btree ("email") WHERE email IS NOT NULL;)
+    end
+
+    it "inlines an Avram::Queryable's where conditions for a unique index" do
+      helper = IndexHelperHarness.new
+      helper.add_index(:age, unique: true, where: UserQuery.new.age(18))
+      helper.added_indexes.last.should eq %(CREATE UNIQUE INDEX users_age_index ON users USING btree ("age") WHERE "users"."age" = '18';)
+    end
+  end
+end

spec/avram/query_builder_spec.cr

@@ -383,6 +383,27 @@ describe Avram::QueryBuilder do
       end
     end
   end
+
+  describe "#to_prepared_where_sql" do
+    it "returns the inlined predicate without the leading WHERE" do
+      query = new_query
+        .where(Avram::Where::Equal.new(:name, "Paul"))
+        .where(Avram::Where::GreaterThan.new(:age, "18"))
+      query.to_prepared_where_sql.should eq "name = 'Paul' AND age > '18'"
+    end
+
+    it "raises when there are no where conditions" do
+      expect_raises(Avram::InvalidQueryError) do
+        new_query.to_prepared_where_sql
+      end
+    end
+
+    it "does not mutate the query's placeholder numbering" do
+      query = new_query.where(Avram::Where::Equal.new(:name, "Paul"))
+      query.to_prepared_where_sql
+      query.statement.should eq "SELECT * FROM users WHERE name = $1"
+    end
+  end
 end
 
 private def new_query

spec/avram/queryable_spec.cr

@@ -1684,6 +1684,19 @@ describe Avram::Queryable do
     end
   end
 
+  describe "#to_prepared_where_sql" do
+    it "delegates to the query builder and inlines values as table-qualified literals" do
+      query = UserQuery.new.name("Don").age.gt(30)
+      query.to_prepared_where_sql.should eq %("users"."name" = 'Don' AND "users"."age" > '30')
+    end
+
+    it "raises when the query has no where conditions" do
+      expect_raises(Avram::InvalidQueryError) do
+        UserQuery.new.to_prepared_where_sql
+      end
+    end
+  end
+
   describe "#reset_limit" do
     it "resets the limit to nil" do
       users = UserQuery.new.limit(10)

src/avram/migrator/create_index_statement.cr

@@ -36,7 +36,7 @@ class Avram::Migrator::CreateIndexStatement
     Brin
   end
 
-  def initialize(@table : TableName, @columns : Columns, using : Symbol = :btree, @unique : Bool = false, @concurrently : Bool = false, @name : String? | Symbol? = nil)
+  def initialize(@table : TableName, @columns : Columns, using : Symbol = :btree, @unique : Bool = false, @concurrently : Bool = false, @name : String? | Symbol? = nil, @where : String? = nil)
     @using = IndexTypes.parse?(using.to_s)
     raise "index type '#{using}' not supported" if @using.nil?
   end
@@ -54,7 +54,11 @@ class Avram::Migrator::CreateIndexStatement
       index << index_name
       index << " ON " << @table
       index << " USING " << @using.to_s.downcase
-      index << " (" << mapped_columns << ");"
+      index << " (" << mapped_columns << ")"
+      if where = @where
+        index << " WHERE " << where
+      end
+      index << ";"
     end
   end
 

src/avram/migrator/index_statement_helpers.cr

@@ -3,12 +3,17 @@ module Avram::Migrator::IndexStatementHelpers
 
   private getter index_statements = [] of String
 
-  # Generates raw sql for adding an index to a table column. Accepts 'unique' and 'using' options.
-  def add_index(column : Symbol, unique = false, using : Symbol = :btree)
-    index = CreateIndexStatement.new(@table_name, column, using, unique).build
+  def add_index(column : Symbol, unique = false, using : Symbol = :btree, where : Avram::Queryable? = nil, where_raw : String? = nil)
+    index = CreateIndexStatement.new(@table_name, column, using, unique, where: index_where_predicate(where, where_raw)).build
     index_statements << index unless index_added?(index, column)
   end
 
+  # `where_raw` is dropped into the index predicate unescaped — the caller owns its safety.
+  private def index_where_predicate(where : Avram::Queryable?, where_raw : String?) : String?
+    raise ArgumentError.new("Pass `where:` or `where_raw:`, not both") if where && where_raw
+    where ? where.to_prepared_where_sql : where_raw
+  end
+
   # Returns false unless matching index exists. Ignores UNIQUE
   def index_added?(index : String, column : Symbol)
     return false unless index_statements.includes?(index) || index_statements.includes?(index.gsub(" UNIQUE", ""))

src/avram/migrator/statement_helpers.cr

@@ -36,8 +36,8 @@ module Avram::Migrator::StatementHelpers
     prepared_statements << DropForeignKeyStatement.new(from, references, column).build
   end
 
-  def create_index(table_name : TableName, columns : Columns, unique = false, concurrently = false, using = :btree, name : String? | Symbol? = nil)
-    prepared_statements << CreateIndexStatement.new(table_name, columns, using, unique, concurrently, name).build
+  def create_index(table_name : TableName, columns : Columns, unique = false, concurrently = false, using = :btree, name : String? | Symbol? = nil, where : Avram::Queryable? = nil, where_raw : String? = nil)
+    prepared_statements << CreateIndexStatement.new(table_name, columns, using, unique, concurrently, name, where: index_where_predicate(where, where_raw)).build
   end
 
   def drop_index(table_name : TableName, columns : Columns? = nil, if_exists = false, on_delete = :do_nothing, name : String? | Symbol? = nil)

src/avram/query_builder.cr

@@ -26,14 +26,20 @@ class Avram::QueryBuilder
   # Prepares the SQL statement by combining the `args` and `statement`
   # in to a single `String`
   def to_prepared_sql : String
+    inline_prepared_args(statement)
+  end
+
+  # Bare WHERE predicate (no "WHERE " prefix), values inlined as literals since
+  # DDL index predicates can't use bind params. Raises if there are no conditions.
+  def to_prepared_where_sql : String
+    predicate = clone.where_predicate_sql!
+    raise Avram::InvalidQueryError.new("Cannot build a partial index predicate: the query has no `where` conditions") unless predicate
+    inline_prepared_args(predicate)
+  end
+
+  private def inline_prepared_args(sql : String) : String
     params = args.map { |arg| "'#{String.new(PQ::Param.encode(arg).slice)}'" }
-    i = 0
-    sql = statement
-    sql.scan(/\$\d+/) do |match|
-      sql = sql.sub(match[0], params[i])
-      i += 1
-    end
-    sql
+    sql.gsub(/\$\d+/) { |placeholder| params[placeholder.lchop('$').to_i - 1] }
   end
 
   # Merges the wheres, joins, and orders from the passed in query
@@ -362,21 +368,26 @@ class Avram::QueryBuilder
   end
 
   private def wheres_sql : String?
-    if !wheres.empty?
-      statements = wheres.flat_map do |sql_clause|
-        clause = sql_clause.prepare(->next_prepared_statement_placeholder)
-
-        [clause, sql_clause.conjunction.to_s]
-      end
+    predicate = where_predicate_sql!
+    "WHERE #{predicate}" if predicate
+  end
 
-      # Remove blank conjunctions
-      statements.reject!(&.blank?)
+  # Mutates the placeholder counter (like `statement!`) — call it on a clone.
+  protected def where_predicate_sql! : String?
+    return if wheres.empty?
 
-      # Remove the last floating conjunction
-      statements.pop
+    statements = wheres.flat_map do |sql_clause|
+      clause = sql_clause.prepare(->next_prepared_statement_placeholder)
 
-      "WHERE #{statements.join(" ")}"
+      [clause, sql_clause.conjunction.to_s]
     end
+
+    statements.reject!(&.blank?)
+
+    # Last element is a dangling conjunction the flat_map appends after each clause
+    statements.pop
+
+    statements.join(" ")
   end
 
   def wheres : Array(Avram::Where::Condition)

src/avram/queryable.cr

@@ -344,4 +344,8 @@ module Avram::Queryable(T)
   def to_prepared_sql
     query.to_prepared_sql
   end
+
+  def to_prepared_where_sql
+    query.to_prepared_where_sql
+  end
 end