From 8d50ed82fd03592a25bfbcaab9902429dae41194 Mon Sep 17 00:00:00 2001
From: Mildred Ki'Lya <mildred-pub.git@mildred.fr>
Date: Tue, 28 Sep 2021 16:46:39 +0200
Subject: [PATCH] Allow to override environment on spring client startup

Spring.env_override can be set to
- false (default) for current behaviour
- true to always override client provided environment variables in the
  application
- an array listing the environment variables to override in the
  application

The after_fork callback can take an ptional argument that contains the
Spring provided environment to help the application detect when there is
a mismatch between application environment and provided environment.
---
 README.md                   | 27 +++++++++++++++++++++++++++
 lib/spring/application.rb   | 26 ++++++++++++++++++++++----
 lib/spring/configuration.rb |  3 ++-
 3 files changed, 51 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index a61acef4..60583abd 100644
--- a/README.md
+++ b/README.md
@@ -358,6 +358,16 @@ end
 If you want to register multiple callbacks you can simply call
 `Spring.after_fork` multiple times with different blocks.
 
+If you accept an argument, it will contain an object with information about the
+Spring client request and server state. You can use it to detect environment
+variable mismatch between the invocation and the spring instance startup:
+
+```ruby
+Spring.after_fork do |cb|
+  raise "error" if cb.env['MASTER'].present? && cb.env['MASTER'] != ENV['MASTER']
+end
+```
+
 ### Watching files and directories
 
 Spring will automatically detect file changes to any file loaded when the server
@@ -386,6 +396,23 @@ a command runs:
 Spring.quiet = true
 ```
 
+### Environment variable overrides
+
+When the spring server starts up a new instance of the application, it does not
+set all environment variables by default. Only the environment variables that
+were set when the spring server started and were not changed by the application
+are propagated.
+
+This can be customized with `Spring.env_override`:
+
+- `Spring.env_override = false` : default behavior noted above
+
+- `Spring.env_override = true` : all variables are applied, regardless of if
+  they were changed by the application or not.
+
+- `Spring.env_override = []` (array of strings) : in addition to the usual
+  behavior, the variable listed here will be propagated
+
 ### Environment variables
 
 The following environment variables are used by Spring:
diff --git a/lib/spring/application.rb b/lib/spring/application.rb
index fe640043..6027c448 100644
--- a/lib/spring/application.rb
+++ b/lib/spring/application.rb
@@ -187,13 +187,27 @@ def serve(client)
         # Delete all env vars which are unchanged from before Spring started
         original_env.each { |k, v| ENV.delete k if ENV[k] == v }
 
+        # if env_override is an array, delete the variables listed
+        if Spring.env_override.is_a? Array
+          Spring.env_override.each { |k| ENV.delete k }
+        end
+
         # Load in the current env vars, except those which *were* changed when Spring started
-        env.each { |k, v| ENV[k] ||= v }
+        # if env_override is true, load in *all* the current vars
+        env.each do |k, v|
+          if Spring.env_override == true
+            ENV[k] = v
+          else
+            ENV[k] ||= v
+          end
+        end
 
         connect_database
         srand
 
-        invoke_after_fork_callbacks
+        invoke_after_fork_callbacks(OpenStruct.new(
+          env: env,
+          original_env: original_env))
         shush_backtraces
 
         command.call
@@ -256,9 +270,13 @@ def setup(command)
       end
     end
 
-    def invoke_after_fork_callbacks
+    def invoke_after_fork_callbacks(arg)
       Spring.after_fork_callbacks.each do |callback|
-        callback.call
+        if callback.arity == 1
+          callback.call(arg)
+        else
+          callback.call
+        end
       end
     end
 
diff --git a/lib/spring/configuration.rb b/lib/spring/configuration.rb
index eba8e9a7..45926d80 100644
--- a/lib/spring/configuration.rb
+++ b/lib/spring/configuration.rb
@@ -2,7 +2,7 @@
 
 module Spring
   class << self
-    attr_accessor :application_root, :quiet
+    attr_accessor :application_root, :quiet, :env_override
 
     def gemfile
       if /\s1.9.[0-9]/ ===  Bundler.ruby_scope.gsub(/[\/\s]+/,'')
@@ -55,4 +55,5 @@ def find_project_root(current_dir)
   end
 
   self.quiet = false
+  self.env_override = false
 end