code_to_query 0.1.0

data/lib/code_to_query/planner.rb

@@ -0,0 +1,289 @@
+# frozen_string_literal: true
+
+module CodeToQuery
+  class Planner
+    def initialize(config)
+      @config = config
+    end
+
+    def plan(prompt:, schema:, allow_tables:)
+      schema ||= load_context_pack
+      allow_tables = Array(allow_tables).compact
+
+      attempt = 0
+      last_error = nil
+      feedback = nil
+      provider = build_provider
+      max_attempts = @config.planner_max_attempts || 1
+
+      while attempt < max_attempts
+        attempt += 1
+        begin
+          intent = provider.extract_intent(
+            prompt: build_prompt_with_feedback(prompt, feedback),
+            schema: schema,
+            allow_tables: allow_tables
+          )
+
+          # Optional schema strictness pass: drop filters referencing unknown columns
+          if @config.planner_feedback_mode.to_s == 'schema_strict'
+            intent = strip_unknown_columns(intent, schema)
+          end
+
+          # Expose provider metrics if available
+          if provider.respond_to?(:metrics) && provider.metrics.is_a?(Hash)
+            intent = intent.merge('_metrics' => provider.metrics)
+          end
+
+          # Heuristic backfill of missing params from the prompt (IDs and enum-like labels)
+          intent = backfill_params_from_prompt(prompt, intent, schema)
+
+          return intent
+        rescue StandardError => e
+          last_error = e
+          feedback = generate_feedback(e)
+          @config.logger.warn("[code_to_query] Planning attempt #{attempt} failed: #{e.message}")
+        end
+      end
+
+      @config.logger.warn("[code_to_query] Query planning failed after #{max_attempts} attempts: #{last_error&.message}")
+      fallback_intent(allow_tables)
+    rescue StandardError => e
+      @config.logger.warn("[code_to_query] Query planning failed: #{e.message}")
+      fallback_intent(allow_tables)
+    end
+
+    private
+
+    def build_provider
+      case @config.provider
+      when :openai
+        Providers::OpenAI.new(@config)
+      when :local
+        Providers::Local.new(@config)
+      else
+        detect_best_provider
+      end
+    end
+
+    def detect_best_provider
+      if @config.openai_api_key.present? && !@config.stub_llm
+        Providers::OpenAI.new(@config)
+      else
+        Providers::Local.new(@config)
+      end
+    end
+
+    def load_context_pack
+      path = @config.context_pack_path
+      unless File.exist?(path)
+        begin
+          # Attempt to bootstrap a context pack automatically if missing
+          if defined?(CodeToQuery::Context::Builder)
+            CodeToQuery::Context::Builder.new(@config).bootstrap!
+          end
+        rescue StandardError => e
+          @config.logger.warn("[code_to_query] Auto-bootstrap of context pack failed: #{e.message}")
+        end
+      end
+
+      return {} unless File.exist?(path)
+
+      JSON.parse(File.read(path))
+    rescue StandardError => e
+      @config.logger.warn("[code_to_query] Failed to load context pack: #{e.message}")
+      {}
+    end
+
+    def fallback_intent(allow_tables)
+      intent = {
+        'type' => 'select',
+        'table' => Array(allow_tables).compact.first || 'main_table',
+        'columns' => ['*'],
+        'filters' => [],
+        'order' => [],
+        'params' => {}
+      }
+      intent['limit'] = @config.default_limit if @config.default_limit
+      intent
+    end
+
+    def build_prompt_with_feedback(prompt, feedback)
+      return prompt if feedback.to_s.strip.empty?
+
+      "#{prompt}\n\nConstraints/feedback: #{feedback}"
+    end
+
+    def generate_feedback(error)
+      return '' unless @config.planner_feedback_mode && @config.planner_feedback_mode != :none
+
+      case error
+      when ArgumentError
+        'Ensure all columns and tables exist in the provided schema and avoid unknown fields. For EXISTS, provide related_table and fk_column.'
+      else
+        'Avoid inventing tables/columns; map business terms to schema; use EXISTS for relationships; set column to "id" in EXISTS filters.'
+      end
+    end
+
+    def strip_unknown_columns(intent, schema)
+      return intent unless intent.is_a?(Hash)
+
+      tables = Array(schema.dig('schema', 'tables') || schema['tables'] || [])
+      table_name = intent['table']
+      table = tables.find { |t| (t['name'] || t[:name]).to_s == table_name.to_s }
+      return intent unless table
+
+      columns = Array(table['columns'] || table[:columns]).map { |c| c['name'] || c[:name] }.compact
+
+      if intent['filters'].is_a?(Array)
+        intent['filters'] = intent['filters'].select do |f|
+          op = f['op'].to_s
+          next true if %w[exists not_exists].include?(op)
+
+          columns.include?(f['column'].to_s)
+        end
+      end
+
+      if intent['order'].is_a?(Array)
+        intent['order'] = intent['order'].select { |o| columns.include?(o['column'].to_s) }
+      end
+
+      if intent['columns'].is_a?(Array)
+        intent['columns'] = intent['columns'].select { |c| c == '*' || columns.include?(c.to_s) }
+      end
+
+      intent
+    rescue StandardError
+      intent
+    end
+
+    # Fill missing intent params with simple values parsed from the prompt.
+    # Example: "student with id 17963" -> maps first integer to first missing *_id param.
+    def backfill_params_from_prompt(prompt, intent, schema)
+      return intent unless prompt.is_a?(String)
+      return intent unless intent.is_a?(Hash)
+
+      numbers = begin
+        prompt.scan(/\b\d+\b/).map { |n| Integer(n) }
+      rescue StandardError
+        []
+      end
+      return intent if numbers.empty?
+
+      params = (intent['params'] || {}).dup
+      remaining = numbers.dup
+
+      # Helper to assign a value for a filter if missing
+      assign_for = lambda do |param_key|
+        return if param_key.nil? || param_key.to_s.strip.empty?
+        return if params.key?(param_key.to_s)
+
+        value = remaining.shift
+        params[param_key.to_s] = value if value
+      end
+
+      # Main table filters first, prioritize *_id columns
+      Array(intent['filters']).each do |f|
+        op = f['op'].to_s
+        next unless ['=', 'like', 'ilike', 'in', 'between'].include?(op) || %w[exists not_exists].include?(op)
+
+        if %w[exists not_exists].include?(op)
+          Array(f['related_filters']).each do |rf|
+            pkey = rf['param'] || rf['column']
+            if rf['column'].to_s.end_with?('_id')
+              assign_for.call(pkey)
+            end
+          end
+        else
+          pkey = f['param'] || f['column']
+          if f['column'].to_s.end_with?('_id')
+            assign_for.call(pkey)
+          end
+        end
+      end
+
+      # Enum-like backfill using model enums from schema (if available)
+      tokens = prompt.scan(/[a-z0-9_]+/i).map(&:downcase)
+      enums_by_table = extract_enums_index(schema)
+
+      # Main table columns
+      table_name = intent['table']
+      if table_name && enums_by_table[table_name]
+        Array(intent['filters']).each do |f|
+          next unless f.is_a?(Hash)
+          next unless f['op'].to_s == '='
+
+          col = f['column']
+          next unless col
+
+          pkey = (f['param'] || col).to_s
+          next if params.key?(pkey)
+
+          mapping = enums_by_table[table_name][col]
+          next unless mapping.is_a?(Hash) && mapping.any?
+
+          label = tokens.find { |tk| mapping.key?(tk) }
+          params[pkey] = label if label
+        end
+      end
+
+      # Related table columns in EXISTS
+      Array(intent['filters']).each do |f|
+        next unless f.is_a?(Hash) && %w[exists not_exists].include?(f['op'].to_s)
+
+        rtable = f['related_table']
+        next unless rtable && enums_by_table[rtable]
+
+        Array(f['related_filters']).each do |rf|
+          next unless rf.is_a?(Hash) && rf['op'].to_s == '='
+
+          rcol = rf['column']
+          next unless rcol
+
+          rpkey = (rf['param'] || rcol).to_s
+          next if params.key?(rpkey)
+
+          rmapping = enums_by_table[rtable][rcol]
+          next unless rmapping.is_a?(Hash) && rmapping.any?
+
+          rlabel = tokens.find { |tk| rmapping.key?(tk) }
+          params[rpkey] = rlabel if rlabel
+        end
+      end
+
+      intent.merge('params' => params)
+    rescue StandardError
+      intent
+    end
+
+    def extract_enums_index(schema)
+      index = Hash.new { |h, k| h[k] = {} }
+      return index unless schema.is_a?(Hash)
+
+      models = schema.dig('models', 'models')
+      return index unless models.is_a?(Hash)
+
+      models.each_value do |meta|
+        next unless meta.is_a?(Hash)
+
+        table = meta['table_name'] || meta[:table_name]
+        next unless table
+
+        enums = meta['enums'] || meta[:enums] || {}
+        next unless enums.is_a?(Hash)
+
+        enums.each do |col, mapping|
+          # Normalize mapping keys to downcased strings
+          next unless mapping.is_a?(Hash)
+
+          norm = {}
+          mapping.each { |k, v| norm[k.to_s.downcase] = v }
+          index[table.to_s][col.to_s] = norm
+        end
+      end
+      index
+    rescue StandardError
+      {}
+    end
+  end
+end

data/lib/code_to_query/policies/pundit_adapter.rb

@@ -0,0 +1,71 @@
+# frozen_string_literal: true
+
+begin
+  require 'pundit'
+rescue LoadError
+end
+
+module CodeToQuery
+  module Policies
+    class PunditAdapter
+      def call(current_user, table:, intent: nil) # rubocop:disable Lint/UnusedMethodArgument
+        return {} unless defined?(Pundit)
+
+        info = {
+          enforced_predicates: inferred_tenant_predicates(current_user, table),
+          allowed_tables: [],
+          allowed_columns: {}
+        }
+
+        model = infer_model_for_table(table)
+        if model
+          begin
+            Pundit.policy_scope!(current_user, model)
+            info[:allowed_tables] << table
+          rescue StandardError
+          end
+
+          policy = Pundit.policy(current_user, model)
+          if policy.respond_to?(:permitted_columns)
+            cols = Array(policy.permitted_columns).map(&:to_s)
+            info[:allowed_columns][table] = cols if cols.any?
+          end
+        end
+
+        info
+      end
+
+      private
+
+      def infer_model_for_table(table_name)
+        return nil unless defined?(ActiveRecord::Base)
+        return nil unless table_name
+
+        candidates = [
+          table_name.to_s.singularize.camelize,
+          table_name.to_s.camelize,
+          table_name.to_s.singularize.camelize.gsub(/s$/, '')
+        ]
+
+        candidates.each do |klass|
+          k = klass.constantize
+          return k if k < ActiveRecord::Base && k.table_name == table_name
+        rescue NameError
+          next
+        end
+        nil
+      end
+
+      def inferred_tenant_predicates(current_user, _table)
+        return {} unless current_user
+        return {} unless current_user.respond_to?(:company_id)
+
+        if current_user.company_id
+          { company_id: current_user.company_id }
+        else
+          {}
+        end
+      end
+    end
+  end
+end

data/lib/code_to_query/providers/base.rb

@@ -0,0 +1,173 @@
+# frozen_string_literal: true
+
+module CodeToQuery
+  module Providers
+    class Base
+      attr_reader :metrics
+
+      def initialize(config)
+        @config = config
+        @metrics = {}
+      end
+
+      def extract_intent(prompt:, schema:, allow_tables:)
+        raise NotImplementedError, 'Subclasses must implement #extract_intent'
+      end
+
+      protected
+
+      # rubocop:disable Metrics/BlockNesting
+      def build_system_context(schema, allow_tables)
+        # Support passing either a raw schema hash ({'tables'=>[...]})
+        # or a full context pack ({'schema'=>{...}, 'models'=>{...}, 'glossary'=>{...}})
+        raw_schema = schema
+        model_defs = nil
+        scopes_map = nil
+        glossary = {}
+        if schema.is_a?(Hash) && schema['schema'].is_a?(Hash)
+          raw_schema = schema['schema']
+          if schema['models'].is_a?(Hash)
+            model_defs = schema['models']['models']
+            scopes_map = schema['models']['scopes']
+          end
+          glossary = schema['glossary'] || {}
+        end
+
+        available_tables = if allow_tables.any?
+                             allow_tables
+                           elsif raw_schema.is_a?(Hash) && raw_schema['tables']
+                             raw_schema['tables'].map { |t| t['name'] || t[:name] }.compact
+                           else
+                             []
+                           end
+
+        schema_info = if raw_schema.is_a?(Hash) && raw_schema['tables']
+                        raw_schema['tables'].map do |table|
+                          table_name = table['name'] || table[:name]
+                          cols = Array(table['columns'] || table[:columns])
+                          col_names = cols.map { |c| c['name'] || c[:name] }.compact
+                          fks = col_names.select { |c| c.end_with?('_id') }
+                          col_summary = col_names.take(10).join(', ')
+                          fk_summary = fks.any? ? " | fks: #{fks.join(', ')}" : ''
+
+                          scope_summary = ''
+                          if model_defs && scopes_map
+                            # find model for this table
+                            model_name, _def = model_defs.find { |_mn, md| (md['table_name'] || md[:table_name]) == table_name }
+                            if model_name
+                              scopes_for_model = scopes_map[model_name] || scopes_map[model_name.to_sym]
+                              if scopes_for_model.is_a?(Hash) && scopes_for_model.any?
+                                # Include compact where summaries when available
+                                pairs = scopes_for_model.to_a.take(4).map do |(sname, meta)|
+                                  w = meta['where'] || meta[:where]
+                                  w ? "#{sname}: #{w}" : sname.to_s
+                                end
+                                scope_summary = pairs.any? ? " | scopes: #{pairs.join('; ')}" : ''
+                              end
+                              # include enum mapping if present
+                              enums = model_defs.dig(model_name, 'enums') || model_defs.dig(model_name.to_sym, :enums) || {}
+                              if enums.is_a?(Hash) && enums.any?
+                                enum_pairs = enums.to_a.take(3).map do |(col, mapping)|
+                                  sample = mapping.is_a?(Hash) ? mapping.to_a.take(2).map { |k, v| "#{k}=#{v}" }.join(', ') : ''
+                                  sample.empty? ? col.to_s : "#{col}(#{sample})"
+                                end
+                                scope_summary += enum_pairs.any? ? " | enums: #{enum_pairs.join('; ')}" : ''
+                              end
+                            end
+                          end
+
+                          "#{table_name}: #{col_summary}#{fk_summary}#{scope_summary}"
+                        end.join("\n")
+                      else
+                        'No schema information available'
+                      end
+
+        # Include glossary information for better business understanding
+        glossary_info = if glossary.any?
+                          "\nBusiness Glossary (user terms -> database concepts):\n" \
+                            "#{glossary.map { |term, synonyms| "#{term}: #{Array(synonyms).join(', ')}" }.join("\n")}"
+                        else
+                          ''
+                        end
+
+        {
+          available_tables: available_tables,
+          schema_info: schema_info + glossary_info,
+          constraints: build_constraints_info
+        }
+      end
+      # rubocop:enable Metrics/BlockNesting
+
+      def build_constraints_info
+        base_constraints = [
+          'Only generate SELECT queries',
+          'All values must be parameterized',
+          'Use standard SQL operators: =, !=, <>, >, <, >=, <=, between, in, like, ilike',
+          'Prefer EXISTS/NOT EXISTS when expressing presence/absence of related rows',
+          'DISTINCT is allowed for unique result sets',
+          'DISTINCT ON (columns) is supported for PostgreSQL',
+          "For 'top N' queries, use ORDER BY with LIMIT"
+        ]
+
+        limit_constraints = build_limit_constraints
+        aggregation_constraints = build_aggregation_constraints
+
+        base_constraints + limit_constraints + aggregation_constraints
+      end
+
+      def build_limit_constraints
+        constraints = []
+
+        if @config.default_limit
+          constraints << "Default LIMIT: #{@config.default_limit} for SELECT queries"
+        end
+
+        constraints << if @config.count_limit
+                         "COUNT queries limited to #{@config.count_limit} rows"
+                       else
+                         'COUNT queries have no automatic LIMIT'
+                       end
+
+        constraints << if @config.aggregation_limit
+                         "Aggregation queries limited to #{@config.aggregation_limit} rows"
+                       else
+                         'Aggregation queries (SUM, AVG, MAX, MIN) have no automatic LIMIT'
+                       end
+
+        if @config.distinct_limit
+          constraints << "DISTINCT queries limited to #{@config.distinct_limit} rows"
+        end
+
+        if @config.exists_limit
+          constraints << "EXISTS checks automatically use LIMIT #{@config.exists_limit}"
+        end
+
+        constraints
+      end
+
+      def build_aggregation_constraints
+        [
+          'COUNT(*) and COUNT(column) are supported',
+          'SUM, AVG, MAX, MIN aggregations are supported',
+          'GROUP BY is supported for aggregations',
+          'Multiple aggregations can be combined in a single query'
+        ]
+      end
+
+      def validate_and_enhance_intent(intent, allow_tables)
+        intent['type'] ||= 'select'
+        intent['columns'] ||= ['*']
+        intent['filters'] ||= []
+        intent['order'] ||= []
+        intent['limit'] ||= @config.default_limit if @config.default_limit
+        intent['params'] ||= {}
+
+        if allow_tables.any? && !allow_tables.include?(intent['table'])
+          raise ArgumentError, "Table '#{intent['table']}' not in allowlist: #{allow_tables.join(', ')}"
+        end
+
+        intent
+      end
+    end
+  end
+end

data/lib/code_to_query/providers/local.rb

@@ -0,0 +1,84 @@
+# frozen_string_literal: true
+
+module CodeToQuery
+  module Providers
+    class Local < Base
+      def extract_intent(prompt:, schema:, allow_tables:)
+        @prompt = prompt.to_s.strip
+        @schema = schema || {}
+        @allow_tables = Array(allow_tables).compact
+        started_at = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+
+        table = select_table
+        table ||= 'main_table' # Back-compat default expected by specs
+
+        intent = {
+          'type' => 'select',
+          'table' => table,
+          'columns' => ['*'],
+          'filters' => [],
+          'order' => [],
+          'params' => {}
+        }
+
+        # Only set limit when configured (avoid nil which fails validation)
+        intent['limit'] = @config.default_limit if @config.default_limit
+
+        result = validate_and_enhance_intent(intent, allow_tables)
+
+        # Lightweight metrics: elapsed and estimated tokens from prompt + schema table names
+        elapsed = Process.clock_gettime(Process::CLOCK_MONOTONIC) - started_at
+        prompt_blob = build_prompt_blob(@prompt, @schema)
+        est = estimate_tokens(prompt_blob)
+        @metrics[:prompt_tokens] = est
+        @metrics[:completion_tokens] = 0
+        @metrics[:total_tokens] = est
+        @metrics[:elapsed_s] = elapsed
+
+        result
+      end
+
+      private
+
+      def build_prompt_blob(prompt, schema)
+        tables = Array(schema['tables'] || schema.dig('schema', 'tables') || [])
+        table_names = tables.map { |t| t['name'] || t[:name] }.compact.join(',')
+        [prompt.to_s, table_names].join("\n")
+      end
+
+      def estimate_tokens(text)
+        (text.to_s.length / 4.0).ceil
+      end
+
+      def select_table
+        return @allow_tables.first if @allow_tables.any?
+
+        tables = extract_schema_tables
+        return tables.first[:name] if tables.any?
+
+        nil
+      end
+
+      def extract_schema_tables
+        return [] unless @schema.is_a?(Hash)
+
+        raw_tables = if @schema['tables'].is_a?(Array)
+                       @schema['tables']
+                     elsif @schema['schema'].is_a?(Hash) && @schema['schema']['tables'].is_a?(Array)
+                       @schema['schema']['tables']
+                     else
+                       []
+                     end
+
+        Array(raw_tables).map do |table|
+          next unless table.is_a?(Hash)
+
+          {
+            name: table['name'] || table[:name],
+            columns: Array(table['columns'] || table[:columns])
+          }
+        end.compact
+      end
+    end
+  end
+end