litestack 0.2.6 → 0.4.1

data/lib/litestack/litesearch/index.rb

@@ -0,0 +1,230 @@
+require 'oj'
+require_relative './schema.rb'
+
+class Litesearch::Index
+  
+  DEFAULT_SEARCH_OPTIONS = {limit: 25, offset: 0}
+
+  def initialize(db, name)
+    @db = db # this index instance will always belong to this db instance
+    @stmts = {}
+    name = name.to_s.downcase.to_sym
+    # if in the db then put in cache and return if no schema is given
+    # if a schema is given then compare the new and the existing schema
+    # if they are the same put in cache and return
+    # if they differ only in weights then set the new weights, update the schema, put in cache and return
+    # if they differ in fields (added/removed/renamed) then update the structure, then rebuild if auto-rebuild is on
+    # if they differ in tokenizer then rebuild if auto-rebuild is on (error otherwise)
+    # if they differ in both then update the structure and rebuild if auto-rebuild is on (error otherwise)
+    load_index(name) if exists?(name)
+    
+    if block_given?
+      schema = Litesearch::Schema.new
+      schema.schema[:name] = name
+      yield schema
+      schema.post_init
+      # now that we have a schema object we need to check if we need to create or modify and existing index
+      @db.transaction(:immediate) do
+        if exists?(name)
+          load_index(name)
+          do_modify(schema)
+        else
+          do_create(schema)
+        end
+        prepare_statements
+      end
+    else
+      if exists?(name)
+        # an index already exists, load it from the database and return the index instance to the caller
+        load_index(name)
+        prepare_statements
+      else
+        raise "index does not exist and no schema was supplied"
+      end
+    end
+  end
+  
+  def load_index(name)
+    # we cannot use get_config_value here since the schema object is not created yet, should we allow something here?
+    @schema = Litesearch::Schema.new(Oj.load(@db.get_first_value("SELECT v from #{name}_config where k = ?", :litesearch_schema.to_s))) rescue nil
+    raise "index configuration not found, either corrupted or not a litesearch index!" if @schema.nil?
+    self
+  end
+          
+  def modify
+    schema = Litesearch::Schema.new
+    yield schema
+    schema.schema[:name] = @schema.schema[:name]
+    do_modify(schema)
+  end
+  
+  def rebuild!
+    @db.transaction(:immediate) do
+      do_rebuild
+    end
+  end
+  
+  def add(document)
+    @stmts[:insert].execute!(document)
+    return @db.last_insert_row_id
+  end
+  
+  def remove(id)
+    @stmts[:delete].execute!(id)
+  end
+  
+  def count(term = nil)
+    if term
+      @stmts[:count].execute!(term)[0][0]
+    else
+      @stmts[:count_all].execute!()[0][0]
+    end
+  end
+  
+  # search options include
+  # limit: how many records to return
+  # offset: start from which record
+  def search(term, options = {})
+    result = []
+    options = DEFAULT_SEARCH_OPTIONS.merge(options)
+    rs = @stmts[:search].execute(term, options[:limit], options[:offset])
+    if @db.results_as_hash
+      rs.each_hash do |hash|
+        result << hash 
+      end
+    else
+      result = rs.to_a
+    end    
+    result
+  end
+      
+  def clear!
+    @stmts[:delete_all].execute!(id)
+  end
+  
+  def drop!
+    if @schema.get(:type) == :backed
+      @db.execute_batch(@schema.sql_for(:drop_primary_triggers))
+      if secondary_triggers_sql = @schema.sql_for(:create_secondary_triggers)
+        @db.execute_batch(@schema.sql_for(:drop_secondary_triggers))
+      end
+    end
+    @db.execute(@schema.sql_for(:drop))
+  end
+  
+  
+  private
+  
+  def exists?(name)
+    @db.get_first_value("SELECT count(*) FROM SQLITE_MASTER WHERE name = ? AND type = 'table' AND (sql like '%fts5%' OR sql like '%FTS5%')", name.to_s) == 1
+  end
+
+  def prepare_statements
+    stmt_names = [:insert, :delete, :delete_all, :drop, :count, :count_all, :search]
+    stmt_names.each do |stmt_name|
+      @stmts[stmt_name] = @db.prepare(@schema.sql_for(stmt_name))
+    end      
+  end
+  
+  def do_create(schema)   
+    @schema = schema   
+    @schema.clean
+    # create index
+    @db.execute(schema.sql_for(:create_index, true))
+    # adjust ranking function
+    @db.execute(schema.sql_for(:ranks, true))
+    # create triggers (if any)
+    if @schema.get(:type) == :backed
+      @db.execute_batch(@schema.sql_for(:create_primary_triggers))
+      if secondary_triggers_sql = @schema.sql_for(:create_secondary_triggers)
+        @db.execute_batch(secondary_triggers_sql)
+      end
+      @db.execute(@schema.sql_for(:rebuild)) if @schema.get(:rebuild_on_create)
+    end
+    set_config_value(:litesearch_schema, @schema.schema)
+  end
+
+  def do_modify(new_schema)
+    changes = @schema.compare(new_schema)
+    # ensure the new schema maintains feild order
+    new_schema.order_fields(@schema)
+    # with the changes object decide what needs to be done to the schema
+    requires_schema_change = false
+    requires_trigger_change = false
+    requires_rebuild = false
+    if changes[:fields] || changes[:table] || changes[:tokenizer] || changes[:filter_column] || changes[:removed_fields_count] > 0# any change here will require a schema change
+      requires_schema_change = true
+      # only a change in tokenizer
+      requires_rebuild = changes[:tokenizer] || new_schema.get(:rebuild_on_modify)
+      requires_trigger_change = (changes[:table] || changes[:fields] || changes[:filter_column]) && @schema.get(:type) == :backed 
+    end
+    if requires_schema_change
+      # 1. enable schema editing
+      @db.execute("PRAGMA WRITABLE_SCHEMA = TRUE")
+      # 2. update the index sql
+      @db.execute(new_schema.sql_for(:update_index), new_schema.sql_for(:create_index))
+      # 3. update the content table sql (if it exists)
+      @db.execute(new_schema.sql_for(:update_content_table), new_schema.sql_for(:create_content_table, new_schema.schema[:fields].count))
+      # adjust shadow tables
+      @db.execute(new_schema.sql_for(:expand_data), changes[:extra_fields_count])
+      @db.execute(new_schema.sql_for(:expand_docsize), changes[:extra_fields_count])
+      @db.execute("PRAGMA WRITABLE_SCHEMA = RESET")
+      # need to reprepare statements        
+    end
+    if requires_trigger_change
+      @db.execute_batch(new_schema.sql_for(:drop_primary_triggers))
+      @db.execute_batch(new_schema.sql_for(:create_primary_triggers))
+      if secondary_triggers_sql = new_schema.sql_for(:create_secondary_triggers)
+        @db.execute_batch(new_schema.sql_for(:drop_secondary_triggers))
+        @db.execute_batch(secondary_triggers_sql)
+      end
+    end
+    if changes[:fields] || changes[:table] || changes[:tokenizer] || changes[:weights] || changes[:filter_column]
+      @schema = new_schema
+      set_config_value(:litesearch_schema, @schema.schema)
+      prepare_statements
+      #save_schema
+    end
+    do_rebuild if requires_rebuild
+    # update the weights if they changed 
+    @db.execute(@schema.sql_for(:ranks)) if changes[:weights]
+  end
+
+  def do_rebuild
+    # remove any zero weight columns
+    if @schema.get(:type) == :backed
+      @db.execute_batch(@schema.sql_for(:drop_primary_triggers))
+      if secondary_triggers_sql = @schema.sql_for(:create_secondary_triggers)
+        @db.execute_batch(@schema.sql_for(:drop_secondary_triggers))
+      end
+      @db.execute(@schema.sql_for(:drop))
+      @db.execute(@schema.sql_for(:create_index, true))
+      @db.execute_batch(@schema.sql_for(:create_primary_triggers))
+      @db.execute_batch(secondary_triggers_sql) if secondary_triggers_sql
+      @db.execute(@schema.sql_for(:rebuild)) 
+    elsif @schema.get(:type) == :standalone
+      removables = []
+      @schema.get(:fields).each_with_index{|f, i| removables << [f[0], i] if f[1][:weight] == 0  }
+      removables.each do |col|
+        @db.execute(@schema.sql_for(:drop_content_col, col[1]))
+        @schema.get(:fields).delete(col[0]) 
+      end
+      @db.execute("PRAGMA WRITABLE_SCHEMA = TRUE")
+      @db.execute(@schema.sql_for(:update_index), @schema.sql_for(:create_index, true))
+      @db.execute(@schema.sql_for(:update_content_table), @schema.sql_for(:create_content_table, @schema.schema[:fields].count)) 
+      @db.execute("PRAGMA WRITABLE_SCHEMA = RESET")
+      @db.execute(@schema.sql_for(:rebuild)) 
+    end
+    set_config_value(:litesearch_schema, @schema.schema)
+    @db.execute(@schema.sql_for(:ranks, true)) 
+ end
+  
+  def get_config_value(key)
+    Oj.load(@db.get_first_value(@schema.sql_for(:get_config_value), key.to_s)) #rescue nil
+  end
+  
+  def set_config_value(key, value)
+    @db.execute(@schema.sql_for(:set_config_value), key.to_s, Oj.dump(value))
+  end
+
+end

data/lib/litestack/litesearch/model.rb

@@ -0,0 +1,178 @@
+module Litesearch::Model
+                      
+  def self.included(klass)
+    klass.include InstanceMethods
+    klass.extend ClassMethods 
+    klass.attribute :search_rank, :float if klass.respond_to? :attribute
+    if defined?(Sequel::Model) != nil && klass.ancestors.include?(Sequel::Model)
+       klass.include Litesearch::Model::SequelInstanceMethods
+       klass.extend Litesearch::Model::SequelClassMethods
+       Sequel::Model.extend Litesearch::Model::BaseClassMethods
+    elsif defined?(ActiveRecord::Base) != nil && klass.ancestors.include?(ActiveRecord::Base)
+       klass.include Litesearch::Model::ActiveRecordInstanceMethods
+       klass.extend Litesearch::Model::ActiveRecordClassMethods
+       ActiveRecord::Base.extend Litesearch::Model::BaseClassMethods
+    end
+  end
+  
+  module BaseClassMethods
+    def search_models
+      @@models ||= {}
+    end 
+  end
+  
+  module InstanceMethods
+  
+  end
+  
+  module ClassMethods  
+     
+    def litesearch
+      idx = get_connection.search_index(index_name) do |schema|
+        schema.type :backed
+        schema.table table_name.to_sym
+        yield schema
+        schema.post_init
+        @schema = schema #save the schema   
+      end
+      if defined?(Sequel::Model) != nil && self.ancestors.include?(Sequel::Model)
+        Sequel::Model.search_models[self.name] = self
+      elsif defined?(ActiveRecord::Base) != nil && self.ancestors.include?(ActiveRecord::Base)
+        ActiveRecord::Base.search_models[self.name] = self
+      end
+      idx
+    end
+    
+    def rebuild_index!
+      get_connection.search_index(index_name).rebuild!
+    end
+    
+    def drop_index!
+      get_connection.search_index(index_name).drop!
+    end
+    
+    def search_all(term, options={})
+      options[:offset] ||= 0
+      options[:limit] ||= 25
+      selects = [] 
+      if models = options[:models]
+        models_hash = {}
+        models.each do |model|
+          models_hash[model.name] = model
+        end
+      else
+        models_hash = search_models
+      end
+      models_hash.each do |name, klass|
+        selects << "SELECT '#{name}' AS model, rowid, -rank AS search_rank FROM #{index_name_for_table(klass.table_name)}('#{term}')"
+      end
+      conn = get_connection
+      sql = selects.join(" UNION ") << " ORDER BY search_rank DESC LIMIT #{options[:limit]} OFFSET #{options[:offset]}"
+      result = []
+      rs = conn.query(sql) #, options[:limit], options[:offset])
+      rs.each_hash do |row|
+        obj = models_hash[row["model"]].fetch_row(row["rowid"])
+        obj.search_rank = row["search_rank"]
+        result << obj
+      end
+      rs.close
+      result
+    end
+          
+    # AR specific
+                     
+    private
+    
+    def index_name
+      "#{table_name}_search_idx"
+    end
+    
+    def index_name_for_table(table)
+      "#{table}_search_idx"
+    end
+    
+    # create a new instance of self with the row as an argument
+    def create_instance(row)
+      self.new(row)
+    end
+    
+          
+  end
+  
+  module ActiveRecordInstanceMethods;end
+    
+  module ActiveRecordClassMethods
+    
+    def get_connection
+      connection.raw_connection
+    end
+    
+    def fetch_row(id)
+      find(id)
+    end
+    
+    def search(term)
+      self.select(
+        "#{table_name}.*"
+      ).joins(
+        "INNER JOIN #{index_name} ON #{table_name}.id = #{index_name}.rowid AND rank != 0 AND #{index_name} MATCH ", Arel.sql("'#{term}'")
+      ).select(
+        "-#{index_name}.rank AS search_rank"
+      ).order(
+        Arel.sql("#{index_name}.rank")
+      )        
+    end
+
+    private
+    
+    def create_instance(row)
+      instantiate(row)
+    end
+  end
+  
+  module SequelInstanceMethods
+  
+    def search_rank
+      @values[:search_rank]
+    end
+    
+    def search_rank=(rank)
+      @values[:search_rank] = rank
+    end
+        
+  end
+  
+  module SequelClassMethods
+        
+    def fetch_row(id)
+      self[id]
+    end
+
+    def get_connection
+      db.instance_variable_get(:@raw_db)
+    end
+
+    def search(term)
+      dataset.select(
+        Sequel.lit("#{table_name}.*, -#{index_name}.rank AS search_rank")
+      ).inner_join(
+        Sequel.lit("#{index_name}('#{term}') ON #{table_name}.id = #{index_name}.rowid AND rank != 0")
+      ).order(
+        Sequel.lit('rank')
+      )
+    end
+
+    private
+    
+    def create_instance(row)
+      # we need to convert keys to symbols first!
+      row.keys.each do |k|
+        next if k.is_a? Symbol 
+        row[k.to_sym] = row[k] 
+        row.delete(k) 
+      end
+      self.call(row)
+    end
+  end
+      
+end

data/lib/litestack/litesearch/schema.rb

@@ -0,0 +1,193 @@
+require_relative './schema_adapters.rb'
+
+class Litesearch::Schema
+
+  TOKENIZERS = { 
+    porter: 'porter unicode61 remove_diacritics 2', 
+    unicode: 'unicode61 remove_diacritics 2', 
+    ascii: 'ascii', 
+    trigram: 'trigram' 
+  }
+  
+  INDEX_TYPES = {
+    standalone: Litesearch::Schema::StandaloneAdapter, 
+    contentless: Litesearch::Schema::ContentlessAdapter, 
+    backed: Litesearch::Schema::BackedAdapter
+  }
+
+  DEFAULT_SCHEMA = {
+    name: nil,
+    type: :standalone,
+    fields: nil, 
+    table: nil,
+    filter_column: nil,
+    tokenizer: :porter,
+    auto_create: true,
+    auto_modify: true,
+    rebuild_on_create: false,      
+    rebuild_on_modify: false      
+  }
+
+  attr_accessor :schema
+
+  def initialize(schema = {})
+    @schema = schema #DEFAULT_SCHEMA.merge(schema)
+    @schema[:fields] = {} unless @schema[:fields]
+  end
+
+  # schema definition API
+  def name(new_name)
+    @schema[:name] = new_name
+  end
+    
+  def type(new_type)
+    raise "Unknown index type" if INDEX_TYPES[new_type].nil?
+    @schema[:type] = new_type
+  end
+
+  def table(table_name)
+    @schema[:table] = table_name
+  end
+
+  def fields(field_names)
+    field_names.each {|f| field f }
+  end
+  
+  def field(name, attributes = {})
+    name = name.to_s.downcase.to_sym
+    attributes = {weight: 1}.merge(attributes).select{|k, v| allowed_attributes.include?(k)} # only allow attributes we know, to ease schema comparison later
+    @schema[:fields][name] = attributes
+  end
+  
+  def tokenizer(new_tokenizer)
+    raise "Unknown tokenizer" if TOKENIZERS[new_tokenizer].nil?
+    @schema[:tokenizer] = new_tokenizer    
+  end
+  
+  def filter_column(filter_column)
+    @schema[:filter_column] = filter_column
+  end
+  
+  def auto_create(boolean)
+    @schema[:auto_create] = boolean
+  end
+  
+  def auto_modify(boolean)
+    @schema[:auto_modify] = boolean
+  end
+  
+  def rebuild_on_create(boolean)
+    @schema[:rebuild_on_create] = boolean  
+  end
+   
+  def rebuild_on_modify(boolean)
+    @schema[:rebuild_on_modify] = boolean  
+  end
+
+  def post_init
+    @schema = DEFAULT_SCHEMA.merge(@schema)
+  end
+
+  # schema sql generation API
+  
+  def sql_for(method, *args)
+    adapter.sql_for(method, *args)
+  end
+    
+  # schema data structure API
+  def get(key)
+    @schema[key]
+  end  
+  
+  def get_field(name)
+    @schema[:fields][name]
+  end
+  
+  def adapter
+    @adapter ||= INDEX_TYPES[@schema[:type]].new(@schema)
+  end
+  
+  def reset_sql
+    adapter.generate_sql
+  end
+  
+  def order_fields(old_schema)
+    adapter.order_fields(old_schema)
+  end
+  
+  # should we do this at the schema objects level?
+  def compare(other_schema)
+    other_schema = other_schema.schema
+    # are the schemas identical?
+    # 1 - same fields?
+    [:type, :tokenizer, :name, :table].each do |key|
+       other_schema[key] = @schema[key] if other_schema[key].nil?  
+    end
+    if @schema[:type] != other_schema[:type]
+      raise Litesearch::SchemaChangeException.new "Cannot change the index type, please drop the index before creating it again with the new type"
+    end
+    changes = { tokenizer: @schema[:tokenizer] != other_schema[:tokenizer], table: @schema[:table] != other_schema[:table], removed_fields_count: 0, filter_column: @schema[:filter_column] != other_schema[:filter_column]  }
+    #check tokenizer changes
+    if changes[:tokenizer] && !other_schema[:rebuild_on_modify]
+      raise Litesearch::SchemaChangeException.new "Cannot change the tokenizer without an index rebuild!"
+    end
+
+    
+
+    # check field changes
+    keys = @schema[:fields].keys.sort
+    other_keys = other_schema[:fields].keys.sort
+
+    extra_keys = other_keys - keys
+    extra_keys.each do |key|
+      if other_schema[:fields][key][:weight] == 0
+        other_schema[:fields].delete(key)
+      end
+    end
+    
+    other_keys = other_schema[:fields].keys.sort
+
+    changes[:fields] = keys != other_keys # only acceptable change is adding extra fields
+    changes[:extra_fields_count] = other_keys.count - keys.count 
+    # check for missing fields (please note that adding fields can work without a rebuild)
+    if keys - other_keys != []
+      raise Litesearch::SchemaChangeException.new "Missing fields from existing schema, they have to exist with weight zero until the next rebuild!" 
+    end
+    
+    # check field weights
+    weights = keys.collect{|key| @schema[:fields][key][:weight] }
+    other_weights = other_keys.collect{|key| other_schema[:fields][key][:weight] }
+    changes[:weights] = weights != other_weights # will always be true if fields are added
+    if (removed_count = other_weights.select{|w| w == 0}.count) > 0
+      changes[:removed_fields_count] = removed_count
+    end
+    # check field attributes, only backed tables have attributes
+    attrs = keys.collect do |key| 
+      f = @schema[:fields][key].dup
+      f.delete(:weight)
+      f.select{|k,v| allowed_attributes.include? k }
+    end
+    other_attrs = other_keys.collect do |key|
+      f = other_schema[:fields][key].dup
+      f.delete(:weight)
+      f.select{|k,v| allowed_attributes.include? k }
+    end
+    changes[:attributes] if other_attrs != attrs # this means that we will need to redefine the triggers if any are there and also the table definition if needed
+    
+    # return the changes
+    changes
+  end
+  
+  def clean
+    removable = @schema[:fields].select{|name, f| f[:weight] == 0 }.collect{|name, f| name}
+    removable.each{|name| @schema[:fields].delete(name)}
+  end
+  
+  def allowed_attributes
+    [:weight, :col, :target]
+  end
+    
+end
+
+class Litesearch::SchemaException < StandardError; end
+class Litesearch::SchemaChangeException < StandardError; end