From d92668fe534e6e3fc8c8c0031a500f86c4be3eab Mon Sep 17 00:00:00 2001
From: Troy Dean <tadean@arizona.edu>
Date: Tue, 5 Nov 2024 16:49:04 -0700
Subject: [PATCH] Fixes #3866 Update AZEnterpriseAttributesMigrationSync

---
 composer.json                                 |  2 +-
 ..._enterprise_attributes_import.services.yml |  2 +-
 .../AZEnterpriseAttributesMigrationSync.php   | 26 +++++++++++++------
 3 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/composer.json b/composer.json
index 416dc07887..0dc5e373fa 100644
--- a/composer.json
+++ b/composer.json
@@ -98,7 +98,7 @@
         "drupal/metatag": "2.0.2",
         "drupal/migrate_plus": "6.0.4",
         "drupal/migrate_queue_importer": "2.1.1",
-        "drupal/migrate_tools": "6.0.4",
+        "drupal/migrate_tools": "6.0.5",
         "drupal/optional_end_date": "1.4.0",
         "drupal/paragraphs": "1.18.0",
         "drupal/paragraphs_admin": "1.6.0",
diff --git a/modules/custom/az_core/az_enterprise_attributes_import/az_enterprise_attributes_import.services.yml b/modules/custom/az_core/az_enterprise_attributes_import/az_enterprise_attributes_import.services.yml
index 88f2ef4193..645f32c7cf 100644
--- a/modules/custom/az_core/az_enterprise_attributes_import/az_enterprise_attributes_import.services.yml
+++ b/modules/custom/az_core/az_enterprise_attributes_import/az_enterprise_attributes_import.services.yml
@@ -7,7 +7,7 @@ services:
     decorates: migrate_tools.migration_sync
     arguments:
       - '@event_dispatcher'
-      - '@state'
+      - '@migrate_tools.migrate_tools'
       - '@entity_type.manager'
       - '@logger.channel.az_enterprise_attributes_import'
   logger.channel.az_enterprise_attributes_import:
diff --git a/modules/custom/az_core/az_enterprise_attributes_import/src/EventSubscriber/AZEnterpriseAttributesMigrationSync.php b/modules/custom/az_core/az_enterprise_attributes_import/src/EventSubscriber/AZEnterpriseAttributesMigrationSync.php
index e53a892ad0..10e1734f06 100644
--- a/modules/custom/az_core/az_enterprise_attributes_import/src/EventSubscriber/AZEnterpriseAttributesMigrationSync.php
+++ b/modules/custom/az_core/az_enterprise_attributes_import/src/EventSubscriber/AZEnterpriseAttributesMigrationSync.php
@@ -6,9 +6,9 @@
 
 use Drupal\Core\Entity\EntityTypeManagerInterface;
 use Drupal\Core\Logger\LoggerChannel;
-use Drupal\Core\State\StateInterface;
 use Drupal\migrate\Event\MigrateImportEvent;
 use Drupal\migrate_tools\EventSubscriber\MigrationImportSync;
+use Drupal\migrate_tools\MigrateTools;
 use Symfony\Component\EventDispatcher\EventDispatcherInterface;
 
 /**
@@ -31,17 +31,16 @@ final class AZEnterpriseAttributesMigrationSync extends MigrationImportSync {
    *
    * @param \Symfony\Component\EventDispatcher\EventDispatcherInterface $dispatcher
    *   The event dispatcher.
-   * @param \Drupal\Core\State\StateInterface $state
-   *   The Key/Value Store to use for tracking synced source rows.
+   * @param \Drupal\migrate_tools\MigrateTools $migrateTools
+   *   The MigrateTools helper for source id tracking.
    * @param \Drupal\Core\Entity\EntityTypeManagerInterface $entityTypeManager
    *   The entity type manager service.
    * @param \Drupal\Core\Logger\LoggerChannel $logger
    *   The logger channel service.
    */
-  public function __construct(EventDispatcherInterface $dispatcher, StateInterface $state, EntityTypeManagerInterface $entityTypeManager, LoggerChannel $logger) {
+  public function __construct(EventDispatcherInterface $dispatcher, MigrateTools $migrateTools, EntityTypeManagerInterface $entityTypeManager, LoggerChannel $logger) {
     $this->dispatcher = $dispatcher;
-    $this->state = $state;
-    $this->state->set('migrate_tools_sync', []);
+    $this->migrateTools = $migrateTools;
     $this->entityTypeManager = $entityTypeManager;
     $this->logger = $logger;
   }
@@ -51,12 +50,18 @@ public function __construct(EventDispatcherInterface $dispatcher, StateInterface
    */
   public function sync(MigrateImportEvent $event): void {
     $migration = $event->getMigration();
+    $migrationId = $migration->getPluginId();
     // If this isn't a migration we're concerned with, use the parent.
-    if ($migration->id() !== 'az_enterprise_attributes_import') {
+    if ($migrationId !== 'az_enterprise_attributes_import') {
       parent::sync($event);
       return;
     }
     if (!empty($migration->syncSource)) {
+      // Clear Sync IDs for this migration before starting preparing rows.
+      $this->migrateTools->clearSyncSourceIds($migrationId);
+      // Activate the syncing state for this migration, so
+      // migrate_tools_migrate_prepare_row() can record all IDs.
+      $this->migrateTools->setMigrationSyncingState($migrationId, TRUE);
 
       // Loop through the source to register existing source ids.
       // @see migrate_tools_migrate_prepare_row().
@@ -68,7 +73,12 @@ public function sync(MigrateImportEvent $event): void {
         $source->next();
       }
 
-      $source_id_values = $this->state->get('migrate_tools_sync', []);
+      // Deactivate the syncing state for this migration, so
+      // migrate_tools_migrate_prepare_row() does not record any further IDs
+      // during the actual migration process.
+      $this->migrateTools->setMigrationSyncingState($migrationId, FALSE);
+
+      $source_id_values = $this->migrateTools->getSyncSourceIds($migrationId);
 
       $id_map = $migration->getIdMap();
       $id_map->rewind();