Capistrano 3 support

examples/Capfile

@@ -0,0 +1,19 @@
+require 'capistrano/setup'
+require 'capistrano/datadog'
+
+set :datadog_api_key, 'my_api_key'
+set :stage, :test
+# set :format, :pretty
+# set :log_level, :info
+# set :pty, true
+
+server "host0", roles: ["thing"]
+server "host1", roles: ["other_thing"]
+
+desc "Hello world"
+task :hello do
+  on roles(:all) do |host|
+    info capture('echo "$(date): Hello from $(whoami)@$(hostname) !"')
+  end
+end
+

lib/capistrano/README.md

@@ -7,7 +7,7 @@ To set up your Capfile:
     require "capistrano/datadog"
     set :datadog_api_key, "my_api_key"
 
-You can find your Datadog API key [here](https://app.datad0g.com/account/settings#api). If you don't have a Datadog account, you can sign up for one [here](http://www.datadoghq.com/).
+You can find your Datadog API key [here](https://app.datadoghq.com/account/settings#api). If you don't have a Datadog account, you can sign up for one [here](http://www.datadoghq.com/).
 
 `capistrano/datadog` will capture each Capistrano task that that Capfile runs, including the roles that the task applies to and any logging output that it emits and submits them as events to Datadog at the end of the execution of all the tasks. If sending to Datadog fails for any reason, your scripts will still succeed.
 

lib/capistrano/datadog.rb

@@ -1,57 +1,44 @@
-require "benchmark"
 require "etc"
 require "digest/md5"
-require "socket"
-require "time"
 require "timeout"
-require "delegate"
 
 require "dogapi"
 
-# Monkeypatch capistrano to collect data about the tasks that it's running
-module Capistrano
-  class Configuration
-    module Execution
-      # Attempts to locate the task at the given fully-qualified path, and
-      # execute it. If no such task exists, a Capistrano::NoSuchTaskError
-      # will be raised.
-      # Also, capture the time the task took to execute, and the logs it
-      # outputted for submission to Datadog
-      def find_and_execute_task(path, hooks = {})
-        task = find_task(path) or raise NoSuchTaskError, "the task `#{path}' does not exist"
-        result = nil
-        reporter = Capistrano::Datadog.reporter
-        timing = Benchmark.measure(task.fully_qualified_name) do
-          # Set the current task so that the logger knows which task to
-          # associate the logs with
-          reporter.current_task = task.fully_qualified_name
-          trigger(hooks[:before], task) if hooks[:before]
-          result = execute_task(task)
-          trigger(hooks[:after], task) if hooks[:after]
-          reporter.current_task = nil
-        end
-        # Collect the timing in a list for later reporting
-        reporter.record_task task, timing
-        result
-      end
-    end
-  end
-
-  class Logger
-    # Make the device attribute writeable so we can swap it out
-    # with something that captures logging out by task
-    attr_accessor   :device
-  end
-end
-
-
 module Capistrano
   module Datadog
     # Singleton method for Reporter
     def self.reporter()
       @reporter || @reporter = Reporter.new
     end
 
+    def self.cap_version()
+      if @cap_version.nil? then
+        if Configuration.respond_to? :instance then
+          @cap_version = :v2
+        else
+          @cap_version = :v3
+        end
+      end
+      @cap_version
+    end
+
+    def self.submit(api_key)
+      begin
+        if api_key
+          dog = Dogapi::Client.new(api_key)
+          reporter.report.each do |event|
+            dog.emit_event event
+          end
+        else
+          puts "No api key set, not submitting to Datadog"
+        end
+      rescue Timeout::Error => e
+        puts "Could not submit to Datadog, request timed out."
+      rescue => e
+        puts "Could not submit to Datadog: #{e.inspect}\n#{e.backtrace.join("\n")}"
+      end
+    end
+
     # Collects info about the tasks that ran in order to submit to Datadog
     class Reporter
       attr_accessor :current_task
@@ -62,15 +49,12 @@ def initialize()
         @logging_output = {}
       end
 
-      def record_task(task, timing)
-        roles = task.options[:roles]
-        if roles.is_a? Proc
-          roles = roles.call
-        end
+      def record_task(task_name, timing, roles, stage=nil)
         @tasks << {
-          :name   => task.fully_qualified_name,
-          :timing => timing.real,
-          :roles  => roles
+          :name   => task_name,
+          :timing => timing,
+          :roles  => roles,
+          :stage  => stage
         }
       end
 
@@ -93,11 +77,18 @@ def report()
           name  = task[:name]
           roles = Array(task[:roles]).map(&:to_s).sort
           tags  = ["#capistrano"] + (roles.map { |t| '#role:' + t })
+          if !task[:stage].nil? and !task[:stage].empty? then
+            tags << "#stage:#{task[:stage]}"
+          end
           title = "%s@%s ran %s on %s with capistrano in %.2f secs" % [user, hostname, name, roles.join(', '), task[:timing]]
           type  = "deploy"
           alert_type = "success"
           source_type = "capistrano"
-          message = "@@@" + "\n" + (@logging_output[name] || []).join('') + "@@@"
+          message_content = (@logging_output[name] || []).join('')
+          message = if !message_content.empty? then
+            # Strip out color control characters
+            message_content = message_content.gsub(/\e\[(\d+)m/, '')
+            "@@@\n#{message_content}@@@" else "" end
 
           Dogapi::Event.new(message,
             :msg_title        => title,
@@ -111,46 +102,14 @@ def report()
       end
     end
 
-    class LogCapture < SimpleDelegator
-      def puts(message)
-        Capistrano::Datadog::reporter.record_log message
-        __getobj__.puts message
-      end
-    end
-  end
-
-  if Configuration.respond_to? :instance then
-    # Capistrano v2
-    Configuration.instance(:must_exist).load do
-      # Wrap the existing logging target with the Datadog capture class
-      logger.device = Datadog::LogCapture.new logger.device
-
-      # Trigger the Datadog submission once all the tasks have run
-      on :exit, "datadog:submit"
-      namespace :datadog do
-        desc "Submit the tasks that have run to Datadog as events"
-        task :submit do |ns|
-          begin
-            api_key = variables[:datadog_api_key]
-            if api_key
-              dog = Dogapi::Client.new(api_key)
-              Datadog::reporter.report.each do |event|
-                dog.emit_event event
-              end
-            else
-              puts "No api key set, not submitting to Datadog"
-            end
-          rescue Timeout::Error => e
-            puts "Could not submit to Datadog, request timed out."
-          rescue => e
-            puts "Could not submit to Datadog: #{e.inspect}\n#{e.backtrace.join("\n")}"
-          end
-        end
-      end
-    end
-  else
-    # No support yet for Capistrano v3
-    puts "No Datadog events will be sent since Datadog currently only supports Capistrano v2. For more info, see https://github.com/DataDog/dogapi-rb/issues/40."
   end
+end
 
+case Capistrano::Datadog::cap_version
+when :v2
+  require 'capistrano/datadog/v2'
+when :v3
+  require 'capistrano/datadog/v3'
+else
+  puts "Unknown version: #{Capistrano::Datadog::cap_version.inspect}"
 end

lib/capistrano/datadog/v2.rb

@@ -0,0 +1,74 @@
+require "benchmark"
+require "delegate"
+
+# Capistrano v2
+
+# Monkeypatch capistrano to collect data about the tasks that it's running
+module Capistrano
+  class Configuration
+    module Execution
+      # Attempts to locate the task at the given fully-qualified path, and
+      # execute it. If no such task exists, a Capistrano::NoSuchTaskError
+      # will be raised.
+      # Also, capture the time the task took to execute, and the logs it
+      # outputted for submission to Datadog
+      def find_and_execute_task(path, hooks = {})
+        task = find_task(path) or raise NoSuchTaskError, "the task `#{path}' does not exist"
+        result = nil
+        reporter = Capistrano::Datadog.reporter
+        task_name = task.fully_qualified_name
+        timing = Benchmark.measure(task_name) do
+          # Set the current task so that the logger knows which task to
+          # associate the logs with
+          reporter.current_task = task_name
+          trigger(hooks[:before], task) if hooks[:before]
+          result = execute_task(task)
+          trigger(hooks[:after], task) if hooks[:after]
+          reporter.current_task = nil
+        end
+
+        # Record the task name, its timing and roles
+        roles = task.options[:roles]
+        if roles.is_a? Proc
+          roles = roles.call
+        end
+        reporter.record_task(task_name, timing.real, roles, task.namespace.variables[:stage])
+
+        # Return the original result
+        result
+      end
+    end
+  end
+
+  class Logger
+    # Make the device attribute writeable so we can swap it out
+    # with something that captures logging out by task
+    attr_accessor   :device
+  end
+end
+
+
+module Capistrano
+  module Datadog
+    class LogCapture < SimpleDelegator
+      def puts(message)
+        Capistrano::Datadog::reporter.record_log message
+        __getobj__.puts message
+      end
+    end
+
+    Configuration.instance(:must_exist).load do
+      # Wrap the existing logging target with the Datadog capture class
+      logger.device = Datadog::LogCapture.new logger.device
+
+      # Trigger the Datadog submission once all the tasks have run
+      on :exit, "datadog:submit"
+      namespace :datadog do
+        desc "Submit the tasks that have run to Datadog as events"
+        task :submit do |ns|
+          Capistrano::Datadog.submit variables[:datadog_api_key]
+        end
+      end
+    end
+  end
+end

lib/capistrano/datadog/v3.rb

@@ -0,0 +1,57 @@
+require "benchmark"
+require "sshkit/formatters/pretty"
+
+# Capistrano v3 uses Rake's DSL instead of its own
+
+module Rake
+  class Task
+    alias old_invoke invoke
+    def invoke(*args)
+      result = nil
+      reporter = Capistrano::Datadog.reporter
+      task_name = name
+      reporter.current_task = task_name
+      timing = Benchmark.measure(task_name) do
+        result = old_invoke(*args)
+      end
+      reporter.record_task(task_name, timing.real, roles,
+        Capistrano::Configuration.env.fetch(:stage))
+      result
+    end
+  end
+end
+
+module Capistrano
+  module Datadog
+    class CaptureIO
+      def initialize(wrapped)
+        @wrapped = wrapped
+      end
+
+      def write(*args)
+        @wrapped.write(*args)
+        args.each { |arg| Capistrano::Datadog.reporter.record_log(arg) }
+      end
+      alias :<< :write
+
+      def close
+        @wrapped.close
+      end
+    end
+  end
+end
+
+module SSHKit
+  module Formatter
+    class Pretty
+      def initialize(oio)
+        super(Capistrano::Datadog::CaptureIO.new(oio))
+      end
+    end
+  end
+end
+
+at_exit do
+  api_key = Capistrano::Configuration.env.fetch :datadog_api_key
+  Capistrano::Datadog.submit api_key
+end