sequelizer 0.1.4 → 0.1.6

data/lib/sequel/extensions/make_readyable.rb

@@ -1,60 +1,101 @@
+require 'pathname'
+
 module Sequel
+
+  # = MakeReadyable
+  #
+  # Sequel extension that provides database readiness functionality,
+  # primarily geared towards Spark SQL-based databases. This extension
+  # allows setting up temporary views and schema configurations to prepare
+  # a database for use.
+  #
+  # @example Basic schema usage
+  #   db.extension :make_readyable
+  #   db.make_ready(use_schema: :my_schema)
+  #
+  # @example Search path with schema precedence
+  #   db.make_ready(search_path: [:schema1, :schema2])
+  #
+  # @example External file sources
+  #   db.make_ready(search_path: [Pathname.new('data.parquet')])
   module MakeReadyable
-    ##
-    # This method is primarily geared towards Spark SQL-based databases.
+
+    # Prepares the database by setting up schemas, views, and external data sources.
     #
+    # This method is primarily geared towards Spark SQL-based databases.
     # Given some options, prepares a set of views to represent a set
     # of tables across a collection of different schemas and external,
     # unmanaged tables.
     #
+    # @param opts [Hash] the options used to prepare the database
+    # @option opts [Symbol] :use_schema The schema to be used as the primary schema
+    # @option opts [Array] :search_path A set of symbols (schemas) or Pathnames (external files)
+    # @option opts [Array] :only Limit view creation to these tables only
+    # @option opts [Array] :except Skip view creation for these tables
+    #
+    # @example Set primary schema
     #   DB.make_ready(use_schema: :schema)
     #   # => USE `schema`
     #
-    # When using search_path, tables from previous schema override tables
-    # from the next schema.  This is analogous to the way Unix searches 
-    # the PATH variable for programs.
-    #
-    # Assuming the following tables: schema1.a, schema2.a, schema2.b
-    #
+    # @example Search path with precedence
+    #   # Assuming tables: schema1.a, schema2.a, schema2.b
     #   DB.make_ready(search_path: [:schema1, :schema2])
     #   # => CREATE TEMPORARY VIEW `a` AS SELECT * FROM `schema1`.`a;`
     #   # => CREATE TEMPORARY VIEW `b` AS SELECT * FROM `schema2`.`b;`
     #
-    # When using Pathnames, the extension on the file becomes the format
-    # to try to read from the file.  
-    #
+    # @example External file sources
     #   DB.make_ready(search_path: [Pathname.new("c.parquet"), Pathname.new("d.orc")])
     #   # => CREATE TEMPORARY VIEW `c` USING parquet OPTIONS ('path'='c.parquet')
     #   # => CREATE TEMPORARY VIEW `d` USING orc OPTIONS ('path'='d.orc')
-    #
-    # @param [Hash] opts the options used to prepare the database
-    # @option opts [String] :use_schema The schema to be used as the primary schema
-    # @option opts [Array] :search_path A set of sympbols (to represent schemas) or Pathnames (to represent externally managed data files)
     def make_ready(opts = {})
       ReadyMaker.new(self, opts).run
     end
+
   end
 
-  private 
+  # = ReadyMaker
+  #
+  # Internal class that handles the actual database preparation logic.
+  # This class processes the make_ready options and executes the necessary
+  # SQL statements to set up schemas, views, and external data sources.
   class ReadyMaker
+
+    # @!attribute [r] db
+    #   @return [Sequel::Database] the database instance
+    # @!attribute [r] opts
+    #   @return [Hash] the preparation options
     attr_reader :db, :opts
 
+    # Creates a new ReadyMaker instance.
+    #
+    # @param db [Sequel::Database] the database to prepare
+    # @param opts [Hash] the preparation options
     def initialize(db, opts)
       @db = db
       @opts = opts
     end
-  
+
+    # Executes the database preparation process.
+    #
+    # This method handles:
+    # 1. Setting the primary schema if specified
+    # 2. Processing the search path to create views
+    # 3. Handling table filtering (only/except options)
     def run
       if opts[:use_schema]
         db.extension :usable
         db.use(opts[:use_schema])
       end
       only_tables = Array(opts[:only])
-      created_views = (Array(opts[:except]) || [])
-      (opts[:search_path] || []).each do |schema|
-        schema = schema.is_a?(Pathname) ? schema : schema.to_sym
+      created_views = Array(opts[:except]) || []
+      (opts[:search_path] || []).flatten.each do |schema|
+        schema = schema.to_sym unless schema.is_a?(Pathname)
         source = get_source(db, schema)
-        tables = source.tables(schema: schema) - created_views
+        tables = if schema.is_a?(Pathname)
+                   source.tables - created_views
+                 else
+                   source.tables(schema: schema) - created_views
+                 end
         tables &= only_tables unless only_tables.empty?
         tables.each do |table|
           create_view(source, table, schema)
@@ -63,46 +104,123 @@ module Sequel
       end
     end
 
+    # Creates a temporary view for the given table.
+    #
+    # @param source [Object] the source (database or FileSourcerer)
+    # @param table [Symbol] the table name
+    # @param schema [Symbol, Pathname] the schema or file path
     def create_view(source, table, schema)
       if schema.to_s =~ %r{/}
         source.create_view(table, temp: true)
       else
+        # For schema-based tables, just create temporary views
+        # This extension is primarily for Spark SQL-based databases
         source.create_view(table, db[Sequel.qualify(schema, table)], temp: true)
       end
     end
 
+    # Gets the appropriate source handler for the schema.
+    #
+    # @param db [Sequel::Database] the database instance
+    # @param schema [Symbol, Pathname] the schema or file path
+    # @return [Sequel::Database, FileSourcerer] the source handler
     def get_source(db, schema)
       if schema.to_s =~ %r{/}
-        FileSourcerer.new(db, Pathname.new(schema))
+        FileSourcerer.new(db, Pathname.new(schema.to_s))
       else
         db
       end
     end
 
+    # = FileSourcerer
+    #
+    # Handles external file sources for the make_ready functionality.
+    # This class creates temporary views that read from external files
+    # like Parquet, ORC, etc.
     class FileSourcerer
+
+      # @!attribute [r] db
+      #   @return [Sequel::Database] the database instance
+      # @!attribute [r] schema
+      #   @return [Pathname] the file path
       attr_reader :db, :schema
+
+      # Creates a new FileSourcerer instance.
+      #
+      # @param db [Sequel::Database] the database instance
+      # @param schema [Pathname] the file path
       def initialize(db, schema)
         @db = db
         @schema = schema
       end
 
-      def tables(opts = {})
-        [schema.basename(".*").to_s.to_sym]
+      # Returns the table name derived from the file name.
+      #
+      # @param _opts [Hash] unused options parameter
+      # @return [Array<Symbol>] array containing the table name
+      def tables(_opts = {})
+        [schema.basename(schema.extname).to_s.to_sym]
       end
 
+      # Creates a temporary view that reads from the external file.
+      #
+      # @param table [Symbol] the table/view name
+      # @param opts [Hash] additional options to merge
       def create_view(table, opts = {})
-        db.create_view(table, {
-          temp: true,
-          using: format,
-          options: { path: schema.expand_path }
-        }.merge(opts))
+        case db.database_type
+        when :spark
+          # Spark SQL uses USING clause for external tables
+          db.create_view(table, {
+            temp: true,
+            using: format,
+            options: { path: schema.expand_path },
+          }.merge(opts))
+        when :duckdb
+          # DuckDB uses direct file reading with read_* functions
+          create_duckdb_view(table, opts)
+        else
+          raise Sequel::Error, "External file sources are not supported on #{db.database_type}"
+        end
       end
 
+      private
+
+      # Creates a view for DuckDB to read external files
+      #
+      # @param table [Symbol] the table/view name
+      # @param _opts [Hash] additional options to merge (currently unused for DuckDB)
+      def create_duckdb_view(table, _opts)
+        file_path = if schema.directory?
+                      schema.expand_path.join('**').join("*.#{format}").to_s
+                    else
+                      schema.expand_path.to_s
+                    end
+        read_function = case format
+                        when 'parquet'
+                          :read_parquet
+                        when 'csv'
+                          :read_csv_auto
+                        when 'json'
+                          :read_json_auto
+                        else
+                          raise Sequel::Error, "Unsupported file format '#{format}' for DuckDB"
+                        end
+
+        # DuckDB doesn't support TEMPORARY views, use regular CREATE VIEW
+        db.create_view(table, db.from(Sequel.function(read_function, file_path)))
+      end
+
+      # Returns the file format based on the file extension.
+      #
+      # @return [String] the file format (e.g., 'parquet', 'orc')
       def format
-        schema.extname[1..-1]
+        schema.extname[1..]
       end
+
     end
+
   end
 
   Database.register_extension(:make_readyable, MakeReadyable)
+
 end

data/lib/sequel/extensions/more_sql.rb

@@ -0,0 +1,76 @@
+# frozen_string_literal: true
+
+module Sequel
+
+  # Provides additional SQL helper methods for database operations.
+  #
+  # The more_sql extension adds convenience methods for SQL operations that
+  # aren't covered by Sequel's core functionality, particularly schema-related
+  # operations like CREATE SCHEMA.
+  #
+  # @example Load the extension
+  #   DB.extension :more_sql
+  #
+  # @example Create a schema
+  #   DB.create_schema(:analytics)
+  #   # Executes: CREATE SCHEMA "analytics"
+  #
+  # @example Create schema if it doesn't exist
+  #   DB.create_schema(:staging, if_not_exists: true)
+  #   # Executes: CREATE SCHEMA IF NOT EXISTS "staging"
+  module MoreSql
+
+    # Creates a database schema.
+    #
+    # Generates and executes a CREATE SCHEMA statement with optional
+    # IF NOT EXISTS clause for idempotent schema creation.
+    #
+    # @param schema_name [Symbol, String] The name of the schema to create
+    # @param opts [Hash] Options for schema creation
+    # @option opts [Boolean] :if_not_exists (false) Only create the schema if it doesn't already exist
+    #
+    # @return [nil]
+    #
+    # @example Basic schema creation
+    #   DB.create_schema(:reports)
+    #   # Executes: CREATE SCHEMA "reports"
+    #
+    # @example Idempotent schema creation
+    #   DB.create_schema(:analytics, if_not_exists: true)
+    #   # Executes: CREATE SCHEMA IF NOT EXISTS "analytics"
+    #
+    # @example With string schema name
+    #   DB.create_schema('user_data')
+    #   # Executes: CREATE SCHEMA "user_data"
+    def create_schema(schema_name, opts = {})
+      run(create_schema_sql(schema_name, opts))
+    end
+
+    private
+
+    # Generates the SQL for creating a schema.
+    #
+    # Builds a CREATE SCHEMA statement with proper identifier quoting
+    # and optional IF NOT EXISTS clause.
+    #
+    # @param schema_name [Symbol, String] The name of the schema to create
+    # @param opts [Hash] Options for schema creation
+    # @option opts [Boolean] :if_not_exists (false) Include IF NOT EXISTS clause
+    #
+    # @return [String] The CREATE SCHEMA SQL statement
+    #
+    # @example
+    #   create_schema_sql(:test, if_not_exists: true)
+    #   # => 'CREATE SCHEMA IF NOT EXISTS "test"'
+    def create_schema_sql(schema_name, opts)
+      sql = 'CREATE SCHEMA '
+      sql += 'IF NOT EXISTS ' if opts[:if_not_exists]
+      sql += literal(schema_name)
+      sql
+    end
+
+  end
+
+  Database.register_extension(:more_sql, MoreSql)
+
+end

data/lib/sequel/extensions/settable.rb

@@ -1,5 +1,56 @@
+# frozen_string_literal: true
+
+#
+# The settable extension adds a convenient +set+ method to database connections
+# for executing SET statements with key-value pairs. This is particularly useful
+# for configuring database session parameters.
+#
+#   DB.extension :settable
+#   DB.set(search_path: 'public', timezone: 'UTC')
+#   # Executes: SET search_path=public
+#   #           SET timezone=UTC
+#
+#   DB.set(work_mem: '256MB')
+#   # Executes: SET work_mem=256MB
+#
+# The extension works with any database adapter and supports various value types
+# including strings, numbers, booleans, and nil values.
+#
+# Related module: Sequel::Settable
+
 module Sequel
+
+  # The Settable module provides database configuration functionality through
+  # SET statements. When loaded as an extension, it adds the +set+ method to
+  # database connections.
   module Settable
+
+    # Execute SET statements for the given options hash.
+    #
+    # Each key-value pair in the options hash is converted to a SET statement
+    # and executed against the database. Multiple options result in multiple
+    # SET statements being executed in sequence.
+    #
+    # @param opts [Hash] Hash of configuration options to set
+    # @option opts [Object] key The configuration parameter name
+    # @option opts [Object] value The value to set for the parameter
+    #
+    # @example Set a single parameter
+    #   DB.set(timezone: 'UTC')
+    #   # Executes: SET timezone=UTC
+    #
+    # @example Set multiple parameters
+    #   DB.set(search_path: 'public', work_mem: '256MB')
+    #   # Executes: SET search_path=public
+    #   #           SET work_mem=256MB
+    #
+    # @example Different value types
+    #   DB.set(port: 5432, autocommit: true, custom_setting: nil)
+    #   # Executes: SET port=5432
+    #   #           SET autocommit=true
+    #   #           SET custom_setting=
+    #
+    # @return [void]
     def set(opts = {})
       set_sql(opts).each do |sql|
         run(sql)
@@ -8,10 +59,23 @@ module Sequel
 
     private
 
+    # Generate SET SQL statements from options hash.
+    #
+    # Converts each key-value pair in the options hash into a SET SQL statement
+    # string. This is a private helper method used internally by the +set+ method.
+    #
+    # @param opts [Hash] Hash of options to convert to SET statements
+    # @return [Array<String>] Array of SET SQL statement strings
+    #
+    # @example
+    #   set_sql(timezone: 'UTC', port: 5432)
+    #   # => ["SET timezone=UTC", "SET port=5432"]
     def set_sql(opts)
       opts.map { |k, v| "SET #{k}=#{v}" }
     end
+
   end
 
   Database.register_extension(:settable, Settable)
+
 end

data/lib/sequel/extensions/sql_recorder.rb

@@ -0,0 +1,85 @@
+# frozen_string_literal: true
+
+# == Overview
+#
+# The sql_recorder extension records each SQL statement sent to the database
+# in a thread-safe array accessible via the +sql_recorder+ method.
+#
+# == Usage
+#
+#   DB.extension :sql_recorder
+#   DB[:users].all
+#   DB[:posts].where(id: 1).first
+#
+#   # Access recorded SQL statements
+#   DB.sql_recorder
+#   # => ["SELECT * FROM users", "SELECT * FROM posts WHERE (id = 1) LIMIT 1"]
+#
+#   # Clear the recorded statements
+#   DB.sql_recorder.clear
+#
+# == Thread Safety
+#
+# The extension is thread-safe and uses a mutex to synchronize access to the
+# SQL recording array when multiple threads are executing queries simultaneously.
+#
+# == Compatibility
+#
+# This extension is designed to work alongside mock databases and other SQL
+# recording mechanisms. It uses the method name +sql_recorder+ to avoid
+# conflicts with existing +sqls+ methods that may be present in test frameworks.
+#
+# Related module: Sequel::SqlRecorder
+
+module Sequel
+
+  # Extension module that adds SQL recording capabilities to Sequel databases.
+  # When included, it provides a +sql_recorder+ method that returns an array
+  # of all SQL statements executed against the database.
+  module SqlRecorder
+
+    # Returns the array of recorded SQL statements.
+    #
+    # The array accumulates all SQL statements sent to the database since the
+    # extension was loaded or since the last time +clear+ was called on the array.
+    #
+    # @return [Array<String>] array of SQL statement strings
+    # @example
+    #   DB.extension :sql_recorder
+    #   DB[:users].all
+    #   DB.sql_recorder #=> ["SELECT * FROM users"]
+    attr_reader :sql_recorder
+
+    # Intercepts SQL execution to record statements.
+    #
+    # This method overrides Sequel's +log_connection_yield+ to capture each SQL
+    # statement in a thread-safe manner before delegating to the parent implementation.
+    #
+    # @param sql [String] the SQL statement being executed
+    # @param conn [Object] the database connection object
+    # @param args [Object] additional arguments (optional)
+    # @return [Object] result from the parent +log_connection_yield+ method
+    def log_connection_yield(sql, conn, args = nil)
+      @sql_recorder_mutex.synchronize { sql_recorder.push(sql) }
+      super
+    end
+
+    # Initializes the SQL recording infrastructure when the extension is loaded.
+    #
+    # Sets up the mutex for thread-safe access and initializes the SQL recording
+    # array. This method is automatically called when the extension is loaded
+    # via +DB.extension :sql_recorder+.
+    #
+    # @param db [Sequel::Database] the database instance being extended
+    def self.extended(db)
+      db.instance_exec do
+        @sql_recorder_mutex ||= Mutex.new
+        @sql_recorder ||= []
+      end
+    end
+
+  end
+
+  Database.register_extension(:sql_recorder, SqlRecorder)
+
+end

data/lib/sequel/extensions/unionize.rb

@@ -0,0 +1,169 @@
+# frozen_string_literal: true
+
+require 'digest'
+
+module Sequel
+
+  # Provides efficient handling of large UNION operations.
+  #
+  # The unionize extension allows combining many datasets through UNION operations
+  # by chunking them into manageable temporary tables or views. This is particularly
+  # useful when dealing with databases that have limitations on the number of UNION
+  # operations in a single query (e.g., Spark SQL, DuckDB).
+  #
+  # @example Load the extension
+  #   DB.extension :unionize
+  #
+  # @example Basic usage
+  #   DB.unionize([dataset1, dataset2, dataset3, dataset4])
+  #
+  # @example With options
+  #   DB.unionize(datasets, chunk_size: 50, all: true, temp_table_prefix: 'my_union')
+  module Unionize
+
+    # Handles the chunking and union of multiple datasets.
+    #
+    # This class manages the process of splitting a large collection of datasets
+    # into smaller chunks, creating temporary tables/views for each chunk, and
+    # then recursively combining them until a single unified dataset is produced.
+    class Unionizer
+
+      # Default number of datasets to combine in each chunk
+      DEFAULT_CHUNK_SIZE = 100
+
+      # Represents a chunk of datasets to be combined via UNION.
+      #
+      # Each chunk handles a subset of datasets, creates a temporary table/view
+      # for the combined result, and provides access to the unified dataset.
+      class Chunk
+
+        # @!attribute [r] db
+        #   @return [Sequel::Database] The database connection
+        # @!attribute [r] dses
+        #   @return [Array<Sequel::Dataset>] The datasets in this chunk
+        # @!attribute [r] opts
+        #   @return [Hash] Options for the union operation
+        attr_reader :db, :dses, :opts
+
+        # Creates a new chunk instance.
+        #
+        # @param db [Sequel::Database] The database connection
+        # @param dses [Array<Sequel::Dataset>] The datasets to combine
+        # @param opts [Hash] Options for the union operation
+        def initialize(db, dses, opts)
+          @db = db
+          @dses = dses
+          @opts = opts
+        end
+
+        # Returns the unified dataset created by combining all datasets in this chunk.
+        #
+        # @return [Sequel::Dataset] The combined dataset
+        def union
+          @union ||= dses.reduce { |a, b| a.union(b, all: opts[:all], from_self: opts[:from_self]) }
+        end
+
+        # Generates a unique name for the temporary table/view.
+        #
+        # The name is based on a hash of the SQL query to ensure uniqueness
+        # and avoid collisions when multiple unionize operations are running.
+        #
+        # @return [Symbol] The temporary table/view name
+        def name
+          @name ||= :"#{opts[:temp_table_prefix]}_#{Digest::SHA1.hexdigest(union.sql)}"
+        end
+
+        # Creates a temporary table or view for this chunk's union result.
+        #
+        # The method used depends on the database type:
+        # - Spark: Creates a temporary view
+        # - DuckDB: Creates a temporary table
+        #
+        # @raise [RuntimeError] If the database type is not supported
+        # @return [void]
+        def create
+          if db.database_type == :spark
+            db.create_view(name, union, temp: true)
+          elsif db.database_type == :duckdb
+            db.create_table(name, temp: true, as: union)
+          else
+            raise "Unsupported database type: #{db.database_type}"
+          end
+        end
+
+      end
+
+      # @!attribute [r] db
+      #   @return [Sequel::Database] The database connection
+      attr_reader :db
+
+      # Creates a new Unionizer instance.
+      #
+      # @param db [Sequel::Database] The database connection
+      # @param ds_set [Array<Sequel::Dataset>] The datasets to combine
+      # @param opts [Hash] Options for the union operation
+      # @option opts [Integer] :chunk_size (100) Number of datasets per chunk
+      # @option opts [String] :temp_table_prefix ('temp_union') Prefix for temporary tables
+      # @option opts [Boolean] :all (false) Use UNION ALL instead of UNION
+      # @option opts [Boolean] :from_self (true) Wrap individual datasets in subqueries
+      def initialize(db, ds_set, opts = {})
+        @db = db
+        @ds_set = ds_set
+        @opts = opts
+        opts[:chunk_size] ||= DEFAULT_CHUNK_SIZE
+        opts[:temp_table_prefix] ||= 'temp_union'
+        opts[:all] ||= false
+        opts[:from_self] = opts.fetch(:from_self, true)
+      end
+
+      # Performs the unionization of datasets.
+      #
+      # This method recursively chunks the datasets, creates temporary tables/views
+      # for each chunk, and then combines them until a single dataset remains.
+      #
+      # @param dses [Array<Sequel::Dataset>] The datasets to combine (defaults to @ds_set)
+      # @return [Sequel::Dataset] The final combined dataset
+      def unionize(dses = @ds_set)
+        chunks = dses.each_slice(@opts[:chunk_size]).map do |chunk_of_dses|
+          Chunk.new(db, chunk_of_dses, @opts)
+        end
+
+        return chunks.first.union if chunks.size == 1
+
+        unionize(chunks.each(&:create).map { |chunk| db[chunk.name] })
+      end
+
+    end
+
+    # Efficiently combines multiple datasets using UNION operations.
+    #
+    # This method handles large numbers of datasets by chunking them into
+    # manageable groups, creating temporary tables/views for intermediate
+    # results, and recursively combining them until a single dataset is produced.
+    #
+    # @param ds_set [Array<Sequel::Dataset>] The datasets to combine via UNION
+    # @param opts [Hash] Options for the union operation
+    # @option opts [Integer] :chunk_size (100) Number of datasets to combine in each chunk
+    # @option opts [String] :temp_table_prefix ('temp_union') Prefix for temporary table names
+    # @option opts [Boolean] :all (false) Use UNION ALL instead of UNION (keeps duplicates)
+    # @option opts [Boolean] :from_self (true) Wrap individual datasets in subqueries
+    #
+    # @return [Sequel::Dataset] The combined dataset
+    #
+    # @example Basic union of datasets
+    #   db.unionize([ds1, ds2, ds3, ds4])
+    #
+    # @example Union all with custom chunk size
+    #   db.unionize(datasets, all: true, chunk_size: 50)
+    #
+    # @example Custom temporary table prefix
+    #   db.unionize(datasets, temp_table_prefix: 'my_union_batch')
+    def unionize(ds_set, opts = {})
+      Unionizer.new(self, ds_set, opts).unionize
+    end
+
+  end
+
+  Database.register_extension(:unionize, Unionize)
+
+end