URL rewrite fix while product website update using mass action

Author: AnshuMishra17

app/code/Magento/CatalogUrlRewrite/Observer/ProductToWebsiteChangeObserver.php

@@ -13,6 +13,8 @@
 use Magento\Store\Model\Store;
 use Magento\UrlRewrite\Model\UrlPersistInterface;
 use Magento\UrlRewrite\Service\V1\Data\UrlRewrite;
+use Magento\Store\Api\StoreWebsiteRelationInterface;
+use Magento\Framework\App\ObjectManager;
 
 /**
  * Observer to assign the products to website
@@ -39,22 +41,31 @@ class ProductToWebsiteChangeObserver implements ObserverInterface
      */
     protected $request;
 
+    /**
+     * @var StoreWebsiteRelationInterface
+     */
+    private $storeWebsiteRelation;
+
     /**
      * @param ProductUrlRewriteGenerator $productUrlRewriteGenerator
      * @param UrlPersistInterface $urlPersist
      * @param ProductRepositoryInterface $productRepository
      * @param RequestInterface $request
+     * @param StoreWebsiteRelationInterface $storeWebsiteRelation
      */
     public function __construct(
         ProductUrlRewriteGenerator $productUrlRewriteGenerator,
         UrlPersistInterface $urlPersist,
         ProductRepositoryInterface $productRepository,
-        RequestInterface $request
+        RequestInterface $request,
+        StoreWebsiteRelationInterface $storeWebsiteRelation = null
     ) {
         $this->productUrlRewriteGenerator = $productUrlRewriteGenerator;
         $this->urlPersist = $urlPersist;
         $this->productRepository = $productRepository;
         $this->request = $request;
+        $this->storeWebsiteRelation = $storeWebsiteRelation ?:
+            ObjectManager::getInstance()->get(StoreWebsiteRelationInterface::class);
     }
 
     /**
@@ -73,10 +84,17 @@ public function execute(\Magento\Framework\Event\Observer $observer)
             );
 
             if (!empty($this->productUrlRewriteGenerator->generate($product))) {
-                $this->urlPersist->deleteByData([
-                    UrlRewrite::ENTITY_ID => $product->getId(),
-                    UrlRewrite::ENTITY_TYPE => ProductUrlRewriteGenerator::ENTITY_TYPE,
-                ]);
+                if ($this->request->getParam('remove_website_ids')) {
+                    foreach ($this->request->getParam('remove_website_ids') as $webId) {
+                        foreach ($this->storeWebsiteRelation->getStoreByWebsiteId($webId) as $storeId) {
+                            $this->urlPersist->deleteByData([
+                                UrlRewrite::ENTITY_ID => $product->getId(),
+                                UrlRewrite::ENTITY_TYPE => ProductUrlRewriteGenerator::ENTITY_TYPE,
+                                UrlRewrite::STORE_ID => $storeId
+                            ]);
+                        }
+                    }
+                }
                 if ($product->getVisibility() != Visibility::VISIBILITY_NOT_VISIBLE) {
                     $this->urlPersist->replace($this->productUrlRewriteGenerator->generate($product));
                 }