[Forwardport] Use batches and direct queries to fix sales address upgrade

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

@@ -8,27 +8,22 @@
 
 use Magento\Eav\Model\Config;
 use Magento\Framework\App\State;
-use Magento\Quote\Model\QuoteFactory;
-use Magento\Sales\Model\OrderFactory;
-use Magento\Sales\Model\ResourceModel\Order\Address\CollectionFactory as AddressCollectionFactory;
-use Magento\Framework\App\ResourceConnection;
-use Magento\Sales\Setup\SalesSetupFactory;
+use Magento\Framework\Setup\ModuleDataSetupInterface;
 use Magento\Framework\Setup\Patch\DataPatchInterface;
 use Magento\Framework\Setup\Patch\PatchVersionInterface;
-use Magento\Framework\Setup\ModuleDataSetupInterface;
+use Magento\Sales\Model\Order\Address;
+use Magento\Sales\Setup\SalesSetupFactory;
 
+/**
+ * Fills quote_address_id in table sales_order_address if it is empty.
+ */
 class FillQuoteAddressIdInSalesOrderAddress implements DataPatchInterface, PatchVersionInterface
 {
     /**
      * @var ModuleDataSetupInterface
      */
     private $moduleDataSetup;
 
-    /**
-     * @var SalesSetupFactory
-     */
-    private $salesSetupFactory;
-
     /**
      * @var State
      */
@@ -40,44 +35,22 @@ class FillQuoteAddressIdInSalesOrderAddress implements DataPatchInterface, Patch
     private $eavConfig;
 
     /**
-     * @var AddressCollectionFactory
-     */
-    private $addressCollectionFactory;
-
-    /**
-     * @var OrderFactory
-     */
-    private $orderFactory;
-
-    /**
-     * @var QuoteFactory
-     */
-    private $quoteFactory;
-
-    /**
-     * PatchInitial constructor.
      * @param ModuleDataSetupInterface $moduleDataSetup
+     * @param State $state
+     * @param Config $eavConfig
      */
     public function __construct(
         ModuleDataSetupInterface $moduleDataSetup,
-        SalesSetupFactory $salesSetupFactory,
         State $state,
-        Config $eavConfig,
-        AddressCollectionFactory $addressCollectionFactory,
-        OrderFactory $orderFactory,
-        QuoteFactory $quoteFactory
+        Config $eavConfig
     ) {
         $this->moduleDataSetup = $moduleDataSetup;
-        $this->salesSetupFactory = $salesSetupFactory;
         $this->state = $state;
         $this->eavConfig = $eavConfig;
-        $this->addressCollectionFactory = $addressCollectionFactory;
-        $this->orderFactory = $orderFactory;
-        $this->quoteFactory = $quoteFactory;
     }
 
     /**
-     * {@inheritdoc}
+     * @inheritdoc
      */
     public function apply()
     {
@@ -96,32 +69,12 @@ 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);
     }
 
     /**
-     * {@inheritdoc}
+     * @inheritdoc
      */
     public static function getDependencies()
     {
@@ -131,18 +84,107 @@ public static function getDependencies()
     }
 
     /**
-     * {@inheritdoc}
+     * @inheritdoc
      */
     public static function getVersion()
     {
         return '2.0.8';
     }
 
     /**
-     * {@inheritdoc}
+     * @inheritdoc
      */
     public function getAliases()
     {
         return [];
     }
+
+    /**
+     * Fill quote_address_id in sales_order_address by type.
+     *
+     * @param ModuleDataSetupInterface $setup
+     * @param string $addressType
+     * @throws \Zend_Db_Statement_Exception
+     */
+    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);
+        }
+    }
+
+    /**
+     * Process filling quote_address_id in sales_order_address in batch.
+     *
+     * @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 = [
+                'entity_id = ?' => $orderAddress['entity_id']
+            ];
+
+            $salesConnection->update($salesOrderAddressTable, $bind, $where);
+        }
+    }
 }