From 598e7ab97fb7c57f0fb934a237ad30719d8934e4 Mon Sep 17 00:00:00 2001
From: David Heinemeier Hansson <david@loudthinking.com>
Date: Sun, 22 Jan 2023 11:27:31 +0100
Subject: [PATCH] Add power to follow logs on app and traefik

---
 lib/mrsk/cli/app.rb          | 28 +++++++++++++++++++---------
 lib/mrsk/cli/traefik.rb      | 20 +++++++++++++++++---
 lib/mrsk/commands/app.rb     | 10 +++++++++-
 lib/mrsk/commands/base.rb    |  4 ++++
 lib/mrsk/commands/traefik.rb | 17 +++++++++++------
 5 files changed, 60 insertions(+), 19 deletions(-)

diff --git a/lib/mrsk/cli/app.rb b/lib/mrsk/cli/app.rb
index b34355a8..9363afd5 100644
--- a/lib/mrsk/cli/app.rb
+++ b/lib/mrsk/cli/app.rb
@@ -81,22 +81,32 @@ class Mrsk::Cli::App < Mrsk::Cli::Base
   option :since, aliases: "-s", desc: "Show logs since timestamp (e.g. 2013-01-02T13:23:37Z) or relative (e.g. 42m for 42 minutes)"
   option :lines, type: :numeric, aliases: "-n", desc: "Number of log lines to pull from each server"
   option :grep, aliases: "-g", desc: "Show lines with grep match only (use this to fetch specific requests by id)"
+  option :follow, aliases: "-f", desc: "Follow logs on primary server (or specific host set by --hosts)"
   def logs
     # FIXME: Catch when app containers aren't running
 
-    since = options[:since]
-    lines = options[:lines]
-    grep  = options[:grep]
+    grep = options[:grep]
 
-    on(MRSK.hosts) do |host|
-      begin
-        puts_by_host host, capture_with_info(*MRSK.app.logs(since: since, lines: lines, grep: grep))
-      rescue SSHKit::Command::Failed
-        puts_by_host host, "Nothing found"
+    if options[:follow]
+      run_locally do
+        info "Following logs on #{MRSK.primary_host}..."
+        info MRSK.app.follow_logs(host: MRSK.primary_host, grep: grep)
+        exec MRSK.app.follow_logs(host: MRSK.primary_host, grep: grep)
+      end
+    else
+      since = options[:since]
+      lines = options[:lines]
+
+      on(MRSK.hosts) do |host|
+        begin
+          puts_by_host host, capture_with_info(*MRSK.app.logs(since: since, lines: lines, grep: grep))
+        rescue SSHKit::Command::Failed
+          puts_by_host host, "Nothing found"
+        end
       end
     end
   end
-  
+
   desc "remove", "Remove app containers and images from servers"
   option :only, default: "", desc: "Use 'containers' or 'images'"
   def remove
diff --git a/lib/mrsk/cli/traefik.rb b/lib/mrsk/cli/traefik.rb
index 002fa19d..975db7d8 100644
--- a/lib/mrsk/cli/traefik.rb
+++ b/lib/mrsk/cli/traefik.rb
@@ -37,11 +37,25 @@ class Mrsk::Cli::Traefik < Mrsk::Cli::Base
   desc "logs", "Show log lines from Traefik on servers"
   option :since, aliases: "-s", desc: "Show logs since timestamp (e.g. 2013-01-02T13:23:37Z) or relative (e.g. 42m for 42 minutes)"
   option :lines, type: :numeric, aliases: "-n", desc: "Number of log lines to pull from each server"
+  option :grep, aliases: "-g", desc: "Show lines with grep match only (use this to fetch specific requests by id)"
+  option :follow, aliases: "-f", desc: "Follow logs on primary server (or specific host set by --hosts)"
   def logs
-    since = options[:since]
-    lines = options[:lines]
+    grep = options[:grep]
 
-    on(MRSK.traefik_hosts) { |host| puts_by_host host, capture(*MRSK.traefik.logs(since: since, lines: lines)), type: "Traefik" }
+    if options[:follow]
+      run_locally do
+        info "Following logs on #{MRSK.primary_host}..."
+        info MRSK.traefik.follow_logs(host: MRSK.primary_host, grep: grep)
+        exec MRSK.traefik.follow_logs(host: MRSK.primary_host, grep: grep)
+      end
+    else
+      since = options[:since]
+      lines = options[:lines]
+
+      on(MRSK.traefik_hosts) do |host|
+        puts_by_host host, capture(*MRSK.traefik.logs(since: since, lines: lines, grep: grep)), type: "Traefik"
+      end
+    end
   end
 
   desc "remove", "Remove Traefik container and image from servers"
diff --git a/lib/mrsk/commands/app.rb b/lib/mrsk/commands/app.rb
index e1faf9dc..245d0e33 100644
--- a/lib/mrsk/commands/app.rb
+++ b/lib/mrsk/commands/app.rb
@@ -57,6 +57,14 @@ class Mrsk::Commands::App < Mrsk::Commands::Base
       *command
   end
 
+  def follow_logs(host:, grep: nil)
+    run_over_ssh pipe(
+      current_container_id,
+      "xargs docker logs -t -n 10 -f 2>&1",
+      ("grep '#{grep}'" if grep)
+    ).join(" "), host: host
+  end
+
   def console(host:)
     exec_over_ssh "bin/rails", "c", host: host
   end
@@ -79,7 +87,7 @@ class Mrsk::Commands::App < Mrsk::Commands::Base
 
   private
     def exec_over_ssh(*command, host:)
-      "ssh -t #{config.ssh_user}@#{host} '#{run_exec(*command, interactive: true).join(" ")}'"
+      run_over_ssh run_exec(*command, interactive: true).join(" "), host: host
     end
 
     def service_filter
diff --git a/lib/mrsk/commands/base.rb b/lib/mrsk/commands/base.rb
index d2a0ae89..e7305701 100644
--- a/lib/mrsk/commands/base.rb
+++ b/lib/mrsk/commands/base.rb
@@ -23,5 +23,9 @@ module Mrsk::Commands
       def docker(*args)
         args.compact.unshift :docker
       end
+
+      def run_over_ssh(command, host:)
+        "ssh -t #{config.ssh_user}@#{host} '#{command}'"
+      end
   end
 end
diff --git a/lib/mrsk/commands/traefik.rb b/lib/mrsk/commands/traefik.rb
index 8ede5c6c..a2d395d7 100644
--- a/lib/mrsk/commands/traefik.rb
+++ b/lib/mrsk/commands/traefik.rb
@@ -24,12 +24,17 @@ class Mrsk::Commands::Traefik < Mrsk::Commands::Base
     docker :ps, "--filter", "name=traefik"
   end
 
-  def logs(since: nil, lines: nil)
-    docker :logs, "traefik",
-      (" --since #{since}" if since),
-      (" -n #{lines}" if lines),
-      "-t",
-      "2>&1"
+  def logs(since: nil, lines: nil, grep: nil)
+    pipe \
+      docker(:logs, "traefik", (" --since #{since}" if since), (" -n #{lines}" if lines), "-t", "2>&1"),
+      ("grep '#{grep}'" if grep)
+  end
+
+  def follow_logs(host:, grep: nil)
+    run_over_ssh pipe(
+      docker(:logs, "traefik", "-t", "-n", "10", "-f", "2>&1"),
+      ("grep '#{grep}'" if grep)
+    ).join(" "), host: host
   end
 
   def remove_container