From edc8d6e5e76560eca7b59feb62eb1c06c4167d9d Mon Sep 17 00:00:00 2001
From: Azat Khuzhin <a3at.mail@gmail.com>
Date: Sun, 25 Oct 2020 14:14:52 +0300
Subject: [PATCH] Fix async Distributed INSERT w/ prefer_localhost_replica=0
 and internal_replication

---
 programs/server/config.xml                    | 16 +++++++++
 src/Interpreters/Cluster.cpp                  | 11 ++++--
 ..._directory_monitor_batch_inserts.reference | 16 +++++++++
 ...ibuted_directory_monitor_batch_inserts.sql | 36 ++++++++++++++++++-
 4 files changed, 75 insertions(+), 4 deletions(-)

diff --git a/programs/server/config.xml b/programs/server/config.xml
index 9850d77abb7..5bdec5377fd 100644
--- a/programs/server/config.xml
+++ b/programs/server/config.xml
@@ -392,6 +392,22 @@
                 </replica>
             </shard>
         </test_cluster_two_shards>
+        <test_cluster_two_shards_internal_replication>
+            <shard>
+                <internal_replication>true</internal_replication>
+                <replica>
+                    <host>127.0.0.1</host>
+                    <port>9000</port>
+                </replica>
+            </shard>
+            <shard>
+                <internal_replication>true</internal_replication>
+                <replica>
+                    <host>127.0.0.2</host>
+                    <port>9000</port>
+                </replica>
+            </shard>
+        </test_cluster_two_shards_internal_replication>
         <test_shard_localhost_secure>
             <shard>
                 <replica>
diff --git a/src/Interpreters/Cluster.cpp b/src/Interpreters/Cluster.cpp
index 8a98e8282a6..9c2766ae7d6 100644
--- a/src/Interpreters/Cluster.cpp
+++ b/src/Interpreters/Cluster.cpp
@@ -614,13 +614,18 @@ const std::string & Cluster::ShardInfo::pathForInsert(bool prefer_localhost_repl
     if (!has_internal_replication)
         throw Exception("internal_replication is not set", ErrorCodes::LOGICAL_ERROR);
 
-    if (dir_name_for_internal_replication.empty() || dir_name_for_internal_replication_with_local.empty())
-        throw Exception("Directory name for async inserts is empty", ErrorCodes::LOGICAL_ERROR);
-
     if (prefer_localhost_replica)
+    {
+        if (dir_name_for_internal_replication.empty())
+            throw Exception("Directory name for async inserts is empty", ErrorCodes::LOGICAL_ERROR);
         return dir_name_for_internal_replication;
+    }
     else
+    {
+        if (dir_name_for_internal_replication_with_local.empty())
+            throw Exception("Directory name for async inserts is empty", ErrorCodes::LOGICAL_ERROR);
         return dir_name_for_internal_replication_with_local;
+    }
 }
 
 bool Cluster::maybeCrossReplication() const
diff --git a/tests/queries/0_stateless/01040_distributed_directory_monitor_batch_inserts.reference b/tests/queries/0_stateless/01040_distributed_directory_monitor_batch_inserts.reference
index 5565ed6787f..03e58c13ff2 100644
--- a/tests/queries/0_stateless/01040_distributed_directory_monitor_batch_inserts.reference
+++ b/tests/queries/0_stateless/01040_distributed_directory_monitor_batch_inserts.reference
@@ -1,4 +1,20 @@
+test_cluster_two_shards prefer_localhost_replica=0
+0
 0
 1
+1
+test_cluster_two_shards prefer_localhost_replica=1
+0
 0
 1
+1
+test_cluster_two_shards_internal_replication prefer_localhost_replica=0
+0
+0
+1
+1
+test_cluster_two_shards_internal_replication prefer_localhost_replica=1
+0
+0
+1
+1
diff --git a/tests/queries/0_stateless/01040_distributed_directory_monitor_batch_inserts.sql b/tests/queries/0_stateless/01040_distributed_directory_monitor_batch_inserts.sql
index dbec319ab76..dec748789c8 100644
--- a/tests/queries/0_stateless/01040_distributed_directory_monitor_batch_inserts.sql
+++ b/tests/queries/0_stateless/01040_distributed_directory_monitor_batch_inserts.sql
@@ -7,6 +7,40 @@ DROP TABLE IF EXISTS dist_test_01040;
 
 CREATE TABLE test_01040 (key UInt64) ENGINE=TinyLog();
 CREATE TABLE dist_test_01040 AS test_01040 Engine=Distributed(test_cluster_two_shards, currentDatabase(), test_01040, key);
+
+-- internal_replication=false
+SELECT 'test_cluster_two_shards prefer_localhost_replica=0';
+SET prefer_localhost_replica=0;
 INSERT INTO dist_test_01040 SELECT toUInt64(number) FROM numbers(2);
 SYSTEM FLUSH DISTRIBUTED dist_test_01040;
-SELECT * FROM dist_test_01040;
+SELECT * FROM dist_test_01040 ORDER BY key;
+TRUNCATE TABLE test_01040;
+
+SELECT 'test_cluster_two_shards prefer_localhost_replica=1';
+SET prefer_localhost_replica=1;
+INSERT INTO dist_test_01040 SELECT toUInt64(number) FROM numbers(2);
+SYSTEM FLUSH DISTRIBUTED dist_test_01040;
+SELECT * FROM dist_test_01040 ORDER BY key;
+TRUNCATE TABLE test_01040;
+
+DROP TABLE dist_test_01040;
+
+-- internal_replication=true
+CREATE TABLE dist_test_01040 AS test_01040 Engine=Distributed(test_cluster_two_shards_internal_replication, currentDatabase(), test_01040, key);
+SELECT 'test_cluster_two_shards_internal_replication prefer_localhost_replica=0';
+SET prefer_localhost_replica=0;
+INSERT INTO dist_test_01040 SELECT toUInt64(number) FROM numbers(2);
+SYSTEM FLUSH DISTRIBUTED dist_test_01040;
+SELECT * FROM dist_test_01040 ORDER BY key;
+TRUNCATE TABLE test_01040;
+
+SELECT 'test_cluster_two_shards_internal_replication prefer_localhost_replica=1';
+SET prefer_localhost_replica=1;
+INSERT INTO dist_test_01040 SELECT toUInt64(number) FROM numbers(2);
+SYSTEM FLUSH DISTRIBUTED dist_test_01040;
+SELECT * FROM dist_test_01040 ORDER BY key;
+TRUNCATE TABLE test_01040;
+
+
+DROP TABLE dist_test_01040;
+DROP TABLE test_01040;