typesensual 0.1.0

data/lib/typesensual/collection.rb

@@ -0,0 +1,202 @@
+# frozen_string_literal: true
+
+require 'typesensual/schema'
+require 'typesensual/state_helpers'
+require 'typesensual/callbacks' if defined?(ActiveRecord)
+
+class Typesensual
+  class Collection
+    include StateHelpers
+
+    # The pattern we use for collection names, `name:env@version`, where `name`
+    # is the name of the index, `env` is the environment, and `version` is the
+    # timestamp of the collection's creation. If a name doesn't follow this
+    # pattern, `name` collects everything.
+    COLLECTION_NAME_PATTERN = /^
+      (?<name>.*?)           # the name of the collection cannot include : or @
+      (?::(?<env>.*?))?      # the env is optional, but also cannot include : or @
+      (?:@(?<version>\d+))?  # the version is also optional but must be an integer
+    $/x.freeze
+
+    include StateHelpers
+
+    # @overload initialize(collection)
+    #   Initialize a new collection from a Typesense collection hash
+    #
+    #   @param collection [Hash] the Typesense collection hash
+    #     * `created_at` [Integer] the timestamp of the collection's creation
+    #     * `default_sorting_field` [String] the default sorting field
+    #     * `enable_nested_fields` [Boolean] whether nested fields are enabled
+    #     * `fields` [Array<Hash>] the fields in the collection
+    #     * `name` [String] the name of the collection
+    #     * `num_documents` [Integer] the number of documents in the collection
+    #     * `symbols_to_index` [String] the symbols to index
+    #     * `token_separators` [String] the token separators
+    # @overload initialize(name)
+    #   Initialize a new collection, loading info from Typesense
+    #
+    #   @param name [String] the name of the collection
+    #   @raise [Typesense::Error::ObjectNotFound] if the collection doesn't exist
+    def initialize(collection_or_name)
+      @collection = if collection_or_name.is_a?(Hash)
+                      collection_or_name.deep_stringify_keys
+                    else
+                      client.collections[collection_or_name].retrieve
+                    end
+    end
+
+    # Reload the underlying collection data from Typesense
+    # @return [self]
+    def reload
+      @collection = client.collections[name].retrieve
+      self
+    end
+
+    # The time the collection was created, as tracked by Typesense
+    # @return [Time] the time the collection was created
+    def created_at
+      @created_at ||= Time.strptime(@collection['created_at'].to_s, '%s')
+    end
+
+    # The default sorting field for the collection
+    # @return [String] the default sorting field
+    def default_sorting_field
+      @collection['default_sorting_field']
+    end
+
+    # Whether the collection has nested fields enabled
+    # @return [Boolean] whether nested fields are enabled
+    def enable_nested_fields?
+      @collection['enable_nested_fields']
+    end
+
+    # The fields in the collection
+    # @return [Array<Field>] the field information
+    def fields
+      @collection['fields'].map do |field|
+        Field.new(field)
+      end
+    end
+
+    # The raw, underlying name of the collection
+    # @return [String] the name of the collection
+    def name
+      @collection['name']
+    end
+
+    # The number of documents in the collection
+    # @return [Integer] the number of documents in the collection
+    def num_documents
+      @collection['num_documents']
+    end
+
+    # Special characters in strings which should be indexed as text
+    # @return [Array<String>] the symbols to index
+    def symbols_to_index
+      @collection['symbols_to_index']
+    end
+
+    # Additional characters to be treated as separators when indexing text
+    # @return [Array<String>] the token separators
+    def token_separators
+      @collection['token_separators']
+    end
+
+    # The name of the index, parsed from the Typesensual collection naming scheme.
+    # @see COLLECTION_NAME_PATTERN
+    # @return [String] the name of the index
+    def index_name
+      parsed_name['name']
+    end
+
+    # The environment the collection is in, parsed from the Typesensual collection
+    # naming scheme.
+    # @see COLLECTION_NAME_PATTERN
+    # @return [String] the environment the collection is in
+    def env
+      parsed_name['env']
+    end
+
+    # The version of the collection, parsed from the Typesensual collection naming
+    # scheme.
+    # @see COLLECTION_NAME_PATTERN
+    # @return [String] the version of the collection
+    def version
+      parsed_name['version']
+    end
+
+    # Creates the collection in Typesense
+    # @return [self]
+    def create!
+      client.collections.create(@collection)
+      self
+    end
+
+    # Deletes the collection in Typesense
+    # @return [void]
+    def delete!
+      typesense_collection.delete
+    end
+
+    # Create a new collection using the given collection hash
+    #
+    # @param collection [Hash] the Typesense collection hash
+    #   * `default_sorting_field` [String] the default sorting field
+    #   * `enable_nested_fields` [Boolean] whether nested fields are enabled
+    #   * `fields` [Array<Hash>] the fields in the collection
+    #   * `name` [String] the name of the collection
+    #   * `symbols_to_index` [String] the symbols to index
+    #   * `token_separators` [String] the token separators
+    # @return [Collection] the created collection
+    def self.create!(collection)
+      new(collection).tap(&:create!)
+    end
+
+    # Insert a single document into typesense
+    #
+    # @param doc [Hash] the document to insert
+    # @return [Boolean] if the document was inserted successfully
+    def insert_one!(doc)
+      typesense_collection.documents.create(doc)
+    end
+
+    # Insert many documents into typesense. Notably, the input can be an enumerable
+    # or enumerator, which will be batched into groups of `batch_size` and inserted
+    # with the ID of any failed rows being provided in the response.
+    #
+    # @param docs [Enumerable<Hash>] the documents to insert
+    # @return [Array<Hash>] any failed insertions
+    def insert_many!(docs, batch_size: 100)
+      docs.lazy.each_slice(batch_size).with_object([]) do |slice, failures|
+        results = typesense_collection.documents.import(slice, return_id: true)
+        failures.push(*results.reject { |result| result['success'] })
+      end
+    end
+
+    # Remove a single document from typesense
+    #
+    # @param id [String] the ID of the document to remove
+    # @return [void]
+    def remove_one!(id)
+      typesense_collection.documents[id.to_s].delete
+    end
+
+    def search(query:, query_by:)
+      Search.new(
+        collection: typesense_collection,
+        query: query,
+        query_by: query_by
+      )
+    end
+
+    def typesense_collection
+      @typesense_collection ||= client.collections[name]
+    end
+
+    private
+
+    def parsed_name
+      @parsed_name ||= name.match(COLLECTION_NAME_PATTERN)
+    end
+  end
+end

data/lib/typesensual/config.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+class Typesensual
+  class Config
+    attr_accessor :nodes, :api_key
+    attr_writer :env, :client
+
+    def initialize(&block)
+      yield self if block
+    end
+
+    def env
+      @env ||= (defined?(Rails) ? Rails.env : nil)
+    end
+
+    def client
+      @client ||= Typesense::Client.new(connection_options)
+    end
+
+    private
+
+    def connection_options
+      { nodes: nodes, api_key: api_key }
+    end
+  end
+end

data/lib/typesensual/field.rb

@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+class Typesensual
+  class Field
+    def initialize(hash)
+      @field = hash.stringify_keys
+    end
+
+    def facet?
+      @field['facet']
+    end
+
+    def index?
+      @field['index']
+    end
+
+    def infix?
+      @field['infix']
+    end
+
+    def locale
+      @field['locale'].presence
+    end
+
+    def name
+      if @field['name'].is_a?(Regexp)
+        @field['name'].source
+      else
+        @field['name'].to_s
+      end
+    end
+
+    def optional?
+      @field['optional']
+    end
+
+    def sort?
+      @field['sort']
+    end
+
+    def type
+      @field['type']
+    end
+
+    def to_h
+      @field.to_h.merge!(
+        'name' => name,
+        'locale' => locale
+      ).compact!
+    end
+  end
+end

data/lib/typesensual/index.rb

@@ -0,0 +1,161 @@
+# frozen_string_literal: true
+
+require 'typesensual/schema'
+require 'typesensual/state_helpers'
+
+class Typesensual
+  # Represents your index as a set of collections in Typesense. Manages
+  # Typesense aliases, but with some extra functionality for managing multiple
+  # environments and versions, and a nice DSL for defining your schema.
+  #
+  # @example Defining an index
+  #   class PostsIndex < Typesensual::Index
+  #     # Optional, default is inferred from the class name
+  #     index_name 'user_posts'
+  #
+  #     schema do
+  #       field 'text', type: 'string[]'
+  #       field 'user', type: 'string', facet: true
+  #     end
+  #   end
+  class Index
+    include StateHelpers
+
+    # Get or set the name for this index
+    #
+    # @overload index_name(value)
+    #   Explicitly set the index name
+    #
+    #   @param value [String] the name to identify this index
+    #   @return [void]
+    #
+    # @overload index_name
+    #   Get the index name (either explicitly set or inferred from the class name)
+    #
+    #   @return [String] the name of this index
+    def self.index_name(value = nil)
+      if value
+        @index_name = value
+      else
+        @index_name ||= name.underscore.sub(/_index$/, '')
+      end
+    end
+
+    # The alias name for this index in the current environment
+    #
+    # @return [String] the alias name
+    def self.alias_name
+      [index_name, env].compact.join(':')
+    end
+
+    # Generate a new collection name for the given version
+    #
+    # @param version [String] the version to generate the collection name for
+    # @return [String] the generated collection name
+    def self.collection_name_for(version: Time.now.strftime('%s'))
+      "#{alias_name}@#{version}"
+    end
+
+    # Create a new collection for this index
+    #
+    # @param version [String] the version to create the collection for
+    # @return [Collection] the newly created collection
+    def self.create!(version: Time.now.strftime('%s'))
+      generated_name = collection_name_for(version: version)
+
+      Collection.create!(@schema.to_h.merge('name' => generated_name))
+    end
+
+    # Get the collections for this index
+    #
+    # @return [Array<Collection>] the collections that match the alias
+    def self.collections
+      Typesensual.collections.filter do |collection|
+        collection.index_name == index_name
+      end
+    end
+
+    def self.collection_for(version:)
+      Typesensual.collections.find do |collection|
+        collection.version == version
+      end
+    end
+
+    # Get the collection that the alias points to
+    #
+    # @return [Collection] the collection that the alias points to
+    def self.collection
+      @collection ||= Collection.new(alias_name)
+    rescue Typesense::Error::ObjectNotFound
+      nil
+    end
+
+    # Define the schema for the collection
+    #
+    # See {Schema} for more information
+    def self.schema(&block)
+      @schema = Typesensual::Schema.new(&block)
+    end
+
+    # Updates the alias to point to the given collection name
+    #
+    # @param name [String, Collection] the collection to point the alias to
+    def self.update_alias!(name_or_collection)
+      name = if name_or_collection.is_a?(Collection)
+               name_or_collection.name
+             else
+               name_or_collection
+             end
+
+      client.aliases.upsert(alias_name, collection_name: name)
+    end
+
+    # Indexes the given records into a collection, then updates the alias to
+    # point to it.
+    #
+    # @param records [Enumerable] the records to index
+    # @param collection [Collection] the collection to index into, defaults to a
+    #   new collection
+    def self.reindex!(ids, collection: create!)
+      index_many(ids, collection: collection)
+
+      update_alias!(collection)
+    end
+
+    # The method to implement to index *one* record.
+    #
+    # @return [Hash] the document to upsert in Typesense
+    def index_one(_id); end
+
+    def self.index_one(id, collection: self.collection)
+      collection.insert_one!(new.index_one(id))
+    end
+
+    # The method to implement to index *many* records
+    # Unlike {#index_one}, this method should yield successive records to index
+    #
+    # @yield [Hash] a document to upsert in Typesense
+    def index_many(ids)
+      ids.each do |id|
+        yield index_one(id)
+      end
+    end
+
+    def self.index_many(ids, collection: self.collection, batch_size: 100)
+      collection.insert_many!(
+        new.enum_for(:index_many, ids),
+        batch_size: batch_size
+      )
+    end
+
+    def self.remove_one(id, collection: self.collection)
+      collection.remove_one!(id)
+    end
+
+    if defined?(ActiveRecord)
+      def self.ar_callbacks
+        Typesensual::Callbacks.new(self)
+      end
+    end
+  end
+end

data/lib/typesensual/railtie.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+class Typesensual
+  class Railtie < Rails::Railtie
+    rake_tasks do
+      load 'tasks/typesensual.rake'
+    end
+  end
+end

data/lib/typesensual/rake_helper.rb

@@ -0,0 +1,103 @@
+# frozen_string_literal: true
+
+require 'paint'
+
+class Typesensual
+  class RakeHelper
+    HEADER_FORMAT = Paint["==> %s\n", :bold]
+    LIST_ROW_FORMAT = "%<prefix>4s %<version>-20s %<created_at>-20s %<documents>-20s\n"
+    LIST_HEADER = Paint[format(
+      LIST_ROW_FORMAT,
+      prefix: '', version: 'Version', created_at: 'Created At', documents: 'Documents'
+    ), :bold, :underline]
+
+    class << self
+      # List the collections in all Index descendants
+      #
+      # @param output [IO] The output stream to write to
+      # @example
+      #   rake typesensual:list
+      def list(output: $stdout)
+        # Build up a hash of indices and their (sorted) collections
+        indices = Index.descendants.to_h do |index|
+          [index, index.collections.sort_by(&:created_at).reverse]
+        end
+
+        indices.each do |index, collections|
+          alias_name = index.collection.name
+
+          output.printf(HEADER_FORMAT, index.name.titleize)
+          output.printf(LIST_HEADER)
+
+          collections.each do |collection|
+            output.printf(LIST_ROW_FORMAT,
+              prefix: collection.name == alias_name ? '->' : '',
+              version: collection.version,
+              created_at: collection.created_at.strftime('%Y-%m-%d %H:%M:%S'),
+              documents: collection.num_documents)
+          end
+
+          output.printf("\n")
+        end
+      end
+
+      # Index all records from a model into an index
+      #
+      # @param index [String] The name of the index to index into
+      # @param model [String] The name of the model to index from
+      # @example
+      #   rake typesensual:index[FooIndex,Foo]
+      def index(index:, model:, output: $stdout)
+        index = index.safe_constantize
+        model = model.safe_constantize
+
+        collection = index.create!
+        output.printf(
+          Paint["==> Indexing %<model>s into %<index>s (Version %<version>s)\n", :bold],
+          model: model.name,
+          index: index.name,
+          version: collection.version
+        )
+        failures = index.index_many(
+          model.ids,
+          collection: collection
+        )
+
+        failures.each do |failure|
+          output.puts(failure.to_json)
+        end
+      end
+
+      # Update the alias for an index to point to a specific version
+      #
+      # @param index [String] The name of the index to update
+      # @param version [String] The version to update the alias to
+      # @example
+      #   rake typesensual:update_alias[FooIndex,1]
+      def update_alias(index:, version:, output: $stdout)
+        index = index.safe_constantize
+        old_coll = index.collection
+        new_coll = index.collection_for(version: version)
+
+        unless new_coll
+          output.puts(Paint["--> No such version #{version} for #{index.name}", :bold])
+          return
+        end
+
+        output.puts(Paint["==> Alias for #{index.name}", :bold])
+        output.printf(
+          "Old: %<version>s (%<created_at>s)\n",
+          version: old_coll&.version || 'None',
+          created_at: old_coll&.created_at&.strftime('%Y-%m-%d %H:%M:%S') || 'N/A'
+        )
+        index.update_alias!(new_coll)
+
+        output.printf(
+          "New: %<version>s (%<created_at>s)\n",
+          version: new_coll.version,
+          created_at: new_coll.created_at.strftime('%Y-%m-%d %H:%M:%S')
+        )
+      end
+    end
+  end
+end

data/lib/typesensual/schema.rb

@@ -0,0 +1,49 @@
+# frozen_string_literal: true
+
+require 'typesensual/field'
+
+class Typesensual
+  class Schema
+    def initialize(&block)
+      instance_eval(&block)
+    end
+
+    def field(name, type: 'auto', locale: nil, facet: nil, index: nil, optional: nil)
+      @fields ||= []
+      @fields << Field.new(
+        name: name,
+        type: type,
+        locale: locale,
+        facet: facet,
+        index: index,
+        optional: optional
+      )
+    end
+
+    def token_separators(*separators)
+      @token_separators = separators
+    end
+
+    def symbols_to_index(*symbols)
+      @symbols_to_index = symbols
+    end
+
+    def default_sorting_field(field_name)
+      @default_sorting_field = field_name.to_s
+    end
+
+    def enable_nested_fields(value = true) # rubocop:disable Style/OptionalBooleanParameter
+      @enable_nested_fields = value
+    end
+
+    def to_h
+      {
+        'fields' => @fields&.map(&:to_h),
+        'token_separators' => @token_separators,
+        'symbols_to_index' => @symbols_to_index,
+        'default_sorting_field' => @default_sorting_field,
+        'enable_nested_fields' => @enable_nested_fields
+      }.compact!
+    end
+  end
+end

data/lib/typesensual/search/hit.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+class Typesensual
+  class Search
+    class Hit
+      # {
+      #   "highlights": [
+      #     {
+      #       "field": "company_name",
+      #       "matched_tokens": ["Stark"],
+      #       "snippet": "<mark>Stark</mark> Industries"
+      #     }
+      #   ],
+      #   "document": {
+      #     "id": "124",
+      #     "company_name": "Stark Industries",
+      #     "num_employees": 5215,
+      #     "country": "USA"
+      #   },
+      #   "text_match": 130916
+      # }
+      # @param collection [Hash] the Typesense hit hash
+      #   * `highlights` [Array<Hash>] the highlights for the hit
+      #   * `document` [Hash] the matching document
+      #   * `text_match` [Integer] the text matching score
+      def initialize(hit)
+        @hit = hit
+      end
+
+      def highlights
+        @hit['highlights']
+      end
+
+      def document
+        @hit['document']
+      end
+
+      def score
+        @hit['text_match']
+      end
+    end
+  end
+end

data/lib/typesensual/search/results.rb

@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+
+class Typesensual
+  class Search
+    class Results
+      def initialize(results)
+        @results = results
+      end
+
+      def hits
+        @results['hits'].map { |hit| Hit.new(hit) }
+      end
+
+      def count
+        @results['found']
+      end
+
+      def out_of
+        @results['out_of']
+      end
+
+      def current_page
+        @results['page']
+      end
+
+      def first_page?
+        current_page == 1
+      end
+
+      def last_page?
+        current_page == total_pages
+      end
+
+      def prev_page
+        current_page - 1 unless first_page?
+      end
+
+      def next_page
+        current_page + 1 unless last_page?
+      end
+
+      def search_time_ms
+        @results['search_time_ms']
+      end
+
+      def total_pages
+        (@results['found'] / @results['per_page'].to_f).ceil
+      end
+    end
+  end
+end