Refactor

lib-injection/auto_inject.rb

@@ -1,67 +1,74 @@
-if !ENV['skip_autoinject']
-  if system 'skip_autoinject=true bundle show ddtrace'
-    puts 'ddtrace already installed... skipping auto-injection'
-  else
-    version = "<DD_TRACE_VERSION_TO_BE_REPLACED>"
-    sha = "<DD_TRACE_SHA_TO_BE_REPLACED>"
 
-    condition =
-      if !version.empty?
-        # For public release
-        "--version '#{version.gsub(/^v/, '')}'"
-      elsif !sha.empty?
-        # For internal testing
-        "--github 'datadog/dd-trace-rb' --ref '#{sha}'"
-      else
-        puts "NO VERSION"
-      end
+return if ENV['skip_autoinject']
 
-    puts "ddtrace is not installed... Perform auto-injection... for dd-trace-rb"
+if system 'skip_autoinject=true bundle show ddtrace'
+  puts '[DATADOG LIB INJECTION] ddtrace already installed... skipping auto-injection'
+  return
+end
 
-    require 'bundler'
-    require 'fileutils'
+begin
+  require 'bundler'
+  # require 'shellwords'
 
-    gemfile  = Bundler::SharedHelpers.default_gemfile
-    lockfile = Bundler::SharedHelpers.default_lockfile
+  if Bundler.frozen_bundle?
+    puts "[DATADOG LIB INJECTION] Cannot inject with frozen Bundler"
+    return
+  end
 
-    puts gemfile
-    puts lockfile
+  version = "<DD_TRACE_VERSION_TO_BE_REPLACED>"
+  sha = "<DD_TRACE_SHA_TO_BE_REPLACED>"
 
-    if gemfile.basename.to_s == 'gems.rb'
-      datadog_gemfile = gemfile.dirname + "datadog-Gemfile"
-      datadog_lockfile = lockfile.dirname + "datadog-Gemfile.lock"
-    else
-      datadog_gemfile = gemfile.dirname + ("datadog-#{gemfile.basename}")
-      datadog_lockfile = lockfile.dirname + ("datadog-#{lockfile.basename}")
-    end
+  # Stronger restrict
+  condition = if !version.empty?
+    # For public release
+    "--version #{version.gsub(/^v/, '')}"
+  elsif !sha.empty?
+    # For internal testing
+    "--github datadog/dd-trace-rb --ref #{sha}"
+  end
+
+  unless condition
+    puts "[DATADOG LIB INJECTION] NO VERSION"
+    return
+  end
+
+  puts "[DATADOG LIB INJECTION] ddtrace is not installed... Perform lib injection for dd-trace-rb."
 
-    puts datadog_gemfile
-    puts datadog_lockfile
+  gemfile   = Bundler::SharedHelpers.default_gemfile
+  lockfile  = Bundler::SharedHelpers.default_lockfile
 
+  if gemfile.basename.to_s == 'gems.rb'
+    datadog_gemfile = gemfile.dirname + "datadog-Gemfile"
+    datadog_lockfile = lockfile.dirname + "datadog-Gemfile.lock"
+  else
+    datadog_gemfile = gemfile.dirname + ("datadog-#{gemfile.basename}")
+    datadog_lockfile = lockfile.dirname + ("datadog-#{lockfile.basename}")
+  end
+
+  require 'fileutils'
+
+  begin
     # Copies for trial
     FileUtils.cp gemfile, datadog_gemfile
     FileUtils.cp lockfile, datadog_lockfile
 
     if system "skip_autoinject=true BUNDLE_GEMFILE=#{datadog_gemfile} bundle add ddtrace #{condition} --require ddtrace/auto_instrument"
-      puts 'ddtrace added to bundle...'
+      puts '[DATADOG LIB INJECTION] ddtrace added to bundle...'
 
       # Trial success, replace the original
       FileUtils.cp datadog_gemfile, gemfile
       FileUtils.cp datadog_lockfile, lockfile
     else
-      puts 'Something went wrong when adding ddtrace to bundle...'
+      puts '[DATADOG LIB INJECTION] Something went wrong when adding ddtrace to bundle...'
     end
 
+  rescue => e
+    puts "[DATADOG LIB INJECTION] #{e}"
+  ensure
     # Remove the copies
     FileUtils.rm datadog_gemfile
     FileUtils.rm datadog_lockfile
   end
-
-  begin
-    require 'ddtrace'
-  rescue LoadError => e
-    puts e
-  else
-    puts 'ddtrace loaded...'
-  end
+rescue LoadError, Bundler::BundlerError => e
+  puts "[DATADOG LIB INJECTION] #{e}"
 end