sql-chatbot-rails 1.0.0

data/lib/sql_chatbot/services/sql_executor.rb

@@ -0,0 +1,81 @@
+module SqlChatbot
+  module Services
+    module SqlExecutor
+      KEYWORD_BLOCKLIST = %w[INSERT UPDATE DELETE DROP ALTER CREATE GRANT TRUNCATE EXECUTE REVOKE COPY INTO].freeze
+      FUNCTION_BLOCKLIST = %w[pg_read_file pg_read_binary_file dblink pg_terminate_backend lo_import lo_export pg_sleep set_config current_setting].freeze
+      CATALOG_BLOCKLIST = %w[pg_shadow pg_roles pg_authid pg_user information_schema].freeze
+      AGGREGATE_PATTERN = /\b(COUNT|SUM|AVG|MIN|MAX)\s*\(/i
+
+      def self.validate_sql(sql)
+        trimmed = sql.gsub(/^\s+/, "")
+        trimmed = trimmed.gsub(/--[^\n]*/, "").gsub(/\/\*[\s\S]*?\*\//, "")
+        trimmed = trimmed.strip
+
+        # Fix missing space after SELECT
+        trimmed = trimmed.sub(/^SELECT(?=[A-Z])/i, "SELECT ")
+
+        # Single statement check
+        parts = trimmed.split(";").select { |p| p.strip.length > 0 }
+        if parts.length > 1
+          return { valid: false, reason: "Only a single statement is allowed" }
+        end
+
+        working_sql = trimmed.sub(/;\s*$/, "").strip
+
+        # Must start with SELECT
+        unless working_sql.match?(/^SELECT\b/i)
+          return { valid: false, reason: "Only SELECT queries are allowed" }
+        end
+
+        # Keyword blocklist
+        KEYWORD_BLOCKLIST.each do |keyword|
+          if working_sql.match?(/\b#{keyword}\b/i)
+            return { valid: false, reason: "Blocked keyword: #{keyword}" }
+          end
+        end
+
+        # Function blocklist
+        FUNCTION_BLOCKLIST.each do |fn|
+          if working_sql.match?(/\b#{fn}\b/i)
+            return { valid: false, reason: "Blocked function: #{fn}" }
+          end
+        end
+
+        # Catalog blocklist
+        CATALOG_BLOCKLIST.each do |catalog|
+          if working_sql.match?(/\b#{catalog}\b/i)
+            return { valid: false, reason: "Blocked system catalog: #{catalog}" }
+          end
+        end
+
+        # Auto-add LIMIT
+        has_limit = working_sql.match?(/\bLIMIT\b/i)
+        is_aggregate = AGGREGATE_PATTERN.match?(working_sql)
+        working_sql += " LIMIT 500" if !has_limit && !is_aggregate
+
+        { valid: true, sql: working_sql }
+      end
+
+      def self.execute_sql(sql)
+        connection = ActiveRecord::Base.connection
+        connection.execute("SET statement_timeout = '10s'")
+        connection.execute("BEGIN")
+        connection.execute("SET TRANSACTION READ ONLY")
+
+        result = connection.execute(sql)
+        rows = result.to_a
+        columns = rows.first&.keys || []
+
+        connection.execute("COMMIT")
+
+        { columns: columns, rows: rows, row_count: rows.length }
+      rescue => e
+        connection.execute("ROLLBACK") rescue nil
+        raise e
+      ensure
+        # Reset statement_timeout on shared AR connection
+        connection.execute("SET statement_timeout = '0'") rescue nil
+      end
+    end
+  end
+end

data/lib/sql_chatbot/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module SqlChatbot
+  VERSION = "1.0.0"
+end

data/lib/sql_chatbot_rails.rb

@@ -0,0 +1,91 @@
+# frozen_string_literal: true
+
+require "sql_chatbot/version"
+require "sql_chatbot/configuration"
+require "sql_chatbot/llm/client"
+require "sql_chatbot/prompts/classify"
+require "sql_chatbot/prompts/generate_sql"
+require "sql_chatbot/prompts/answer"
+require "sql_chatbot/services/sql_executor"
+require "sql_chatbot/services/schema_service"
+require "sql_chatbot/services/code_indexer"
+require "sql_chatbot/services/orchestrator"
+require "sql_chatbot/services/model_introspector"
+require "sql_chatbot/services/route_introspector"
+require "sql_chatbot/auth/jwt"
+require "sql_chatbot/auth/cors"
+require "sql_chatbot/engine" if defined?(Rails)
+
+module SqlChatbot
+  class << self
+    attr_accessor :config, :schema_service, :code_indexer, :orchestrator, :registry
+
+    def configure
+      self.config ||= Configuration.new
+      yield(config) if block_given?
+    end
+
+    def reset!
+      self.config = Configuration.new
+      @schema_service = nil
+      @code_indexer = nil
+      @orchestrator = nil
+      @registry = nil
+      @initialized = false
+      @init_mutex = Mutex.new
+    end
+
+    def ensure_initialized!
+      return if @initialized
+      @init_mutex ||= Mutex.new
+      @init_mutex.synchronize do
+        return if @initialized
+        cfg = config || Configuration.new
+
+        @schema_service = Services::SchemaService.new
+        @schema_service.discover
+
+        # Introspect Rails models for enums, non-standard FKs, and soft delete gems
+        introspector = Services::ModelIntrospector.new
+        introspection = introspector.introspect
+
+        # Apply soft delete annotations conditionally (gem-based vs enum-based)
+        @schema_service.apply_soft_delete_annotations(
+          soft_delete_tables: introspection.soft_delete_tables,
+          enum_soft_delete_tables: introspection.enum_soft_delete_tables,
+        )
+
+        # Inject model annotations (enums, FKs)
+        @schema_service.append_model_annotations(introspection.annotations)
+
+        # Move lookup values from referenced tables to FK columns
+        @schema_service.relocate_lookup_annotations
+
+        # Introspect Rails routes for navigation context
+        route_introspector = Services::RouteIntrospector.new
+        route_data = route_introspector.introspect
+
+        @code_indexer = Services::CodeIndexer.new
+        @code_indexer.index(cfg.code_paths)
+
+        llm_client = LLM::Client.new(
+          api_key: cfg.resolved_api_key,
+          base_url: cfg.resolved_base_url,
+          model: cfg.resolved_model,
+        )
+
+        @orchestrator = Services::Orchestrator.new(
+          llm_client: llm_client,
+          schema_service: @schema_service,
+          code_indexer: @code_indexer,
+          route_introspector_data: route_data,
+        )
+
+        @initialized = true
+      end
+    rescue => e
+      @initialized = false
+      raise e
+    end
+  end
+end