arfi 0.5.0 → 1.0.0

data/lib/arfi/sql_function_loader.rb

@@ -1,152 +1,355 @@
 # frozen_string_literal: true
 
 module Arfi
-  # +Arfi::SqlFunctionLoader+ is a class which loads user defined SQL functions into database.
+  # Loads user-defined SQL functions into the database by executing SQL files located under `db/functions`.
+  #
+  # Supported directory layout (Option B: explicit public):
+  #
+  #   db/functions/public/*.sql                       (generic public)
+  #   db/functions/postgresql/public/*.sql            (postgres public)
+  #   db/functions/postgresql/<schema>/*.sql          (postgres schema)
+  #   db/functions/mysql/public/*.sql                 (mysql/trilogy public)
+  #
+  # Backward-compatible legacy aliases:
+  #
+  #   db/functions/*.sql                              (generic public legacy)
+  #   db/functions/postgresql/*.sql                   (postgres public legacy)
+  #   db/functions/mysql/*.sql                        (mysql public legacy)
+  #
+  # Rules:
+  # - underscore-prefixed files are ignored (_shared.sql)
+  # - adapter-specific overrides generic when schema+filename matches
+  #
+  # @api public
   class SqlFunctionLoader
     class << self
-      # +Arfi::SqlFunctionLoader.load!+                        -> (nil | void)
+      # Load all SQL function files into the database.
       #
-      # Loads user defined SQL functions into database.
+      # Handles both single-DB and multi-DB setups. Uses the given connection or infers it.
       #
-      # @param task_name [String|nil] Name of the task.
-      # @return [nil] if there is no `db/functions` directory.
-      # @return [void] if there is no errors.
-      def load!(task_name: nil)
-        self.task_name = task_name[/([^:]+$)/] if task_name
-        return puts 'No SQL files found. Skipping db population with ARFI' unless sql_files.any?
+      # @param [String?] task_name Optional task name for logging (e.g. 'db:migrate')
+      # @param [ActiveRecord::ConnectionAdapters::AbstractAdapter?] connection Specific connection to load into
+      # @param [Boolean] clear_active_connections Whether to clear connections after loading
+      # @param [Boolean] verbose Whether to log each loaded file
+      # @return [void]
+      def load!(task_name: nil, connection: nil, clear_active_connections: true, verbose: true)
+        task_short = task_name ? task_name[/([^:]+$)/] : nil
+        conn = connection || default_connection
+
+        raise_unless_supported_adapter(conn)
 
-        raise_unless_supported_adapter
-        handle_db_population
-        conn.close
+        if connection.nil? && multi_db? && task_name.nil?
+          populate_multiple_db(verbose: verbose, task_name: task_short)
+        else
+          populate_db(conn, verbose: verbose, task_name: task_short)
+        end
+      ensure
+        clear_active_connections_if_needed(clear_active_connections)
       end
 
       private
 
-      attr_accessor :task_name
-
-      # +Arfi::SqlFunctionLoader#raise_unless_supported_adapter+ -> void
-      #
-      # Checks if the database adapter is supported.
+      # Raise unless the connection adapter is one of the supported types.
       #
-      # @!visibility private
       # @private
+      # @param [ActiveRecord::ConnectionAdapters::AbstractAdapter] conn Database connection
+      # @raise [Arfi::Errors::AdapterNotSupported] If adapter is not PostgreSQL, Mysql2, or Trilogy
       # @return [void]
-      # @raise [Arfi::Errors::AdapterNotSupported]
-      def raise_unless_supported_adapter
-        allowed = %w[ActiveRecord::ConnectionAdapters::PostgreSQLAdapter
-                     ActiveRecord::ConnectionAdapters::Mysql2Adapter].freeze
-        return if allowed.include?(conn.class.to_s) # steep:ignore ArgumentTypeMismatch
+      def raise_unless_supported_adapter(conn) # rubocop:disable SortedMethodsByCall/Waterfall
+        allowed = %w[
+          ActiveRecord::ConnectionAdapters::PostgreSQLAdapter
+          ActiveRecord::ConnectionAdapters::Mysql2Adapter
+          ActiveRecord::ConnectionAdapters::TrilogyAdapter
+        ].freeze
 
-        raise Arfi::Errors::AdapterNotSupported
+        raise Arfi::Errors::AdapterNotSupported unless allowed.include?(conn.class.name)
       end
 
-      # +Arfi::SqlFunctionLoader#handle_db_population+         -> void
+      # Check whether the Rails app has multiple database configurations for the current environment.
       #
-      # Loads user defined SQL functions into database. This conditional branch was written this way because if we
-      # call db:migrate:db_name, then task_name will not be nil, but it will be zero if we call db:migrate. Then we
-      # check that the application has been configured to work with multiple databases in order to populate all
-      # databases, and only after this check can we populate the database in case the db:migrate (or any other) task
-      # has been called for configuration with a single database. Go to `lib/arfi/tasks/db.rake` for additional info.
-      #
-      # @!visibility private
       # @private
-      # @return [void]
-      def handle_db_population
-        if task_name || (task_name && multi_db?) || task_name.nil?
-          populate_db
-        elsif multi_db?
-          populate_multiple_db
-        end
-      end
-
-      # +Arfi::SqlFunctionLoader#multi_db?+                       -> Boolean
-      #
-      # Checks if the application has been configured to work with multiple databases.
-      #
-      # @return [Boolean]
+      # @return [Boolean] Whether multi-DB is configured
       def multi_db?
         ActiveRecord::Base.configurations.configurations.count { _1.env_name == Rails.env } > 1 # steep:ignore NoMethod
       end
 
-      # +Arfi::SqlFunctionLoader#populate_multiple_db+          -> void
+      # Load functions into all databases in a multi-DB setup.
       #
-      # Loads user defined SQL functions into all databases.
+      # Saves the original connection and restores it after iterating all configs,
+      # matching the pattern used in `run_with_connection_switch` (db.rake).
       #
-      # @!visibility private
       # @private
+      # @param [Boolean] verbose Whether to log each loaded file
+      # @param [String?] task_name Optional task name for logging
       # @return [void]
-      # @see Arfi::SqlFunctionLoader#multi_db?
-      # @see Arfi::SqlFunctionLoader#populate_db
-      def populate_multiple_db
+      def populate_multiple_db(verbose:, task_name: nil)
+        original = current_db_config
         # steep:ignore:start
         ActiveRecord::Base.configurations.configurations.select { _1.env_name == Rails.env }.each do |config|
           ActiveRecord::Base.establish_connection(config)
-          populate_db
+          populate_db(default_connection, verbose: verbose, task_name: task_name)
         end
         # steep:ignore:end
+      ensure
+        ActiveRecord::Base.establish_connection(original) if original
       end
 
-      # +Arfi::SqlFunctionLoader#populate_db+                   -> void
+      # Get the default database connection, handling Rails version differences.
       #
-      # Loads user defined SQL functions into database.
+      # @private
+      # @return [ActiveRecord::ConnectionAdapters::AbstractAdapter] Default connection
+      def default_connection
+        if Rails::VERSION::MAJOR < 7 || (Rails::VERSION::MAJOR == 7 && Rails::VERSION::MINOR < 2)
+          ActiveRecord::Base.connection
+        else
+          ActiveRecord::Base.lease_connection
+        end
+      end
+
+      # Get the current database config, or nil if no connection is established.
+      #
+      # @private
+      # @raise [StandardError]
+      # @return [ActiveRecord::DatabaseConfig, nil] Current database config, or nil if no connection is established
+      def current_db_config
+        ActiveRecord::Base.connection_db_config # steep:ignore NoMethod
+      rescue StandardError
+        nil
+      end
+
+      # Load SQL function files into a single database connection.
       #
-      # @!visibility private
       # @private
+      # @param [ActiveRecord::ConnectionAdapters::AbstractAdapter] conn Database connection
+      # @param [Boolean] verbose Whether to log each loaded file
+      # @param [String?] task_name Optional task name for logging
       # @return [void]
-      def populate_db
-        sql_files.each do |file|
-          sql = File.read(file).strip
+      def populate_db(conn, verbose:, task_name:)
+        files = sql_files(conn)
+        if files.empty?
+          log(conn, "No SQL files found for adapter #{conn.class}. Skipping db population with ARFI")
+          return
+        end
+
+        files.each { |file| load_sql_file(conn, file, verbose, task_name) }
+      end
+
+      # Clear all active database connections if requested, handling Rails version differences.
+      #
+      # @private
+      # @param [Boolean] clear Whether to clear connections
+      # @return [void]
+      def clear_active_connections_if_needed(clear)
+        return unless clear && defined?(ActiveRecord::Base)
+
+        if ActiveRecord::Base.respond_to?(:connection_handler) &&
+           ActiveRecord::Base.connection_handler.respond_to?(:clear_active_connections!)
+          ActiveRecord::Base.connection_handler.clear_active_connections!
+        elsif ActiveRecord::Base.respond_to?(:clear_active_connections!)
+          ActiveRecord::Base.clear_active_connections!
+        end
+      end
+
+      # Load a single SQL file into the database, wrapping errors with file context.
+      #
+      # @private
+      # @param [ActiveRecord::ConnectionAdapters::AbstractAdapter] conn Database connection
+      # @param [Pathname, String] file Path to the SQL file
+      # @param [Boolean] verbose Whether to log the loaded file
+      # @param [String?] task_name Optional task name for logging
+      # @raise [StandardError] If the SQL execution fails (re-raised with file context)
+      # @return [void]
+      def load_sql_file(conn, file, verbose, task_name)
+        sql = File.read(file.to_s).strip
+        return if sql.empty?
+
+        begin
           conn.execute(sql)
-          puts "[ARFI] Loaded: #{File.basename(file)} into #{conn.pool.db_config.env_name} #{conn.pool.db_config.name}"
+        rescue StandardError => e
+          raise e.class, "#{e.message}\n[ARFI] while loading #{file}", e.backtrace
         end
+        return unless verbose
+
+        log_sql_load(conn, file, task_name)
       end
 
-      # +Arfi::SqlFunctionLoader#sql_files+                     -> Array<String>
+      # Log that a SQL file was successfully loaded into the database.
       #
-      # Helper method to get list of SQL files. Here we check if we need to populate all databases or just one.
+      # @private
+      # @param [ActiveRecord::ConnectionAdapters::AbstractAdapter] conn Database connection
+      # @param [Pathname, String] file Path to the loaded SQL file
+      # @param [String?] task_name Optional task name for logging
+      # @return [void]
+      def log_sql_load(conn, file, task_name)
+        label = "[ARFI] Loaded: #{File.basename(file)} into #{safe_db_env(conn)} #{safe_db_name(conn)}"
+        label += " (#{task_name})" if task_name
+        log(conn, label)
+      end
+
+      # Get the database environment name safely, returning empty string on error.
       #
-      # @!visibility private
       # @private
-      # @return [Array<String>] List of SQL files.
-      # @see Arfi::SqlFunctionLoader#load!
-      # @see Arfi::SqlFunctionLoader#multi_db?
-      # @see Arfi::SqlFunctionLoader#sql_functions_by_adapter
-      def sql_files
-        if task_name || multi_db?
-          sql_functions_by_adapter
+      # @param [ActiveRecord::ConnectionAdapters::AbstractAdapter] conn Database connection
+      # @raise [StandardError]
+      # @return [String] Environment name, or empty string on error
+      # @return [String] if StandardError
+      def safe_db_env(conn)
+        conn.pool&.db_config&.env_name.to_s
+      rescue StandardError
+        ''
+      end
+
+      # Get the database name safely, returning empty string on error.
+      #
+      # @private
+      # @param [ActiveRecord::ConnectionAdapters::AbstractAdapter] conn Database connection
+      # @raise [StandardError]
+      # @return [String] Database name, or empty string on error
+      # @return [String] if StandardError
+      def safe_db_name(conn)
+        conn.pool&.db_config&.name.to_s
+      rescue StandardError
+        ''
+      end
+
+      # Log a message via Rails.logger or stdout.
+      #
+      # @private
+      # @param [ActiveRecord::ConnectionAdapters::AbstractAdapter] _conn Database connection (unused)
+      # @param [String] msg Message to log
+      # @return [void]
+      def log(_conn, msg)
+        if defined?(Rails) && Rails.respond_to?(:logger) && Rails.logger
+          Rails.logger.info(msg)
         else
-          Dir.glob(Rails.root.join('db', 'functions').join('*.sql'))
+          $stdout.puts(msg)
         end
       end
 
-      # +Arfi::SqlFunctionLoader#sql_functions_by_adapter+      -> Array<String>
+      # Collect all SQL files to load for the given connection, applying override resolution.
       #
-      # Helper method to get list of SQL files for specific database adapter.
+      # @private
+      # @param [ActiveRecord::ConnectionAdapters::AbstractAdapter] conn Database connection
+      # @return [Array<String>] Ordered list of SQL file paths
+      def sql_files(conn)
+        root = Rails.root.join('db', 'functions')
+        return [] unless root.directory?
+
+        generic = [
+          collect_sql(glob: root.join('*.sql'), schema: 'public', priority: 1),
+          collect_sql(glob: root.join('public', '*.sql'), schema: 'public', priority: 2)
+        ].flatten
+
+        adapter_root = adapter_root_for(conn, root)
+        return finalize_items(generic) unless adapter_root&.directory?
+
+        finalize_items(generic + collect_adapter_sql_files(conn, adapter_root))
+      end
+
+      # Collect adapter-specific SQL files, dispatching to the correct strategy based on adapter type.
       #
-      # @!visibility private
       # @private
-      # @return [Array<String>] List of SQL files.
-      # @raise [Arfi::Errors::AdapterNotSupported] if database adapter is not supported.
-      def sql_functions_by_adapter
-        case conn
-        when ActiveRecord::ConnectionAdapters::PostgreSQLAdapter
-          Dir.glob(Rails.root.join('db', 'functions', 'postgresql').join('*.sql'))
-        when ActiveRecord::ConnectionAdapters::Mysql2Adapter
-          Dir.glob(Rails.root.join('db', 'functions', 'mysql').join('*.sql'))
+      # @param [ActiveRecord::ConnectionAdapters::AbstractAdapter] conn Database connection
+      # @param [Pathname] adapter_root Adapter root directory (e.g. db/functions/postgresql)
+      # @raise [Arfi::Errors::AdapterNotSupported] If adapter is not supported
+      # @return [Array<Arfi::sql_file_item>] List of SQL file items
+      def collect_adapter_sql_files(conn, adapter_root)
+        case conn.class.name
+        when 'ActiveRecord::ConnectionAdapters::PostgreSQLAdapter'
+          collect_postgresql_sql_files(adapter_root)
+        when 'ActiveRecord::ConnectionAdapters::Mysql2Adapter', 'ActiveRecord::ConnectionAdapters::TrilogyAdapter'
+          collect_adapter_public_sql_files(adapter_root)
         else
           raise Arfi::Errors::AdapterNotSupported
         end
       end
 
-      # +Arfi::SqlFunctionLoader#conn+                          -> ActiveRecord::ConnectionAdapters::AbstractAdapter
+      # Collect SQL files for PostgreSQL, including schema subdirectories.
+      #
+      # @private
+      # @param [Pathname] adapter_root PostgreSQL adapter root directory
+      # @return [Array<Arfi::sql_file_item>] List of SQL file items
+      def collect_postgresql_sql_files(adapter_root)
+        items = collect_adapter_public_sql_files(adapter_root)
+
+        Dir.children(adapter_root).sort.each do |child|
+          next if child.start_with?('_')
+          next if child == 'public'
+
+          dir = adapter_root.join(child)
+          next unless dir.directory?
+
+          items.concat collect_sql(glob: dir.join('*.sql'), schema: child, priority: 10)
+        end
+
+        items
+      end
+
+      # Collect SQL files from the adapter's public directory (legacy + explicit).
+      #
+      # @private
+      # @param [Pathname] adapter_root Adapter root directory
+      # @return [Array<Arfi::sql_file_item>] List of SQL file items
+      def collect_adapter_public_sql_files(adapter_root)
+        items = [] # steep:ignore
+        items.concat collect_sql(glob: adapter_root.join('*.sql'), schema: 'public', priority: 8)
+        items.concat collect_sql(glob: adapter_root.join('public', '*.sql'), schema: 'public', priority: 9)
+        items
+      end
+
+      # Collect SQL files matching a glob pattern, skipping underscore-prefixed files.
       #
-      # Helper method to get database connection.
+      # @private
+      # @param [Pathname, String] glob Glob pattern to match SQL files
+      # @param [String] schema Schema name to assign to matched files
+      # @param [Integer] priority Priority for override resolution (higher = preferred)
+      # @return [Array<Arfi::sql_file_item>] List of SQL file items
+      def collect_sql(glob:, schema:, priority:)
+        Dir.glob(glob.to_s).filter_map do |path|
+          base = File.basename(path)
+          next if base.start_with?('_')
+
+          { schema: schema, base: base, path: path, priority: priority }
+        end
+      end
+
+      # Finalize the file list by selecting the highest-priority item per schema/filename key.
       #
-      # @!visibility private
       # @private
-      # @return [ActiveRecord::ConnectionAdapters::AbstractAdapter] Database connection.
-      def conn
-        ActiveRecord::Base.lease_connection
+      # @param [Array<Arfi::sql_file_item>] items All collected SQL file items
+      # @return [Array<String>] Ordered list of chosen file paths
+      def finalize_items(items)
+        # @type var chosen: Hash[String, { schema: String, base: String, path: String, priority: Integer }]
+        chosen = {}
+
+        items.each do |item|
+          key = "#{item[:schema]}/#{item[:base]}"
+          prev = chosen[key]
+          chosen[key] = item if prev.nil? || item[:priority] > prev[:priority]
+        end
+
+        chosen.values
+              .sort_by { |item| [item[:schema], item[:base]] }
+              .map { |item| item[:path] }
+      end
+
+      # Resolve the adapter-specific root directory for the given connection.
+      #
+      # @private
+      # @param [ActiveRecord::ConnectionAdapters::AbstractAdapter] conn Database connection
+      # @param [Pathname] root Project root (db/functions)
+      # @raise [Arfi::Errors::AdapterNotSupported] If adapter is not supported
+      # @return [Pathname] Adapter root directory
+      def adapter_root_for(conn, root)
+        case conn.class.name
+        when 'ActiveRecord::ConnectionAdapters::PostgreSQLAdapter'
+          root.join('postgresql')
+        when 'ActiveRecord::ConnectionAdapters::Mysql2Adapter',
+          'ActiveRecord::ConnectionAdapters::TrilogyAdapter'
+          root.join('mysql')
+        else
+          raise Arfi::Errors::AdapterNotSupported
+        end
       end
     end
   end

data/lib/arfi/tasks/db.rake

@@ -1,38 +1,80 @@
 # frozen_string_literal: true
 
 require 'rake'
+require 'active_record'
 require 'arfi/sql_function_loader'
 
 namespace :_db do
   task :arfi_enhance do
-    Arfi::SqlFunctionLoader.load!
+    # For non-suffixed tasks (db:migrate / db:prepare etc),
+    # loader will populate all DBs when multi-db and task_name is nil.
+    Arfi::SqlFunctionLoader.load!(verbose: true)
   end
 end
 
 Rake::Task.define_task(:environment) unless Rake::Task.task_defined?(:environment)
 
-# Enhancing single db tasks
-%w[db:migrate db:schema:load db:setup].each do |task|
+# Enhance common single-db tasks (and multi-db "default" tasks)
+%w[
+  db:migrate
+  db:schema:load
+  db:setup
+  db:prepare
+  db:test:prepare
+].each do |task|
   Rake::Task[task].enhance(['_db:arfi_enhance']) if Rake::Task.task_defined?(task)
 end
 
-# We remove user defined keys in database.yml and use only the ones about RDBMS connections.
-# Here we try to find tasks like db:migrate:your_db_name and enhance them as well as tasks for single db connections.
-rdbms_configs = Rails.configuration.database_configuration[Rails.env].select { |_k, v| v.is_a?(Hash) }.keys
-possible_tasks = Rake::Task.tasks.select do |task|
-  task.name.match?(/^(db:migrate:|db:schema:load:|db:setup:|db:prepare:|db:test:prepare:)(\w+)$/)
+def define_arfi_enhance_for(task, db_name)
+  Rake::Task.define_task("_db:arfi_enhance:#{task.name}") do
+    run_arfi_loader_for(task, arfi_db_config_for(db_name))
+  end
+  task.enhance(["_db:arfi_enhance:#{task.name}"])
+end
+
+def run_arfi_loader_for(task, config_hash)
+  if config_hash
+    run_with_connection_switch(task, config_hash)
+  else
+    run_arfi_loader(task)
+  end
+end
+
+def run_with_connection_switch(task, config_hash)
+  original = current_db_config
+  ActiveRecord::Base.establish_connection(config_hash)
+  run_arfi_loader(task)
+ensure
+  ActiveRecord::Base.establish_connection(original) if original
+end
+
+def current_db_config
+  ActiveRecord::Base.connection_db_config
+rescue StandardError
+  nil
+end
+
+def run_arfi_loader(task)
+  Arfi::SqlFunctionLoader.load!(
+    task_name: task.name,
+    connection: ActiveRecord::Base.connection,
+    verbose: true
+  )
 end
-possible_tasks = possible_tasks.select do |task|
-  rdbms_configs.any? { |n| task.name.include?(n) }
+
+def arfi_db_config_for(name)
+  return unless ActiveRecord::Base.respond_to?(:configurations) && ActiveRecord::Base.configurations
+
+  list = ActiveRecord::Base.configurations.configs_for(env_name: Rails.env, name: name)
+  cfg = Array(list).first
+  cfg&.configuration_hash
 end
 
-# In general, this is a small trick due to the fact that Rake does not allow you to view parent tasks from a task that
-# was called for enhancing. Moreover, the utility does not provide for passing parameters to the called task.
-# To get around this limitation, we dynamically create a task with the name we need, and then pass this name as
-# an argument to the method.
+# Enhance suffixed tasks: db:migrate:<db>, db:schema:load:<db>, etc.
+pattern = /^(db:migrate:|db:schema:load:|db:setup:|db:prepare:|db:test:prepare:)([^:]+)$/
+possible_tasks = Rake::Task.tasks.select { |t| t.name.match?(pattern) }
+
 possible_tasks.each do |task|
-  Rake::Task.define_task("_db:arfi_enhance:#{task.name}") do
-    Arfi::SqlFunctionLoader.load!(task_name: task.name)
-  end
-  task.enhance(["_db:arfi_enhance:#{task.name}"])
+  db_name = task.name.match(pattern)[2]
+  define_arfi_enhance_for(task, db_name)
 end

data/lib/arfi/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Arfi
-  VERSION = '0.5.0'
+  VERSION = '1.0.0'
 end

data/lib/arfi.rb

@@ -2,6 +2,8 @@
 
 require_relative 'arfi/version'
 require_relative 'arfi/errors'
+require_relative 'arfi/cli'
+require_relative 'arfi/commands/f_idx'
 require 'arfi/extensions/extensions'
 require 'rails' if defined?(Rails)