rails-ai-context 5.35.0 → 5.35.2

data/lib/rails_ai_context/database_family/postgresql_plan.rb

@@ -0,0 +1,124 @@
+# frozen_string_literal: true
+
+module RailsAiContext
+  module DatabaseFamily
+    # A PostgreSQL EXPLAIN (VERBOSE, FORMAT JSON) plan, read for the relations
+    # and base columns its expressions name.
+    class PostgresqlPlan
+      # A plain qualified Var, either part bare or quoted. A set operation's "*SELECT* 1"
+      # is no relation's alias, so only a real relation's column is read.
+      NAME = /(?:([A-Za-z_]\w*)|"((?:""|[^"])+)")/
+      BASE_VAR = /\A#{NAME}\.#{NAME}\z/
+      # Nodes whose every read of a child's column is an expression of their own that names it.
+      READS_BY_NAME = [ "Aggregate", "Hash Join", "Merge Join", "Nested Loop" ].freeze
+
+      # Catalogs and TOAST tables: the never-read list and the catalog credentials hold what is secret there.
+      SYSTEM_SCHEMA = /\A(?:pg_catalog|information_schema|pg_toast(?:_temp_\d+)?)\z/
+
+      APPENDS = [ "Append", "Merge Append" ].freeze
+      REMOTE_NAME = /"(?:""|[^"])+"|[^\s()."]+/
+      REMOTE_RELATION = /\((?:(#{REMOTE_NAME})\.)?(#{REMOTE_NAME})(?: (#{REMOTE_NAME}))?\)/
+
+      # aliases: alias => the relations it stands for, each as `schema.relation` where the plan gives a schema that is no system one.
+      attr_reader :nodes, :aliases
+
+      # @param top [Hash] the plan's top node, as Family#plan returns it
+      def initialize(top)
+        @top = top
+        @nodes = []
+        @parents = {}.compare_by_identity
+        @appended = {}.compare_by_identity
+        collect(top)
+        # The alias defaults to the relation name when the query gave none, which is how `users.*` resolves.
+        # Keyed as EXPLAIN prints it: `t` and "T" are two aliases.
+        @names = {}
+        @aliases = @nodes.each_with_object({}) do |node, map|
+          scanned(node).each do |schema, rel, as|
+            key = [ (schema unless schema.to_s.match?(SYSTEM_SCHEMA)), rel ].compact.join(".").downcase
+            @names[key] = rel.downcase
+            # The members of an Append over partitions are `name_1`, `name_2`, and a node over them names their columns by `name`.
+            # Anywhere else `name_1` is the plan's own name for a second use of `name`, which stands for another relation.
+            [ as, (as[/\A(.+)_\d+\z/m, 1] if @appended[node]) ].compact.each { |name| (map[name] ||= []) << key }
+          end
+        end
+        @aliases.each_value(&:uniq!)
+      end
+
+      # The relation's own name, with no schema, for an answer.
+      def name(relation)
+        @names.fetch(relation, relation)
+      end
+
+      # @return [Array<Array(String, String)>] [schema, `schema.relation`] for each relation the plan reads
+      #   outside the system schemas
+      def relations
+        @nodes.flat_map { |node| scanned(node) }.filter_map do |schema, rel, _as|
+          [ schema, "#{schema}.#{rel}".downcase ] if schema && !schema.match?(SYSTEM_SCHEMA)
+        end.uniq
+      end
+
+      # @return [Array<Array<String>>] the expressions that can stand at result column i, with any sort or group key
+      #   past the result width. An Append prints no Output, and each of its members hands up a row of its own.
+      def output
+        return Array(@top["Output"]).map { |expr| [ expr ] } if @top.key?("Output") || !APPENDS.include?(@top["Node Type"])
+
+        members = Array(@top["Plans"])
+        outputs = members.map { |member| Array(member["Output"]) }
+        # Read only where each member is a relation that hands up the same columns in the same places.
+        columns = outputs.map { |exprs| exprs.map { |expr| expr.to_s.match(BASE_VAR)&.then { |m| m[3] || m[4] } } }.uniq
+        members.all? { |member| member["Relation Name"] } && columns.size == 1 ? outputs.transpose : []
+      end
+
+      # @return [Array<Array(String, String)>] [relation, column] for a plain Var of a base relation, the column
+      #   folded as the secret column rule folds a name: one for each relation its alias stands for
+      def base_columns(expr)
+        m = expr.to_s.strip.match(BASE_VAR) or return []
+        column = SecretColumns.fold(m[3] || m[4].gsub('""', '"'))
+        Array(aliases[m[1] || m[2].gsub('""', '"')]).map { |rel| [ rel, column ] }
+      end
+
+      # The node that takes this node's rows as its input and reads their columns by name, or nil. A scan
+      # under one hands up every column of its table, read or not. A subplan's columns are read by position.
+      def reader(node)
+        parent = @parents[node]
+        parent if parent && READS_BY_NAME.include?(parent["Node Type"]) && %w[Outer Inner].include?(node["Parent Relationship"])
+      end
+
+      # Every expression the node itself holds, under whatever key.
+      def expressions(node)
+        strings(node.except("Plans"))
+      end
+
+      private
+
+      # [schema, relation, alias] of each relation the node scans. A join or an aggregate that a foreign server runs has
+      # no relation of its own: postgres_fdw names its tables there, each with its alias, as `(schema.table alias)`.
+      def scanned(node)
+        rel = node["Relation Name"]
+        return [ [ node["Schema"], rel.to_s, (node["Alias"] || rel).to_s ] ] if rel
+
+        node["Relations"].to_s.scan(REMOTE_RELATION).map do |parts|
+          schema, name, as = parts.map { |part| part&.delete_prefix('"')&.delete_suffix('"')&.gsub('""', '"') }
+          [ schema, name, as || name ]
+        end
+      end
+
+      def collect(node, parent = nil)
+        return unless node.is_a?(Hash)
+
+        @nodes << node
+        @parents[node] = parent
+        @appended[node] = true if parent && (@appended[parent] || APPENDS.include?(parent["Node Type"]))
+        Array(node["Plans"]).each { |child| collect(child, node) }
+      end
+
+      def strings(value)
+        case value
+        when Hash then value.values.flat_map { |held| strings(held) }
+        when Array then value.flat_map { |held| strings(held) }
+        else [ value.to_s ]
+        end
+      end
+    end
+  end
+end

data/lib/rails_ai_context/database_family/postgresql_runtime.rb

@@ -0,0 +1,211 @@
+# frozen_string_literal: true
+
+require "json"
+
+module RailsAiContext
+  module DatabaseFamily
+    # PostgreSQL bounds a statement with SET LOCAL statement_timeout inside a read-only transaction.
+    module PostgresqlRuntime
+      module_function
+
+      def read_only(conn, sql, timeout)
+        result = bounded(conn, timeout) { conn.select_all(sql) }
+        Rows.new(result.columns, result.rows, false)
+      end
+
+      # statement_timeout ends each statement in the block, a wait on a lock included. Read-only, rolled back.
+      # Inside a transaction that is already open it is a savepoint, whose rollback puts both settings back.
+      def bounded(conn, timeout)
+        result = nil
+        DatabaseFamily.bounding do
+          conn.transaction(requires_new: true) do
+            conn.execute("SET TRANSACTION READ ONLY")
+            conn.execute("SET LOCAL statement_timeout = '#{[ (timeout * 1000).to_i, 1 ].max}'")
+            result = yield
+            raise ActiveRecord::Rollback
+          end
+        end
+        result
+      end
+
+      def explain(conn, sql, timeout)
+        raw = read_only(conn, "EXPLAIN (FORMAT JSON, ANALYZE) #{sql}", timeout).rows.map { |r| r.first.to_s }.join("\n")
+        scans = []
+        warnings = []
+        begin
+          data = JSON.parse(raw)
+          plan = data.first["Plan"] if data.is_a?(Array) && data.first.is_a?(Hash)
+          scan_nodes(plan, scans, warnings) if plan
+        rescue JSON::ParserError
+          nil
+        end
+        { scan_types: scans, warnings: warnings, raw: raw, unbounded: false }
+      end
+
+      # Planning only: EXPLAIN without ANALYZE runs nothing, not even a VOLATILE function.
+      # A timeout and an SQL error are raised, so the caller never runs what it could not plan.
+      def plan(conn, sql, timeout)
+        json = read_only(conn, "EXPLAIN (VERBOSE, FORMAT JSON) #{sql}", timeout).rows.first&.first
+        data = JSON.parse(json.to_s)
+        data.is_a?(Array) && data.first.is_a?(Hash) ? data.first["Plan"] : nil
+      rescue JSON::ParserError
+        nil
+      end
+
+      # The described statement's ftable/ftablecol, then the plan's top Output,
+      # which also sees through a view, CTE or subquery that renamed a column.
+      def column_sources(conn, sql, plan)
+        raise ArgumentError, "PostgreSQL column sources are read from the statement's plan, and none was given" unless plan
+
+        described_sources(conn, sql) + planned_sources(plan)
+      end
+
+      def described_sources(conn, sql)
+        raw = conn.raw_connection
+        stmt = "rac_prov_#{Process.pid}_#{rand(1 << 30)}"
+        begin
+          raw.prepare(stmt, sql)
+          desc = raw.describe_prepared(stmt)
+        ensure
+          raw.exec("DEALLOCATE #{stmt}") rescue nil
+        end
+
+        relations = {}
+        desc.nfields.times.filter_map do |i|
+          oid = desc.ftable(i).to_i
+          col = desc.ftablecol(i).to_i
+          next if oid.zero? || col.zero?
+
+          relation, names = relations[oid] ||= [ relation_name(conn, oid), attribute_names(conn, oid) ]
+          # A relation whose columns could not be read still names itself, for the caller to fail closed on.
+          next [ i, nil, relation ] unless names
+          [ i, names[col], relation ] if names[col]
+        end
+      rescue StandardError => e
+        # The driver raises its own class here. Half a trace is no answer, so a timeout is raised as the call's.
+        raise TimedOut, e.message if defined?(PG::QueryCanceled) && e.is_a?(PG::QueryCanceled)
+        raise if e.is_a?(ActiveRecord::StatementInvalid) && DatabaseFamily.timed_out?(e)
+
+        []
+      end
+
+      def planned_sources(plan)
+        read = PostgresqlPlan.new(plan)
+        read.output.each_with_index.flat_map do |exprs, i|
+          exprs.flat_map { |expr| read.base_columns(expr) }.map { |rel, col| [ i, col, rel ] }.uniq { |_, col, _| col }
+        end
+      end
+
+      def unqualified_schemas(_conn)
+        []
+      end
+
+      # A bare name folds to lower case and a quoted one does not, so both spellings are asked for.
+      def schemas_named(conn, names)
+        asked = (names | names.map(&:downcase)).map { |name| conn.quote(name) }.join(", ")
+        conn.select_values(
+          "SELECT nspname FROM pg_namespace WHERE nspname IN (#{asked}) AND nspname !~ '^pg_' " \
+          "AND nspname <> 'information_schema' AND nspname <> ALL (current_schemas(false))"
+        )
+      end
+
+      # A plan names the session's temporary schema pg_temp, whatever number its own name carries.
+      SCHEMA = "CASE WHEN n.oid = pg_my_temp_schema() THEN 'pg_temp' ELSE n.nspname END"
+
+      def in_schemas(conn, schemas)
+        schemas.empty? ? "n.nspname = ANY (current_schemas(false))" : "#{SCHEMA} IN (#{schemas.map { |name| conn.quote(name) }.join(", ")})"
+      end
+
+      # Every relation a plan can name that stores rows: tables, partitioned tables, materialized views, foreign tables and sequences.
+      # One read that takes no lock on any of them.
+      def relation_columns(conn, schemas = [])
+        rows = conn.select_rows(
+          "SELECT #{SCHEMA}, c.relname, a.attname FROM pg_attribute a JOIN pg_class c ON c.oid = a.attrelid " \
+          "JOIN pg_namespace n ON n.oid = c.relnamespace WHERE c.relkind IN ('r', 'p', 'm', 'f', 'S') AND a.attnum > 0 " \
+          "AND NOT a.attisdropped AND #{in_schemas(conn, schemas)} ORDER BY c.relname, a.attnum"
+        )
+        rows.each_with_object({}) do |(schema, relation, column), found|
+          (found[relation] ||= []) << column
+          (found["#{schema}.#{relation}"] ||= []) << column
+        end
+      end
+
+      VIEWS = "FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace WHERE c.relkind"
+
+      # A materialized view stores its rows, so its result columns have no table column to be traced to. The
+      # views are read with it for what a materialized view built on one reads. Each is also under `schema.view`.
+      def view_definitions(conn, schemas = [])
+        rows = conn.select_rows("SELECT #{SCHEMA}, c.relname, pg_get_viewdef(c.oid) #{VIEWS} IN ('v', 'm') AND #{in_schemas(conn, schemas)}")
+        DatabaseFamily.by_view(rows.flat_map { |schema, view, definition| [ [ view, definition ], [ "#{schema}.#{view}", definition ] ] })
+      end
+
+      def traced_views(conn, schemas = [])
+        conn.select_rows("SELECT #{SCHEMA}, c.relname #{VIEWS} = 'v' AND #{in_schemas(conn, schemas)}")
+          .flat_map { |schema, view| [ view, "#{schema}.#{view}" ] }
+      end
+
+      # A stored generated column is a plain column to the plan, so its expression is read here.
+      def generated_columns(conn, _schemas = [])
+        return {} if conn.database_version < 12_00_00
+
+        DatabaseFamily.keep_expressions({}, conn.select_rows(
+          "SELECT c.relname, a.attname, pg_get_expr(d.adbin, d.adrelid) FROM pg_attribute a " \
+          "JOIN pg_class c ON c.oid = a.attrelid JOIN pg_attrdef d ON d.adrelid = a.attrelid AND d.adnum = a.attnum " \
+          "WHERE a.attgenerated <> '' AND NOT a.attisdropped"
+        ))
+      end
+
+      def relation_name(conn, oid)
+        conn.select_value("SELECT relname FROM pg_class WHERE oid = #{oid.to_i}")&.to_s
+      rescue StandardError
+        nil
+      end
+
+      # attnum => column name for one relation oid, or nil when they cannot be read.
+      def attribute_names(conn, oid)
+        conn.select_rows(
+          "SELECT attnum, attname FROM pg_attribute WHERE attrelid = #{oid.to_i} AND attnum > 0 AND NOT attisdropped"
+        ).to_h { |num, name| [ num.to_i, name.to_s ] }
+      rescue StandardError
+        nil
+      end
+
+      def index_tables(_conn, _schemas = [])
+        {}
+      end
+
+      def table_sizes(conn)
+        Introspectors::PgPartitions.table_bytes(conn).sort_by { |r| [ -r[:bytes], r[:name] ] }
+      end
+
+      def row_estimates(conn)
+        Introspectors::PgPartitions.table_rows(conn).sort_by { |t| [ -t[:rows], t[:name] ] }.map do |t|
+          entry = { table: t[:name], approximate_rows: t[:rows] }
+          entry[:dead_rows] = t[:dead_rows] if t[:dead_rows] > 0
+          entry
+        end
+      end
+
+      def index_usage(conn)
+        Introspectors::PgPartitions.index_stats(conn).sort_by { |i| [ i[:scans], i[:table], i[:index] ] }
+      end
+
+      def scan_nodes(node, scans, warnings)
+        return unless node.is_a?(Hash)
+
+        node_type = node["Node Type"].to_s
+        table = node["Relation Name"]
+        rows = node["Actual Rows"] || node["Plan Rows"]
+
+        if node_type.include?("Seq Scan")
+          scans << { table: table, type: "sequential scan", index: nil }
+          warnings << "Sequential scan on #{table} (#{rows} rows)" if rows.to_i > 1000
+        elsif node_type.include?("Index")
+          scans << { table: table, type: node_type.downcase, index: node["Index Name"] }
+        end
+
+        (node["Plans"] || []).each { |child| scan_nodes(child, scans, warnings) }
+      end
+    end
+  end
+end