diff --git a/lib/internal/Magento/Framework/Setup/Declaration/Schema/OperationsExecutor.php b/lib/internal/Magento/Framework/Setup/Declaration/Schema/OperationsExecutor.php
index ed1baff1d3323..3ed7c9d69fc41 100644
--- a/lib/internal/Magento/Framework/Setup/Declaration/Schema/OperationsExecutor.php
+++ b/lib/internal/Magento/Framework/Setup/Declaration/Schema/OperationsExecutor.php
@@ -132,10 +132,10 @@ public function getDestructiveOperations()
     private function startSetupForAllConnections()
     {
         foreach ($this->sharding->getResources() as $resource) {
-            $this->resourceConnection->getConnection($resource)
-                ->startSetup();
-            $this->resourceConnection->getConnection($resource)
-                ->query('SET UNIQUE_CHECKS=0');
+            $connection = $this->resourceConnection->getConnection($resource);
+
+            $connection->startSetup();
+            $connection->query('SET @OLD_UNIQUE_CHECKS=@@UNIQUE_CHECKS, UNIQUE_CHECKS=0');
         }
     }
 
@@ -148,8 +148,10 @@ private function startSetupForAllConnections()
     private function endSetupForAllConnections()
     {
         foreach ($this->sharding->getResources() as $resource) {
-            $this->resourceConnection->getConnection($resource)
-                ->endSetup();
+            $connection = $this->resourceConnection->getConnection($resource);
+
+            $connection->query('SET UNIQUE_CHECKS=IF(@OLD_UNIQUE_CHECKS=0, 0, 1)');
+            $connection->endSetup();
         }
     }
 
diff --git a/lib/internal/Magento/Framework/Setup/Test/Unit/Declaration/Schema/OperationsExecutorTest.php b/lib/internal/Magento/Framework/Setup/Test/Unit/Declaration/Schema/OperationsExecutorTest.php
index 1090f2e04e190..ced07a1fc7650 100644
--- a/lib/internal/Magento/Framework/Setup/Test/Unit/Declaration/Schema/OperationsExecutorTest.php
+++ b/lib/internal/Magento/Framework/Setup/Test/Unit/Declaration/Schema/OperationsExecutorTest.php
@@ -160,7 +160,7 @@ public function testExecute()
         $connectionMock = $this->getMockBuilder(Mysql::class)
             ->disableOriginalConstructor()
             ->getMock();
-        $this->resourceConnectionMock->expects(self::exactly(3))
+        $this->resourceConnectionMock->expects(self::exactly(2))
             ->method('getConnection')
             ->with('default')
             ->willReturn($connectionMock);