2.2.6 Use batches and direct queries to fix sales address upgrade

app/code/Magento/Sales/Setup/UpgradeData.php

@@ -15,6 +15,7 @@
 use Magento\Framework\Setup\ModuleDataSetupInterface;
 use Magento\Framework\Setup\UpgradeDataInterface;
 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;
 
@@ -42,27 +43,14 @@ class UpgradeData implements UpgradeDataInterface
      */
     private $aggregatedFieldConverter;
 
-    /**
-     * @var AddressCollectionFactory
-     */
-    private $addressCollectionFactory;
-
-    /**
-     * @var OrderFactory
-     */
-    private $orderFactory;
-
-    /**
-     * @var QuoteFactory
-     */
-    private $quoteFactory;
-
     /**
      * @var State
      */
     private $state;
 
     /**
+     * @SuppressWarnings(PHPMD.UnusedFormalParameter)
+     *
      * @param SalesSetupFactory $salesSetupFactory
      * @param Config $eavConfig
      * @param AggregatedFieldDataConverter $aggregatedFieldConverter
@@ -83,9 +71,6 @@ public function __construct(
         $this->salesSetupFactory = $salesSetupFactory;
         $this->eavConfig = $eavConfig;
         $this->aggregatedFieldConverter = $aggregatedFieldConverter;
-        $this->addressCollectionFactory = $addressCollFactory;
-        $this->orderFactory = $orderFactory;
-        $this->quoteFactory = $quoteFactory;
         $this->state = $state;
     }
 
@@ -125,6 +110,7 @@ public function upgrade(ModuleDataSetupInterface $setup, ModuleContextInterface
      * @param string $setupVersion
      * @param SalesSetup $salesSetup
      * @return void
+     * @throws \Magento\Framework\DB\FieldDataConversionException
      */
     private function convertSerializedDataToJson($setupVersion, SalesSetup $salesSetup)
     {
@@ -173,32 +159,95 @@ private function convertSerializedDataToJson($setupVersion, SalesSetup $salesSet
 
     /**
      * Fill quote_address_id in table sales_order_address if it is empty.
-     *
      * @param ModuleDataSetupInterface $setup
      */
     public function fillQuoteAddressIdInSalesOrderAddress(ModuleDataSetupInterface $setup)
     {
-        $addressTable = $setup->getTable('sales_order_address');
-        $updateOrderAddress = $setup->getConnection()
+        $this->fillQuoteAddressIdInSalesOrderAddressByType($setup, Address::TYPE_SHIPPING);
+        $this->fillQuoteAddressIdInSalesOrderAddressByType($setup, Address::TYPE_BILLING);
+    }
+
+    /**
+     * @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' => $setup->getTable('sales_order')],
-                $addressTable . '.parent_id = sales_order.entity_id',
-                ['quote_address_id' => 'quote_address.address_id']
+                ['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_address' => $setup->getTable('quote_address')],
-                'sales_order.quote_id = quote_address.quote_id 
-                AND ' . $addressTable . '.address_type = quote_address.address_type',
+                ['quote' => $quoteTable],
+                'quote_address.quote_id = quote.entity_id',
                 []
             )
-            ->where(
-                $addressTable . '.quote_address_id IS NULL'
-            );
-        $updateOrderAddress = $setup->getConnection()->updateFromSelect(
-            $updateOrderAddress,
-            $addressTable
-        );
-        $setup->getConnection()->query($updateOrderAddress);
+            ->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 = [
+                'entity_id = ?' => $orderAddress['entity_id']
+            ];
+
+            $salesConnection->update($salesOrderAddressTable, $bind, $where);
+        }
     }
 }