Partially cherry-picked 9a3adc7 for Context::getDefaultOrAuxiliaryZooKeeper

al13n321 · Enmk · commit 0659951883d1 · 2024-11-25T13:14:05.000+01:00
diff --git a/src/Common/ZooKeeper/ZooKeeper.cpp b/src/Common/ZooKeeper/ZooKeeper.cpp
@@ -1686,11 +1686,10 @@ std::string normalizeZooKeeperPath(std::string zookeeper_path, bool check_starts
 
 String extractZooKeeperName(const String & path)
 {
-    static constexpr auto default_zookeeper_name = "default";
     if (path.empty())
         throw DB::Exception(DB::ErrorCodes::BAD_ARGUMENTS, "ZooKeeper path should not be empty");
     if (path[0] == '/')
-        return default_zookeeper_name;
+        return String(DEFAULT_ZOOKEEPER_NAME);
     auto pos = path.find(":/");
     if (pos != String::npos && pos < path.find('/'))
     {
@@ -1699,7 +1698,7 @@ String extractZooKeeperName(const String & path)
             throw DB::Exception(DB::ErrorCodes::BAD_ARGUMENTS, "Zookeeper path should start with '/' or '<auxiliary_zookeeper_name>:/'");
         return zookeeper_name;
     }
-    return default_zookeeper_name;
+    return String(DEFAULT_ZOOKEEPER_NAME);
 }
 
 String extractZooKeeperPath(const String & path, bool check_starts_with_slash, LoggerPtr log)
diff --git a/src/Common/ZooKeeper/ZooKeeper.h b/src/Common/ZooKeeper/ZooKeeper.h
@@ -47,6 +47,10 @@ namespace zkutil
 /// Preferred size of multi() command (in number of ops)
 constexpr size_t MULTI_BATCH_SIZE = 100;
 
+/// Path "default:/foo" refers to znode "/foo" in the default zookeeper,
+/// path "other:/foo" refers to znode "/foo" in auxiliary zookeeper named "other".
+constexpr std::string_view DEFAULT_ZOOKEEPER_NAME = "default";
+
 struct ShuffleHost
 {
     enum AvailabilityZoneInfo
diff --git a/src/Interpreters/Context.cpp b/src/Interpreters/Context.cpp
@@ -3746,6 +3746,11 @@ zkutil::ZooKeeperPtr Context::getAuxiliaryZooKeeper(const String & name) const
     return zookeeper->second;
 }
 
+std::shared_ptr<zkutil::ZooKeeper> Context::getDefaultOrAuxiliaryZooKeeper(const String & name) const
+{
+    return name == zkutil::DEFAULT_ZOOKEEPER_NAME ? getZooKeeper() : getAuxiliaryZooKeeper(name);
+}
+
 
 std::map<String, zkutil::ZooKeeperPtr> Context::getAuxiliaryZooKeepers() const
 {
diff --git a/src/Interpreters/Context.h b/src/Interpreters/Context.h
@@ -997,6 +997,8 @@ class Context: public ContextData, public std::enable_shared_from_this<Context>
     std::shared_ptr<zkutil::ZooKeeper> getZooKeeper() const;
     /// Same as above but return a zookeeper connection from auxiliary_zookeepers configuration entry.
     std::shared_ptr<zkutil::ZooKeeper> getAuxiliaryZooKeeper(const String & name) const;
+    /// If name == "default", same as getZooKeeper(), otherwise same as getAuxiliaryZooKeeper().
+    std::shared_ptr<zkutil::ZooKeeper> getDefaultOrAuxiliaryZooKeeper(const String & name) const;
     /// return Auxiliary Zookeeper map
     std::map<String, zkutil::ZooKeeperPtr> getAuxiliaryZooKeepers() const;
 

Original file line number	Diff line number	Diff line change
`@@ -1686,11 +1686,10 @@ std::string normalizeZooKeeperPath(std::string zookeeper_path, bool check_starts`
`1686`	`1686`
`1687`	`1687`	`String extractZooKeeperName(const String & path)`
`1688`	`1688`	`{`
`1689`		`- static constexpr auto default_zookeeper_name = "default";`
`1690`	`1689`	`if (path.empty())`
`1691`	`1690`	`throw DB::Exception(DB::ErrorCodes::BAD_ARGUMENTS, "ZooKeeper path should not be empty");`
`1692`	`1691`	`if (path[0] == '/')`
`1693`		`- return default_zookeeper_name;`
	`1692`	`+ return String(DEFAULT_ZOOKEEPER_NAME);`
`1694`	`1693`	`auto pos = path.find(":/");`
`1695`	`1694`	`if (pos != String::npos && pos < path.find('/'))`
`1696`	`1695`	`{`
`@@ -1699,7 +1698,7 @@ String extractZooKeeperName(const String & path)`
`1699`	`1698`	`throw DB::Exception(DB::ErrorCodes::BAD_ARGUMENTS, "Zookeeper path should start with '/' or '<auxiliary_zookeeper_name>:/'");`
`1700`	`1699`	`return zookeeper_name;`
`1701`	`1700`	`}`
`1702`		`- return default_zookeeper_name;`
	`1701`	`+ return String(DEFAULT_ZOOKEEPER_NAME);`
`1703`	`1702`	`}`
`1704`	`1703`
`1705`	`1704`	`String extractZooKeeperPath(const String & path, bool check_starts_with_slash, LoggerPtr log)`
Original file line number	Diff line number	Diff line change
`@@ -3746,6 +3746,11 @@ zkutil::ZooKeeperPtr Context::getAuxiliaryZooKeeper(const String & name) const`
`3746`	`3746`	`return zookeeper->second;`
`3747`	`3747`	`}`
`3748`	`3748`
	`3749`	`+std::shared_ptr<zkutil::ZooKeeper> Context::getDefaultOrAuxiliaryZooKeeper(const String & name) const`
	`3750`	`+{`
	`3751`	`+ return name == zkutil::DEFAULT_ZOOKEEPER_NAME ? getZooKeeper() : getAuxiliaryZooKeeper(name);`
	`3752`	`+}`
	`3753`	`+`
`3749`	`3754`
`3750`	`3755`	`std::map<String, zkutil::ZooKeeperPtr> Context::getAuxiliaryZooKeepers() const`
`3751`	`3756`	`{`