switchman 3.6.4 → 3.6.6

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: fe8b34e257391bacd4982df3b99510262ebd3ed88d9f00c943a018fdec7d015e
-  data.tar.gz: d55e7363f9595127992f2f874f39e1c017ac9d6e757d8b6b8e0a30fd4bd778b0
+  metadata.gz: 8c9938f088ef1bcce6ba0647f9e46db87d777a8b654a832bc82dbb24be2bea0a
+  data.tar.gz: 22087e848ed37bf0551999ea9923248688a5194501be7d02d07e2a2448fe701d
 SHA512:
-  metadata.gz: 61ef5a93ffd16bf1cca7ac8cd064230e62257c7d12567fb52f0f2e5e7ceac13ab2b34b61077cea8d62cb2390f4224229a3444183a52b877fca883cdfc50e5b10
-  data.tar.gz: 17c1326e8daf6ad3febbaacc48a6e8642befe6c0a0847de83e55cbc40e687ee8cb9850ebd885ace4e5847993171e28a236bfed9b9690ada52b23c2e0d7687f4e
+  metadata.gz: 21931a6cedf1dfb74539e19657a9df58becb94e098a94423a059c3961ee042f36a66b547c40bc62ca66b17498235330743f12f8bc3d75a1c6f5172b25df15d37
+  data.tar.gz: 0220b63454657b7cddd89100a2d1ebe940369d6c0749345153912676741eba50f48d6ec1ae5e68f1eed08c46f4fa4e6c77032a0b0839f46645e7a4b376ef9ac0

data/lib/switchman/active_record/base.rb

@@ -22,13 +22,59 @@ module Switchman
           @integral_id
         end
 
-        %w[transaction insert_all upsert_all].each do |method|
+        def transaction(**)
+          if self != ::ActiveRecord::Base && current_scope
+            current_scope.activate do
+              db = Shard.current(connection_class_for_self).database_server
+              db.unguard { super }
+            end
+          else
+            db = Shard.current(connection_class_for_self).database_server
+            db.unguard { super }
+          end
+        end
+
+        # NOTE: `returning` values are _not_ transposed back to the current shard
+        %w[insert_all upsert_all].each do |method|
           class_eval <<-RUBY, __FILE__, __LINE__ + 1
-            def #{method}(*, **)
-              if self != ::ActiveRecord::Base && current_scope
-                current_scope.activate do
+            def #{method}(attributes, returning: nil, **)
+              scope = self != ::ActiveRecord::Base && current_scope
+              if (target_shard = scope&.primary_shard) == (current_shard = Shard.current(connection_class_for_self))
+                scope = nil
+              end
+              if scope
+                dupped = false
+                attributes.each_with_index do |hash, i|
+                  if dupped || hash.any? { |k, v| sharded_column?(k) }
+                    unless dupped
+                      attributes = attributes.dup
+                      dupped = true
+                    end
+                    attributes[i] = hash.to_h do |k, v|
+                      if sharded_column?(k)
+                        [k, Shard.relative_id_for(v, current_shard, target_shard)]
+                      else
+                        [k, v]
+                      end
+                    end
+                  end
+                end
+              end
+
+              if scope
+                scope.activate do
                   db = Shard.current(connection_class_for_self).database_server
-                  db.unguard { super }
+                  result = db.unguard { super }
+                  if result&.columns&.any? { |c| sharded_column?(c) }
+                    transposed_rows = result.rows.map do |row|
+                      row.map.with_index do |value, i|
+                        sharded_column?(result.columns[i]) ? Shard.relative_id_for(value, target_shard, current_shard) : value
+                      end
+                    end
+                    result = ::ActiveRecord::Result.new(result.columns, transposed_rows, result.column_types)
+                  end
+
+                  result
                 end
               else
                 db = Shard.current(connection_class_for_self).database_server

data/lib/switchman/parallel.rb

@@ -50,16 +50,16 @@ module Switchman
       end
     end
 
-    class PrefixingIO
+    class TransformingIO
       delegate_missing_to :@original_io
 
-      def initialize(prefix, original_io)
-        @prefix = prefix
+      def initialize(transformer, original_io)
+        @transformer = transformer
         @original_io = original_io
       end
 
       def puts(*args)
-        args.flatten.each { |arg| @original_io.puts "#{@prefix}: #{arg}" }
+        args.flatten.each { |arg| @original_io.puts @transformer.call(arg) }
       end
     end
   end

data/lib/switchman/shard.rb

@@ -165,7 +165,7 @@ module Switchman
       #                forking.
       #    exception: - :ignore, :raise, :defer (wait until the end and raise the first
       #                error), or a proc
-      #    output: - :simple, :decorated (with database_server_id:shard_name)
+      #    output: - :simple, :decorated (with database_server_id:shard_name), custom lambda transformer
       def with_each_shard(*args, parallel: false, exception: :raise, output: :simple)
         raise ArgumentError, "wrong number of arguments (#{args.length} for 0...2)" if args.length > 2
 
@@ -181,6 +181,13 @@ module Switchman
           scope, classes = args
         end
 
+        output = if output == :decorated
+                   ->(arg) { "#{Shard.current.description}: #{arg}" }
+                 elsif output == :simple
+                   nil
+                 else
+                   output
+                 end
         parallel = [Environment.cpu_count || 2, 2].min if parallel == true
         parallel = 0 if parallel == false || parallel.nil?
 
@@ -224,7 +231,10 @@ module Switchman
               new_title = [short_parent_name, name].join(" ")
               Process.setproctitle(new_title)
               Switchman.config[:on_fork_proc]&.call
-              with_each_shard(subscope, classes, exception: exception, output: :decorated) do
+              with_each_shard(subscope,
+                              classes,
+                              exception: exception,
+                              output: output || :decorated) do
                 last_description = Shard.current.description
                 Parallel::ResultWrapper.new(yield)
               end
@@ -258,9 +268,9 @@ module Switchman
           next unless shard.database_server
 
           shard.activate(*classes) do
-            if output == :decorated
-              $stdout = Parallel::PrefixingIO.new(shard.description, $stdout)
-              $stderr = Parallel::PrefixingIO.new(shard.description, $stderr)
+            if output
+              $stdout = Parallel::TransformingIO.new(output, $stdout)
+              $stderr = Parallel::TransformingIO.new(output, $stderr)
             end
 
             result.concat Array.wrap(yield)

data/lib/switchman/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Switchman
-  VERSION = "3.6.4"
+  VERSION = "3.6.6"
 end

data/lib/tasks/switchman.rake

@@ -71,12 +71,13 @@ module Switchman
     end
 
     def self.options
-      { parallel: ENV["PARALLEL"].to_i }
+      { exception: (ENV["FAIL_FAST"] == "0") ? :defer : :raise, parallel: ENV["PARALLEL"].to_i }
     end
 
     # classes - an array or proc, to activate as the current shard during the
     # task.
     def self.shardify_task(task_name, classes: [::ActiveRecord::Base])
+      log_format = ENV.fetch("LOG_FORMAT", nil)
       old_task = ::Rake::Task[task_name]
       old_actions = old_task.actions.dup
       old_task.actions.clear
@@ -90,14 +91,56 @@ module Switchman
         ::GuardRail.activate(:deploy) do
           Shard.default.database_server.unguard do
             classes = classes.call if classes.respond_to?(:call)
-            Shard.with_each_shard(scope, classes, **options) do
+
+            # We don't want the shard status messages to be wrapped using a custom log transfomer
+            original_stderr = $stderr
+            original_stdout = $stdout
+            output = if log_format == "json"
+                       lambda { |msg|
+                         JSON.dump(shard: Shard.current.id,
+                                   database_server: Shard.current.database_server.id,
+                                   type: "log",
+                                   message: msg)
+                       }
+                     else
+                       nil
+                     end
+            Shard.with_each_shard(scope, classes, output: output, **options) do
               shard = Shard.current
-              puts "#{shard.id}: #{shard.description}"
+
+              if log_format == "json"
+                original_stdout.puts JSON.dump(
+                  shard: shard.id,
+                  database_server: shard.database_server.id,
+                  type: "started"
+                )
+              else
+                original_stdout.puts "#{shard.id}: #{shard.description}"
+              end
 
               shard.database_server.unguard do
                 old_actions.each { |action| action.call(*task_args) }
               end
+
+              if log_format == "json"
+                original_stdout.puts JSON.dump(
+                  shard: shard.id,
+                  database_server: shard.database_server.id,
+                  type: "completed"
+                )
+              end
               nil
+            rescue => e
+              if log_format == "json"
+                original_stderr.puts JSON.dump(
+                  shard: shard.id,
+                  database_server: shard.database_server.id,
+                  type: "failed",
+                  message: e.full_message
+                )
+              end
+
+              raise
             end
           rescue => e
             if options[:parallel] != 0

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: switchman
 version: !ruby/object:Gem::Version
-  version: 3.6.4
+  version: 3.6.6
 platform: ruby
 authors:
 - Cody Cutrer
@@ -10,7 +10,7 @@ authors:
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2024-05-10 00:00:00.000000000 Z
+date: 2024-08-20 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: activerecord