From f72dc826bf05d7e8dc2a3feb728734aa2b9d6df2 Mon Sep 17 00:00:00 2001
From: Pere Miquel Brull <peremiquelbrull@gmail.com>
Date: Sat, 27 Sep 2025 16:47:03 +0200
Subject: [PATCH] FIX - Cleanup orphaned data contracts (#23596)

---
 .../migration/mysql/v1100/Migration.java      |  1 +
 .../migration/postgres/v1100/Migration.java   |  1 +
 .../migration/utils/v1100/MigrationUtil.java  | 65 +++++++++++++++++++
 3 files changed, 67 insertions(+)

diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/migration/mysql/v1100/Migration.java b/openmetadata-service/src/main/java/org/openmetadata/service/migration/mysql/v1100/Migration.java
index 852d0bf2d22..bab7dde84c4 100644
--- a/openmetadata-service/src/main/java/org/openmetadata/service/migration/mysql/v1100/Migration.java
+++ b/openmetadata-service/src/main/java/org/openmetadata/service/migration/mysql/v1100/Migration.java
@@ -17,5 +17,6 @@ public class Migration extends MigrationProcessImpl {
   public void runDataMigration() {
     MigrationUtil migrationUtil = new MigrationUtil(handle, ConnectionType.MYSQL);
     migrationUtil.migrateEntityStatusForExistingEntities();
+    migrationUtil.cleanupOrphanedDataContracts();
   }
 }
diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/migration/postgres/v1100/Migration.java b/openmetadata-service/src/main/java/org/openmetadata/service/migration/postgres/v1100/Migration.java
index 8686a462ea3..3b2f664f4ee 100644
--- a/openmetadata-service/src/main/java/org/openmetadata/service/migration/postgres/v1100/Migration.java
+++ b/openmetadata-service/src/main/java/org/openmetadata/service/migration/postgres/v1100/Migration.java
@@ -17,5 +17,6 @@ public class Migration extends MigrationProcessImpl {
   public void runDataMigration() {
     MigrationUtil migrationUtil = new MigrationUtil(handle, ConnectionType.POSTGRES);
     migrationUtil.migrateEntityStatusForExistingEntities();
+    migrationUtil.cleanupOrphanedDataContracts();
   }
 }
diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/migration/utils/v1100/MigrationUtil.java b/openmetadata-service/src/main/java/org/openmetadata/service/migration/utils/v1100/MigrationUtil.java
index 4f2a150e098..cc10edf3844 100644
--- a/openmetadata-service/src/main/java/org/openmetadata/service/migration/utils/v1100/MigrationUtil.java
+++ b/openmetadata-service/src/main/java/org/openmetadata/service/migration/utils/v1100/MigrationUtil.java
@@ -1,7 +1,14 @@
 package org.openmetadata.service.migration.utils.v1100;
 
+import java.util.List;
 import lombok.extern.slf4j.Slf4j;
 import org.jdbi.v3.core.Handle;
+import org.openmetadata.schema.entity.data.DataContract;
+import org.openmetadata.schema.type.Include;
+import org.openmetadata.service.Entity;
+import org.openmetadata.service.exception.EntityNotFoundException;
+import org.openmetadata.service.jdbi3.DataContractRepository;
+import org.openmetadata.service.jdbi3.ListFilter;
 import org.openmetadata.service.jdbi3.locator.ConnectionType;
 
 @Slf4j
@@ -304,4 +311,62 @@ public class MigrationUtil {
 
     return totalMigrated;
   }
+
+  public void cleanupOrphanedDataContracts() {
+    LOG.info("Starting cleanup of orphaned data contracts...");
+
+    try {
+      DataContractRepository dataContractRepository =
+          (DataContractRepository) Entity.getEntityRepository(Entity.DATA_CONTRACT);
+
+      List<DataContract> allDataContracts =
+          dataContractRepository.listAll(
+              dataContractRepository.getFields("id,entity"), new ListFilter(Include.ALL));
+
+      if (allDataContracts.isEmpty()) {
+        LOG.info("✓ No data contracts found - cleanup complete");
+        return;
+      }
+
+      int deletedCount = 0;
+      int totalContracts = allDataContracts.size();
+
+      LOG.info("Found {} data contracts to validate", totalContracts);
+
+      for (DataContract dataContract : allDataContracts) {
+        try {
+          // Try to get the associated entity
+          Entity.getEntityReferenceById(
+              dataContract.getEntity().getType(),
+              dataContract.getEntity().getId(),
+              Include.NON_DELETED);
+
+        } catch (EntityNotFoundException e) {
+          LOG.info(
+              "Deleting orphaned data contract '{}' - associated {} entity with ID {} not found",
+              dataContract.getFullyQualifiedName(),
+              dataContract.getEntity().getType(),
+              dataContract.getEntity().getId());
+
+          try {
+            dataContractRepository.delete(Entity.ADMIN_USER_NAME, dataContract.getId(), true, true);
+            deletedCount++;
+          } catch (Exception deleteException) {
+            LOG.warn(
+                "Failed to delete orphaned data contract '{}': {}",
+                dataContract.getFullyQualifiedName(),
+                deleteException.getMessage());
+          }
+        }
+      }
+
+      LOG.info(
+          "✓ Cleanup complete: {} orphaned data contracts deleted out of {} total",
+          deletedCount,
+          totalContracts);
+
+    } catch (Exception e) {
+      LOG.error("✗ FAILED cleanup of orphaned data contracts: {}", e.getMessage(), e);
+    }
+  }
 }