tweak: Drop external foreign key constraints during GTFS import

lib/arrow/gtfs.ex

@@ -7,6 +7,7 @@ defmodule Arrow.Gtfs do
   alias Arrow.Gtfs.Importable
   alias Arrow.Gtfs.JobHelper
   alias Arrow.Repo
+  alias Arrow.Repo.ForeignKeyConstraint
 
   @import_timeout_ms :timer.minutes(10)
 
@@ -60,9 +61,14 @@ defmodule Arrow.Gtfs do
 
   defp import_transaction(unzip, dry_run?) do
     transaction = fn ->
-      _ = truncate_all()
+      external_fkeys = get_external_fkeys()
+      drop_external_fkeys(external_fkeys)
+
+      truncate_all()
       import_all(unzip)
 
+      add_external_fkeys(external_fkeys)
+
       if dry_run? do
         # Set any deferred constraints to run now, instead of on transaction commit,
         # since we don't actually commit the transaction in this case.
@@ -81,9 +87,11 @@ defmodule Arrow.Gtfs do
     result
   end
 
+  @spec truncate_all() :: :ok
   defp truncate_all do
     tables = Enum.map_join(importable_schemas(), ", ", & &1.__schema__(:source))
-    Repo.query!("TRUNCATE #{tables}")
+    _ = Repo.query!("TRUNCATE #{tables}")
+    :ok
   end
 
   defp import_all(unzip) do
@@ -141,4 +149,46 @@ defmodule Arrow.Gtfs do
     |> Enum.flat_map(& &1.filenames())
     |> MapSet.new()
   end
+
+  defp get_external_fkeys do
+    importable_schemas()
+    |> Enum.map(& &1.__schema__(:source))
+    |> ForeignKeyConstraint.external_constraints_referencing_tables()
+  end
+
+  @spec drop_external_fkeys(list(ForeignKeyConstraint.t())) :: :ok
+  defp drop_external_fkeys(external_fkeys) do
+    # To allow all GTFS tables to be truncated, we first need to
+    # temporarily drop all foreign key constraints referencing them
+    # from non-GTFS tables.
+    fkey_names = Enum.map_join(external_fkeys, ",", & &1.name)
+
+    Logger.info(
+      "temporarily dropping external foreign keys referencing GTFS tables fkey_names=#{fkey_names}"
+    )
+
+    Enum.each(external_fkeys, &ForeignKeyConstraint.drop/1)
+
+    Logger.info("finished dropping external foreign keys referencing GTFS tables")
+
+    :ok
+  end
+
+  @spec add_external_fkeys(list(ForeignKeyConstraint.t())) :: :ok
+  defp add_external_fkeys(external_fkeys) do
+    fkey_names = Enum.map_join(external_fkeys, ",", & &1.name)
+
+    Logger.info(
+      "re-adding external foreign keys referencing GTFS tables fkey_names=#{fkey_names}"
+    )
+
+    Enum.each(external_fkeys, fn fkey ->
+      Logger.info("re-adding foreign key fkey_name=#{fkey.name}")
+      ForeignKeyConstraint.add(fkey)
+    end)
+
+    Logger.info("finished re-adding external foreign keys referencing GTFS tables")
+
+    :ok
+  end
 end

lib/arrow/repo/foreign_key_constraint.ex

@@ -0,0 +1,98 @@
+defmodule Arrow.Repo.ForeignKeyConstraint do
+  @moduledoc """
+  Schema allowing Arrow to introspect its DB's foreign key constraints.
+  """
+  use Ecto.Schema
+  import Ecto.Query
+  alias Arrow.Repo
+
+  @type t :: %__MODULE__{
+          name: String.t(),
+          origin_table: String.t(),
+          referenced_table: String.t(),
+          definition: String.t()
+        }
+
+  @primary_key false
+
+  schema "foreign_key_constraints" do
+    field :name
+    field :origin_table
+    field :referenced_table
+    field :definition
+  end
+
+  @doc """
+  Returns foreign key constraints that reference any table in `tables`,
+  and originate in any table _not_ in `tables`.
+
+  For example, given the following foreign key relations:
+
+      foo.bar_id -> bar.id
+      foo.baz_id -> baz.id
+      baz.bar_id -> bar.id
+
+  Calling this:
+
+      external_constraints_referencing_tables(["bar", "baz"])
+
+  Would produce this:
+
+      [
+        %ForeignKeyConstraint{origin_table: "foo", referenced_table: "bar"},
+        %ForeignKeyConstraint{origin_table: "foo", referenced_table: "baz"},
+      ]
+  """
+  @spec external_constraints_referencing_tables(list(String.t() | atom)) :: list(t())
+  def external_constraints_referencing_tables(tables) when is_list(tables) do
+    from(fk in __MODULE__,
+      where: fk.referenced_table in ^tables,
+      where: fk.origin_table not in ^tables
+    )
+    |> Repo.all()
+  end
+
+  @doc """
+  Drops a foreign key constraint.
+
+  This function should not be used to permanently drop a constraint--
+  use Ecto's migration utilities to do that.
+  """
+  @spec drop(t()) :: :ok
+  def drop(%__MODULE__{} = fk) do
+    if Repo.in_transaction?() do
+      _ =
+        Repo.query!("""
+        ALTER TABLE "#{fk.origin_table}"
+          DROP CONSTRAINT "#{fk.name}"
+        """)
+
+      :ok
+    else
+      raise "must be in a transaction"
+    end
+  end
+
+  @doc """
+  Adds a foreign key constraint.
+
+  This function should not be used to permanently add a new constraint--
+  use Ecto's migration utilities to do that.
+
+  This is intended only for re-adding a previously, temporarily dropped constraint.
+  """
+  @spec add(t()) :: :ok
+  def add(%__MODULE__{} = fk) do
+    if Repo.in_transaction?() do
+      _ =
+        Repo.query!("""
+        ALTER TABLE "#{fk.origin_table}"
+          ADD CONSTRAINT "#{fk.name}" #{fk.definition}
+        """)
+
+      :ok
+    else
+      raise "must be in a transaction"
+    end
+  end
+end

priv/repo/migrations/20241030181351_add_foreign_keys_view.exs

@@ -0,0 +1,22 @@
+defmodule Arrow.Repo.Migrations.AddForeignKeysView do
+  use Ecto.Migration
+
+  def up do
+    execute("""
+    CREATE VIEW "foreign_key_constraints" AS
+    SELECT
+      pgc.conname AS name,
+      pgc.conrelid::regclass::text AS origin_table,
+      pgc.confrelid::regclass::text AS referenced_table,
+      pg_get_constraintdef(pgc.oid, true) AS definition
+    FROM pg_catalog.pg_constraint pgc
+    WHERE pgc.contype = 'f'
+    """)
+  end
+
+  def down do
+    execute("""
+    DROP VIEW "foreign_key_constraints"
+    """)
+  end
+end