<commit_before>module Dexter
  class Indexer
    attr_reader :client

    def initialize(client)
      @client = client

      select_all("SET client_min_messages = warning")
      select_all("CREATE EXTENSION IF NOT EXISTS hypopg")
    end

    def process_queries(queries)
      # narrow down queries and tables
      tables, queries = narrow_queries(queries)
      return [] if tables.empty?

      # get ready for hypothetical indexes
      select_all("SELECT hypopg_reset()")

      # ensure tables have recently been analyzed
      analyze_tables(tables)

      # get initial plans
      initial_plans = {}
      queries.each do |query|
        begin
          initial_plans[query] = plan(query)
        rescue PG::Error
          # do nothing
        end
      end
      # TODO log skipping
      queries.select! { |q| initial_plans[q] }

      # get existing indexes
      index_set = Set.new
      indexes(tables).each do |index|
        # TODO make sure btree
        index_set << [index["table"], index["columns"]]
      end

      # create hypothetical indexes
      candidates = {}
      columns(tables).each do |col|
        unless index_set.include?([col[:table], [col[:column]]])
          candidates[col] = select_all("SELECT * FROM hypopg_create_index('CREATE INDEX ON #{col[:table]} (#{[col[:column]].join(", ")})');").first["indexname"]
        end
      end

      queries_by_index = {}

      if client.options[:log_level].start_with?("debug")
        # TODO don't generate fingerprints again
        fingerprints = {}
        queries.each do |query|
          fingerprints[query] = PgQuery.fingerprint(query)
        end
      end

      new_indexes = []
      queries.each do |query|
        starting_cost = initial_plans[query]["Total Cost"]
        plan2 = plan(query)
        cost2 = plan2["Total Cost"]
        best_indexes = []
        found_indexes = []

        candidates.each do |col, index_name|
          if plan2.inspect.include?(index_name)
            best_index = {
              table: col[:table],
              columns: [col[:column]]
            }
            found_indexes << best_index
            if cost2 < starting_cost * 0.5
              best_indexes << best_index
              (queries_by_index[best_index] ||= []) << {
                starting_cost: starting_cost,
                final_cost: cost2,
                query: query
              }
            end
          end
        end

        new_indexes.concat(best_indexes)

        if client.options[:log_level] == "debug2"
          log "Processed #{fingerprints[query]}"
          log "Cost: #{starting_cost} -> #{cost2}"

          index_str = found_indexes.any? ? found_indexes.map { |i| "#{i[:table]} (#{i[:columns].join(", ")})" }.join(", ") : "None"
          log "Indexes: #{index_str}"

          if found_indexes != best_indexes
            log "Need 50% cost savings to suggest index"
          end

          puts
          puts query
          puts
        end
      end

      new_indexes = new_indexes.uniq.sort_by(&:to_a)

      # create indexes
      if new_indexes.any?
        new_indexes.each do |index|
          index[:queries] = queries_by_index[index]

          log "Index found: #{index[:table]} (#{index[:columns].join(", ")})"

          if client.options[:log_level].start_with?("debug")
            index[:queries].sort_by { |q| fingerprints[q[:query]] }.each do |query|
              log "Query #{fingerprints[query[:query]]} (Cost: #{query[:starting_cost]} -> #{query[:final_cost]})"
              puts
              puts query[:query]
              puts
            end
          end
        end

        new_indexes.each do |index|
          statement = "CREATE INDEX CONCURRENTLY ON #{index[:table]} (#{index[:columns].join(", ")})"
          if client.options[:create]
            log "Creating index: #{statement}"
            started_at = Time.now
            select_all(statement)
            log "Index created: #{((Time.now - started_at) * 1000).to_i} ms"
          end
        end
      else
        log "No indexes found"
      end

      new_indexes
    end

    def conn
      @conn ||= begin
        uri = URI.parse(client.arguments[0])
        config = {
          host: uri.host,
          port: uri.port,
          dbname: uri.path.sub(/\A\//, ""),
          user: uri.user,
          password: <PASSWORD>,
          connect_timeout: 3
        }.reject { |_, value| value.to_s.empty? }
        PG::Connection.new(config)
      end
    rescue PG::ConnectionBad
      abort "Bad database url"
    end

    def select_all(query)
      conn.exec(query).to_a
    end

    def plan(query)
      JSON.parse(select_all("EXPLAIN (FORMAT JSON) #{query}").first["QUERY PLAN"]).first["Plan"]
    end

    def narrow_queries(queries)
      result = select_all <<-SQL
        SELECT
          table_name
        FROM
          information_schema.tables
        WHERE
          table_catalog = current_database() AND
          table_schema NOT IN ('pg_catalog', 'information_schema')
      SQL
      possible_tables = Set.new(result.map { |r| r["table_name"] })

      query_tables = {}
      queries.each do |query|
        query_tables[query] = PgQuery.parse(query).tables
      end

      tables = queries.flat_map { |q| query_tables[q] }.uniq.select { |t| possible_tables.include?(t) }

      new_queries = queries.select { |q| query_tables[q].all? { |t| possible_tables.include?(t) } }

      if client.options[:log_level] == "debug2"
        (queries - new_queries).each do |query|
          log "Processed #{PgQuery.fingerprint(query)}"
          if query_tables[query].empty?
            log "No tables"
          else
            log "Some tables missing in current database"
          end
          puts
          puts query<commit_msg>Make sure all queries are logged<commit_after>module Dexter
  class Indexer
    attr_reader :client

    def initialize(client)
      @client = client

      select_all("SET client_min_messages = warning")
      select_all("CREATE EXTENSION IF NOT EXISTS hypopg")
    end

    def process_queries(queries)
      # narrow down queries and tables
      tables, queries = narrow_queries(queries)
      return [] if tables.empty?

      # get ready for hypothetical indexes
      select_all("SELECT hypopg_reset()")

      # ensure tables have recently been analyzed
      analyze_tables(tables)

      # get initial plans
      initial_plans = {}
      queries.each do |query|
        begin
          initial_plans[query] = plan(query)
        rescue PG::Error
          # do nothing
        end
      end

      # get existing indexes
      index_set = Set.new
      indexes(tables).each do |index|
        # TODO make sure btree
        index_set << [index["table"], index["columns"]]
      end

      # create hypothetical indexes
      candidates = {}
      columns(tables).each do |col|
        unless index_set.include?([col[:table], [col[:column]]])
          candidates[col] = select_all("SELECT * FROM hypopg_create_index('CREATE INDEX ON #{col[:table]} (#{[col[:column]].join(", ")})');").first["indexname"]
        end
      end

      queries_by_index = {}

      if client.options[:log_level].start_with?("debug")
        # TODO don't generate fingerprints again
        fingerprints = {}
        queries.each do |query|
          fingerprints[query] = PgQuery.fingerprint(query)
        end
      end

      new_indexes = []
      queries.each do |query|
        if initial_plans[query]
          starting_cost = initial_plans[query]["Total Cost"]
          plan2 = plan(query)
          cost2 = plan2["Total Cost"]
          best_indexes = []
          found_indexes = []

          candidates.each do |col, index_name|
            if plan2.inspect.include?(index_name)
              best_index = {
                table: col[:table],
                columns: [col[:column]]
              }
              found_indexes << best_index
              if cost2 < starting_cost * 0.5
                best_indexes << best_index
                (queries_by_index[best_index] ||= []) << {
                  starting_cost: starting_cost,
                  final_cost: cost2,
                  query: query
                }
              end
            end
          end

          new_indexes.concat(best_indexes)
        end

        if client.options[:log_level] == "debug2"
          log "Processed #{fingerprints[query]}"
          if initial_plans[query]
            log "Cost: #{starting_cost} -> #{cost2}"

            index_str = found_indexes.any? ? found_indexes.map { |i| "#{i[:table]} (#{i[:columns].join(", ")})" }.join(", ") : "None"
            log "Indexes: #{index_str}"

            if found_indexes != best_indexes
              log "Need 50% cost savings to suggest index"
            end
          else
            log "Could not run explain"
          end

          puts
          puts query
          puts
        end
      end

      new_indexes = new_indexes.uniq.sort_by(&:to_a)

      # create indexes
      if new_indexes.any?
        new_indexes.each do |index|
          index[:queries] = queries_by_index[index]

          log "Index found: #{index[:table]} (#{index[:columns].join(", ")})"

          if client.options[:log_level].start_with?("debug")
            index[:queries].sort_by { |q| fingerprints[q[:query]] }.each do |query|
              log "Query #{fingerprints[query[:query]]} (Cost: #{query[:starting_cost]} -> #{query[:final_cost]})"
              puts
              puts query[:query]
              puts
            end
          end
        end

        new_indexes.each do |index|
          statement = "CREATE INDEX CONCURRENTLY ON #{index[:table]} (#{index[:columns].join(", ")})"
          if client.options[:create]
            log "Creating index: #{statement}"
            started_at = Time.now
            select_all(statement)
            log "Index created: #{((Time.now - started_at) * 1000).to_i} ms"
          end
        end
      else
        log "No indexes found"
      end

      new_indexes
    end

    def conn
      @conn ||= begin
        uri = URI.parse(client.arguments[0])
        config = {
          host: uri.host,
          port: uri.port,
          dbname: uri.path.sub(/\A\//, ""),
          user: uri.user,
          password: <PASSWORD>,
          connect_timeout: 3
        }.reject { |_, value| value.to_s.empty? }
        PG::Connection.new(config)
      end
    rescue PG::ConnectionBad
      abort "Bad database url"
    end

    def select_all(query)
      conn.exec(query).to_a
    end

    def plan(query)
      JSON.parse(select_all("EXPLAIN (FORMAT JSON) #{query}").first["QUERY PLAN"]).first["Plan"]
    end

    def narrow_queries(queries)
      result = select_all <<-SQL
        SELECT
          table_name
        FROM
          information_schema.tables
        WHERE
          table_catalog = current_database() AND
          table_schema NOT IN ('pg_catalog', 'information_schema')
      SQL
      possible_tables = Set.new(result.map { |r| r["table_name"] })

      query_tables = {}
      queries.each do |query|
        query_tables[query] = PgQuery.parse(query).tables
      end

      tables = queries.flat_map { |q| query_tables[q] }.uniq.select { |t| possible_tables.include?(t) }

      new_queries = queries.select { |q| query_tables[q].any? && query_tables[q].all? { |t| possible_tables.include?(t) } }

      if client.options[:log_level] == "debug2"
        (queries - new_queries).each do |query|
          log "Processed #{PgQuery.fingerprint(query)}"
          if query_tables[query].empty?
            log "No tables"
          else
            log "Some tables missing in current database"
          end
          puts
          puts query
{'fraction_non_alphanumeric': 0.11648089849400153, 'fraction_numerical': 0.003233217050965711, 'mean_word_length': 1.91734921816828, 'pattern_counts': {'":': 0, '<': 27, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '620527', 'n_tokens_mistral': 3495, 'n_tokens_neox': 3409, 'n_words': 793}