appquery 0.4.0 → 0.6.0.alpha

data/lib/app_query.rb

@@ -1,6 +1,9 @@
 # frozen_string_literal: true
 
 require_relative "app_query/version"
+require_relative "app_query/base_query"
+require_relative "app_query/paginatable"
+require_relative "app_query/mappable"
 require_relative "app_query/tokenizer"
 require_relative "app_query/render_helpers"
 require "active_record"
@@ -22,7 +25,7 @@ require "active_record"
 #   end
 #
 # @example CTE manipulation
-#   AppQuery(<<~SQL).select_all(select: "select * from articles where id = 1")
+#   AppQuery(<<~SQL).select_all("select * from articles where id = 1")
 #     WITH articles AS(...)
 #     SELECT * FROM articles
 #     ORDER BY id
@@ -72,21 +75,44 @@ module AppQuery
 
   # Loads a query from a file in the configured query path.
   #
+  # When no extension is provided, tries `.sql` first, then `.sql.erb`.
+  # Raises an error if both files exist (ambiguous).
+  #
   # @param query_name [String, Symbol] the query name or path (without extension)
   # @param opts [Hash] additional options passed to {Q#initialize}
   # @return [Q] a new query object loaded from the file
   #
-  # @example Load a simple query
+  # @example Load a .sql file
   #   AppQuery[:invoices]  # loads app/queries/invoices.sql
   #
+  # @example Load a .sql.erb file (when .sql doesn't exist)
+  #   AppQuery[:dynamic_report]  # loads app/queries/dynamic_report.sql.erb
+  #
   # @example Load from a subdirectory
   #   AppQuery["reports/weekly"]  # loads app/queries/reports/weekly.sql
   #
   # @example Load with explicit extension
   #   AppQuery["invoices.sql.erb"]  # loads app/queries/invoices.sql.erb
+  #
+  # @raise [Error] if both `.sql` and `.sql.erb` files exist for the same name
   def self.[](query_name, **opts)
-    filename = File.extname(query_name.to_s).empty? ? "#{query_name}.sql" : query_name.to_s
-    full_path = (Pathname.new(configuration.query_path) / filename).expand_path
+    base = Pathname.new(configuration.query_path) / query_name.to_s
+
+    full_path = if File.extname(query_name.to_s).empty?
+      sql_path = base.sub_ext(".sql").expand_path
+      erb_path = base.sub_ext(".sql.erb").expand_path
+      sql_exists = sql_path.exist?
+      erb_exists = erb_path.exist?
+
+      if sql_exists && erb_exists
+        raise Error, "Ambiguous query name #{query_name.inspect}: both #{sql_path} and #{erb_path} exist"
+      end
+
+      sql_exists ? sql_path : erb_path
+    else
+      base.expand_path
+    end
+
     Q.new(full_path.read, name: "AppQuery #{query_name}", filename: full_path.to_s, **opts)
   end
 
@@ -114,6 +140,48 @@ module AppQuery
       count
     end
 
+    private
+
+    # Override to provide indifferent access (string or symbol keys).
+    def hash_rows
+      @hash_rows ||= rows.map do |row|
+        columns.zip(row).to_h.with_indifferent_access
+      end
+    end
+
+    public
+
+    # Transforms each record in-place using the provided block.
+    #
+    # @yield [Hash] each record as a hash with indifferent access
+    # @yieldreturn [Hash] the transformed record
+    # @return [self] the result object for chaining
+    #
+    # @example Add a computed field
+    #   result = AppQuery[:users].select_all
+    #   result.transform! { |r| r.merge("full_name" => "#{r['first']} #{r['last']}") }
+    def transform!
+      @hash_rows = hash_rows.map { |r| yield(r) } unless empty?
+      self
+    end
+
+    # Resolves a cast type value, converting symbols to ActiveRecord types.
+    #
+    # @param value [Symbol, Object] the cast type (symbol shorthand or type instance)
+    # @return [Object] the resolved type instance
+    #
+    # @example
+    #   resolve_cast_type(:date)  #=> ActiveRecord::Type::Date instance
+    #   resolve_cast_type(ActiveRecord::Type::Json.new)  #=> returns as-is
+    def self.resolve_cast_type(value)
+      case value
+      when Symbol
+        ActiveRecord::Type.lookup(value)
+      else
+        value
+      end
+    end
+
     def self.from_ar_result(r, cast = nil)
       if r.empty?
         EMPTY
@@ -122,7 +190,7 @@ module AppQuery
         when Array
           r.columns.zip(cast).to_h
         when Hash
-          cast
+          cast.transform_keys(&:to_s).transform_values { |v| resolve_cast_type(v) }
         else
           {}
         end
@@ -186,14 +254,14 @@ module AppQuery
     # @return [String] the SQL string
     # @return [Array, Hash] bind parameters
     # @return [Boolean, Hash, Array] casting configuration
-    attr_reader :name, :sql, :binds, :cast
+    attr_reader :sql, :name, :filename, :binds, :cast
 
     # Creates a new query object.
     #
     # @param sql [String] the SQL query string (may contain ERB)
     # @param name [String, nil] optional name for logging
     # @param filename [String, nil] optional filename for ERB error reporting
-    # @param binds [Array, Hash] bind parameters for the query
+    # @param binds [Hash, nil] bind parameters for the query
     # @param cast [Boolean, Hash, Array] type casting configuration
     #
     # @example Simple query
@@ -201,28 +269,38 @@ module AppQuery
     #
     # @example With ERB and binds
     #   Q.new("SELECT * FROM users WHERE id = :id", binds: {id: 1})
-    def initialize(sql, name: nil, filename: nil, binds: [], cast: true)
+    def initialize(sql, name: nil, filename: nil, binds: {}, cast: true, cte_depth: 0)
       @sql = sql
       @name = name
       @filename = filename
       @binds = binds
       @cast = cast
+      @cte_depth = cte_depth
+      @binds = binds_with_defaults(sql, binds)
     end
 
-    # @private
-    def deep_dup
-      super.send(:reset!)
+    attr_reader :cte_depth
+
+    def to_arel
+      if binds.presence
+        Arel::Nodes::BoundSqlLiteral.new sql, [], binds
+      else
+        # TODO: add retryable? available from >=7.1
+        Arel::Nodes::SqlLiteral.new(sql)
+      end
     end
-    private :deep_dup
 
-    # @private
-    def reset!
-      (instance_variables - %i[@sql @filename @name @binds @cast]).each do
-        instance_variable_set(_1, nil)
+    private def binds_with_defaults(sql, binds)
+      if (named_binds = sql.scan(/:(?<!::)([a-zA-Z]\w*)/).flatten.map(&:to_sym).uniq.presence)
+        named_binds.zip(Array.new(named_binds.count)).to_h.merge(binds.to_h)
+      else
+        binds.to_h
       end
-      self
     end
-    private :reset!
+
+    def deep_dup(sql: self.sql, name: self.name, filename: self.filename, binds: self.binds.dup, cast: self.cast, cte_depth: self.cte_depth)
+      self.class.new(sql, name:, filename:, binds:, cast:, cte_depth:)
+    end
 
     # @!group Rendering
 
@@ -267,12 +345,7 @@ module AppQuery
       sql = to_erb.result(helper.get_binding)
       collected = helper.collected_binds
 
-      with_sql(sql).tap do |q|
-        # Merge collected binds with existing binds (convert array to hash if needed)
-        existing = @binds.is_a?(Hash) ? @binds : {}
-        new_binds = existing.merge(collected)
-        q.instance_variable_set(:@binds, new_binds) if new_binds.any?
-      end
+      with_sql(sql).add_binds(**collected)
     end
 
     def to_erb
@@ -306,56 +379,35 @@ module AppQuery
 
     # Executes the query and returns all matching rows.
     #
-    # @param binds [Array, Hash, nil] bind parameters (positional or named)
     # @param select [String, nil] override the SELECT clause
+    # @param binds [Hash, nil] bind parameters to add
     # @param cast [Boolean, Hash, Array] type casting configuration
     # @return [Result] the query results with optional type casting
     #
-    # @example Simple query with positional binds
-    #   AppQuery("SELECT * FROM users WHERE id = $1").select_all(binds: [1])
-    #
-    # @example Named binds
+    # @example (Named) binds
     #   AppQuery("SELECT * FROM users WHERE id = :id").select_all(binds: {id: 1})
     #
-    # @example With type casting
-    #   AppQuery("SELECT created_at FROM users")
-    #     .select_all(cast: {created_at: ActiveRecord::Type::DateTime.new})
+    # @example With type casting (shorthand)
+    #   AppQuery("SELECT published_on FROM articles")
+    #     .select_all(cast: {"published_on" => :date})
+    #
+    # @example With type casting (explicit)
+    #   AppQuery("SELECT metadata FROM products")
+    #     .select_all(cast: {"metadata" => ActiveRecord::Type::Json.new})
     #
     # @example Override SELECT clause
-    #   AppQuery("SELECT * FROM users").select_all(select: "COUNT(*)")
+    #   AppQuery("SELECT * FROM users").select_all("COUNT(*)")
     #
     # @raise [UnrenderedQueryError] if the query contains unrendered ERB
-    # @raise [ArgumentError] if mixing positional binds with collected named binds
-    #
-    # TODO: have aliases for common casts: select_all(cast: {"today" => :date})
-    def select_all(binds: nil, select: nil, cast: self.cast)
-      with_select(select).render({}).then do |aq|
-        # Support both positional (array) and named (hash) binds
-        if binds.is_a?(Array)
-          if @binds.is_a?(Hash) && @binds.any?
-            raise ArgumentError, "Cannot use positional binds (Array) when query has collected named binds from values()/bind() helpers. Use named binds (Hash) instead."
-          end
-          # Positional binds using $1, $2, etc.
-          ActiveRecord::Base.connection.select_all(aq.to_s, name, binds).then do |result|
-            Result.from_ar_result(result, cast)
-          end
+    def select_all(s = nil, binds: {}, cast: self.cast)
+      add_binds(**binds).with_select(s).render({}).then do |aq|
+        sql = if ActiveRecord::VERSION::STRING.to_f >= 7.1
+          aq.to_arel
         else
-          # Named binds - merge collected binds with explicitly passed binds
-          merged_binds = (@binds.is_a?(Hash) ? @binds : {}).merge(binds || {})
-          if merged_binds.any?
-            sql = if ActiveRecord::VERSION::STRING.to_f >= 7.1
-              Arel.sql(aq.to_s, **merged_binds)
-            else
-              ActiveRecord::Base.sanitize_sql_array([aq.to_s, **merged_binds])
-            end
-            ActiveRecord::Base.connection.select_all(sql, name).then do |result|
-              Result.from_ar_result(result, cast)
-            end
-          else
-            ActiveRecord::Base.connection.select_all(aq.to_s, name).then do |result|
-              Result.from_ar_result(result, cast)
-            end
-          end
+          ActiveRecord::Base.sanitize_sql_array([aq.to_s, aq.binds])
+        end
+        ActiveRecord::Base.connection.select_all(sql, aq.name).then do |result|
+          Result.from_ar_result(result, cast)
         end
       end
     rescue NameError => e
@@ -366,23 +418,24 @@ module AppQuery
 
     # Executes the query and returns the first row.
     #
-    # @param binds [Array, Hash, nil] bind parameters (positional or named)
+    # @param binds [Hash, nil] bind parameters to add
     # @param select [String, nil] override the SELECT clause
     # @param cast [Boolean, Hash, Array] type casting configuration
     # @return [Hash, nil] the first row as a hash, or nil if no results
     #
     # @example
-    #   AppQuery("SELECT * FROM users WHERE id = $1").select_one(binds: [1])
+    #   AppQuery("SELECT * FROM users WHERE id = :id").select_one(binds: {id: 1})
     #   # => {"id" => 1, "name" => "Alice"}
     #
     # @see #select_all
-    def select_one(binds: nil, select: nil, cast: self.cast)
-      select_all(binds:, select:, cast:).first
+    def select_one(s = nil, binds: {}, cast: self.cast)
+      with_select(s).select_all("SELECT * FROM :_ LIMIT 1", binds:, cast:).first
     end
+    alias_method :first, :select_one
 
     # Executes the query and returns the first value of the first row.
     #
-    # @param binds [Array, Hash, nil] bind parameters (positional or named)
+    # @param binds [Hash, nil] named bind parameters
     # @param select [String, nil] override the SELECT clause
     # @param cast [Boolean, Hash, Array] type casting configuration
     # @return [Object, nil] the first value, or nil if no results
@@ -392,13 +445,129 @@ module AppQuery
     #   # => 42
     #
     # @see #select_one
-    def select_value(binds: nil, select: nil, cast: self.cast)
-      select_one(binds:, select:, cast:)&.values&.first
+    def select_value(s = nil, binds: {}, cast: self.cast)
+      select_one(s, binds:, cast:)&.values&.first
+    end
+
+    # Returns the count of rows from the query.
+    #
+    # Wraps the query in a CTE and selects only the count, which is more
+    # efficient than fetching all rows via `select_all.count`.
+    #
+    # @param s [String, nil] optional SELECT to apply before counting
+    # @param binds [Hash, nil] bind parameters to add
+    # @return [Integer] the count of rows
+    #
+    # @example Simple count
+    #   AppQuery("SELECT * FROM users").count
+    #   # => 42
+    #
+    # @example Count with filtering
+    #   AppQuery("SELECT * FROM users")
+    #     .with_select("SELECT * FROM :_ WHERE active")
+    #     .count
+    #   # => 10
+    def count(s = nil, binds: {})
+      with_select(s).select_all("SELECT COUNT(*) c FROM :_", binds:).column("c").first
+    end
+
+    # Returns whether any rows exist in the query result.
+    #
+    # Uses `EXISTS` which stops at the first matching row, making it more
+    # efficient than `count > 0` for large result sets.
+    #
+    # @param s [String, nil] optional SELECT to apply before checking
+    # @param binds [Hash, nil] bind parameters to add
+    # @return [Boolean] true if at least one row exists
+    #
+    # @example Check if query has results
+    #   AppQuery("SELECT * FROM users").any?
+    #   # => true
+    #
+    # @example Check with filtering
+    #   AppQuery("SELECT * FROM users").any?("SELECT * FROM :_ WHERE admin")
+    #   # => false
+    def any?(s = nil, binds: {})
+      with_select(s).select_all("SELECT EXISTS(SELECT 1 FROM :_) e", binds:).column("e").first
+    end
+
+    # Returns whether no rows exist in the query result.
+    #
+    # Inverse of {#any?}. Uses `EXISTS` for efficiency.
+    #
+    # @param s [String, nil] optional SELECT to apply before checking
+    # @param binds [Hash, nil] bind parameters to add
+    # @return [Boolean] true if no rows exist
+    #
+    # @example Check if query is empty
+    #   AppQuery("SELECT * FROM users WHERE admin").none?
+    #   # => true
+    def none?(s = nil, binds: {})
+      !any?(s, binds:)
+    end
+
+    # Returns an array of values for a single column.
+    #
+    # Wraps the query in a CTE and selects only the specified column, which is
+    # more efficient than fetching all columns via `select_all.column(name)`.
+    # The column name is safely quoted, making this method safe for user input.
+    #
+    # @param c [String, Symbol] the column name to extract
+    # @param s [String, nil] optional SELECT to apply before extracting
+    # @param binds [Hash, nil] bind parameters to add
+    # @return [Array] the column values
+    #
+    # @example Extract a single column
+    #   AppQuery("SELECT id, name FROM users").column(:name)
+    #   # => ["Alice", "Bob", "Charlie"]
+    #
+    # @example With additional filtering
+    #   AppQuery("SELECT * FROM users").column(:email, "SELECT * FROM :_ WHERE active")
+    #   # => ["alice@example.com", "bob@example.com"]
+    def column(c, s = nil, binds: {})
+      quoted_column = ActiveRecord::Base.connection.quote_column_name(c)
+      with_select(s).select_all("SELECT #{quoted_column} AS column FROM :_", binds:).column("column")
+    end
+
+    # Returns an array of id values from the query.
+    #
+    # Convenience method equivalent to `column(:id)`. More efficient than
+    # fetching all columns via `select_all.column("id")`.
+    #
+    # @param s [String, nil] optional SELECT to apply before extracting
+    # @param binds [Hash, nil] bind parameters to add
+    # @return [Array] the id values
+    #
+    # @example Get all user IDs
+    #   AppQuery("SELECT * FROM users").ids
+    #   # => [1, 2, 3]
+    #
+    # @example With filtering
+    #   AppQuery("SELECT * FROM users").ids("SELECT * FROM :_ WHERE active")
+    #   # => [1, 3]
+    def ids(s = nil, binds: {})
+      column(:id, s, binds:)
+    end
+
+    # Executes the query and returns results as an Array of Hashes.
+    #
+    # Shorthand for `select_all(...).entries`. Accepts the same arguments as
+    # {#select_all}.
+    #
+    # @return [Array<Hash>] the query results as an array
+    #
+    # @example
+    #   AppQuery("SELECT * FROM users").entries
+    #   # => [{"id" => 1, "name" => "Alice"}, {"id" => 2, "name" => "Bob"}]
+    #
+    # @see #select_all
+    def entries(...)
+      select_all(...).entries
     end
 
     # Executes an INSERT query.
     #
-    # @param binds [Array, Hash] bind parameters for the query
+    # @param binds [Hash, nil] bind parameters for the query
     # @param returning [String, nil] columns to return (Rails 7.1+ only)
     # @return [Integer, Object] the inserted ID or returning value
     #
@@ -419,30 +588,22 @@ module AppQuery
     #
     # @raise [UnrenderedQueryError] if the query contains unrendered ERB
     # @raise [ArgumentError] if returning is used with Rails < 7.1
-    def insert(binds: [], returning: nil)
+    def insert(binds: {}, returning: nil)
       # ActiveRecord::Base.connection.insert(sql, name, _pk = nil, _id_value = nil, _sequence_name = nil, binds, returning: nil)
       if returning && ActiveRecord::VERSION::STRING.to_f < 7.1
         raise ArgumentError, "The 'returning' option requires Rails 7.1+. Current version: #{ActiveRecord::VERSION::STRING}"
       end
 
-      binds = binds.presence || @binds
-      render({}).then do |aq|
-        if binds.is_a?(Hash)
-          sql = if ActiveRecord::VERSION::STRING.to_f >= 7.1
-            Arel.sql(aq.to_s, **binds)
-          else
-            ActiveRecord::Base.sanitize_sql_array([aq.to_s, **binds])
-          end
-          if ActiveRecord::VERSION::STRING.to_f >= 7.1
-            ActiveRecord::Base.connection.insert(sql, name, returning:)
-          else
-            ActiveRecord::Base.connection.insert(sql, name)
-          end
-        elsif ActiveRecord::VERSION::STRING.to_f >= 7.1
-          # pk is the less flexible returning
-          ActiveRecord::Base.connection.insert(aq.to_s, name, _pk = nil, _id_value = nil, _sequence_name = nil, binds, returning:)
+      with_binds(**binds).render({}).then do |aq|
+        sql = if ActiveRecord::VERSION::STRING.to_f >= 7.1
+          aq.to_arel
+        else
+          ActiveRecord::Base.sanitize_sql_array([aq.to_s, **aq.binds])
+        end
+        if ActiveRecord::VERSION::STRING.to_f >= 7.1
+          ActiveRecord::Base.connection.insert(sql, name, returning:)
         else
-          ActiveRecord::Base.connection.insert(aq.to_s, name, _pk = nil, _id_value = nil, _sequence_name = nil, binds)
+          ActiveRecord::Base.connection.insert(sql, name)
         end
       end
     rescue NameError => e
@@ -453,7 +614,7 @@ module AppQuery
 
     # Executes an UPDATE query.
     #
-    # @param binds [Array, Hash] bind parameters for the query
+    # @param binds [Hash, nil] bind parameters for the query
     # @return [Integer] the number of affected rows
     #
     # @example With named binds
@@ -465,19 +626,14 @@ module AppQuery
     #     .update(binds: ["New Title", 1])
     #
     # @raise [UnrenderedQueryError] if the query contains unrendered ERB
-    def update(binds: [])
-      binds = binds.presence || @binds
-      render({}).then do |aq|
-        if binds.is_a?(Hash)
-          sql = if ActiveRecord::VERSION::STRING.to_f >= 7.1
-            Arel.sql(aq.to_s, **binds)
-          else
-            ActiveRecord::Base.sanitize_sql_array([aq.to_s, **binds])
-          end
-          ActiveRecord::Base.connection.update(sql, name)
+    def update(binds: {})
+      with_binds(**binds).render({}).then do |aq|
+        sql = if ActiveRecord::VERSION::STRING.to_f >= 7.1
+          aq.to_arel
         else
-          ActiveRecord::Base.connection.update(aq.to_s, name, binds)
+          ActiveRecord::Base.sanitize_sql_array([aq.to_s, **aq.binds])
         end
+        ActiveRecord::Base.connection.update(sql, name)
       end
     rescue NameError => e
       raise e unless e.instance_of?(NameError)
@@ -486,7 +642,7 @@ module AppQuery
 
     # Executes a DELETE query.
     #
-    # @param binds [Array, Hash] bind parameters for the query
+    # @param binds [Hash, nil] bind parameters for the query
     # @return [Integer] the number of deleted rows
     #
     # @example With named binds
@@ -496,19 +652,14 @@ module AppQuery
     #   AppQuery("DELETE FROM videos WHERE id = $1").delete(binds: [1])
     #
     # @raise [UnrenderedQueryError] if the query contains unrendered ERB
-    def delete(binds: [])
-      binds = binds.presence || @binds
-      render({}).then do |aq|
-        if binds.is_a?(Hash)
-          sql = if ActiveRecord::VERSION::STRING.to_f >= 7.1
-            Arel.sql(aq.to_s, **binds)
-          else
-            ActiveRecord::Base.sanitize_sql_array([aq.to_s, **binds])
-          end
-          ActiveRecord::Base.connection.delete(sql, name)
+    def delete(binds: {})
+      with_binds(**binds).render({}).then do |aq|
+        sql = if ActiveRecord::VERSION::STRING.to_f >= 7.1
+          aq.to_arel
         else
-          ActiveRecord::Base.connection.delete(aq.to_s, name, binds)
+          ActiveRecord::Base.sanitize_sql_array([aq.to_s, **aq.binds])
         end
+        ActiveRecord::Base.connection.delete(sql, name)
       end
     rescue NameError => e
       raise e unless e.instance_of?(NameError)
@@ -547,16 +698,29 @@ module AppQuery
 
     # Returns a new query with different bind parameters.
     #
-    # @param binds [Array, Hash] the new bind parameters
-    # @return [Q] a new query object with the specified binds
+    # @param binds [Hash, nil] the bind parameters
+    # @return [Q] a new query object with the binds replaced
     #
     # @example
-    #   query = AppQuery("SELECT * FROM users WHERE id = :id")
-    #   query.with_binds(id: 1).select_one
-    def with_binds(binds)
-      deep_dup.tap do
-        _1.instance_variable_set(:@binds, binds)
-      end
+    #   query = AppQuery("SELECT :foo, :bar", binds: {foo: 1})
+    #   query.with_binds(bar: 2).binds
+    #   # => {foo: nil, bar: 2}
+    def with_binds(**binds)
+      deep_dup(binds:)
+    end
+    alias_method :replace_binds, :with_binds
+
+    # Returns a new query with binds added.
+    #
+    # @param binds [Hash, nil] the bind parameters to add
+    # @return [Q] a new query object with the added binds
+    #
+    # @example
+    #   query = AppQuery("SELECT :foo, :bar", binds: {foo: 1})
+    #   query.add_binds(bar: 2).binds
+    #   # => {foo: 1, bar: 2}
+    def add_binds(**binds)
+      deep_dup(binds: self.binds.merge(binds))
     end
 
     # Returns a new query with different cast settings.
@@ -568,9 +732,7 @@ module AppQuery
     #   query = AppQuery("SELECT created_at FROM users")
     #   query.with_cast(false).select_all  # disable casting
     def with_cast(cast)
-      deep_dup.tap do
-        _1.instance_variable_set(:@cast, cast)
-      end
+      deep_dup(cast:)
     end
 
     # Returns a new query with different SQL.
@@ -578,31 +740,48 @@ module AppQuery
     # @param sql [String] the new SQL string
     # @return [Q] a new query object with the specified SQL
     def with_sql(sql)
-      deep_dup.tap do
-        _1.instance_variable_set(:@sql, sql)
-      end
+      deep_dup(sql:)
     end
 
     # Returns a new query with a modified SELECT statement.
     #
-    # If the query has a CTE named `"_"`, replaces the SELECT statement.
-    # Otherwise, wraps the original query in a `"_"` CTE and uses the new SELECT.
+    # Wraps the current SELECT in a numbered CTE and applies the new SELECT.
+    # CTEs are named `_`, `_1`, `_2`, etc. Use `:_` in the new SELECT to
+    # reference the previous result.
     #
     # @param sql [String, nil] the new SELECT statement (nil returns self)
     # @return [Q] a new query object with the modified SELECT
     #
-    # @example
-    #   AppQuery("SELECT id, name FROM users").with_select("SELECT COUNT(*) FROM _")
-    #   # => "WITH _ AS (\n  SELECT id, name FROM users\n)\nSELECT COUNT(*) FROM _"
+    # @example Single transformation
+    #   AppQuery("SELECT * FROM users").with_select("SELECT COUNT(*) FROM :_")
+    #   # => "WITH _ AS (\n  SELECT * FROM users\n)\nSELECT COUNT(*) FROM _"
+    #
+    # @example Chained transformations
+    #   AppQuery("SELECT * FROM users")
+    #     .with_select("SELECT * FROM :_ WHERE active")
+    #     .with_select("SELECT COUNT(*) FROM :_")
+    #   # => WITH _ AS (SELECT * FROM users),
+    #   #         _1 AS (SELECT * FROM _ WHERE active)
+    #   #    SELECT COUNT(*) FROM _1
     def with_select(sql)
       return self if sql.nil?
-      if cte_names.include?("_")
-        with_sql(tokens.each_with_object([]) do |token, acc|
-          v = (token[:t] == "SELECT") ? sql : token[:v]
+
+      # First CTE is "_", then "_1", "_2", etc.
+      current_cte = (cte_depth == 0) ? "_" : "_#{cte_depth}"
+
+      # Replace :_ with the current CTE name
+      processed_sql = sql.gsub(/:_\b/, current_cte)
+
+      # Wrap current SELECT in numbered CTE
+      new_cte = "#{current_cte} AS (\n  #{select}\n)"
+
+      append_cte(new_cte).then do |q|
+        # Replace the SELECT token with processed_sql and increment depth
+        new_sql = q.tokens.each_with_object([]) do |token, acc|
+          v = (token[:t] == "SELECT") ? processed_sql : token[:v]
           acc << v
-        end.join)
-      else
-        append_cte("_ as (\n  #{select}\n)").with_select(sql)
+        end.join
+        q.deep_dup(sql: new_sql, cte_depth: cte_depth + 1)
       end
     end
 

data/mise.toml

@@ -3,4 +3,4 @@ ruby = "3.4"
 
 [env]
 _.path = ["bin"]
-PROJECT_ROOT = "{{config_root}}"
+PROJECT_ROOT = "{{config_root}}"