From 40d94279e69c368f80f651a97b0469a1fb723d14 Mon Sep 17 00:00:00 2001
From: John Davis <jdavcs@gmail.com>
Date: Thu, 11 May 2023 18:53:06 -0400
Subject: [PATCH] Add create_foreign_key migration utility

---
 lib/galaxy/model/migrations/util.py | 57 +++++++++++++++++++++++++++++
 1 file changed, 57 insertions(+)

diff --git a/lib/galaxy/model/migrations/util.py b/lib/galaxy/model/migrations/util.py
index d27a50d84ddc..da15b720ed52 100644
--- a/lib/galaxy/model/migrations/util.py
+++ b/lib/galaxy/model/migrations/util.py
@@ -223,6 +223,43 @@ def log_check_not_passed(self) -> None:
         self._log_object_does_not_exist_message(name)
 
 
+class CreateForeignKey(DDLAlterOperation):
+    """Wraps alembic's create_foreign_key directive."""
+
+    def __init__(
+        self,
+        foreign_key_name: str,
+        table_name: str,
+        referent_table: str,
+        local_cols: List[str],
+        remote_cols: List[str],
+        **kw: Any,
+    ) -> None:
+        super().__init__(table_name)
+        self.foreign_key_name = foreign_key_name
+        self.referent_table = referent_table
+        self.local_cols = local_cols
+        self.remote_cols = remote_cols
+        self.kw = kw
+
+    def batch_execute(self, batch_op) -> None:
+        batch_op.create_foreign_key(
+            self.foreign_key_name, self.referent_table, self.local_cols, self.remote_cols, **self.kw
+        )
+
+    def non_batch_execute(self) -> None:
+        op.create_foreign_key(
+            self.foreign_key_name, self.table_name, self.referent_table, self.local_cols, self.remote_cols, **self.kw
+        )
+
+    def pre_execute_check(self) -> bool:
+        return not foreign_key_exists(self.foreign_key_name, self.table_name, False)
+
+    def log_check_not_passed(self) -> None:
+        name = _table_object_description(self.foreign_key_name, self.table_name)
+        self._log_object_exists_message(name)
+
+
 class CreateUniqueConstraint(DDLAlterOperation):
     """Wraps alembic's create_unique_constraint directive."""
 
@@ -294,6 +331,17 @@ def drop_index(index_name, table_name) -> None:
     DropIndex(index_name, table_name).run()
 
 
+def create_foreign_key(
+    foreign_key_name: str,
+    table_name: str,
+    referent_table: str,
+    local_cols: List[str],
+    remote_cols: List[str],
+    **kw: Any,
+) -> None:
+    CreateForeignKey(foreign_key_name, table_name, referent_table, local_cols, remote_cols, **kw).run()
+
+
 def create_unique_constraint(constraint_name: str, table_name: str, columns: List[str]) -> None:
     CreateUniqueConstraint(constraint_name, table_name, columns).run()
 
@@ -328,6 +376,15 @@ def index_exists(index_name: str, table_name: str, default: bool) -> bool:
     return any(index["name"] == index_name for index in indexes)
 
 
+def foreign_key_exists(constraint_name: str, table_name: str, default: bool) -> bool:
+    """Check if unique constraint exists. If running in offline mode, return default."""
+    if context.is_offline_mode():
+        _log_offline_mode_message(foreign_key_exists.__name__, default)
+        return default
+    constraints = _inspector().get_foreign_keys(table_name)
+    return any(c["name"] == constraint_name for c in constraints)
+
+
 def unique_constraint_exists(constraint_name: str, table_name: str, default: bool) -> bool:
     """Check if unique constraint exists. If running in offline mode, return default."""
     if context.is_offline_mode():