From d968c8240406bb982920c9314fdbe2c9e8aa4090 Mon Sep 17 00:00:00 2001
From: Ivan Habunek <ivan@habunek.com>
Date: Wed, 18 Jan 2023 16:27:32 +0100
Subject: [PATCH] Enable specifying operators to load in pt_update

---
 lib/mix/tasks/triglav/pt_update.ex |  3 ++-
 lib/triglav/public_transport.ex    | 33 +++++++++++++++++++-----------
 2 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/lib/mix/tasks/triglav/pt_update.ex b/lib/mix/tasks/triglav/pt_update.ex
index 2fa9792..e9bb833 100644
--- a/lib/mix/tasks/triglav/pt_update.ex
+++ b/lib/mix/tasks/triglav/pt_update.ex
@@ -11,7 +11,8 @@ defmodule Mix.Tasks.Triglav.PtUpdate do
   def run(args) do
     Application.put_env(:triglav, :repo_only, true)
     {:ok, _} = Application.ensure_all_started(:triglav)
-    {opts, _} = OptionParser.parse!(args, strict: [force: :boolean])
+    {opts, operators} = OptionParser.parse!(args, strict: [force: :boolean])
+    opts = if operators == [], do: opts, else: Keyword.put(opts, :operators, operators)
 
     {time, _} = :timer.tc(fn -> PublicTransport.update_all(opts) end)
     time_seconds = :erlang.float_to_binary(time / 1_000_000, decimals: 2)
diff --git a/lib/triglav/public_transport.ex b/lib/triglav/public_transport.ex
index e6d3dde..e334ecf 100644
--- a/lib/triglav/public_transport.ex
+++ b/lib/triglav/public_transport.ex
@@ -12,28 +12,37 @@ defmodule Triglav.PublicTransport do
 
   defdelegate generate_mappings(feed_id), to: Mappings, as: :generate
   defdelegate import_autotrolej(opts \\ []), to: Triglav.Autotrolej, as: :import
+  defdelegate import_gpp_osijek(opts \\ []), to: Triglav.PublicTransport.ImportGtfs
   defdelegate import_hzpp(opts \\ []), to: Triglav.PublicTransport.ImportGtfs
+  defdelegate import_jadrolinija(opts \\ []), to: Triglav.PublicTransport.ImportGtfs
   defdelegate import_zet(opts \\ []), to: Triglav.PublicTransport.ImportGtfs
-  defdelegate import_gpp_osijek(opts \\ []), to: Triglav.PublicTransport.ImportGtfs
 
   @doc """
   For each operator, check if a new feed is available. If so, update the feed,
   generate new mappings and update errors.
 
   Options:
-  - force - if set to `true` will re-load the feeds even if a newer feed is not
-    found. Default is `false.
+  - `:operators` - A list of operator slugs to load
+  - `:force` - If set to `true` will re-load the feeds even if a newer feed is
+               not found. Default is `false`.
   """
   def update_all(opts \\ []) do
-    Operator
-    |> Repo.all()
-    |> Enum.each(fn operator ->
-      try do
-        update_feed(operator, opts)
-      rescue
-        e -> Triglav.Tooling.log_error("Failed loading #{operator.name}", e, __STACKTRACE__)
-      end
-    end)
+    slugs = Keyword.get(opts, :operators)
+
+    operators =
+      if slugs,
+        do: Repo.all(from(o in Operator, where: o.slug in ^slugs)),
+        else: Repo.all(Operator)
+
+    Enum.each(operators, &update_operator(&1, opts))
+  end
+
+  defp update_operator(operator, opts) do
+    try do
+      update_feed(operator, opts)
+    rescue
+      e -> Triglav.Tooling.log_error("Failed loading #{operator.name}", e, __STACKTRACE__)
+    end
   end
 
   @doc """
-- 
2.38.5