asktive_record 0.1.7 → 0.2.0

data/lib/asktive_record/query.rb

@@ -1,5 +1,8 @@
 # frozen_string_literal: true
 
+require "asktive_record/sql_sanitizer"
+require "asktive_record/log"
+
 module AsktiveRecord
   # The Query class encapsulates a natural language question, its corresponding SQL,
   # and provides methods for sanitization, execution, and generating answers using LLMs.
@@ -12,25 +15,27 @@ module AsktiveRecord
       @model_class = model_class
       @natural_question = natural_question
       @sanitized_sql = raw_sql # Initially, sanitized SQL is the same as raw SQL
+
+      AsktiveRecord::Log.info("Received question: \"#{natural_question}\"")
+      AsktiveRecord::Log.info("Generated SQL: #{raw_sql}")
     end
 
-    # Placeholder for sanitization logic
-    # In a real scenario, this would involve more sophisticated checks,
-    # potentially allowing only SELECT statements or using a whitelist of allowed SQL patterns.
+    # Sanitizes the SQL using the SqlSanitizer.
+    # Raises AsktiveRecord::SanitizationError if the query is unsafe.
+    #
+    # @param allow_only_select [Boolean] restrict to SELECT-only (default: true)
+    # @return [self] for chaining
     def sanitize!(allow_only_select: true)
-      if allow_only_select && !@sanitized_sql.strip.downcase.start_with?("select")
-        raise SanitizationError, "Query sanitization failed: Only SELECT statements are allowed by default."
-      end
-
-      # Add more sanitization rules here as needed
+      @sanitized_sql = SqlSanitizer.sanitize!(@sanitized_sql, allow_only_select: allow_only_select)
+      AsktiveRecord::Log.info("Sanitized SQL: #{@sanitized_sql}")
       self # Return self for chaining
     end
 
     def answer
       response = execute
       llm = AsktiveRecord::LlmService.new(AsktiveRecord.configuration)
-      response = response.inspect if response.respond_to?(:inspect)
-      llm.answer(@natural_question, @sanitized_sql, response)
+      response_text = response.respond_to?(:inspect) ? response.inspect : response.to_s
+      llm.answer(@natural_question, @sanitized_sql, response_text)
     end
 
     def execute
@@ -39,8 +44,15 @@ module AsktiveRecord
               "Cannot execute raw SQL. Call sanitize! first or work with sanitized_sql."
       end
 
+      # Auto-sanitize before execution if not already done
+      validate_before_execution!
+
       result = execute_query
-      extract_count_if_present(result)
+      AsktiveRecord::Log.info("Query executed successfully.")
+      AsktiveRecord::Log.debug("Query results: #{result.inspect}")
+      result
+    rescue QueryExecutionError
+      raise
     rescue StandardError => e
       raise QueryExecutionError, "Failed to execute SQL query: #{e.message}"
     end
@@ -51,27 +63,32 @@ module AsktiveRecord
 
     private
 
+    def validate_before_execution!
+      # Always validate the SQL is safe before execution
+      SqlSanitizer.sanitize!(@sanitized_sql, allow_only_select: read_only_mode?)
+    end
+
+    def read_only_mode?
+      AsktiveRecord.configuration&.read_only != false
+    end
+
     def execute_query
       if active_record_model?
-        result = model_class.find_by_sql(@sanitized_sql)
-        return result[0].count if result[0].respond_to?(:count)
-
-        result
+        AsktiveRecord::Log.info("Executing SQL via #{model_class.name}.find_by_sql")
+        model_class.find_by_sql(@sanitized_sql)
       else
         execute_raw_sql
       end
     end
 
     def active_record_model?
-      model_class.respond_to?(:find_by_sql) && model_class.table_name.present?
+      model_class.respond_to?(:find_by_sql) && model_class.respond_to?(:table_name) && model_class.table_name.present?
     end
 
     def execute_raw_sql
       return unless defined?(ActiveRecord::Base)
 
-      if ActiveRecord::Base.connection.respond_to?(:exec_query)
-        # no-op here unless you're logging or observing
-      end
+      AsktiveRecord::Log.info("Executing SQL via ActiveRecord::Base.connection")
 
       if select_query?
         ActiveRecord::Base.connection.select_all(@sanitized_sql)
@@ -81,13 +98,7 @@ module AsktiveRecord
     end
 
     def select_query?
-      @sanitized_sql.strip.downcase.start_with?("select")
-    end
-
-    def extract_count_if_present(result)
-      return result unless result.is_a?(Array) && result[0].is_a?(Hash) && result[0].key?("count")
-
-      result[0]["count"]
+      @sanitized_sql.strip.match?(/\ASELECT\b/i)
     end
   end
 end

data/lib/asktive_record/schema_loader.rb

@@ -0,0 +1,63 @@
+# frozen_string_literal: true
+
+require "asktive_record/log"
+
+module AsktiveRecord
+  # SchemaLoader provides shared schema loading functionality used by both
+  # Model::ClassMethods and Service::ClassMethods. It handles reading schema files,
+  # falling back to structure.sql, and providing clear error messages.
+  module SchemaLoader
+    # Loads the database schema content from the configured path.
+    # Falls back to db/structure.sql if the primary schema file is not found.
+    #
+    # @return [String] the schema content
+    # @raise [ConfigurationError] if no schema file is found or content is empty
+    def load_schema_content
+      path = AsktiveRecord.configuration.db_schema_path
+      return File.read(path) if File.exist?(path)
+
+      attempt_schema_fallback(path)
+    rescue SystemCallError => e
+      raise ConfigurationError, "Error reading schema file at #{path}: #{e.message}"
+    end
+
+    # Validates that the schema content is not empty.
+    #
+    # @param schema_content [String] the schema content to validate
+    # @raise [ConfigurationError] if schema content is blank
+    def ensure_schema_is_not_empty!(schema_content)
+      return unless schema_content.to_s.strip.empty?
+
+      raise ConfigurationError,
+            "Schema content is empty. Cannot proceed without database schema context."
+    end
+
+    # Validates that the LLM API key is configured.
+    #
+    # @raise [ConfigurationError] if the API key is not set
+    def validate_llm_api_key!
+      return if AsktiveRecord.configuration&.llm_api_key
+
+      raise ConfigurationError, "LLM API key is not configured for AsktiveRecord."
+    end
+
+    private
+
+    def attempt_schema_fallback(path)
+      AsktiveRecord::Log.warn(
+        "Schema file not found at #{path}. " \
+        "Run 'bundle exec rails generate asktive_record:setup' for robust schema handling."
+      )
+
+      alt_path = "db/structure.sql"
+      if File.exist?(alt_path)
+        AsktiveRecord::Log.info("Using schema from #{alt_path}")
+        return File.read(alt_path)
+      end
+
+      raise ConfigurationError,
+            "Database schema file not found at #{path}. AsktiveRecord needs schema context. " \
+            "Ensure the schema file is present or configure the correct path."
+    end
+  end
+end

data/lib/asktive_record/service.rb

@@ -1,6 +1,8 @@
 # frozen_string_literal: true
 
 require "asktive_record/llm_service"
+require "asktive_record/schema_loader"
+require "asktive_record/log"
 
 module AsktiveRecord
   module Service
@@ -9,6 +11,8 @@ module AsktiveRecord
     # It allows for more complex queries that may involve multiple tables or relationships,
     # without requiring the user to specify a target table.
     module ClassMethods
+      include AsktiveRecord::SchemaLoader
+
       def ask(natural_language_query, options = {})
         validate_llm_api_key!
         schema_content = load_schema_content
@@ -22,56 +26,6 @@ module AsktiveRecord
 
       private
 
-      def ensure_schema_is_not_empty!(schema_content)
-        return unless schema_content.to_s.strip.empty?
-
-        raise ConfigurationError,
-              "Schema content is empty. Cannot proceed without database schema context."
-      end
-
-      def validate_llm_api_key!
-        return if AsktiveRecord.configuration&.llm_api_key
-
-        raise ConfigurationError, "LLM API key is not configured for AsktiveRecord."
-      end
-
-      def load_schema_content
-        path = AsktiveRecord.configuration.db_schema_path
-        return File.read(path) if File.exist?(path)
-
-        attempt_schema_fallback(path)
-      rescue SystemCallError => e
-        raise ConfigurationError, "Error reading schema file at #{path}: #{e.message}"
-      end
-
-      def attempt_schema_fallback(path)
-        puts "Schema file not found at #{path}. Attempting to generate it. " \
-             "Run 'bundle exec asktive_record:setup' for robust schema handling."
-
-        return fallback_schema_from_rails(path) if defined?(Rails)
-
-        raise ConfigurationError,
-              "Database schema file not found at #{path}. AsktiveRecord needs schema context. " \
-              "Run in a Rails environment or ensure the schema file is present."
-      end
-
-      def fallback_schema_from_rails(path)
-        system("bin/rails db:schema:dump") unless AsktiveRecord.configuration.skip_dump_schema
-        return File.read(path) if File.exist?(path)
-
-        alt_path = "db/structure.sql"
-        return use_alternative_schema(alt_path) if File.exist?(alt_path)
-
-        raise ConfigurationError,
-              "Database schema file not found at #{path} or #{alt_path} even after attempting to dump. " \
-              "Please run asktive_record:setup or configure the correct path."
-      end
-
-      def use_alternative_schema(path)
-        puts "Using schema from #{path}"
-        File.read(path)
-      end
-
       def resolve_model(provided_model)
         return provided_model if provided_model
         return ApplicationRecord if defined?(Rails) && defined?(ApplicationRecord)

data/lib/asktive_record/sql_sanitizer.rb

@@ -0,0 +1,92 @@
+# frozen_string_literal: true
+
+module AsktiveRecord
+  # SqlSanitizer provides robust SQL sanitization to prevent injection attacks
+  # from LLM-generated queries. It validates that queries are safe SELECT statements
+  # and blocks dangerous patterns like DDL, DML, and injection techniques.
+  class SqlSanitizer
+    # SQL keywords that indicate dangerous operations
+    DANGEROUS_KEYWORDS = %w[
+      INSERT UPDATE DELETE DROP ALTER CREATE TRUNCATE
+      REPLACE MERGE GRANT REVOKE EXEC EXECUTE
+      CALL RENAME LOAD COPY INTO OUTFILE DUMPFILE
+    ].freeze
+
+    # Patterns that indicate SQL injection attempts
+    INJECTION_PATTERNS = [
+      /;\s*\S/i,                    # Semicolon followed by another statement
+      /--\s/,                       # SQL line comment
+      %r{/\*.*?\*/}m,               # SQL block comment
+      /\bUNION\b.*\bSELECT\b/i, # UNION-based injection
+      /\bINTO\s+OUTFILE\b/i,        # File write attempt
+      /\bINTO\s+DUMPFILE\b/i,       # File dump attempt
+      /\bLOAD_FILE\b/i,             # File read attempt
+      /\bBENCHMARK\b/i,             # Time-based injection
+      /\bSLEEP\b\s*\(/i,           # Time-based injection
+      /\bIF\b\s*\(/i,              # Conditional injection (IF function)
+      /\bCASE\s+WHEN\b.*\bTHEN\b.*\bWAITFOR\b/i, # Conditional time-based
+      /\bWAITFOR\s+DELAY\b/i,      # MSSQL time-based injection
+      /\bpg_sleep\b/i,             # PostgreSQL time-based injection
+      /\bDBMS_LOCK\.SLEEP\b/i      # Oracle time-based injection
+    ].freeze
+
+    class << self
+      # Sanitizes the given SQL string.
+      # Raises AsktiveRecord::SanitizationError if the query is unsafe.
+      #
+      # @param sql [String] the SQL string to sanitize
+      # @param allow_only_select [Boolean] whether to restrict to SELECT-only (default: true)
+      # @return [String] the sanitized SQL string
+      def sanitize!(sql, allow_only_select: true)
+        raise SanitizationError, "SQL query cannot be nil or empty." if sql.nil? || sql.strip.empty?
+
+        cleaned = clean_sql(sql)
+
+        validate_select_only!(cleaned) if allow_only_select
+
+        check_dangerous_keywords!(cleaned)
+        check_injection_patterns!(cleaned)
+
+        cleaned
+      end
+
+      private
+
+      # Strips trailing semicolons and extra whitespace
+      def clean_sql(sql)
+        sql.strip.gsub(/;\s*\z/, "")
+      end
+
+      # Validates the query starts with SELECT
+      def validate_select_only!(sql)
+        return if sql.strip.match?(/\ASELECT\b/i)
+
+        raise SanitizationError,
+              "Query sanitization failed: Only SELECT statements are allowed. Got: #{sql.split.first}"
+      end
+
+      # Checks for dangerous DDL/DML keywords in the query
+      def check_dangerous_keywords!(sql)
+        # Remove string literals to avoid false positives on quoted content
+        sql_without_strings = sql.gsub(/'[^']*'/, "''").gsub(/"[^"]*"/, '""')
+
+        DANGEROUS_KEYWORDS.each do |keyword|
+          next unless sql_without_strings.match?(/\b#{keyword}\b/i)
+
+          raise SanitizationError,
+                "Query sanitization failed: Dangerous SQL keyword '#{keyword}' detected."
+        end
+      end
+
+      # Checks for common SQL injection patterns
+      def check_injection_patterns!(sql)
+        INJECTION_PATTERNS.each do |pattern|
+          next unless sql.match?(pattern)
+
+          raise SanitizationError,
+                "Query sanitization failed: Potentially dangerous SQL pattern detected."
+        end
+      end
+    end
+  end
+end

data/lib/asktive_record/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module AsktiveRecord
-  VERSION = "0.1.7"
+  VERSION = "0.2.0"
 end

data/lib/asktive_record.rb

@@ -2,7 +2,10 @@
 
 require "asktive_record/version"
 require "asktive_record/error"
+require "asktive_record/log"
 require "asktive_record/configuration"
+require "asktive_record/sql_sanitizer"
+require "asktive_record/schema_loader"
 require "asktive_record/model"
 require "asktive_record/service"
 require "asktive_record/query"
@@ -13,6 +16,12 @@ module AsktiveRecord
   extend Service::ClassMethods
   class << self
     attr_accessor :configuration
+
+    # Provides access to the logger instance
+    # @return [AsktiveRecord::Log]
+    def logger
+      AsktiveRecord::Log
+    end
   end
 
   def self.configure
@@ -34,7 +43,32 @@ module AsktiveRecord
 
   # Class method to allow direct querying from AsktiveRecord module
   def self.ask(natural_language_query, options = {})
-    # Delegate to the Service module's implementation
-    Service::ClassMethods.instance_method(:ask).bind(self).call(natural_language_query, options)
+    validate_configuration!
+    schema_content = load_schema_for_module
+    llm_service = AsktiveRecord::LlmService.new(configuration)
+    target_table = options[:table_name] || "any"
+    raw_sql = llm_service.generate_sql_for_service(natural_language_query, schema_content, target_table)
+    target_model = options[:model] || self
+    AsktiveRecord::Query.new(natural_language_query, raw_sql, target_model)
+  end
+
+  # @api private
+  def self.validate_configuration!
+    return if configuration&.llm_api_key
+
+    raise ConfigurationError, "LLM API key is not configured for AsktiveRecord."
   end
+
+  # @api private
+  def self.load_schema_for_module
+    path = configuration.db_schema_path
+    raise ConfigurationError, "Schema file not found at #{path}." unless File.exist?(path)
+
+    content = File.read(path)
+    raise ConfigurationError, "Schema content is empty." if content.strip.empty?
+
+    content
+  end
+
+  private_class_method :validate_configuration!, :load_schema_for_module
 end

data/lib/generators/asktive_record/templates/asktive_record_initializer.rb

@@ -10,25 +10,48 @@ AsktiveRecord.configure do |config|
   # === LLM API Key ===
   # Set your API key for the chosen LLM provider.
   # It is strongly recommended to use environment variables for sensitive data.
-  # For example, for OpenAI:
-  # config.llm_api_key = ENV["OPENAI_API_KEY"]
-  config.llm_api_key = "YOUR_OPENAI_API_KEY_HERE"
+  # For OpenAI:
+  config.llm_api_key = ENV.fetch("OPENAI_API_KEY", nil)
 
   # === LLM Model Name ===
   # Specify the model name for the LLM provider if applicable.
-  # For OpenAI, default is "gpt-3.5-turbo". Other models like "gpt-4" can be used.
-  # config.llm_model_name = "gpt-3.5-turbo"
+  # For OpenAI, default is "gpt-4o-mini". Other models like "gpt-4o" can be used.
+  # config.llm_model_name = "gpt-4o-mini"
 
   # === Database Schema Path ===
   # Path to your Rails application's schema file (usually schema.rb or structure.sql).
   # This is used by the `asktive_record:setup` command to provide context to the LLM.
   # Default is "db/schema.rb".
   # config.db_schema_path = "db/schema.rb"
-  #
+
   # === Skip dump schema ===
   # If set to true, the schema will not be dumped when running the
   # `asktive_record:setup` command.
   # This is useful if you want to manage schema dumps manually
   # or if you are using a different schema management strategy.
   # config.skip_dump_schema = false
+
+  # === Read-Only Mode ===
+  # When true (default), only SELECT queries are allowed to execute.
+  # Set to false if you want to allow other query types (not recommended).
+  # config.read_only = true
+
+  # === LLM Temperature ===
+  # Controls randomness in LLM responses. Lower = more deterministic.
+  # Default is 0.2 (low randomness for reliable SQL generation).
+  # config.temperature = 0.2
+
+  # === LLM Max Tokens ===
+  # Maximum number of tokens in the LLM response.
+  # Default is 250.
+  # config.max_tokens = 250
+
+  # === Custom Adapter ===
+  # Provide a custom LLM adapter instead of using the built-in provider.
+  # The adapter must inherit from AsktiveRecord::Adapters::Base and implement #chat.
+  # config.adapter = MyCustomAdapter.new(api_key: ENV["MY_LLM_KEY"])
+
+  # === Custom Logger ===
+  # Set a custom logger. Defaults to Rails.logger when available.
+  # config.logger = Logger.new($stdout)
 end

data/sig/asktive_record.rbs

@@ -1,4 +1,180 @@
 module AsktiveRecord
   VERSION: String
-  # See the writing guide of rbs: https://github.com/ruby/rbs#guides
+
+  def self.configure: () { (Configuration) -> void } -> void
+  def self.configuration: () -> Configuration?
+  def self.configuration=: (Configuration?) -> void
+  def self.logger: () -> untyped
+  def self.logger=: (untyped) -> void
+  def self.ask: (String, ?Hash[Symbol, untyped]) -> Query
+  def self.validate_configuration!: () -> void
+  def self.included: (Module) -> void
+
+  class Configuration
+    attr_accessor llm_provider: Symbol
+    attr_accessor llm_api_key: String?
+    attr_accessor llm_model_name: String
+    attr_accessor db_schema_path: String
+    attr_accessor skip_dump_schema: bool
+    attr_accessor logger: untyped
+    attr_accessor read_only: bool
+    attr_accessor adapter: Adapters::Base?
+    attr_accessor temperature: Float
+    attr_accessor max_tokens: Integer
+    attr_accessor cache_enabled: bool
+
+    def initialize: () -> void
+    def resolved_adapter: () -> Adapters::Base
+
+    private
+
+    def build_adapter_from_provider: () -> Adapters::Base
+  end
+
+  class LlmService
+    attr_reader configuration: Configuration
+
+    def initialize: (Configuration) -> void
+    def upload_schema: (String) -> bool
+    def answer: (String, String, untyped) -> String?
+    def generate_sql: (String, String, String) -> String
+    def generate_sql_for_service: (String, String, ?String) -> String
+
+    private
+
+    def answer_as_human: (String, String, untyped) -> String?
+    def adapter: () -> Adapters::Base
+    def llm_options: () -> Hash[Symbol, untyped]
+    def generate_and_validate_sql: (String) -> String
+    def validate_sql_response!: (String?) -> void
+    def clean_sql: (String) -> String
+  end
+
+  class Query
+    attr_reader raw_sql: String
+    attr_reader model_class: untyped
+    attr_reader natural_question: String
+
+    def initialize: (String, String, untyped) -> void
+    def sanitize!: (?allow_only_select: bool) -> self
+    def answer: () -> String?
+    def execute: () -> untyped
+    def to_s: () -> String
+
+    private
+
+    def validate_before_execution!: () -> void
+    def read_only_mode?: () -> bool
+    def execute_query: () -> untyped
+    def execute_raw_sql: () -> untyped
+  end
+
+  class Prompt
+    PROMPT_INJECTION_PATTERNS: Array[Regexp]
+
+    def self.as_human_answerer: (String, String, untyped) -> String
+    def self.as_sql_generator: (String, String) -> String
+    def self.as_sql_generator_for_model: (String, String, String) -> String
+    def self.escape_input: (String) -> String
+  end
+
+  class SqlSanitizer
+    DANGEROUS_KEYWORDS: Array[String]
+    INJECTION_PATTERNS: Array[Regexp]
+
+    def self.sanitize!: (String?, ?allow_only_select: bool) -> String
+    def self.dangerous_keyword?: (String) -> bool
+    def self.injection_pattern?: (String) -> bool
+
+    private
+
+    def self.strip_quoted_strings: (String) -> String
+  end
+
+  module Log
+    PREFIX: String
+
+    def self.logger: () -> untyped
+    def self.logger=: (untyped) -> void
+    def self.info: (String) -> void
+    def self.debug: (String) -> void
+    def self.warn: (String) -> void
+    def self.error: (String) -> void
+
+    private
+
+    def self.default_logger: () -> untyped
+  end
+
+  module SchemaLoader
+    def load_schema_content: () -> String
+    def ensure_schema_is_not_empty!: (String) -> void
+    def validate_llm_api_key!: () -> void
+  end
+
+  module Model
+    module ClassMethods
+      include SchemaLoader
+
+      def asktive_record: () -> void
+      def ask: (String) -> Query
+    end
+  end
+
+  module Service
+    module ClassMethods
+      include SchemaLoader
+
+      def ask: (String, ?Hash[Symbol, untyped]) -> Query
+
+      private
+
+      def resolve_model: (untyped) -> untyped
+    end
+  end
+
+  module Adapters
+    class Base
+      attr_reader api_key: String
+      attr_reader model_name: String?
+
+      def initialize: (api_key: String, ?model_name: String?) -> void
+      def chat: (String, ?Hash[Symbol, untyped]) -> String?
+      def default_model_name: () -> String
+      def resolved_model_name: () -> String
+    end
+
+    class OpenAI < Base
+      DEFAULT_MODEL: String
+      DEFAULT_TEMPERATURE: Float
+      DEFAULT_MAX_TOKENS: Integer
+
+      def initialize: (api_key: String, ?model_name: String?) -> void
+      def chat: (String, ?Hash[Symbol, untyped]) -> String?
+      def default_model_name: () -> String
+
+      private
+
+      def client: () -> untyped
+    end
+  end
+
+  # Error classes
+  class Error < StandardError
+  end
+
+  class ConfigurationError < Error
+  end
+
+  class QueryGenerationError < Error
+  end
+
+  class QueryExecutionError < Error
+  end
+
+  class SanitizationError < Error
+  end
+
+  class ApiError < Error
+  end
 end