atlas_engine 0.1.0

data/app/models/atlas_engine/country_profile.rb

@@ -0,0 +1,270 @@
+# typed: strict
+# frozen_string_literal: true
+
+module AtlasEngine
+  class CountryProfile < FrozenRecord::Base
+    extend T::Sig
+
+    class CountryNotFoundError < StandardError; end
+
+    module Backend
+      extend FrozenRecord::Backends::Yaml
+
+      class << self
+        extend T::Sig
+
+        sig { params(_file_path: String).returns(T::Array[T.untyped]) }
+        def load(_file_path)
+          # FrozenRecord's default is to operate on a single YAML file containing all the records.
+          # A custom backend like ours, that uses separate files, must load all of them and return an array.
+          country_profiles = load_and_merge_fragments(CountryProfile.country_paths)
+          locale_profiles = load_and_merge_fragments(CountryProfile.locale_paths)
+
+          # These hashes are not complete country profiles, but rather fragments that will be merged
+          # onto the default profile template.
+          country_profiles + localize_profiles(country_profiles, locale_profiles)
+        end
+
+        private
+
+        sig { params(path_patterns: T::Array[String]).returns(T::Array[T::Hash[String, T.untyped]]) }
+        def load_and_merge_fragments(path_patterns)
+          path_patterns.flat_map do |directory_pattern|
+            Dir[directory_pattern]
+          end.map do |profile_path|
+            FrozenRecord::Backends::Yaml.load(profile_path)
+          end.group_by do |profile|
+            profile["id"]
+          end.transform_values do |profile_fragments|
+            profile_fragments.inject({}) do |memo, fragment|
+              memo.deep_merge(fragment)
+            end
+          end.values
+        end
+
+        sig do
+          params(country_profiles: T::Array[T.untyped], locale_profiles: T::Array[T.untyped])
+            .returns(T::Array[T.untyped])
+        end
+        def localize_profiles(country_profiles, locale_profiles)
+          locale_profiles.map do |locale_profile|
+            base_profile_id = country_id_from_locale_id(locale_profile["id"])
+
+            base_profile = country_profiles.find do |country_profile|
+              country_profile["id"] == base_profile_id
+            end
+
+            (base_profile || {}).deep_merge(locale_profile)
+          end
+        end
+
+        sig { params(locale_id: String).returns(String) }
+        def country_id_from_locale_id(locale_id)
+          unless locale_id.match?(/\A[A-Z]{2}_[A-Z]{2}\z/)
+            raise "Invalid id for localized country profile: #{locale_id}"
+          end
+
+          T.must(locale_id.split("_").first)
+        end
+      end
+    end
+
+    DEFAULT_PROFILE = "DEFAULT"
+
+    add_index :id, unique: true
+    self.base_path = ""
+    self.backend = Backend
+
+    # rubocop:disable Style/ClassVars
+    @@default_paths = T.let(
+      [
+        File.join(AtlasEngine::Engine.root, "db/data/country_profiles/default.yml"),
+      ],
+      T::Array[String],
+    )
+
+    @@country_paths = T.let(
+      [
+        File.join(AtlasEngine::Engine.root, "app/countries/atlas_engine/*/country_profile.yml"),
+      ],
+      T::Array[String],
+    )
+
+    @@locale_paths = T.let(
+      [
+        File.join(AtlasEngine::Engine.root, "app/countries/atlas_engine/*/locales/*/country_profile.yml"),
+      ],
+      T::Array[String],
+    )
+
+    @attributes = T.let([], T::Array[T.untyped])
+    @records = T.let(nil, T.nilable(T::Array[T.untyped]))
+
+    class << self
+      extend T::Sig
+
+      sig { returns(T::Array[String]) }
+      def default_paths
+        @@default_paths
+      end
+
+      sig { params(paths: T::Array[String]).void }
+      def default_paths=(paths)
+        @@default_paths = paths
+      end
+
+      sig { returns(T::Array[String]) }
+      def country_paths
+        @@country_paths
+      end
+
+      sig { params(paths: T::Array[String]).void }
+      def country_paths=(paths)
+        @@country_paths = paths
+      end
+
+      sig { returns(T::Array[String]) }
+      def locale_paths
+        @@locale_paths
+      end
+
+      sig { params(paths: T::Array[String]).void }
+      def locale_paths=(paths)
+        @@locale_paths = paths
+      end
+
+      sig { params(paths: T.any(String, T::Array[String])).void }
+      def add_default_paths(paths)
+        T.unsafe(@@default_paths).append(*Array(paths))
+      end
+
+      sig { params(paths: T.any(String, T::Array[String])).void }
+      def add_country_paths(paths)
+        T.unsafe(@@country_paths).append(*Array(paths))
+      end
+
+      sig { params(paths: T.any(String, T::Array[String])).void }
+      def add_locale_paths(paths)
+        T.unsafe(@@locale_paths).append(*Array(paths))
+      end
+
+      sig { void }
+      def reset!
+        unload!
+        @@default_paths = []
+        @@country_paths = []
+        @@locale_paths = []
+        @default_attributes = nil
+      end
+      # rubocop:enable Style/ClassVars
+
+      # Overriding (load_records) from FrozenRecord::Base
+      # so that we only create attribute methods that are not already defined
+      sig { params(force: T::Boolean).returns(T::Array[T.untyped]) }
+      def load_records(force: false)
+        if force || (auto_reloading && file_changed?)
+          unload!
+        end
+
+        @records ||= begin
+          records = backend.load(file_path)
+          if attribute_deserializers.any? || default_attributes
+            records = records.map { |r| assign_defaults!(deserialize_attributes!(r.dup)).freeze }.freeze
+          end
+          @attributes = list_attributes(records).freeze
+          define_attribute_methods(methods_to_be_created)
+          records = FrozenRecord.ignore_max_records_scan { records.map { |r| load(r) }.freeze }
+          index_definitions.values.each { |index| index.build(records) }
+          records
+        end
+      end
+
+      sig { params(country_code: String, locale: T.nilable(String)).returns(CountryProfile) }
+      def for(country_code, locale = nil)
+        raise CountryNotFoundError if country_code.blank?
+
+        unless country_code == DEFAULT_PROFILE || Worldwide.region(code: country_code).country?
+          raise CountryNotFoundError
+        end
+
+        ids = [country_code.upcase]
+        ids.push("#{country_code.upcase}_#{locale.upcase}") if locale.present?
+
+        # if a localized profile is not found, fall back to the
+        # country's profile before falling back to the default profile
+        begin
+          id = ids.pop
+          find(id)
+        rescue FrozenRecord::RecordNotFound
+          if ids.present?
+            retry
+          else
+            new(default_attributes.merge("id" => id))
+          end
+        end
+      end
+
+      sig { returns(T::Hash[String, T.untyped]) }
+      def default_attributes
+        @default_attributes ||= T.let(
+          default_paths.each_with_object({}) do |path, hash|
+            hash.deep_merge!(YAML.load_file(path))
+          end,
+          T.nilable(T::Hash[String, T.untyped]),
+        )
+      end
+
+      sig { override.params(record: T::Hash[String, T.untyped]).returns(T::Hash[String, T.untyped]) }
+      def assign_defaults!(record)
+        default_attributes.deep_merge(record)
+      end
+
+      sig { returns(T::Set[String]) }
+      def partial_zip_allowed_countries
+        @partial_zip_allowed_countries ||= T.let(
+          where.not(id: DEFAULT_PROFILE).filter do |country_profile|
+            country_profile.ingestion.allow_partial_zip?
+          end.map(&:id).to_set,
+          T.nilable(T::Set[String]),
+        )
+      end
+
+      sig { returns(T::Set[String]) }
+      def validation_enabled_countries
+        @validation_enabled_countries ||= T.let(
+          where.not(id: DEFAULT_PROFILE).filter do |country_profile|
+            country_profile.validation.enabled
+          end.pluck(:id).to_set,
+          T.nilable(T::Set[String]),
+        )
+      end
+
+      private
+
+      sig { returns(T::Array[T.untyped]) }
+      def methods_to_be_created
+        @attributes.to_a.flatten.reject do |attribute_name|
+          instance_methods.include?(attribute_name.to_sym)
+        end
+      end
+    end
+
+    sig { returns(CountryProfileValidationSubset).checked(:tests) }
+    def validation
+      CountryProfileValidationSubset.new(hash: attributes["validation"] || {})
+    end
+
+    sig { returns(CountryProfileIngestionSubset).checked(:tests) }
+    def ingestion
+      CountryProfileIngestionSubset.new(hash: attributes["ingestion"])
+    end
+
+    sig { returns(T::Hash[Symbol, T.untyped]) }
+    def open_address = (attributes["open_address"] || {}).with_indifferent_access
+
+    sig { params(field: Symbol).returns(T::Array[String]) }
+    def decompounding_patterns(field)
+      attributes.dig("decompounding_patterns", field.to_s) || []
+    end
+  end
+end

data/app/models/atlas_engine/country_profile_ingestion_subset.rb

@@ -0,0 +1,42 @@
+# typed: true
+# frozen_string_literal: true
+
+module AtlasEngine
+  class CountryProfileIngestionSubset < CountryProfileSubsetBase
+    sig { params(source: String).returns(T::Array[T::Class[T.anything]]) }
+    def correctors(source:)
+      corrector_names = attributes.dig("correctors", source) || []
+      corrector_names.map(&:constantize)
+    end
+
+    sig { returns(T.nilable(String)) }
+    def settings_number_of_shards
+      attributes.dig("settings", "number_of_shards")
+    end
+
+    sig { returns(T.nilable(String)) }
+    def settings_number_of_replicas
+      attributes.dig("settings", "number_of_replicas")
+    end
+
+    sig { returns(T.nilable(String)) }
+    def settings_min_zip_edge_ngram
+      attributes.dig("settings", "min_zip_edge_ngram")
+    end
+
+    sig { returns(T::Class[AddressImporter::OpenAddress::DefaultMapper]) }
+    def open_address_feature_mapper
+      attributes.dig("open_address", "feature_mapper").constantize
+    end
+
+    sig { returns(T.nilable(String)) }
+    def settings_max_zip_edge_ngram
+      attributes.dig("settings", "max_zip_edge_ngram")
+    end
+
+    sig { returns(T::Class[AddressValidation::Es::DataMappers::DefaultDataMapper]) }
+    def data_mapper
+      attributes.dig("data_mapper").constantize
+    end
+  end
+end

data/app/models/atlas_engine/country_profile_subset_base.rb

@@ -0,0 +1,22 @@
+# typed: true
+# frozen_string_literal: true
+
+module AtlasEngine
+  class CountryProfileSubsetBase
+    extend T::Sig
+
+    sig { returns(T::Hash[T.untyped, T.untyped]) }
+    attr_reader :attributes
+
+    sig { params(hash: T::Hash[T.untyped, T.untyped]).void }
+    def initialize(hash:)
+      @attributes = hash
+
+      attributes.keys.each do |key|
+        define_singleton_method(key.to_s) do
+          @attributes[key]
+        end unless respond_to?(key.to_s)
+      end
+    end
+  end
+end

data/app/models/atlas_engine/country_profile_validation_subset.rb

@@ -0,0 +1,48 @@
+# typed: true
+# frozen_string_literal: true
+
+module AtlasEngine
+  class CountryProfileValidationSubset < CountryProfileSubsetBase
+    sig { returns(T::Boolean) }
+    def enabled
+      !!attributes.dig("enabled")
+    end
+
+    sig { returns(T.nilable(String)) }
+    def default_matching_strategy
+      attributes.dig("default_matching_strategy")
+    end
+
+    sig { returns(T::Array[String]) }
+    def index_locales
+      attributes.dig("index_locales") || []
+    end
+
+    sig { returns(T::Boolean) }
+    def multi_locale?
+      index_locales.size > 1
+    end
+
+    sig do
+      params(component: String)
+        .returns(T::Array[T.class_of(AddressValidation::Validators::FullAddress::Exclusions::ExclusionBase)])
+    end
+    def validation_exclusions(component:)
+      validation_exclusions = attributes.dig("exclusions", component) || []
+      validation_exclusions.map(&:constantize)
+    end
+
+    sig { params(length: Integer).returns(T.nilable(T::Range[T.untyped])) }
+    def partial_postal_code_range(length)
+      range = attributes.dig("partial_postal_code_range_for_length", length)
+      return unless range
+
+      Range.new(*range.split("..").map(&:to_i))
+    end
+
+    sig { returns(T::Class[ValidationTranscriber::AddressParserBase]) }
+    def address_parser
+      attributes.dig("address_parser").constantize
+    end
+  end
+end

data/app/models/atlas_engine/country_repository.rb

@@ -0,0 +1,110 @@
+# typed: strict
+# frozen_string_literal: true
+
+module AtlasEngine
+  class CountryRepository
+    include LogHelper
+    extend T::Sig
+
+    PostAddressData = T.type_alias { T.any(PostAddress, T::Hash[Symbol, T.untyped]) }
+
+    delegate :active_alias,
+      :archived_alias,
+      :new_alias,
+      :read_alias_name,
+      :record_source,
+      :create_next_index,
+      :switch_to_next_index,
+      :save_records_backfill,
+      to: :repository
+
+    sig do
+      params(
+        country_code: T.any(String, Symbol),
+        repository_class: T::Class[Elasticsearch::RepositoryInterface],
+        locale: T.nilable(String),
+        index_configuration: T.nilable(IndexConfigurationFactory::IndexConfigurations),
+      ).void
+    end
+    def initialize(country_code:, repository_class:, locale: nil, index_configuration: nil)
+      @country_code = T.let(country_code.to_s.downcase, String)
+      @country_profile = T.let(CountryProfile.for(@country_code, locale), CountryProfile)
+
+      @repository = T.let(
+        repository_class.new(
+          index_base_name: index_name(country_code: @country_code, locale: locale),
+          index_mappings: index_configuration.present? ? index_configuration["mappings"] : {},
+          index_settings: index_configuration.present? ? index_configuration["settings"] : {},
+          mapper_callable: mapper_callable,
+        ),
+        Elasticsearch::RepositoryInterface,
+      )
+    end
+
+    sig { params(id: T.any(String, Integer)).returns(T::Hash[String, T.untyped]) }
+    def find(id)
+      repository.find(id).deep_stringify_keys
+    end
+
+    sig { params(query: T::Hash[String, T.untyped]).returns(T::Array[T::Hash[String, T.untyped]]) }
+    def search(query)
+      repository.search(query)["hits"]["hits"]
+    end
+
+    sig { params(query: T::Hash[String, T.untyped]).returns(T::Array[T::Hash[String, T.untyped]]) }
+    def analyze(query)
+      build_analyze_result(repository.analyze(query))
+    end
+
+    sig { params(query: T::Hash[String, T.untyped]).returns(T::Array[T::Hash[String, T.untyped]]) }
+    def term_vectors(query)
+      build_term_vectors(repository.term_vectors(query))
+    end
+
+    private
+
+    sig { params(country_code: T.any(String, Symbol), locale: T.nilable(String)).returns(String) }
+    def index_name(country_code:, locale: nil)
+      if country_profile.validation.multi_locale?
+        if country_profile.validation.index_locales.include?(locale)
+          "#{country_code}_#{locale}"
+        else
+          raise ArgumentError, "#{country_code} is a multi-locale country and requires a locale"
+        end
+      else
+        country_code.to_s.downcase
+      end
+    end
+
+    sig { returns(CountryProfile) }
+    attr_reader :country_profile
+
+    sig { returns(Elasticsearch::RepositoryInterface) }
+    attr_reader :repository
+
+    sig { params(es_response: T::Hash[String, T.untyped]).returns(T::Array[T.untyped]) }
+    def build_analyze_result(es_response)
+      Array(es_response["tokens"])
+    end
+
+    sig { params(es_response: T::Hash[String, T.untyped]).returns(T::Array[T.untyped]) }
+    def build_term_vectors(es_response)
+      Array(es_response["docs"])
+    end
+
+    sig { returns(T.nilable(T.proc.params(arg0: T.untyped).returns(T.untyped))) }
+    def mapper_callable
+      # Caching because there's only one per country
+      return @mapper_callable if @mapper_callable
+
+      mapper_class = country_profile.ingestion.data_mapper
+
+      @mapper_callable = T.let(
+        ->(address) {
+          mapper_class.new(post_address: address, country_profile:).map_data
+        },
+        T.nilable(T.proc.params(arg0: T.untyped).returns(T.untyped)),
+      )
+    end
+  end
+end

data/app/models/atlas_engine/elasticsearch/client.rb

@@ -0,0 +1,116 @@
+# typed: strict
+# frozen_string_literal: true
+
+module AtlasEngine
+  module Elasticsearch
+    class Client
+      extend T::Sig
+      include ClientInterface
+
+      DEFAULT_OPTIONS = T.let(
+        {
+          read_timeout: 1,
+          open_timeout: 1,
+          keep_alive_timeout: 60,
+          retry_on_failure: false,
+          headers: {},
+        },
+        ConfigType,
+      )
+
+      sig { override.returns(ConfigType) }
+      attr_reader :config
+
+      sig { params(config: ConfigType).void }
+      def initialize(config = {})
+        @config = T.let(DEFAULT_OPTIONS.merge(config).freeze, ConfigType)
+
+        options = {
+          url: @config[:url] || ENV["ELASTICSEARCH_URL"],
+          retry_on_failure: false,
+        }.compact
+
+        @client = T.let(
+          Elastic::Transport::Client.new(options) do |faraday_connection|
+            faraday_connection.options.timeout = read_timeout
+            faraday_connection.options.open_timeout = open_timeout
+            @config[:headers][:"Content-Type"] = "application/json"
+
+            if ENV["ELASTICSEARCH_API_KEY"].present?
+              @config[:headers][:Authorization] = "ApiKey #{ENV["ELASTICSEARCH_API_KEY"]}"
+            end
+            faraday_connection.headers = @config[:headers] if @config[:headers].present?
+
+            if ENV["ELASTICSEARCH_CLIENT_CERT"] && ENV["ELASTICSEARCH_CLIENT_KEY"]
+              faraday_connection.ssl.client_cert = ENV["ELASTICSEARCH_CLIENT_CERT"]
+              faraday_connection.ssl.client_key = ENV["ELASTICSEARCH_CLIENT_KEY"]
+              faraday_connection.ssl.verify = true
+            end
+
+            if ENV["ELASTICSEARCH_CLIENT_CA_CERT"]
+              faraday_connection.ssl.ca_file = ENV["ELASTICSEARCH_CLIENT_CA_CERT"]
+              faraday_connection.ssl.verify = true
+            end
+
+            if ENV["ELASTICSEARCH_INSECURE_NO_VERIFY_SERVER"]
+              faraday_connection.ssl.verify = false
+            end
+          end,
+          Elastic::Transport::Client,
+        )
+      end
+
+      sig do
+        override.params(
+          method: T.any(Symbol, String),
+          path: String,
+          body: T.nilable(BodyType),
+          options: ConfigType,
+        ).returns(AtlasEngine::Elasticsearch::Response)
+      end
+      def request(method, path, body = nil, options = config.dup)
+        raw_request(method, path, body, options)
+      end
+
+      sig { override.params(alias_name: String).returns(T.nilable(String)) }
+      def find_index_by(alias_name:)
+        get("_alias/#{alias_name}").body.keys.first
+      rescue Elastic::Transport::Transport::Errors::NotFound
+        nil
+      end
+
+      sig { override.params(name: String).returns(T::Boolean) }
+      def index_or_alias_exists?(name)
+        head(name).status == 200
+      rescue Elastic::Transport::Transport::Errors::NotFound
+        false
+      end
+
+      private
+
+      sig { returns(Integer) }
+      def read_timeout
+        @config[:read_timeout] # Value is in seconds
+      end
+
+      sig { returns(Integer) }
+      def open_timeout
+        @config[:open_timeout] # Value is in seconds
+      end
+
+      sig do
+        params(
+          method: T.any(Symbol, String),
+          path: String,
+          body: T.nilable(BodyType),
+          options: ConfigType,
+        ).returns(Response)
+      end
+      def raw_request(method, path, body, options)
+        params = options[:params] ||= {}
+        headers = options[:headers] ||= {}
+        Response.new(@client.transport.perform_request(method, path, params, body, headers))
+      end
+    end
+  end
+end

data/app/models/atlas_engine/elasticsearch/client_interface.rb

@@ -0,0 +1,89 @@
+# typed: strict
+# frozen_string_literal: true
+
+module AtlasEngine
+  module Elasticsearch
+    module ClientInterface
+      extend T::Sig
+      extend T::Helpers
+      abstract!
+
+      ConfigType = T.type_alias { T::Hash[Symbol, T.untyped] }
+      BodyType = T.type_alias { T.any(String, T::Hash[T.untyped, T.untyped]) }
+
+      sig { abstract.returns(ConfigType) }
+      def config; end
+
+      sig do
+        params(
+          path: String,
+          body: T.nilable(BodyType),
+          options: ConfigType,
+        ).returns(Response)
+      end
+      def get(path, body = nil, options = config.dup)
+        request(:get, path, nil, options)
+      end
+
+      sig do
+        params(
+          path: String,
+          body: T.nilable(BodyType),
+          options: ConfigType,
+        ).returns(Response)
+      end
+      def head(path, body = nil, options = config.dup)
+        request(:head, path, body, options)
+      end
+
+      sig do
+        params(
+          path: String,
+          body: T.nilable(BodyType),
+          options: ConfigType,
+        ).returns(Response)
+      end
+      def post(path, body = nil, options = config.dup)
+        request(:post, path, body, options)
+      end
+
+      sig do
+        params(
+          path: String,
+          body: T.nilable(BodyType),
+          options: ConfigType,
+        ).returns(Response)
+      end
+      def put(path, body = nil, options = config.dup)
+        request(:put, path, body, options)
+      end
+
+      sig do
+        params(
+          path: String,
+          body: T.nilable(BodyType),
+          options: ConfigType,
+        ).returns(Response)
+      end
+      def delete(path, body = nil, options = config.dup)
+        request(:delete, path, body, options)
+      end
+
+      sig do
+        abstract.params(
+          method: T.any(Symbol, String),
+          path: String,
+          body: T.nilable(BodyType),
+          options: T::Hash[Symbol, T.untyped],
+        ).returns(Response)
+      end
+      def request(method, path, body = nil, options = {}); end
+
+      sig { abstract.params(alias_name: String).returns(T.nilable(String)) }
+      def find_index_by(alias_name:); end
+
+      sig { abstract.params(name: String).returns(T::Boolean) }
+      def index_or_alias_exists?(name); end
+    end
+  end
+end