Use batches and direct queries to fix sales address upgrade

Author: rikwillems

app/code/Magento/Sales/Setup/Patch/Data/FillQuoteAddressIdInSalesOrderAddress.php

@@ -9,6 +9,7 @@
 use Magento\Eav\Model\Config;
 use Magento\Framework\App\State;
 use Magento\Quote\Model\QuoteFactory;
+use Magento\Sales\Model\Order\Address;
 use Magento\Sales\Model\OrderFactory;
 use Magento\Sales\Model\ResourceModel\Order\Address\CollectionFactory as AddressCollectionFactory;
 use Magento\Framework\App\ResourceConnection;
@@ -39,23 +40,11 @@ class FillQuoteAddressIdInSalesOrderAddress implements DataPatchInterface, Patch
      */
     private $eavConfig;
 
-    /**
-     * @var AddressCollectionFactory
-     */
-    private $addressCollectionFactory;
-
-    /**
-     * @var OrderFactory
-     */
-    private $orderFactory;
-
-    /**
-     * @var QuoteFactory
-     */
-    private $quoteFactory;
-
     /**
      * PatchInitial constructor.
+     *
+     * @SuppressWarnings(PHPMD.UnusedFormalParameter)
+     *
      * @param ModuleDataSetupInterface $moduleDataSetup
      */
     public function __construct(
@@ -71,9 +60,6 @@ public function __construct(
         $this->salesSetupFactory = $salesSetupFactory;
         $this->state = $state;
         $this->eavConfig = $eavConfig;
-        $this->addressCollectionFactory = $addressCollectionFactory;
-        $this->orderFactory = $orderFactory;
-        $this->quoteFactory = $quoteFactory;
     }
 
     /**
@@ -96,28 +82,8 @@ public function apply()
      */
     public function fillQuoteAddressIdInSalesOrderAddress(ModuleDataSetupInterface $setup)
     {
-        $addressTable = $setup->getTable('sales_order_address');
-        $updateOrderAddress = $setup->getConnection()
-            ->select()
-            ->joinInner(
-                ['sales_order' => $setup->getTable('sales_order')],
-                $addressTable . '.parent_id = sales_order.entity_id',
-                ['quote_address_id' => 'quote_address.address_id']
-            )
-            ->joinInner(
-                ['quote_address' => $setup->getTable('quote_address')],
-                'sales_order.quote_id = quote_address.quote_id 
-                AND ' . $addressTable . '.address_type = quote_address.address_type',
-                []
-            )
-            ->where(
-                $addressTable . '.quote_address_id IS NULL'
-            );
-        $updateOrderAddress = $setup->getConnection()->updateFromSelect(
-            $updateOrderAddress,
-            $addressTable
-        );
-        $setup->getConnection()->query($updateOrderAddress);
+        $this->fillQuoteAddressIdInSalesOrderAddressByType($setup, Address::TYPE_SHIPPING);
+        $this->fillQuoteAddressIdInSalesOrderAddressByType($setup, Address::TYPE_BILLING);
     }
 
     /**
@@ -145,4 +111,88 @@ public function getAliases()
     {
         return [];
     }
+
+    /**
+     * @param ModuleDataSetupInterface $setup
+     * @param string $addressType
+     */
+    private function fillQuoteAddressIdInSalesOrderAddressByType(ModuleDataSetupInterface $setup, $addressType)
+    {
+        $salesConnection = $setup->getConnection('sales');
+
+        $orderTable = $setup->getTable('sales_order', 'sales');
+        $orderAddressTable = $setup->getTable('sales_order_address', 'sales');
+
+        $query = $salesConnection
+            ->select()
+            ->from(
+                ['sales_order_address' => $orderAddressTable],
+                ['entity_id', 'address_type']
+            )
+            ->joinInner(
+                ['sales_order' => $orderTable],
+                'sales_order_address.parent_id = sales_order.entity_id',
+                ['quote_id' => 'sales_order.quote_id']
+            )
+            ->where('sales_order_address.quote_address_id IS NULL')
+            ->where('sales_order_address.address_type = ?', $addressType)
+            ->order('sales_order_address.entity_id');
+
+        $batchSize = 5000;
+        $result = $salesConnection->query($query);
+        $count = $result->rowCount();
+        $batches = ceil($count / $batchSize);
+
+        for ($batch = $batches; $batch > 0; $batch--) {
+            $query->limitPage($batch, $batchSize);
+            $result = $salesConnection->fetchAssoc($query);
+
+            $this->fillQuoteAddressIdInSalesOrderAddressProcessBatch($setup, $result, $addressType);
+        }
+    }
+
+    /**
+     * @param ModuleDataSetupInterface $setup
+     * @param array $orderAddresses
+     * @param string $addressType
+     */
+    private function fillQuoteAddressIdInSalesOrderAddressProcessBatch(
+        ModuleDataSetupInterface $setup,
+        array $orderAddresses,
+        $addressType
+    ) {
+        $salesConnection = $setup->getConnection('sales');
+        $quoteConnection = $setup->getConnection('checkout');
+
+        $quoteAddressTable = $setup->getTable('quote_address', 'checkout');
+        $quoteTable = $setup->getTable('quote', 'checkout');
+        $salesOrderAddressTable = $setup->getTable('sales_order_address', 'sales');
+
+        $query = $quoteConnection
+            ->select()
+            ->from(
+                ['quote_address' => $quoteAddressTable],
+                ['quote_id', 'address_id']
+            )
+            ->joinInner(
+                ['quote' => $quoteTable],
+                'quote_address.quote_id = quote.entity_id',
+                []
+            )
+            ->where('quote.entity_id in (?)', array_column($orderAddresses, 'quote_id'))
+            ->where('address_type = ?', $addressType);
+
+        $quoteAddresses = $quoteConnection->fetchAssoc($query);
+
+        foreach ($orderAddresses as $orderAddress) {
+            $bind = [
+                'quote_address_id' => $quoteAddresses[$orderAddress['quote_id']]['address_id'] ?? null,
+            ];
+            $where = [
+                'orderAddressId' => $orderAddress['entity_id']
+            ];
+
+            $salesConnection->update($salesOrderAddressTable, $bind, $where);
+        }
+    }
 }