Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Handle thread safety for Dotenv.restore, Add Dotenv.modify #475

Merged
merged 1 commit into from
Jan 25, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
97 changes: 71 additions & 26 deletions lib/dotenv.rb
Original file line number Diff line number Diff line change
Expand Up @@ -3,22 +3,21 @@
require "dotenv/missing_keys"
require "dotenv/diff"

# The top level Dotenv module. The entrypoint for the application logic.
# Shim to load environment variables from `.env files into `ENV`.
module Dotenv
class << self
attr_accessor :instrumenter
end
extend self

# An internal monitor to synchronize access to ENV in multi-threaded environments.
SEMAPHORE = Monitor.new
private_constant :SEMAPHORE

module_function
attr_accessor :instrumenter

# Loads environment variables from one or more `.env` files. See `#parse` for more details.
def load(*filenames, **kwargs)
parse(*filenames, **kwargs) do |env|
def load(*filenames, overwrite: false, ignore: true)
parse(*filenames, overwrite: overwrite, ignore: ignore) do |env|
instrument(:load, env: env) do |payload|
env_before = ENV.to_h
env.apply
payload[:diff] = Dotenv::Diff.new(env_before, ENV.to_h)
env
update(env, overwrite: overwrite)
end
end
end
Expand All @@ -33,14 +32,12 @@ def overwrite(*filenames)
load(*filenames, overwrite: true)
end
alias_method :overload, :overwrite
module_function :overload

# same as `#overwrite`, but raises Errno::ENOENT if any files don't exist
def overwrite!(*filenames)
load(*filenames, overwrite: true, ignore: false)
end
alias_method :overload!, :overwrite!
module_function :overload!

# Parses the given files, yielding for each file if a block is given.
#
Expand All @@ -65,11 +62,60 @@ def parse(*filenames, overwrite: false, ignore: true, &block)
end
end

def instrument(name, payload = {}, &block)
if instrumenter
instrumenter.instrument("#{name}.dotenv", payload, &block)
else
block&.call payload
# Save the current `ENV` to be restored later
def save
instrument(:save) do |payload|
@diff = payload[:diff] = Dotenv::Diff.new
end
end

# Restore `ENV` to a given state
#
# @param env [Hash] Hash of keys and values to restore, defaults to the last saved state
# @param safe [Boolean] Is it safe to modify `ENV`? Defaults to `true` in the main thread, otherwise raises an error.
def restore(env = @diff&.a, safe: Thread.current == Thread.main)
diff = Dotenv::Diff.new(b: env)
return unless diff.any?

unless safe
raise ThreadError, <<~EOE.tr("\n", " ")
Dotenv.restore is not thread safe. Use `Dotenv.modify { }` to update ENV for the duration
of the block in a thread safe manner, or call `Dotenv.restore(safe: true)` to ignore
this error.
EOE
end
instrument(:restore, diff: diff) { ENV.replace(env) }
end

# Update `ENV` with the given hash of keys and values
#
# @param env [Hash] Hash of keys and values to set in `ENV`
# @param overwrite [Boolean] Overwrite existing `ENV` values
def update(env = {}, overwrite: false)
instrument(:update) do |payload|
diff = payload[:diff] = Dotenv::Diff.new do
ENV.update(env.transform_keys(&:to_s)) do |key, old_value, new_value|
# This block is called when a key exists. Return the new value if overwrite is true.
overwrite ? new_value : old_value
end
end
diff.env
end
end

# Modify `ENV` for the block and restore it to its previous state afterwards.
#
# Note that the block is synchronized to prevent concurrent modifications to `ENV`,
# so multiple threads will be executed serially.
#
# @param env [Hash] Hash of keys and values to set in `ENV`
def modify(env = {}, &block)
SEMAPHORE.synchronize do
diff = Dotenv::Diff.new
update(env, overwrite: true)
block.call
ensure
restore(diff.a, safe: true)
end
end

Expand All @@ -79,15 +125,14 @@ def require_keys(*keys)
raise MissingKeys, missing_keys
end

# Save a snapshot of the current `ENV` to be restored later
def save
@snapshot = ENV.to_h.freeze
instrument(:save, snapshot: @snapshot)
end
private

# Restore the previous snapshot of `ENV`
def restore
instrument(:restore, diff: Dotenv::Diff.new(ENV.to_h, @snapshot)) { ENV.replace(@snapshot) }
def instrument(name, payload = {}, &block)
if instrumenter
instrumenter.instrument("#{name}.dotenv", payload, &block)
else
block&.call payload
end
end
end

Expand Down
8 changes: 7 additions & 1 deletion lib/dotenv/autorestore.rb
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,13 @@
setup { Dotenv.save }

# Restore ENV after each test
teardown { Dotenv.restore }
teardown do
Dotenv.restore
rescue ThreadError => e
# Restore will fail if running tests in parallel.
warn e.message
warn "Set `config.dotenv.autorestore = false` in `config/initializers/test.rb`" if defined?(Dotenv::Rails)
end
end
end
end
33 changes: 30 additions & 3 deletions lib/dotenv/diff.rb
Original file line number Diff line number Diff line change
@@ -1,10 +1,28 @@
module Dotenv
# Compare two hashes and return the differences
# A diff between multiple states of ENV.
class Diff
attr_reader :a, :b
attr_reader :a

def initialize(a, b)
# Create a new diff. The initial state defaults to a clone of current ENV. If given a block,
# the state of ENV will be preserved as the final state for comparison. Otherwise, the current
# ENV will be the current state.
def initialize(a: current_env, b: nil, &block)
@a, @b = a, b
if block
begin
block.call self
ensure
@b = current_env
end
end
end

def b
@b || current_env
end

def current_env
ENV.to_h.freeze
end

# Return a Hash of keys added with their new values
Expand All @@ -23,5 +41,14 @@ def changed
[k, [a[k], v]]
end.to_h
end

# Returns a Hash of all added, changed, and removed keys and their new values
def env
@env ||= b.slice(*(added.keys + changed.keys)).merge(removed.transform_values { |v| nil })
end

def any?
[added, removed, changed].any?(&:any?)
end
end
end
12 changes: 1 addition & 11 deletions lib/dotenv/environment.rb
Original file line number Diff line number Diff line change
Expand Up @@ -12,21 +12,11 @@ def initialize(filename, overwrite: false)
end

def load
update Parser.call(read, overwrite: @overwrite)
update Parser.call(read, overwrite: overwrite)
end

def read
File.open(@filename, "rb:bom|utf-8", &:read)
end

def apply
each do |k, v|
if @overwrite
ENV[k] = v
else
ENV[k] ||= v
end
end
end
end
end
10 changes: 6 additions & 4 deletions lib/dotenv/log_subscriber.rb
Original file line number Diff line number Diff line change
Expand Up @@ -9,13 +9,15 @@ def logger
end

def load(event)
diff = event.payload[:diff]
env = event.payload[:env]

# Only show the keys that were added or changed
changed = env.slice(*(diff.added.keys + diff.changed.keys)).keys.map { |key| color_var(key) }
info "Loaded #{color_filename(env.filename)}"
end

info "Set #{changed.to_sentence} from #{color_filename(env.filename)}" if changed.any?
def update(event)
diff = event.payload[:diff]
changed = diff.env.keys.map { |key| color_var(key) }
debug "Set #{changed.to_sentence}" if diff.any?
end

def save(event)
Expand Down
10 changes: 9 additions & 1 deletion spec/dotenv/diff_spec.rb
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
describe Dotenv::Diff do
let(:before) { {} }
let(:after) { {} }
subject { Dotenv::Diff.new(before, after) }
subject { Dotenv::Diff.new(a: before, b: after) }

context "no changes" do
let(:before) { {"A" => 1} }
Expand All @@ -12,6 +12,8 @@
it { expect(subject.added).to eq({}) }
it { expect(subject.removed).to eq({}) }
it { expect(subject.changed).to eq({}) }
it { expect(subject.any?).to eq(false) }
it { expect(subject.env).to eq({}) }
end

context "key added" do
Expand All @@ -20,6 +22,8 @@
it { expect(subject.added).to eq("A" => 1) }
it { expect(subject.removed).to eq({}) }
it { expect(subject.changed).to eq({}) }
it { expect(subject.any?).to eq(true) }
it { expect(subject.env).to eq("A" => 1) }
end

context "key removed" do
Expand All @@ -28,6 +32,8 @@
it { expect(subject.added).to eq({}) }
it { expect(subject.removed).to eq("A" => 1) }
it { expect(subject.changed).to eq({}) }
it { expect(subject.any?).to eq(true) }
it { expect(subject.env).to eq("A" => nil) }
end

context "key changed" do
Expand All @@ -37,5 +43,7 @@
it { expect(subject.added).to eq({}) }
it { expect(subject.removed).to eq({}) }
it { expect(subject.changed).to eq("A" => [1, 2]) }
it { expect(subject.any?).to eq(true) }
it { expect(subject.env).to eq("A" => 2) }
end
end
28 changes: 0 additions & 28 deletions spec/dotenv/environment_spec.rb
Original file line number Diff line number Diff line change
Expand Up @@ -16,34 +16,6 @@
end
end

describe "apply" do
it "sets variables in ENV" do
subject.apply
expect(ENV["OPTION_A"]).to eq("1")
end

it "does not overwrite defined variables" do
ENV["OPTION_A"] = "predefined"
subject.apply
expect(ENV["OPTION_A"]).to eq("predefined")
end

context "with overwrite: true" do
subject { env("OPTION_A=1\nOPTION_B=2", overwrite: true) }

it "sets variables in the ENV" do
subject.apply
expect(ENV["OPTION_A"]).to eq("1")
end

it "overwrites defined variables" do
ENV["OPTION_A"] = "predefined"
subject.apply
expect(ENV["OPTION_A"]).to eq("1")
end
end
end

require "tempfile"
def env(text, ...)
file = Tempfile.new("dotenv")
Expand Down
30 changes: 17 additions & 13 deletions spec/dotenv/log_subscriber_spec.rb
Original file line number Diff line number Diff line change
Expand Up @@ -11,31 +11,35 @@
Dotenv::Rails.logger = Logger.new(logs)
end

context "set" do
it "logs when a new instance variable is set" do
describe "load" do
it "logs when a file is loaded" do
Dotenv.load(fixture_path("plain.env"))
expect(logs.string).to match(/Set.*PLAIN.*from.*plain.env/)
expect(logs.string).to match(/Loaded.*plain.env/)
expect(logs.string).to match(/Set.*PLAIN/)
end
end

context "update" do
it "logs when a new instance variable is set" do
Dotenv.update({"PLAIN" => "true"})
expect(logs.string).to match(/Set.*PLAIN/)
end

it "logs when an instance variable is overwritten" do
ENV["PLAIN"] = "nope"
Dotenv.load(fixture_path("plain.env"), overwrite: true)
expect(logs.string).to match(/Set.*PLAIN.*from.*plain.env/)
Dotenv.update({"PLAIN" => "true"}, overwrite: true)
expect(logs.string).to match(/Set.*PLAIN/)
end

it "does not log when an instance variable is not overwritten" do
# load everything once and clear the logs
Dotenv.load(fixture_path("plain.env"))
logs.truncate(0)

# load again
Dotenv.load(fixture_path("plain.env"))
expect(logs.string).not_to match(/Set.*plain.env/i)
ENV["FOO"] = "existing"
Dotenv.update({"FOO" => "new"})
expect(logs.string).not_to match(/FOO/)
end

it "does not log when an instance variable is unchanged" do
ENV["PLAIN"] = "true"
Dotenv.load(fixture_path("plain.env"), overwrite: true)
Dotenv.update({"PLAIN" => "true"}, overwrite: true)
expect(logs.string).not_to match(/PLAIN/)
end
end
Expand Down
Loading