diamond-orm 0.1.0

data/lib/diamond/query_object.rb

@@ -0,0 +1,463 @@
+require_relative 'struct_factory'
+require_relative 'compiler/base'
+require_relative 'cursor'
+
+module Diamond
+  class QueryObject
+    include Enumerable
+
+    # Phase 7 query modes. :struct is the default — returns Diamond's
+    # frozen Structs. :hash returns native Extralite hash rows. :array
+    # returns positional arrays (SoA-friendly, zero allocation). :splat
+    # returns a single column's values (the ultimate pluck).
+    VALID_MODES = %i[struct hash array splat].freeze
+
+    attr_reader :table, :ast, :mode
+
+    def initialize(table, ast = [])
+      @table = table
+      @ast = ast
+      @cached_result = nil
+      @mode = :struct
+    end
+
+    def materialize
+      return @cached_result if @cached_result
+
+      sql, params, transform_spec = Diamond::Compiler::Base.compile(@table, @ast)
+
+      # eager-loaded queries: hand the row layout to Extralite::Transform
+      # so the result is deduplicated and nested per the join graph.
+      if transform_spec
+        transform = Extralite::Transform.new(transform_spec)
+        @cached_result = []
+        Diamond.engine.db.query(transform, sql, *params).each do |row|
+          @cached_result << Diamond::StructFactory.create_eager(@table, row)
+        end
+        return @cached_result
+      end
+
+      # Phase 7: hash/array/splat modes bypass the StructFactory and
+      # hand back the raw Extralite rows. Much cheaper on the GC.
+      case @mode
+      when :hash
+        @cached_result = Diamond.engine.db.query_hash(sql, *params)
+        return @cached_result
+      when :array
+        @cached_result = Diamond.engine.db.query_array(sql, *params)
+        return @cached_result
+      when :splat
+        @cached_result = Diamond.engine.db.query_splat(sql, *params)
+        return @cached_result
+      end
+
+      # Phase 7: use prepared statement cache for :struct mode.
+      # Array fetch mode: rows arrive positionally, zero Hash
+      # allocations per row (see StructFactory.create_from_array).
+      stmt = Diamond.engine.prepared(sql, :array)
+      begin
+        stmt.bind(*params)
+
+        # union results carry no top-level projection; members come from
+        # the left side (sqlite names union output columns after it).
+        union_node = @ast.find { |n| n.is_a?(AST::Union) }
+        scope_ast = union_node ? union_node.left.ast : @ast
+        projection_node = scope_ast.find { |n| n.is_a?(AST::Projection) }
+        projected_columns = projection_node&.columns
+
+        @cached_result = []
+        raw_rows = []
+        stmt.each { |row| raw_rows << row }
+        @cached_result = Diamond::StructFactory.create_many(@table, raw_rows, projected_columns)
+      ensure
+        # Note: we don't close the stmt here — it's cached for reuse.
+        # The stmt will be closed when the engine is frozen (which doesn't
+        # happen) or explicitly closed by the user. Extralite manages
+        # prepared statement lifecycle.
+      end
+
+      @cached_result
+    end
+
+    # --- Phase 7 query mode terminals ---
+
+    def as_hash
+      with_mode(:hash)
+    end
+
+    def as_array
+      with_mode(:array)
+    end
+
+    def as_splat
+      with_mode(:splat)
+    end
+
+    def first(n = 1)
+      scope = has_order? ? self : order(resolve_pk!)
+      results = scope.limit(n).materialize
+      n == 1 ? results.first : results
+    end
+
+    # without a block, returns an Enumerator so .lazy and Enumerable chains work.
+    # Early-return before compiling/preparing: the enumerator re-enters here
+    # with a block on iteration, so doing work now would leak one statement.
+    def each(&block)
+      return enum_for(:each) unless block_given?
+
+      compiled_sql, compiled_params, transform_spec = Diamond::Compiler::Base.compile(@table, @ast)
+
+      if transform_spec
+        transform = Extralite::Transform.new(transform_spec)
+        Diamond.engine.db.query(transform, compiled_sql, *compiled_params).each do |row|
+          yield Diamond::StructFactory.create_eager(@table, row)
+        end
+        return
+      end
+
+      stmt = Diamond.engine.db.prepare_array(compiled_sql)
+      stmt.bind(*compiled_params)
+
+      union_node = @ast.find { |n| n.is_a?(AST::Union) }
+      scope_ast = union_node ? union_node.left.ast : @ast
+      projection_node   = scope_ast.find { |n| n.is_a?(AST::Projection) }
+      projected_columns = projection_node&.columns
+
+      cursor = Diamond::Cursor.new(@table, stmt, projected_columns)
+
+      cursor.each(&block)
+    end
+
+    def last(n = 1)
+      scope = has_order? ? self : order([resolve_pk!, :desc])
+      results = scope.limit(n).materialize
+      n == 1 ? results.first : results.reverse
+    end
+
+    # skip the structs entirely - read values straight off the cursor.
+    def pluck(*columns)
+      raise ArgumentError, "pluck requires at least one column" if columns.empty?
+      validate_pluck_columns!(columns)
+      rows = pluck_rows(columns)
+      if columns.size == 1
+        rows.map { |row| row[0] }
+      else
+        rows.map(&:dup)
+      end
+    end
+
+    # first row's values only: single value for one column, tuple for
+    # several, nil on miss. `Users.pick(:name)` ~ `SELECT name ... LIMIT 1`.
+    def pick(*columns)
+      raise ArgumentError, "pick requires at least one column" if columns.empty?
+      validate_pluck_columns!(columns)
+      q = projected_query(columns).limit(1)
+      sql, params, _transform = Diamond::Compiler::Base.compile(q.table, q.ast)
+      row = fetch_first_array_row(sql, params)
+      return nil if row.nil?
+      columns.size == 1 ? row[0] : row.dup
+    end
+
+    # primary-key values. `Users.ids` ~ `SELECT id FROM users`.
+    def ids
+      pluck(resolve_pk!)
+    end
+
+    # Zero-translation JSON fast path. Pushes serialization into
+    # SQLite's JSON1 C engine: the normally-compiled query becomes a
+    # subquery whose rows are encoded by json_object and aggregated by
+    # json_group_array — Ruby receives ONE finished string. No
+    # StructFactory, no per-row Hash, no Ruby JSON encoding at all.
+    # Columns default to the table's schema columns; pass an explicit
+    # list to project (which also fixes key order). Returns a
+    # Diamond::JsonString so Rutile passes it to the Rack body with
+    # zero further processing.
+    #
+    # The inner query compiles through the normal path, so WHERE,
+    # JOIN, GROUP BY, ORDER BY, LIMIT, OFFSET, DISTINCT and UNION ALL
+    # all keep working — the wrapper only replaces the SELECT list
+    # with a single aggregate expression. One-shot query (no prepared
+    # cache): a single row crosses the Ruby/C boundary.
+    #
+    # Keys are validated identifiers (schema membership, same as
+    # pluck), so they interpolate safely as single-quoted literals;
+    # value refs are bare subquery-output names, mirroring the
+    # compiler's own unquoted column rendering. COALESCE keeps the
+    # empty set as '[]' on SQLite builds predating native empty-array
+    # aggregation.
+    #
+    # Explicit limits (fail loud): eager joins (object-graph transform
+    # is meaningless for a single string) and WITH/CTE chains (WITH
+    # cannot lead inside a subquery). Guards run before column
+    # validation: a WITH query's table is a NullTable (empty schema),
+    # so validation could never pass — the structural error is the
+    # meaningful one.
+    def to_json_array(*columns)
+      if @ast.any? { |n| n.is_a?(AST::Join) && n.eager }
+        raise ArgumentError, "to_json_array does not support eager joins"
+      end
+      if @ast.any? { |n| n.is_a?(AST::With) }
+        raise ArgumentError, "to_json_array does not support WITH/CTE chains"
+      end
+      columns = @table.schema[:columns] if columns.empty?
+      columns.each do |c|
+        raise Diamond::UnknownColumnError.build(@table.schema, c) unless @table.schema[:columns].include?(c)
+      end
+      # Union chains compile standalone (sides carry their own shape;
+      # appending a Projection would trip compile's "union terminates
+      # the chain" rule). The json_object refs resolve against the
+      # union output, named after the left side — sides must project
+      # those names (SELECT * on one table always qualifies).
+      has_union = @ast.any? { |n| n.is_a?(AST::Union) }
+      filtered = if has_union
+                   @ast
+                 else
+                   new_nodes = columns.map { |c| AST::Column.new(c) }
+                   @ast.reject { |n| n.is_a?(AST::Projection) } + [AST::Projection.new(new_nodes)]
+                 end
+      q = Diamond::QueryObject.new(@table, filtered)
+      inner_sql, params, _transform = Diamond::Compiler::Base.compile(q.table, q.ast)
+      pairs = columns.map { |c| "'#{c}', #{c}" }.join(', ')
+      sql = "SELECT COALESCE(json_group_array(json_object(#{pairs})), '[]') FROM (#{inner_sql})"
+      Diamond::JsonString.new(Diamond.engine.db.query_single_splat(sql, *params))
+    end
+
+    def exists?(filter = nil)
+      unless filter.nil?
+        return filter.is_a?(Hash) ? where(filter).exists? : find(filter).exists?
+      end
+      pk  = resolve_pk!
+      col = AST::Column.new(pk)
+      filtered = @ast.reject { |n| n.is_a?(AST::Projection) } + [AST::Projection.new([col])]
+      q = Diamond::QueryObject.new(@table, filtered).limit(1)
+      sql, params, _transform = Diamond::Compiler::Base.compile(q.table, q.ast)
+      stmt = Diamond.engine.db.prepare_array(sql)
+      begin
+        stmt.bind(*params)
+        found = false
+        stmt.each { |_row| found = true; break }
+        found
+      ensure
+        begin
+          stmt.close unless stmt.closed?
+        rescue StandardError
+          # ensure must not raise
+        end
+      end
+    end
+
+    # SQL-backed existence checks. With a block (or pattern args for
+    # any?) they fall through to Enumerable — same results as today,
+    # just without the SQL fast path. NOTE: never bare `super` here.
+    # The DSL modules are mixed into QueryObject *after* Enumerable,
+    # so `super` would land on the DSL delegate and recurse forever;
+    # bind Enumerable explicitly instead.
+    def any?(*args, &block)
+      return enumerable_fallback(:any?, *args, &block) if block || !args.empty?
+      exists?
+    end
+
+    def none?(*args, &block)
+      return enumerable_fallback(:none?, *args, &block) if block || !args.empty?
+      !exists?
+    end
+
+    def empty?
+      !exists?
+    end
+
+    def count(column = nil)
+      col = column.nil? ? resolve_pk! : validate_aggregate_column!(column)
+      nodes = [AST::Function.new(:count, [AST::Column.new(col)])]
+      filtered = @ast.reject { |n| n.is_a?(AST::Projection) } + [AST::Projection.new(nodes)]
+      Diamond::QueryObject.new(@table, filtered).materialize.first.public_send(:"count_#{col}")
+    end
+
+    # Scalar aggregate terminals. Each builds one Function projection
+    # (exactly like count) and reads the auto-named member
+    # (SUM(age) -> sum_age). Empty scope yields one NULL row, so these
+    # return nil — never NoMethodError.
+    #
+    # Block form (or a non-column arg like an Numeric init) falls through
+    # to Enumerable#sum — same pattern as min/max below.
+    def sum(column = nil, &block)
+      if block
+        return column.nil? ? enumerable_fallback(:sum, &block) : enumerable_fallback(:sum, column, &block)
+      end
+      return enumerable_fallback(:sum) if column.nil?
+      return enumerable_fallback(:sum, column) unless column.is_a?(Symbol) || column.is_a?(String)
+
+      aggregate(:sum, column)
+    end
+
+    def minimum(column)
+      aggregate(:min, column)
+    end
+
+    def maximum(column)
+      aggregate(:max, column)
+    end
+
+    def average(column)
+      aggregate(:avg, column)
+    end
+
+    alias avg average
+
+    # min/max with a column go to SQL; bare min/max, min(n), and
+    # block forms keep today's Enumerable behavior (see the super note
+    # on any? above for why this binds explicitly).
+    def min(column = nil, &block)
+      return minimum(column) if !block && !column.nil? && !column.is_a?(Integer)
+      column.nil? ? enumerable_fallback(:min, &block) : enumerable_fallback(:min, column, &block)
+    end
+
+    def max(column = nil, &block)
+      return maximum(column) if !block && !column.nil? && !column.is_a?(Integer)
+      column.nil? ? enumerable_fallback(:max, &block) : enumerable_fallback(:max, column, &block)
+    end
+
+    def method_missing(name, *args, &block)
+      if @table.schema[:columns].include?(name)
+        record = first
+        raise Diamond::RecordNotFound, "No record found for AST: #{@ast.inspect}" if record.nil?
+        record.public_send(name, *args, &block)
+      else
+        super
+      end
+    end
+
+    def respond_to_missing?(name, include_private = false)
+      @table.schema[:columns].include?(name) || super
+    end
+
+    def inspect
+      if @cached_result
+        "#<Diamond::QueryObject materialized: #{@cached_result.size} records>"
+      else
+        "#<Diamond::QueryObject table=#{@table.name} ast=[#{@ast.map(&:class).map(&:name).join(', ')}]>"
+      end
+    end
+
+    # the SQL this chain compiles to, without running it. returns
+    # [sql, params] — params stay separate so you can bind them yourself.
+    def to_sql
+      sql, params, _transform = Diamond::Compiler::Base.compile(@table, @ast)
+      [sql, params]
+    end
+
+    # sqlite's query plan for this chain. read-only, runs on the
+    # caller's own connection so it works from worker ractors too.
+    # returns extralite rows (:selectid, :order, :from, :detail).
+    def explain
+      sql, params, _transform = Diamond::Compiler::Base.compile(@table, @ast)
+      Diamond.engine.db.query("EXPLAIN QUERY PLAN #{sql}", *params)
+    end
+
+    # the chain's AST as an indented tree. for staring at what a block
+    # actually became.
+    def ast_tree
+      @ast.map { |n| AST.dump(n) }.join("\n")
+    end
+
+    private
+
+    # Bypass for Enumerable fallbacks. See the note on any?.
+    def enumerable_fallback(method_name, *args, &block)
+      Enumerable.instance_method(method_name).bind_call(self, *args, &block)
+    end
+
+    def validate_pluck_columns!(columns)
+      columns.each do |c|
+        raise Diamond::UnknownColumnError.build(@table.schema, c) unless @table.schema[:columns].include?(c)
+      end
+    end
+
+    # fresh QueryObject projecting exactly `columns`, replacing any
+    # existing projection (same rule pluck has always used).
+    def projected_query(columns)
+      new_nodes = columns.map { |c| AST::Column.new(c) }
+      filtered = @ast.reject { |n| n.is_a?(AST::Projection) } + [AST::Projection.new(new_nodes)]
+      Diamond::QueryObject.new(@table, filtered)
+    end
+
+    # raw positional rows for a column projection (Array mode: row[i]
+    # lines up with columns[i]).
+    def pluck_rows(columns)
+      q = projected_query(columns)
+      sql, params, _transform = Diamond::Compiler::Base.compile(q.table, q.ast)
+      stmt = Diamond.engine.db.prepare_array(sql)
+      begin
+        stmt.bind(*params)
+        result = []
+        # dup: the cursor may reuse its row buffer across iterations
+        # (the old pluck body duped for the same reason).
+        stmt.each { |row| result << row.dup }
+        result
+      ensure
+        begin
+          stmt.close unless stmt.closed?
+        rescue StandardError
+          # ensure must not raise
+        end
+      end
+    end
+
+    def fetch_first_array_row(sql, params)
+      stmt = Diamond.engine.db.prepare_array(sql)
+      begin
+        stmt.bind(*params)
+        found = nil
+        stmt.each { |row| found = row; break }
+        found
+      ensure
+        begin
+          stmt.close unless stmt.closed?
+        rescue StandardError
+          # ensure must not raise
+        end
+      end
+    end
+
+    # one Function projection (SUM(age) -> sum_age), scalar back.
+    # Empty scope yields a single NULL row, so this returns nil there.
+    def aggregate(func, column)
+      col = validate_aggregate_column!(column)
+      nodes = [AST::Function.new(func, [AST::Column.new(col)])]
+      filtered = @ast.reject { |n| n.is_a?(AST::Projection) } + [AST::Projection.new(nodes)]
+      row = Diamond::QueryObject.new(@table, filtered).materialize.first
+      row&.public_send(:"#{func.to_s.downcase}_#{col}")
+    end
+
+    def validate_aggregate_column!(column)
+      col = column.to_sym
+      unless @table.schema[:columns].include?(col)
+        raise Diamond::UnknownColumnError.build(@table.schema, col)
+      end
+      col
+    end
+
+    def with_mode(new_mode)
+      raise ArgumentError, "Unknown mode #{new_mode.inspect}" unless VALID_MODES.include?(new_mode)
+      q = Diamond::QueryObject.new(@table, @ast.dup)
+      q.instance_variable_set(:@mode, new_mode)
+      q
+    end
+
+    def has_order?
+      @ast.any? { |n| n.is_a?(AST::Order) }
+    end
+
+    # pk or :id, and it better exist. fail here with a column error
+    # instead of letting sqlite complain about COUNT(missing). Also
+    # the implicit ORDER BY column for first/last.
+    def resolve_pk!
+      pk = @table.schema[:primary_key] || :id
+      unless @table.schema[:columns].include?(pk)
+        raise Diamond::UnknownColumnError.build(@table.schema, pk)
+      end
+      pk
+    end
+
+  end
+end