From bdee57a0571da734d9a1c4b5c35b3bbdb26fda8b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jean-S=C3=A9bastien=20Herbaux?=
 <jean-sebastien.herbaux@epitech.eu>
Date: Mon, 24 Jun 2024 09:41:22 +0200
Subject: [PATCH] Handle non-localized content-types in discard-draft migration
 (#20422)

---
 .../database/5.0.0-discard-drafts.ts          | 36 +++++++++++++++----
 .../core/core/src/migrations/draft-publish.ts | 31 ++++++++++++----
 2 files changed, 54 insertions(+), 13 deletions(-)

diff --git a/packages/core/core/src/migrations/database/5.0.0-discard-drafts.ts b/packages/core/core/src/migrations/database/5.0.0-discard-drafts.ts
index ff6c4e7019..36f8043dea 100644
--- a/packages/core/core/src/migrations/database/5.0.0-discard-drafts.ts
+++ b/packages/core/core/src/migrations/database/5.0.0-discard-drafts.ts
@@ -1,8 +1,10 @@
 /* eslint-disable no-continue */
-import type { UID } from '@strapi/types';
+import type { UID, Modules } from '@strapi/types';
 import type { Database, Migration } from '@strapi/database';
 import { async, contentTypes } from '@strapi/utils';
 
+type DiscardDraftParams<TSchemaUID extends UID.ContentType> =
+  Modules.Documents.ServiceParams<TSchemaUID>['discardDraft'];
 type DocumentVersion = { documentId: string; locale: string };
 type Knex = Parameters<Migration['up']>[0];
 
@@ -16,21 +18,29 @@ export async function* getBatchToDiscard({
   db,
   trx,
   uid,
+  isLocalized,
   batchSize = 1000,
 }: {
   db: Database;
   trx: Knex;
   uid: string;
+  isLocalized: boolean;
   batchSize?: number;
 }) {
   let offset = 0;
   let hasMore = true;
 
+  const fields = ['id', 'documentId'];
+
+  if (isLocalized) {
+    fields.push('locale');
+  }
+
   while (hasMore) {
     // Look for the published entries to discard
     const batch: DocumentVersion[] = await db
       .queryBuilder(uid)
-      .select(['id', 'documentId', 'locale'])
+      .select(fields)
       .where({ publishedAt: { $ne: null } })
       .limit(batchSize)
       .offset(offset)
@@ -57,22 +67,36 @@ const migrateUp = async (trx: Knex, db: Database) => {
 
     const uid = meta.uid as UID.ContentType;
     const model = strapi.getModel(uid);
+
     const hasDP = contentTypes.hasDraftAndPublish(model);
+    const isLocalized = strapi
+      .plugin('i18n')
+      .service('content-types')
+      .isLocalizedContentType(model);
+
     if (!hasDP) {
       continue;
     }
 
-    const discardDraft = async (entry: DocumentVersion) =>
+    const discardDraft = async (entry: DocumentVersion) => {
+      const params: DiscardDraftParams<typeof uid> = { documentId: entry.documentId };
+
+      // Only add the locale param if the model is localized
+      if (isLocalized) {
+        params.locale = entry.locale;
+      }
+
       strapi
         .documents(uid)
-        // Discard draft by referencing the documentId and locale
-        .discardDraft({ documentId: entry.documentId, locale: entry.locale });
+        // Discard draft by referencing the documentId (and locale if the model is localized)
+        .discardDraft(params);
+    };
 
     /**
      * Load a batch of entries (batched to prevent loading millions of rows at once ),
      * and discard them using the document service.
      */
-    for await (const batch of getBatchToDiscard({ db, trx, uid: meta.uid })) {
+    for await (const batch of getBatchToDiscard({ db, trx, uid: meta.uid, isLocalized })) {
       await async.map(batch, discardDraft, { concurrency: 10 });
     }
   }
diff --git a/packages/core/core/src/migrations/draft-publish.ts b/packages/core/core/src/migrations/draft-publish.ts
index d317b08609..66b20bfd12 100644
--- a/packages/core/core/src/migrations/draft-publish.ts
+++ b/packages/core/core/src/migrations/draft-publish.ts
@@ -1,8 +1,10 @@
 import { contentTypes as contentTypesUtils, async } from '@strapi/utils';
-import { Schema } from '@strapi/types';
+import type { Modules, Schema } from '@strapi/types';
 
 import { getBatchToDiscard } from './database/5.0.0-discard-drafts';
 
+type DiscardDraftParams = Modules.Documents.ServiceParams['discardDraft'];
+
 interface Input {
   oldContentTypes: Record<string, Schema.ContentType>;
   contentTypes: Record<string, Schema.ContentType>;
@@ -31,22 +33,37 @@ const enableDraftAndPublish = async ({ oldContentTypes, contentTypes }: Input) =
       const oldContentType = oldContentTypes[uid];
       const contentType = contentTypes[uid];
 
+      const isLocalized = strapi
+        .plugin('i18n')
+        .service('content-types')
+        .isLocalizedContentType(contentType);
+
       // if d&p was enabled set publishedAt to eq createdAt
       if (
         !contentTypesUtils.hasDraftAndPublish(oldContentType) &&
         contentTypesUtils.hasDraftAndPublish(contentType)
       ) {
-        const discardDraft = async (entry: { documentId: string; locale: string }) =>
-          strapi
-            .documents(uid as any)
-            // Discard draft by referencing the documentId and locale
-            .discardDraft({ documentId: entry.documentId, locale: entry.locale });
+        const discardDraft = async (entry: { documentId: string; locale: string }) => {
+          const params: DiscardDraftParams = { documentId: entry.documentId };
+
+          // Only add the locale param if the old content-type is localized
+          if (isLocalized) {
+            params.locale = entry.locale;
+          }
+
+          return (
+            strapi
+              .documents(uid as any)
+              // Discard draft by referencing the documentId (and locale if the model is localized)
+              .discardDraft(params)
+          );
+        };
 
         /**
          * Load a batch of entries (batched to prevent loading millions of rows at once ),
          * and discard them using the document service.
          */
-        for await (const batch of getBatchToDiscard({ db: strapi.db, trx, uid })) {
+        for await (const batch of getBatchToDiscard({ db: strapi.db, trx, uid, isLocalized })) {
           await async.map(batch, discardDraft, { concurrency: 10 });
         }
       }