From 2553b31605b9daa45477d7dc357537b16fd58074 Mon Sep 17 00:00:00 2001
From: Mostafa Ahangarha <ahangarha@riseup.net>
Date: Sun, 12 Nov 2023 18:24:12 +0330
Subject: [PATCH] Add location option to apply-image, ps, run, and run:detached

---
 lib/command/apply_template.rb | 3 ++-
 lib/command/ps.rb             | 5 ++++-
 lib/command/run.rb            | 3 ++-
 lib/command/run_detached.rb   | 3 ++-
 4 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/lib/command/apply_template.rb b/lib/command/apply_template.rb
index abe4a357..0ca6bf28 100644
--- a/lib/command/apply_template.rb
+++ b/lib/command/apply_template.rb
@@ -7,6 +7,7 @@ class ApplyTemplate < Base # rubocop:disable Metrics/ClassLength
     REQUIRES_ARGS = true
     OPTIONS = [
       app_option(required: true),
+      location_option,
       skip_confirm_option
     ].freeze
     DESCRIPTION = "Applies application-specific configs from templates"
@@ -126,7 +127,7 @@ def confirm_workload(template)
     def apply_template(filename)
       data = File.read(filename)
                  .gsub("APP_GVC", config.app)
-                 .gsub("APP_LOCATION", config[:default_location])
+                 .gsub("APP_LOCATION", config.options[:location] || config[:default_location])
                  .gsub("APP_ORG", config.org)
                  .gsub("APP_IMAGE", latest_image)
 
diff --git a/lib/command/ps.rb b/lib/command/ps.rb
index f0bb61ed..6e81dd99 100644
--- a/lib/command/ps.rb
+++ b/lib/command/ps.rb
@@ -5,6 +5,7 @@ class Ps < Base
     NAME = "ps"
     OPTIONS = [
       app_option(required: true),
+      location_option,
       workload_option
     ].freeze
     DESCRIPTION = "Shows running replicas in app"
@@ -24,12 +25,14 @@ class Ps < Base
     def call
       cp.fetch_gvc!
 
+      location = config.options["location"] || config[:default_location]
+
       workloads = [config.options[:workload]] if config.options[:workload]
       workloads ||= config[:app_workloads] + config[:additional_workloads]
       workloads.each do |workload|
         cp.fetch_workload!(workload)
 
-        result = cp.workload_get_replicas(workload, location: config[:default_location])
+        result = cp.workload_get_replicas(workload, location: location)
         result["items"].each { |replica| puts replica }
       end
     end
diff --git a/lib/command/run.rb b/lib/command/run.rb
index b4c8fffd..789e680f 100644
--- a/lib/command/run.rb
+++ b/lib/command/run.rb
@@ -10,6 +10,7 @@ class Run < Base
       app_option(required: true),
       image_option,
       workload_option,
+      location_option,
       use_local_token_option,
       terminal_size_option
     ].freeze
@@ -56,7 +57,7 @@ class Run < Base
     attr_reader :location, :workload, :one_off, :container
 
     def call # rubocop:disable Metrics/MethodLength
-      @location = config[:default_location]
+      @location = config.options["location"] || config[:default_location]
       @workload = config.options["workload"] || config[:one_off_workload]
       @one_off = "#{workload}-run-#{rand(1000..9999)}"
 
diff --git a/lib/command/run_detached.rb b/lib/command/run_detached.rb
index 8fea7cdf..c6b2e2e5 100644
--- a/lib/command/run_detached.rb
+++ b/lib/command/run_detached.rb
@@ -9,6 +9,7 @@ class RunDetached < Base # rubocop:disable Metrics/ClassLength
       app_option(required: true),
       image_option,
       workload_option,
+      location_option,
       use_local_token_option
     ].freeze
     DESCRIPTION = "Runs one-off **_non-interactive_** replicas (close analog of `heroku run:detached`)"
@@ -47,7 +48,7 @@ class RunDetached < Base # rubocop:disable Metrics/ClassLength
     attr_reader :location, :workload, :one_off, :container
 
     def call # rubocop:disable Metrics/MethodLength
-      @location = config[:default_location]
+      @location = config.options["location"] || config[:default_location]
       @workload = config.options["workload"] || config[:one_off_workload]
       @one_off = "#{workload}-runner-#{rand(1000..9999)}"