refactor: Replace getAllConnectors with ConnectorRegistry APIs

presto-native-execution/presto_cpp/main/CMakeLists.txt

@@ -88,6 +88,7 @@ target_link_libraries(
   velox_functions_prestosql
   velox_gcs
   velox_hdfs
+  velox_connector_registry
   velox_hive_connector
   velox_hive_iceberg_splitreader
   velox_hive_partition_function

presto-native-execution/presto_cpp/main/PeriodicTaskManager.cpp

@@ -240,9 +240,8 @@ PeriodicTaskManager::PeriodicTaskManager(
     TaskManager* taskManager,
     const velox::memory::MemoryAllocator* memoryAllocator,
     const velox::cache::AsyncDataCache* asyncDataCache,
-    const std::unordered_map<
-        std::string,
-        std::shared_ptr<velox::connector::Connector>>& connectors,
+    std::vector<std::shared_ptr<velox::connector::hive::HiveConnector>>
+        hiveConnectors,
     PrestoServer* server)
     : driverCPUExecutor_(driverCPUExecutor),
       spillerExecutor_(spillerExecutor),
@@ -254,7 +253,7 @@ PeriodicTaskManager::PeriodicTaskManager(
       memoryAllocator_(memoryAllocator),
       asyncDataCache_(asyncDataCache),
       arbitrator_(velox::memory::memoryManager()->arbitrator()),
-      connectors_(connectors),
+      hiveConnectors_(std::move(hiveConnectors)),
       server_(server) {}
 
 void PeriodicTaskManager::start() {
@@ -416,12 +415,8 @@ void PeriodicTaskManager::addPrestoExchangeSourceMemoryStatsTask() {
 
 void PeriodicTaskManager::addConnectorStatsTask() {
   std::vector<HiveConnectorStatsReporter> reporters;
-  for (const auto& itr : connectors_) {
-    if (auto hiveConnector =
-            std::dynamic_pointer_cast<velox::connector::hive::HiveConnector>(
-                itr.second)) {
-      reporters.emplace_back(std::move(hiveConnector));
-    }
+  for (const auto& hiveConnector : hiveConnectors_) {
+    reporters.emplace_back(hiveConnector);
   }
   addTask(
       [reporters = std::move(reporters)]() mutable {

presto-native-execution/presto_cpp/main/PeriodicTaskManager.h

@@ -22,8 +22,8 @@ class CPUThreadPoolExecutor;
 class IOThreadPoolExecutor;
 } // namespace folly
 
-namespace facebook::velox::connector {
-class Connector;
+namespace facebook::velox::connector::hive {
+class HiveConnector;
 }
 
 namespace facebook::velox::cache {
@@ -50,9 +50,8 @@ class PeriodicTaskManager {
       TaskManager* taskManager,
       const velox::memory::MemoryAllocator* memoryAllocator,
       const velox::cache::AsyncDataCache* asyncDataCache,
-      const std::unordered_map<
-          std::string,
-          std::shared_ptr<velox::connector::Connector>>& connectors,
+      std::vector<std::shared_ptr<velox::connector::hive::HiveConnector>>
+          hiveConnectors,
       PrestoServer* server);
 
   /// Invoked to start all registered, and fundamental periodic tasks running at
@@ -137,9 +136,8 @@ class PeriodicTaskManager {
   const velox::memory::MemoryAllocator* memoryAllocator_;
   const velox::cache::AsyncDataCache* asyncDataCache_;
   const velox::memory::MemoryArbitrator* arbitrator_;
-  const std::unordered_map<
-      std::string,
-      std::shared_ptr<velox::connector::Connector>>& connectors_;
+  std::vector<std::shared_ptr<velox::connector::hive::HiveConnector>>
+      hiveConnectors_;
   PrestoServer* server_;
 
   int64_t lastHttpClientNumConnectionsCreated_{0};

presto-native-execution/presto_cpp/main/PrestoServer.cpp

@@ -56,6 +56,7 @@
 #include "velox/common/file/FileSystems.h"
 #include "velox/common/memory/SharedArbitrator.h"
 #include "velox/connectors/Connector.h"
+#include "velox/connectors/ConnectorRegistry.h"
 #include "velox/connectors/hive/HiveConnector.h"
 #include "velox/connectors/hive/storage_adapters/abfs/RegisterAbfsFileSystem.h"
 #include "velox/connectors/hive/storage_adapters/gcs/RegisterGcsFileSystem.h"
@@ -340,7 +341,8 @@ void PrestoServer::run() {
       taskManager_.get(),
       memoryAllocator,
       asyncDataCache,
-      velox::connector::getAllConnectors(),
+      velox::connector::ConnectorRegistry::findAll<
+          velox::connector::hive::HiveConnector>(),
       this);
   addServerPeriodicTasks();
   addAdditionalPeriodicTasks();
@@ -1448,27 +1450,8 @@ void PrestoServer::registerSystemConnector() {
 
 void PrestoServer::unregisterConnectors() {
   PRESTO_SHUTDOWN_LOG(INFO) << "Unregistering connectors";
-  auto connectors = velox::connector::getAllConnectors();
-  if (connectors.empty()) {
-    PRESTO_SHUTDOWN_LOG(INFO) << "No connectors to unregister";
-    return;
-  }
-
-  PRESTO_SHUTDOWN_LOG(INFO)
-      << "Unregistering " << connectors.size() << " connectors";
-  for (const auto& connectorEntry : connectors) {
-    if (velox::connector::unregisterConnector(connectorEntry.first)) {
-      PRESTO_SHUTDOWN_LOG(INFO)
-          << "Unregistered connector: " << connectorEntry.first;
-    } else {
-      PRESTO_SHUTDOWN_LOG(INFO)
-          << "Unable to unregister connector: " << connectorEntry.first;
-    }
-  }
-
-  velox::connector::unregisterConnector("$system@system");
-  PRESTO_SHUTDOWN_LOG(INFO)
-      << "Unregistered " << connectors.size() << " connectors";
+  velox::connector::ConnectorRegistry::unregisterAll();
+  PRESTO_SHUTDOWN_LOG(INFO) << "Unregistered all connectors";
 }
 
 void PrestoServer::registerShuffleInterfaceFactories() {