From 8bf4ab167cbac8f3a26c2df559124d78f363917a Mon Sep 17 00:00:00 2001 From: Jialiang Tan Date: Mon, 18 Aug 2025 23:34:45 -0700 Subject: [PATCH 001/113] [native][pos] Make system config free form --- .../spark/PrestoSparkInjectorFactory.java | 15 +- .../presto/spark/PrestoSparkModule.java | 2 - .../nativeprocess/NativeExecutionModule.java | 9 +- .../nativeprocess/NativeExecutionProcess.java | 7 +- .../property/NativeExecutionConfigModule.java | 51 + .../property/NativeExecutionSystemConfig.java | 893 ++++-------------- .../execution/TestNativeExecutionProcess.java | 3 +- .../http/TestPrestoSparkHttpClient.java | 2 +- .../TestNativeExecutionSystemConfig.java | 237 +++-- 9 files changed, 398 insertions(+), 821 deletions(-) create mode 100644 presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/NativeExecutionConfigModule.java diff --git a/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkInjectorFactory.java b/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkInjectorFactory.java index 9608adf93f843..bfabba20b632e 100644 --- a/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkInjectorFactory.java +++ b/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkInjectorFactory.java @@ -30,6 +30,7 @@ import com.facebook.presto.server.security.PrestoAuthenticatorManager; import com.facebook.presto.spark.classloader_interface.PrestoSparkBootstrapTimer; import com.facebook.presto.spark.classloader_interface.SparkProcessType; +import com.facebook.presto.spark.execution.property.NativeExecutionConfigModule; import com.facebook.presto.spi.security.AccessControl; import com.facebook.presto.sql.analyzer.FeaturesConfig; import com.facebook.presto.sql.parser.SqlParserOptions; @@ -165,6 +166,11 @@ public Injector create(PrestoSparkBootstrapTimer bootstrapTimer) modules.add(new TempStorageModule()); } + Map nativeWorkerConfigs = new HashMap<>( + this.nativeWorkerConfigProperties.orElse(ImmutableMap.of())); + nativeWorkerConfigs.put("node.environment", "spark"); + modules.add(new NativeExecutionConfigModule(nativeWorkerConfigs)); + modules.addAll(additionalModules); Bootstrap app = new Bootstrap(modules.build()); @@ -172,11 +178,10 @@ public Injector create(PrestoSparkBootstrapTimer bootstrapTimer) // Stream redirect doesn't work well with spark logging app.doNotInitializeLogging(); - Map requiredProperties = new HashMap<>(); - requiredProperties.put("node.environment", "spark"); - requiredProperties.putAll(configProperties); - - app.setRequiredConfigurationProperties(ImmutableMap.copyOf(requiredProperties)); + Map requiredConfigProperties = new HashMap<>(); + requiredConfigProperties.put("node.environment", "spark"); + requiredConfigProperties.putAll(configProperties); + app.setRequiredConfigurationProperties(ImmutableMap.copyOf(requiredConfigProperties)); bootstrapTimer.beginInjectorInitialization(); Injector injector = app.initialize(); diff --git a/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkModule.java b/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkModule.java index 650d533492476..82716e8a840c2 100644 --- a/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkModule.java +++ b/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkModule.java @@ -122,7 +122,6 @@ import com.facebook.presto.spark.execution.http.BatchTaskUpdateRequest; import com.facebook.presto.spark.execution.property.NativeExecutionConnectorConfig; import com.facebook.presto.spark.execution.property.NativeExecutionNodeConfig; -import com.facebook.presto.spark.execution.property.NativeExecutionSystemConfig; import com.facebook.presto.spark.execution.shuffle.PrestoSparkLocalShuffleReadInfo; import com.facebook.presto.spark.execution.shuffle.PrestoSparkLocalShuffleWriteInfo; import com.facebook.presto.spark.execution.task.PrestoSparkNativeTaskExecutorFactory; @@ -284,7 +283,6 @@ protected void setup(Binder binder) configBinder(binder).bindConfig(SessionPropertyProviderConfig.class); configBinder(binder).bindConfig(PrestoSparkConfig.class); configBinder(binder).bindConfig(TracingConfig.class); - configBinder(binder).bindConfig(NativeExecutionSystemConfig.class); configBinder(binder).bindConfig(NativeExecutionNodeConfig.class); configBinder(binder).bindConfig(NativeExecutionConnectorConfig.class); configBinder(binder).bindConfig(PlanCheckerProviderManagerConfig.class); diff --git a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/nativeprocess/NativeExecutionModule.java b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/nativeprocess/NativeExecutionModule.java index d265938edf999..d3e9fb1bc3786 100644 --- a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/nativeprocess/NativeExecutionModule.java +++ b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/nativeprocess/NativeExecutionModule.java @@ -23,6 +23,7 @@ import com.facebook.presto.spark.execution.task.ForNativeExecutionTask; import com.facebook.presto.spark.execution.task.NativeExecutionTaskFactory; import com.google.common.annotations.VisibleForTesting; +import com.google.common.collect.ImmutableMap; import com.google.inject.Binder; import com.google.inject.Module; import com.google.inject.Provides; @@ -73,9 +74,13 @@ protected void bindShuffle(Binder binder) protected void bindWorkerProperties(Binder binder) { - newOptionalBinder(binder, new TypeLiteral>() {}).setDefault().to(PrestoSparkWorkerProperty.class).in(Scopes.SINGLETON); + newOptionalBinder(binder, new TypeLiteral>() { + }).setDefault().to(PrestoSparkWorkerProperty.class).in(Scopes.SINGLETON); if (connectorConfig.isPresent()) { - binder.bind(PrestoSparkWorkerProperty.class).toInstance(new PrestoSparkWorkerProperty(connectorConfig.get(), new NativeExecutionNodeConfig(), new NativeExecutionSystemConfig())); + binder.bind(PrestoSparkWorkerProperty.class).toInstance( + new PrestoSparkWorkerProperty(connectorConfig.get(), + new NativeExecutionNodeConfig(), new NativeExecutionSystemConfig( + ImmutableMap.of()))); } else { binder.bind(PrestoSparkWorkerProperty.class).in(Scopes.SINGLETON); diff --git a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/nativeprocess/NativeExecutionProcess.java b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/nativeprocess/NativeExecutionProcess.java index b1d104cdbe9a7..794574b46c7f3 100644 --- a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/nativeprocess/NativeExecutionProcess.java +++ b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/nativeprocess/NativeExecutionProcess.java @@ -22,6 +22,7 @@ import com.facebook.presto.spark.execution.http.PrestoSparkHttpServerClient; import com.facebook.presto.spark.execution.http.server.RequestErrorTracker; import com.facebook.presto.spark.execution.http.server.smile.BaseResponse; +import com.facebook.presto.spark.execution.property.NativeExecutionSystemConfig; import com.facebook.presto.spark.execution.property.WorkerProperty; import com.facebook.presto.spi.PrestoException; import com.google.common.annotations.VisibleForTesting; @@ -337,11 +338,13 @@ private void populateConfigurationFiles(String configBasePath) // there is no port isolation among all the containers running on the same host, so we have // to pick unique port per worker to avoid port collision. This config will be passed down to // the native execution process eventually for process initialization. - workerProperty.getSystemConfig().setHttpServerPort(port); + workerProperty.getSystemConfig() + .update(NativeExecutionSystemConfig.HTTP_SERVER_HTTP_PORT, String.valueOf(port)); workerProperty.populateAllProperties( Paths.get(configBasePath, WORKER_CONFIG_FILE), Paths.get(configBasePath, WORKER_NODE_CONFIG_FILE), - Paths.get(configBasePath, format("%s%s.properties", WORKER_CONNECTOR_CONFIG_FILE, getNativeExecutionCatalogName(session)))); + Paths.get(configBasePath, format("%s%s.properties", WORKER_CONNECTOR_CONFIG_FILE, + getNativeExecutionCatalogName(session)))); } private void doGetServerInfo(SettableFuture future) diff --git a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/NativeExecutionConfigModule.java b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/NativeExecutionConfigModule.java new file mode 100644 index 0000000000000..38f6b365bdf79 --- /dev/null +++ b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/NativeExecutionConfigModule.java @@ -0,0 +1,51 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.spark.execution.property; + +import com.google.common.collect.ImmutableMap; +import com.google.inject.AbstractModule; +import com.google.inject.Scopes; +import com.google.inject.TypeLiteral; +import com.google.inject.name.Names; + +import java.util.Map; + +import static java.util.Objects.requireNonNull; + +/** + * Native configuration module that allows the population of config properties without config + * annotation checks. + */ +public class NativeExecutionConfigModule + extends AbstractModule +{ + private final Map systemConfigs; + + public NativeExecutionConfigModule(Map systemConfigs) + { + this.systemConfigs = ImmutableMap.copyOf( + requireNonNull(systemConfigs, "systemConfigs is null")); + } + + @Override + protected void configure() + { + bind(new TypeLiteral>() {}) + .annotatedWith( + Names.named(NativeExecutionSystemConfig.NATIVE_EXECUTION_SYSTEM_CONFIG)) + .toInstance(systemConfigs); + + bind(NativeExecutionSystemConfig.class).in(Scopes.SINGLETON); + } +} diff --git a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/NativeExecutionSystemConfig.java b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/NativeExecutionSystemConfig.java index f05eb2fa915c4..a0c757ab58595 100644 --- a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/NativeExecutionSystemConfig.java +++ b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/NativeExecutionSystemConfig.java @@ -13,740 +13,209 @@ */ package com.facebook.presto.spark.execution.property; -import com.facebook.airlift.configuration.Config; -import com.facebook.airlift.units.DataSize; -import com.facebook.airlift.units.Duration; import com.google.common.collect.ImmutableMap; +import org.apache.spark.SparkEnv$; +import org.apache.spark.SparkFiles; +import javax.inject.Inject; +import javax.inject.Named; + +import java.io.File; +import java.util.HashMap; import java.util.Map; -import java.util.concurrent.TimeUnit; +import static java.lang.String.format; import static java.util.Objects.requireNonNull; /** - * This config class corresponds to config.properties for native execution process. Properties inside will be used in Configs::SystemConfig in Configs.h/cpp + * This config class corresponds to config.properties for native execution process. Properties + * inside will be used in Configs::SystemConfig in Configs.h/cpp */ public class NativeExecutionSystemConfig { - private static final String CONCURRENT_LIFESPANS_PER_TASK = "concurrent-lifespans-per-task"; - private static final String ENABLE_SERIALIZED_PAGE_CHECKSUM = "enable-serialized-page-checksum"; - private static final String ENABLE_VELOX_EXPRESSION_LOGGING = "enable_velox_expression_logging"; - private static final String ENABLE_VELOX_TASK_LOGGING = "enable_velox_task_logging"; - // Port on which presto-native http server should run - private static final String HTTP_SERVER_HTTP_PORT = "http-server.http.port"; - private static final String HTTP_SERVER_REUSE_PORT = "http-server.reuse-port"; - private static final String HTTP_SERVER_BIND_TO_NODE_INTERNAL_ADDRESS_ONLY_ENABLED = "http-server.bind-to-node-internal-address-only-enabled"; - // Number of I/O thread to use for serving http request on presto-native (proxygen server) - // this excludes worker thread used by velox - private static final String HTTP_SERVER_HTTPS_PORT = "http-server.https.port"; - private static final String HTTP_SERVER_HTTPS_ENABLED = "http-server.https.enabled"; - - // This config control what cipher suites are supported by Native workers for server and client. - // Note Java and folly::SSLContext use different names to refer to the same cipher. - // (guess for different name, Java specific authentication,key exchange and cipher together and folly just cipher). - // For e.g. TLS_RSA_WITH_AES_256_GCM_SHA384 in Java and AES256-GCM-SHA384 in folly::SSLContext. - // The ciphers need to enable worker to worker, worker to coordinator and coordinator to worker communication. - // Have at least one cipher suite that is shared for the above 3, otherwise weird failures will result. - private static final String HTTPS_CIPHERS = "https-supported-ciphers"; - - // Note: Java packages cert and key in combined JKS file. But CPP requires them separately. - // The HTTPS provides integrity and not security(authentication/authorization). - // But the HTTPS will protect against data corruption by bad router and man in middle attacks. - - // The cert path for the https server - private static final String HTTPS_CERT_PATH = "https-cert-path"; - // The key path for the https server - private static final String HTTPS_KEY_PATH = "https-key-path"; - - // TODO: others use "-" separator and this property use _ separator. Fix them. - private static final String HTTP_SERVER_NUM_IO_THREADS_HW_MULTIPLIER = "http-server.num-io-threads-hw-multiplier"; - private static final String EXCHANGE_HTTP_CLIENT_NUM_IO_THREADS_HW_MULTIPLIER = "exchange.http-client.num-io-threads-hw-multiplier"; - private static final String ASYNC_DATA_CACHE_ENABLED = "async-data-cache-enabled"; - private static final String ASYNC_CACHE_SSD_GB = "async-cache-ssd-gb"; - private static final String CONNECTOR_NUM_IO_THREADS_HW_MULTIPLIER = "connector.num-io-threads-hw-multiplier"; - private static final String PRESTO_VERSION = "presto.version"; - private static final String SHUTDOWN_ONSET_SEC = "shutdown-onset-sec"; - // Memory related configurations. - private static final String SYSTEM_MEMORY_GB = "system-memory-gb"; - private static final String QUERY_MEMORY_GB = "query.max-memory-per-node"; - private static final String USE_MMAP_ALLOCATOR = "use-mmap-allocator"; - // Memory arbitration related configurations. - // Set the memory arbitrator kind. If it is empty, then there is no memory - // arbitration, when a query runs out of its capacity, the query will fail. - // If it set to "SHARED" (default), the shared memory arbitrator will be - // used to conduct arbitration and try to trigger disk spilling to reclaim - // memory so the query can run through completion. - private static final String MEMORY_ARBITRATOR_KIND = "memory-arbitrator-kind"; - // Set memory arbitrator capacity to the same as per-query memory capacity - // as there is only one query running at Presto-on-Spark at a time. - private static final String MEMORY_ARBITRATOR_CAPACITY_GB = "query-memory-gb"; - // Set memory arbitrator reserved capacity. Since there is only one query - // running at Presto-on-Spark at a time, then we shall set this to zero. - private static final String SHARED_ARBITRATOR_RESERVED_CAPACITY = "shared-arbitrator.reserved-capacity"; - // Set the initial memory capacity when we create a query memory pool. For - // Presto-on-Spark, we set it to 'query-memory-gb' to allocate all the - // memory arbitrator capacity to the query memory pool on its creation as - // there is only one query running at a time. - private static final String SHARED_ARBITRATOR_MEMORY_POOL_INITIAL_CAPACITY = "shared-arbitrator.memory-pool-initial-capacity"; - // Set the reserved memory capacity when we create a query memory pool. For - // Presto-on-Spark, we set this to zero as there is only one query running - // at a time. - private static final String SHARED_ARBITRATOR_MEMORY_POOL_RESERVED_CAPACITY = "shared-arbitrator.memory-pool-reserved-capacity"; - private static final String SHARED_ARBITRATOR_MAX_MEMORY_ARBITRATION_TIME = "shared-arbitrator.max-memory-arbitration-time"; - // Spilling related configs. - private static final String SPILLER_SPILL_PATH = "experimental.spiller-spill-path"; - private static final String TASK_MAX_DRIVERS_PER_TASK = "task.max-drivers-per-task"; - // Tasks are considered old, when they are in not-running state and it ended more than - // OLD_TASK_CLEANUP_MS ago or last heartbeat was more than OLD_TASK_CLEANUP_MS ago. - // For Presto-On-Spark, this is not relevant as it runs tasks serially, and spark's speculative - // execution takes care of zombie tasks. - private static final String ENABLE_OLD_TASK_CLEANUP = "enable-old-task-cleanup"; - // Name of exchange client to use - private static final String SHUFFLE_NAME = "shuffle.name"; - // Feature flag for access log on presto-native http server - private static final String HTTP_SERVER_ACCESS_LOGS = "http-server.enable-access-log"; - // Terminates the native process and generates a core file on an allocation failure - private static final String CORE_ON_ALLOCATION_FAILURE_ENABLED = "core-on-allocation-failure-enabled"; - // Spill related properties - private static final String SPILL_ENABLED = "spill-enabled"; - private static final String AGGREGATION_SPILL_ENABLED = "aggregation-spill-enabled"; - private static final String JOIN_SPILL_ENABLED = "join-spill-enabled"; - private static final String ORDER_BY_SPILL_ENABLED = "order-by-spill-enabled"; - private static final String MAX_SPILL_BYTES = "max-spill-bytes"; - - private boolean enableSerializedPageChecksum = true; - private boolean enableVeloxExpressionLogging; - private boolean enableVeloxTaskLogging = true; - private boolean httpServerReusePort = true; - private boolean httpServerBindToNodeInternalAddressOnlyEnabled = true; - private int httpServerPort = 7777; - private double httpServerNumIoThreadsHwMultiplier = 1.0; - private int httpsServerPort = 7778; - private boolean enableHttpsCommunication; - private String httpsCiphers = "AES128-SHA,AES128-SHA256,AES256-GCM-SHA384"; - private String httpsCertPath = ""; - private String httpsKeyPath = ""; - private double exchangeHttpClientNumIoThreadsHwMultiplier = 1.0; - private boolean asyncDataCacheEnabled; // false - private int asyncCacheSsdGb; // 0 - private double connectorNumIoThreadsHwMultiplier; // 0.0 - private int shutdownOnsetSec = 10; - private int systemMemoryGb = 10; - // Reserve 2GB from system memory for system operations such as disk - // spilling and cache prefetch. - private DataSize queryMemoryGb = new DataSize(8, DataSize.Unit.GIGABYTE); - - private boolean useMmapAllocator = true; - private String memoryArbitratorKind = "SHARED"; - private int memoryArbitratorCapacityGb = 8; - private DataSize sharedArbitratorReservedCapacity = new DataSize(0, DataSize.Unit.GIGABYTE); - private DataSize sharedArbitratorMemoryPoolInitialCapacity = new DataSize(4, DataSize.Unit.GIGABYTE); - private DataSize sharedArbitratorMemoryPoolReservedCapacity = new DataSize(32, DataSize.Unit.MEGABYTE); - private Duration sharedArbitratorMaxMemoryArbitrationTime = new Duration(5, TimeUnit.MINUTES); - private String spillerSpillPath = ""; - private int concurrentLifespansPerTask = 5; - private int maxDriversPerTask = 15; - private boolean enableOldTaskCleanUp; // false; - private String prestoVersion = "dummy.presto.version"; - private String shuffleName = "local"; - private boolean enableHttpServerAccessLog = true; - private boolean coreOnAllocationFailureEnabled; - private boolean spillEnabled = true; - private boolean aggregationSpillEnabled = true; - private boolean joinSpillEnabled = true; - private boolean orderBySpillEnabled = true; - private Long maxSpillBytes = 600L << 30; - - // TODO: Deprecate following configs - private static final String REGISTER_TEST_FUNCTIONS = "register-test-functions"; - private static final String MEMORY_POOL_INIT_CAPACITY = "memory-pool-init-capacity"; - private static final String MEMORY_POOL_RESERVED_CAPACITY = "memory-pool-reserved-capacity"; - private static final String MEMORY_POOL_TRANSFER_CAPACITY = "memory-pool-transfer-capacity"; - private static final String MEMORY_RECLAIM_WAIT_MS = "memory-reclaim-wait-ms"; - private boolean registerTestFunctions; - private long memoryPoolInitCapacity; - private long memoryPoolReservedCapacity; - private long memoryPoolTransferCapacity; - private long memoryReclaimWaitMs; - - public Map getAllProperties() - { - ImmutableMap.Builder builder = ImmutableMap.builder(); - return builder.put(CONCURRENT_LIFESPANS_PER_TASK, String.valueOf(getConcurrentLifespansPerTask())) - .put(ENABLE_SERIALIZED_PAGE_CHECKSUM, String.valueOf(isEnableSerializedPageChecksum())) - .put(ENABLE_VELOX_EXPRESSION_LOGGING, String.valueOf(isEnableVeloxExpressionLogging())) - .put(ENABLE_VELOX_TASK_LOGGING, String.valueOf(isEnableVeloxTaskLogging())) - .put(HTTP_SERVER_HTTP_PORT, String.valueOf(getHttpServerPort())) - .put(HTTP_SERVER_REUSE_PORT, String.valueOf(isHttpServerReusePort())) - .put(HTTP_SERVER_BIND_TO_NODE_INTERNAL_ADDRESS_ONLY_ENABLED, String.valueOf(isHttpServerBindToNodeInternalAddressOnlyEnabled())) - .put(HTTP_SERVER_HTTPS_PORT, String.valueOf(getHttpsServerPort())) - .put(HTTP_SERVER_HTTPS_ENABLED, String.valueOf(isEnableHttpsCommunication())) - .put(HTTPS_CIPHERS, String.valueOf(getHttpsCiphers())) - .put(HTTPS_CERT_PATH, String.valueOf(getHttpsCertPath())) - .put(HTTPS_KEY_PATH, String.valueOf(getHttpsKeyPath())) - .put(HTTP_SERVER_NUM_IO_THREADS_HW_MULTIPLIER, String.valueOf(getHttpServerNumIoThreadsHwMultiplier())) - .put(EXCHANGE_HTTP_CLIENT_NUM_IO_THREADS_HW_MULTIPLIER, String.valueOf(getExchangeHttpClientNumIoThreadsHwMultiplier())) - .put(ASYNC_DATA_CACHE_ENABLED, String.valueOf(getAsyncDataCacheEnabled())) - .put(ASYNC_CACHE_SSD_GB, String.valueOf(getAsyncCacheSsdGb())) - .put(CONNECTOR_NUM_IO_THREADS_HW_MULTIPLIER, String.valueOf(getConnectorNumIoThreadsHwMultiplier())) - .put(PRESTO_VERSION, getPrestoVersion()) - .put(SHUTDOWN_ONSET_SEC, String.valueOf(getShutdownOnsetSec())) - .put(SYSTEM_MEMORY_GB, String.valueOf(getSystemMemoryGb())) - .put(QUERY_MEMORY_GB, String.valueOf(getQueryMemoryGb())) - .put(USE_MMAP_ALLOCATOR, String.valueOf(getUseMmapAllocator())) - .put(MEMORY_ARBITRATOR_KIND, String.valueOf(getMemoryArbitratorKind())) - .put(MEMORY_ARBITRATOR_CAPACITY_GB, String.valueOf(getMemoryArbitratorCapacityGb())) - .put(SHARED_ARBITRATOR_RESERVED_CAPACITY, String.valueOf(getSharedArbitratorReservedCapacity())) - .put(SHARED_ARBITRATOR_MEMORY_POOL_INITIAL_CAPACITY, String.valueOf(getSharedArbitratorMemoryPoolInitialCapacity())) - .put(SHARED_ARBITRATOR_MEMORY_POOL_RESERVED_CAPACITY, String.valueOf(getSharedArbitratorMemoryPoolReservedCapacity())) - .put(SHARED_ARBITRATOR_MAX_MEMORY_ARBITRATION_TIME, String.valueOf(getSharedArbitratorMaxMemoryArbitrationTime())) - .put(SPILLER_SPILL_PATH, String.valueOf(getSpillerSpillPath())) - .put(TASK_MAX_DRIVERS_PER_TASK, String.valueOf(getMaxDriversPerTask())) - .put(ENABLE_OLD_TASK_CLEANUP, String.valueOf(getOldTaskCleanupMs())) - .put(SHUFFLE_NAME, getShuffleName()) - .put(HTTP_SERVER_ACCESS_LOGS, String.valueOf(isEnableHttpServerAccessLog())) - .put(CORE_ON_ALLOCATION_FAILURE_ENABLED, String.valueOf(isCoreOnAllocationFailureEnabled())) - .put(SPILL_ENABLED, String.valueOf(getSpillEnabled())) - .put(AGGREGATION_SPILL_ENABLED, String.valueOf(getAggregationSpillEnabled())) - .put(JOIN_SPILL_ENABLED, String.valueOf(getJoinSpillEnabled())) - .put(ORDER_BY_SPILL_ENABLED, String.valueOf(getOrderBySpillEnabled())) - .put(MAX_SPILL_BYTES, String.valueOf(getMaxSpillBytes())) - .put(REGISTER_TEST_FUNCTIONS, String.valueOf(registerTestFunctions)) - .put(MEMORY_POOL_INIT_CAPACITY, String.valueOf(memoryPoolInitCapacity)) - .put(MEMORY_POOL_RESERVED_CAPACITY, String.valueOf(memoryPoolReservedCapacity)) - .put(MEMORY_POOL_TRANSFER_CAPACITY, String.valueOf(memoryPoolTransferCapacity)) - .put(MEMORY_RECLAIM_WAIT_MS, String.valueOf(memoryReclaimWaitMs)) + public static final String NATIVE_EXECUTION_SYSTEM_CONFIG = "native-execution-system-config"; + + public static final String CONCURRENT_LIFESPANS_PER_TASK = "concurrent-lifespans-per-task"; + public static final String ENABLE_SERIALIZED_PAGE_CHECKSUM = "enable-serialized-page-checksum"; + public static final String ENABLE_VELOX_EXPRESSION_LOGGING = "enable_velox_expression_logging"; + public static final String ENABLE_VELOX_TASK_LOGGING = "enable_velox_task_logging"; + public static final String HTTP_SERVER_HTTP_PORT = "http-server.http.port"; + public static final String HTTP_SERVER_REUSE_PORT = "http-server.reuse-port"; + public static final String HTTP_SERVER_BIND_TO_NODE_INTERNAL_ADDRESS_ONLY_ENABLED = "http-server.bind-to-node-internal-address-only-enabled"; + public static final String HTTP_SERVER_HTTPS_PORT = "http-server.https.port"; + public static final String HTTP_SERVER_HTTPS_ENABLED = "http-server.https.enabled"; + public static final String HTTPS_SUPPORTED_CIPHERS = "https-supported-ciphers"; + public static final String HTTPS_CERT_PATH = "https-cert-path"; + public static final String HTTPS_KEY_PATH = "https-key-path"; + public static final String HTTP_SERVER_NUM_IO_THREADS_HW_MULTIPLIER = "http-server.num-io-threads-hw-multiplier"; + public static final String EXCHANGE_HTTP_CLIENT_NUM_IO_THREADS_HW_MULTIPLIER = "exchange.http-client.num-io-threads-hw-multiplier"; + public static final String ASYNC_DATA_CACHE_ENABLED = "async-data-cache-enabled"; + public static final String ASYNC_CACHE_SSD_GB = "async-cache-ssd-gb"; + public static final String CONNECTOR_NUM_IO_THREADS_HW_MULTIPLIER = "connector.num-io-threads-hw-multiplier"; + public static final String PRESTO_VERSION = "presto.version"; + public static final String SHUTDOWN_ONSET_SEC = "shutdown-onset-sec"; + public static final String SYSTEM_MEMORY_GB = "system-memory-gb"; + public static final String QUERY_MEMORY_GB = "query-memory-gb"; + public static final String QUERY_MAX_MEMORY_PER_NODE = "query.max-memory-per-node"; + public static final String USE_MMAP_ALLOCATOR = "use-mmap-allocator"; + public static final String MEMORY_ARBITRATOR_KIND = "memory-arbitrator-kind"; + public static final String SHARED_ARBITRATOR_RESERVED_CAPACITY = "shared-arbitrator.reserved-capacity"; + public static final String SHARED_ARBITRATOR_MEMORY_POOL_INITIAL_CAPACITY = "shared-arbitrator.memory-pool-initial-capacity"; + public static final String SHARED_ARBITRATOR_MAX_MEMORY_ARBITRATION_TIME = "shared-arbitrator.max-memory-arbitration-time"; + public static final String EXPERIMENTAL_SPILLER_SPILL_PATH = "experimental.spiller-spill-path"; + public static final String TASK_MAX_DRIVERS_PER_TASK = "task.max-drivers-per-task"; + public static final String ENABLE_OLD_TASK_CLEANUP = "enable-old-task-cleanup"; + public static final String SHUFFLE_NAME = "shuffle.name"; + public static final String HTTP_SERVER_ENABLE_ACCESS_LOG = "http-server.enable-access-log"; + public static final String CORE_ON_ALLOCATION_FAILURE_ENABLED = "core-on-allocation-failure-enabled"; + public static final String SPILL_ENABLED = "spill-enabled"; + public static final String AGGREGATION_SPILL_ENABLED = "aggregation-spill-enabled"; + public static final String JOIN_SPILL_ENABLED = "join-spill-enabled"; + public static final String ORDER_BY_SPILL_ENABLED = "order-by-spill-enabled"; + public static final String MAX_SPILL_BYTES = "max-spill-bytes"; + public static final String REMOTE_FUNCTION_SERVER_THRIFT_UDS_PATH = "remote-function-server.thrift.uds-path"; + public static final String REMOTE_FUNCTION_SERVER_SIGNATURE_FILES_DIRECTORY_PATH = "remote-function-server.signature.files.directory.path"; + public static final String REMOTE_FUNCTION_SERVER_SERDE = "remote-function-server.serde"; + public static final String REMOTE_FUNCTION_SERVER_CATALOG_NAME = "remote-function-server.catalog-name"; + + private final String remoteFunctionServerSignatureFilesDirectoryPathDefault = "./functions/spark/"; + private final String remoteFunctionServerSerdeDefault = "presto_page"; + private final String remoteFunctionServerCatalogNameDefault = ""; + private final String concurrentLifespansPerTaskDefault = "5"; + private final String enableSerializedPageChecksumDefault = "true"; + private final String enableVeloxExpressionLoggingDefault = "false"; + private final String enableVeloxTaskLoggingDefault = "true"; + private final String httpServerHttpPortDefault = "7777"; + private final String httpServerReusePortDefault = "true"; + private final String httpServerBindToNodeInternalAddressOnlyEnabledDefault = "true"; + private final String httpServerHttpsPortDefault = "7778"; + private final String httpServerHttpsEnabledDefault = "false"; + private final String httpsSupportedCiphersDefault = "AES128-SHA,AES128-SHA256,AES256-GCM-SHA384"; + private final String httpsCertPathDefault = ""; + private final String httpsKeyPathDefault = ""; + private final String httpServerNumIoThreadsHwMultiplierDefault = "1.0"; + private final String exchangeHttpClientNumIoThreadsHwMultiplierDefault = "1.0"; + private final String asyncDataCacheEnabledDefault = "false"; + private final String asyncCacheSsdGbDefault = "0"; + private final String connectorNumIoThreadsHwMultiplierDefault = "0"; + private final String prestoVersionDefault = "dummy.presto.version"; + private final String shutdownOnsetSecDefault = "10"; + private final String systemMemoryGbDefault = "10"; + private final String queryMemoryGbDefault = "8"; + private final String queryMaxMemoryPerNodeDefault = "8GB"; + private final String useMmapAllocatorDefault = "true"; + private final String memoryArbitratorKindDefault = "SHARED"; + private final String sharedArbitratorReservedCapacityDefault = "0GB"; + private final String sharedArbitratorMemoryPoolInitialCapacityDefault = "4GB"; + private final String sharedArbitratorMaxMemoryArbitrationTimeDefault = "5m"; + private final String experimentalSpillerSpillPathDefault = ""; + private final String taskMaxDriversPerTaskDefault = "15"; + private final String enableOldTaskCleanupDefault = "false"; + private final String shuffleNameDefault = "local"; + private final String httpServerEnableAccessLogDefault = "true"; + private final String coreOnAllocationFailureEnabledDefault = "false"; + private final String spillEnabledDefault = "true"; + private final String aggregationSpillEnabledDefault = "true"; + private final String joinSpillEnabledDefault = "true"; + private final String orderBySpillEnabledDefault = "true"; + private final String maxSpillBytesDefault = String.valueOf(600L << 30); + + private final Map systemConfigs; + private final Map defaultSystemConfigs; + + @Inject + public NativeExecutionSystemConfig( + @Named(NATIVE_EXECUTION_SYSTEM_CONFIG) Map systemConfigs) + { + this.systemConfigs = new HashMap<>( + requireNonNull(systemConfigs, "systemConfigs is null")); + + ImmutableMap.Builder defaultSystemConfigsBuilder = ImmutableMap.builder(); + + String remoteFunctionServerThriftUdsPath = System.getProperty(REMOTE_FUNCTION_SERVER_THRIFT_UDS_PATH); + if (remoteFunctionServerThriftUdsPath != null) { + defaultSystemConfigsBuilder.put(REMOTE_FUNCTION_SERVER_THRIFT_UDS_PATH, remoteFunctionServerThriftUdsPath); + defaultSystemConfigsBuilder.put(REMOTE_FUNCTION_SERVER_SIGNATURE_FILES_DIRECTORY_PATH, getAbsolutePath(remoteFunctionServerSignatureFilesDirectoryPathDefault)); + defaultSystemConfigsBuilder.put(REMOTE_FUNCTION_SERVER_SERDE, remoteFunctionServerSerdeDefault); + defaultSystemConfigsBuilder.put(REMOTE_FUNCTION_SERVER_CATALOG_NAME, remoteFunctionServerCatalogNameDefault); + } + + defaultSystemConfigs = defaultSystemConfigsBuilder + .put(CONCURRENT_LIFESPANS_PER_TASK, concurrentLifespansPerTaskDefault) + .put(ENABLE_SERIALIZED_PAGE_CHECKSUM, enableSerializedPageChecksumDefault) + .put(ENABLE_VELOX_EXPRESSION_LOGGING, enableVeloxExpressionLoggingDefault) + .put(ENABLE_VELOX_TASK_LOGGING, enableVeloxTaskLoggingDefault) + .put(HTTP_SERVER_HTTP_PORT, httpServerHttpPortDefault) + .put(HTTP_SERVER_REUSE_PORT, httpServerReusePortDefault) + .put(HTTP_SERVER_BIND_TO_NODE_INTERNAL_ADDRESS_ONLY_ENABLED, + httpServerBindToNodeInternalAddressOnlyEnabledDefault) + .put(HTTP_SERVER_HTTPS_PORT, httpServerHttpsPortDefault) + .put(HTTP_SERVER_HTTPS_ENABLED, httpServerHttpsEnabledDefault) + .put(HTTPS_SUPPORTED_CIPHERS, httpsSupportedCiphersDefault) + .put(HTTPS_CERT_PATH, httpsCertPathDefault) + .put(HTTPS_KEY_PATH, httpsKeyPathDefault) + .put(HTTP_SERVER_NUM_IO_THREADS_HW_MULTIPLIER, + httpServerNumIoThreadsHwMultiplierDefault) + .put(EXCHANGE_HTTP_CLIENT_NUM_IO_THREADS_HW_MULTIPLIER, + exchangeHttpClientNumIoThreadsHwMultiplierDefault) + .put(ASYNC_DATA_CACHE_ENABLED, asyncDataCacheEnabledDefault) + .put(ASYNC_CACHE_SSD_GB, asyncCacheSsdGbDefault) + .put(CONNECTOR_NUM_IO_THREADS_HW_MULTIPLIER, connectorNumIoThreadsHwMultiplierDefault) + .put(PRESTO_VERSION, prestoVersionDefault) + .put(SHUTDOWN_ONSET_SEC, shutdownOnsetSecDefault) + .put(SYSTEM_MEMORY_GB, systemMemoryGbDefault) + .put(QUERY_MEMORY_GB, queryMemoryGbDefault) + .put(QUERY_MAX_MEMORY_PER_NODE, queryMaxMemoryPerNodeDefault) + .put(USE_MMAP_ALLOCATOR, useMmapAllocatorDefault) + .put(MEMORY_ARBITRATOR_KIND, memoryArbitratorKindDefault) + .put(SHARED_ARBITRATOR_RESERVED_CAPACITY, sharedArbitratorReservedCapacityDefault) + .put(SHARED_ARBITRATOR_MEMORY_POOL_INITIAL_CAPACITY, + sharedArbitratorMemoryPoolInitialCapacityDefault) + .put(SHARED_ARBITRATOR_MAX_MEMORY_ARBITRATION_TIME, + sharedArbitratorMaxMemoryArbitrationTimeDefault) + .put(EXPERIMENTAL_SPILLER_SPILL_PATH, experimentalSpillerSpillPathDefault) + .put(TASK_MAX_DRIVERS_PER_TASK, taskMaxDriversPerTaskDefault) + .put(ENABLE_OLD_TASK_CLEANUP, enableOldTaskCleanupDefault) + .put(SHUFFLE_NAME, shuffleNameDefault) + .put(HTTP_SERVER_ENABLE_ACCESS_LOG, httpServerEnableAccessLogDefault) + .put(CORE_ON_ALLOCATION_FAILURE_ENABLED, coreOnAllocationFailureEnabledDefault) + .put(SPILL_ENABLED, spillEnabledDefault) + .put(AGGREGATION_SPILL_ENABLED, aggregationSpillEnabledDefault) + .put(JOIN_SPILL_ENABLED, joinSpillEnabledDefault) + .put(ORDER_BY_SPILL_ENABLED, orderBySpillEnabledDefault) + .put(MAX_SPILL_BYTES, maxSpillBytesDefault) .build(); } - @Config(SHUFFLE_NAME) - public NativeExecutionSystemConfig setShuffleName(String shuffleName) - { - this.shuffleName = requireNonNull(shuffleName); - return this; - } - - public String getShuffleName() - { - return shuffleName; - } - - @Config(ENABLE_SERIALIZED_PAGE_CHECKSUM) - public NativeExecutionSystemConfig setEnableSerializedPageChecksum(boolean enableSerializedPageChecksum) - { - this.enableSerializedPageChecksum = enableSerializedPageChecksum; - return this; - } - - public boolean isEnableSerializedPageChecksum() - { - return enableSerializedPageChecksum; - } - - @Config(ENABLE_VELOX_EXPRESSION_LOGGING) - public NativeExecutionSystemConfig setEnableVeloxExpressionLogging(boolean enableVeloxExpressionLogging) - { - this.enableVeloxExpressionLogging = enableVeloxExpressionLogging; - return this; - } - - public boolean isEnableVeloxExpressionLogging() - { - return enableVeloxExpressionLogging; - } - - @Config(ENABLE_VELOX_TASK_LOGGING) - public NativeExecutionSystemConfig setEnableVeloxTaskLogging(boolean enableVeloxTaskLogging) - { - this.enableVeloxTaskLogging = enableVeloxTaskLogging; - return this; - } - - public boolean isEnableVeloxTaskLogging() - { - return enableVeloxTaskLogging; - } - - @Config(HTTP_SERVER_HTTP_PORT) - public NativeExecutionSystemConfig setHttpServerPort(int httpServerPort) - { - this.httpServerPort = httpServerPort; - return this; - } - - public int getHttpServerPort() - { - return httpServerPort; - } - - @Config(HTTP_SERVER_REUSE_PORT) - public NativeExecutionSystemConfig setHttpServerReusePort(boolean httpServerReusePort) - { - this.httpServerReusePort = httpServerReusePort; - return this; - } - - public boolean isHttpServerReusePort() - { - return httpServerReusePort; - } - - public boolean isHttpServerBindToNodeInternalAddressOnlyEnabled() - { - return httpServerBindToNodeInternalAddressOnlyEnabled; - } - - @Config(HTTP_SERVER_BIND_TO_NODE_INTERNAL_ADDRESS_ONLY_ENABLED) - public NativeExecutionSystemConfig setHttpServerBindToNodeInternalAddressOnlyEnabled(boolean httpServerBindToNodeInternalAddressOnlyEnabled) - { - this.httpServerBindToNodeInternalAddressOnlyEnabled = httpServerBindToNodeInternalAddressOnlyEnabled; - return this; - } - - @Config(HTTP_SERVER_NUM_IO_THREADS_HW_MULTIPLIER) - public NativeExecutionSystemConfig setHttpServerNumIoThreadsHwMultiplier(double httpServerNumIoThreadsHwMultiplier) - { - this.httpServerNumIoThreadsHwMultiplier = httpServerNumIoThreadsHwMultiplier; - return this; - } - - public double getHttpServerNumIoThreadsHwMultiplier() - { - return httpServerNumIoThreadsHwMultiplier; - } - - public int getHttpsServerPort() - { - return httpsServerPort; - } - - @Config(HTTP_SERVER_HTTPS_PORT) - public NativeExecutionSystemConfig setHttpsServerPort(int httpsServerPort) - { - this.httpsServerPort = httpsServerPort; - return this; - } - - public boolean isEnableHttpsCommunication() - { - return enableHttpsCommunication; - } - - @Config(HTTP_SERVER_HTTPS_ENABLED) - public NativeExecutionSystemConfig setEnableHttpsCommunication(boolean enableHttpsCommunication) - { - this.enableHttpsCommunication = enableHttpsCommunication; - return this; - } - - public String getHttpsCiphers() - { - return httpsCiphers; - } - - @Config(HTTPS_CIPHERS) - public NativeExecutionSystemConfig setHttpsCiphers(String httpsCiphers) - { - this.httpsCiphers = httpsCiphers; - return this; - } - - public String getHttpsCertPath() - { - return httpsCertPath; - } - - @Config(HTTPS_CERT_PATH) - public NativeExecutionSystemConfig setHttpsCertPath(String httpsCertPath) - { - this.httpsCertPath = httpsCertPath; - return this; - } - - public String getHttpsKeyPath() - { - return httpsKeyPath; - } - - @Config(HTTPS_KEY_PATH) - public NativeExecutionSystemConfig setHttpsKeyPath(String httpsKeyPath) - { - this.httpsKeyPath = httpsKeyPath; - return this; - } - - @Config(EXCHANGE_HTTP_CLIENT_NUM_IO_THREADS_HW_MULTIPLIER) - public NativeExecutionSystemConfig setExchangeHttpClientNumIoThreadsHwMultiplier(double exchangeHttpClientNumIoThreadsHwMultiplier) - { - this.exchangeHttpClientNumIoThreadsHwMultiplier = exchangeHttpClientNumIoThreadsHwMultiplier; - return this; - } - - public double getExchangeHttpClientNumIoThreadsHwMultiplier() - { - return exchangeHttpClientNumIoThreadsHwMultiplier; - } - - @Config(ASYNC_DATA_CACHE_ENABLED) - public NativeExecutionSystemConfig setAsyncDataCacheEnabled(boolean asyncDataCacheEnabled) - { - this.asyncDataCacheEnabled = asyncDataCacheEnabled; - return this; - } - - public boolean getAsyncDataCacheEnabled() - { - return asyncDataCacheEnabled; - } - - @Config(ASYNC_CACHE_SSD_GB) - public NativeExecutionSystemConfig setAsyncCacheSsdGb(int asyncCacheSsdGb) - { - this.asyncCacheSsdGb = asyncCacheSsdGb; - return this; - } - - public int getAsyncCacheSsdGb() - { - return asyncCacheSsdGb; - } - - @Config(CONNECTOR_NUM_IO_THREADS_HW_MULTIPLIER) - public NativeExecutionSystemConfig setConnectorNumIoThreadsHwMultiplier(double connectorNumIoThreadsHwMultiplier) - { - this.connectorNumIoThreadsHwMultiplier = connectorNumIoThreadsHwMultiplier; - return this; - } - - public double getConnectorNumIoThreadsHwMultiplier() - { - return connectorNumIoThreadsHwMultiplier; - } - - @Config(SHUTDOWN_ONSET_SEC) - public NativeExecutionSystemConfig setShutdownOnsetSec(int shutdownOnsetSec) - { - this.shutdownOnsetSec = shutdownOnsetSec; - return this; - } - - public int getShutdownOnsetSec() - { - return shutdownOnsetSec; - } - - @Config(SYSTEM_MEMORY_GB) - public NativeExecutionSystemConfig setSystemMemoryGb(int systemMemoryGb) - { - this.systemMemoryGb = systemMemoryGb; - return this; - } - - public int getSystemMemoryGb() - { - return systemMemoryGb; - } - - @Config(QUERY_MEMORY_GB) - public NativeExecutionSystemConfig setQueryMemoryGb(DataSize queryMemoryGb) - { - this.queryMemoryGb = queryMemoryGb; - return this; - } - - public DataSize getQueryMemoryGb() - { - return queryMemoryGb; - } - - @Config(USE_MMAP_ALLOCATOR) - public NativeExecutionSystemConfig setUseMmapAllocator(boolean useMmapAllocator) - { - this.useMmapAllocator = useMmapAllocator; - return this; - } - - public boolean getUseMmapAllocator() - { - return useMmapAllocator; - } - - @Config(MEMORY_ARBITRATOR_KIND) - public NativeExecutionSystemConfig setMemoryArbitratorKind(String memoryArbitratorKind) - { - this.memoryArbitratorKind = memoryArbitratorKind; - return this; - } - - public String getMemoryArbitratorKind() - { - return memoryArbitratorKind; - } - - @Config(MEMORY_ARBITRATOR_CAPACITY_GB) - public NativeExecutionSystemConfig setMemoryArbitratorCapacityGb(int memoryArbitratorCapacityGb) - { - this.memoryArbitratorCapacityGb = memoryArbitratorCapacityGb; - return this; - } - - public int getMemoryArbitratorCapacityGb() - { - return memoryArbitratorCapacityGb; - } - - @Config(SHARED_ARBITRATOR_RESERVED_CAPACITY) - public NativeExecutionSystemConfig setSharedArbitratorReservedCapacity(DataSize sharedArbitratorReservedCapacity) - { - this.sharedArbitratorReservedCapacity = sharedArbitratorReservedCapacity; - return this; - } - - public DataSize getSharedArbitratorReservedCapacity() - { - return sharedArbitratorReservedCapacity; - } - - @Config(SHARED_ARBITRATOR_MEMORY_POOL_INITIAL_CAPACITY) - public NativeExecutionSystemConfig setSharedArbitratorMemoryPoolInitialCapacity(DataSize sharedArbitratorMemoryPoolInitialCapacity) - { - this.sharedArbitratorMemoryPoolInitialCapacity = sharedArbitratorMemoryPoolInitialCapacity; - return this; - } - - public DataSize getSharedArbitratorMemoryPoolInitialCapacity() - { - return sharedArbitratorMemoryPoolInitialCapacity; - } - - @Config(SHARED_ARBITRATOR_MEMORY_POOL_RESERVED_CAPACITY) - public NativeExecutionSystemConfig setSharedArbitratorMemoryPoolReservedCapacity(DataSize sharedArbitratorMemoryPoolReservedCapacity) - { - this.sharedArbitratorMemoryPoolReservedCapacity = sharedArbitratorMemoryPoolReservedCapacity; - return this; - } - - public DataSize getSharedArbitratorMemoryPoolReservedCapacity() - { - return sharedArbitratorMemoryPoolReservedCapacity; - } - - @Config(SHARED_ARBITRATOR_MAX_MEMORY_ARBITRATION_TIME) - public NativeExecutionSystemConfig setSharedArbitratorMaxMemoryArbitrationTime(Duration sharedArbitratorMaxMemoryArbitrationTime) - { - this.sharedArbitratorMaxMemoryArbitrationTime = sharedArbitratorMaxMemoryArbitrationTime; - return this; - } - - public Duration getSharedArbitratorMaxMemoryArbitrationTime() - { - return sharedArbitratorMaxMemoryArbitrationTime; - } - - @Config(SPILLER_SPILL_PATH) - public NativeExecutionSystemConfig setSpillerSpillPath(String spillerSpillPath) - { - this.spillerSpillPath = spillerSpillPath; - return this; - } - - public String getSpillerSpillPath() - { - return spillerSpillPath; - } - - @Config(CONCURRENT_LIFESPANS_PER_TASK) - public NativeExecutionSystemConfig setConcurrentLifespansPerTask(int concurrentLifespansPerTask) - { - this.concurrentLifespansPerTask = concurrentLifespansPerTask; - return this; - } - - public int getConcurrentLifespansPerTask() - { - return concurrentLifespansPerTask; - } - - @Config(TASK_MAX_DRIVERS_PER_TASK) - public NativeExecutionSystemConfig setMaxDriversPerTask(int maxDriversPerTask) - { - this.maxDriversPerTask = maxDriversPerTask; - return this; - } - - public int getMaxDriversPerTask() - { - return maxDriversPerTask; - } - - public boolean getOldTaskCleanupMs() - { - return enableOldTaskCleanUp; - } - - @Config(ENABLE_OLD_TASK_CLEANUP) - public NativeExecutionSystemConfig setOldTaskCleanupMs(boolean enableOldTaskCleanUp) - { - this.enableOldTaskCleanUp = enableOldTaskCleanUp; - return this; - } - - @Config(PRESTO_VERSION) - public NativeExecutionSystemConfig setPrestoVersion(String prestoVersion) - { - this.prestoVersion = prestoVersion; - return this; - } - - public String getPrestoVersion() - { - return prestoVersion; - } - - @Config(HTTP_SERVER_ACCESS_LOGS) - public NativeExecutionSystemConfig setEnableHttpServerAccessLog(boolean enableHttpServerAccessLog) - { - this.enableHttpServerAccessLog = enableHttpServerAccessLog; - return this; - } - - public boolean isEnableHttpServerAccessLog() - { - return enableHttpServerAccessLog; - } - - public boolean isCoreOnAllocationFailureEnabled() - { - return coreOnAllocationFailureEnabled; - } - - @Config(CORE_ON_ALLOCATION_FAILURE_ENABLED) - public NativeExecutionSystemConfig setCoreOnAllocationFailureEnabled(boolean coreOnAllocationFailureEnabled) - { - this.coreOnAllocationFailureEnabled = coreOnAllocationFailureEnabled; - return this; - } - - public boolean getSpillEnabled() - { - return spillEnabled; - } - - @Config(SPILL_ENABLED) - public NativeExecutionSystemConfig setSpillEnabled(boolean spillEnabled) - { - this.spillEnabled = spillEnabled; - return this; - } - - public boolean getAggregationSpillEnabled() - { - return aggregationSpillEnabled; - } - - @Config(AGGREGATION_SPILL_ENABLED) - public NativeExecutionSystemConfig setAggregationSpillEnabled(boolean aggregationSpillEnabled) - { - this.aggregationSpillEnabled = aggregationSpillEnabled; - return this; - } - - public boolean getJoinSpillEnabled() - { - return joinSpillEnabled; - } - - @Config(JOIN_SPILL_ENABLED) - public NativeExecutionSystemConfig setJoinSpillEnabled(boolean joinSpillEnabled) - { - this.joinSpillEnabled = joinSpillEnabled; - return this; - } - - public boolean getOrderBySpillEnabled() - { - return orderBySpillEnabled; - } - - @Config(ORDER_BY_SPILL_ENABLED) - public NativeExecutionSystemConfig setOrderBySpillEnabled(boolean orderBySpillEnabled) - { - this.orderBySpillEnabled = orderBySpillEnabled; - return this; - } - - public Long getMaxSpillBytes() - { - return maxSpillBytes; - } - - @Config(MAX_SPILL_BYTES) - public NativeExecutionSystemConfig setMaxSpillBytes(Long maxSpillBytes) - { - this.maxSpillBytes = maxSpillBytes; - return this; - } - - // TODO: All following configs are deprecated and will be removed after references from all - // other systems are cleared. - - @Config(REGISTER_TEST_FUNCTIONS) - public NativeExecutionSystemConfig setRegisterTestFunctions(boolean registerTestFunctions) - { - this.registerTestFunctions = registerTestFunctions; - return this; - } - - public boolean isRegisterTestFunctions() - { - return registerTestFunctions; - } - - @Config(MEMORY_POOL_INIT_CAPACITY) - public NativeExecutionSystemConfig setMemoryPoolInitCapacity(long memoryPoolInitCapacity) - { - this.memoryPoolInitCapacity = memoryPoolInitCapacity; - return this; - } - - public long getMemoryPoolInitCapacity() - { - return memoryPoolInitCapacity; - } - - @Config(MEMORY_POOL_RESERVED_CAPACITY) - public NativeExecutionSystemConfig setMemoryPoolReservedCapacity(long memoryPoolReservedCapacity) - { - this.memoryPoolReservedCapacity = memoryPoolReservedCapacity; - return this; - } - - public long getMemoryPoolReservedCapacity() + public Map getAllProperties() { - return memoryPoolReservedCapacity; + ImmutableMap.Builder builder = ImmutableMap.builder(); + builder.putAll(systemConfigs); + defaultSystemConfigs.entrySet().stream() + .filter(entry -> !systemConfigs.containsKey(entry.getKey())) + .forEach(entry -> builder.put(entry.getKey(), entry.getValue())); + return builder.build(); } - @Config(MEMORY_POOL_TRANSFER_CAPACITY) - public NativeExecutionSystemConfig setMemoryPoolTransferCapacity(long memoryPoolTransferCapacity) + public NativeExecutionSystemConfig update(String key, String value) { - this.memoryPoolTransferCapacity = memoryPoolTransferCapacity; + systemConfigs.put(key, value); return this; } - public long getMemoryPoolTransferCapacity() + private String getAbsolutePath(String path) { - return memoryPoolTransferCapacity; - } + File absolutePath = new File(path); + if (!absolutePath.isAbsolute()) { + // In the case of SparkEnv is not initialed (e.g. unit test), we just use current location instead of calling SparkFiles.getRootDirectory() to avoid error. + String rootDirectory = SparkEnv$.MODULE$.get() != null ? SparkFiles.getRootDirectory() : "."; + absolutePath = new File(rootDirectory, path); + } - @Config(MEMORY_RECLAIM_WAIT_MS) - public NativeExecutionSystemConfig setMemoryReclaimWaitMs(long memoryReclaimWaitMs) - { - this.memoryReclaimWaitMs = memoryReclaimWaitMs; - return this; - } + if (!absolutePath.exists()) { + throw new IllegalArgumentException(format("File doesn't exist %s", absolutePath.getAbsolutePath())); + } - public long getMemoryReclaimWaitMs() - { - return memoryReclaimWaitMs; + return absolutePath.getAbsolutePath(); } } diff --git a/presto-spark-base/src/test/java/com/facebook/presto/spark/execution/TestNativeExecutionProcess.java b/presto-spark-base/src/test/java/com/facebook/presto/spark/execution/TestNativeExecutionProcess.java index 13897a3929a9f..1483270d325e2 100644 --- a/presto-spark-base/src/test/java/com/facebook/presto/spark/execution/TestNativeExecutionProcess.java +++ b/presto-spark-base/src/test/java/com/facebook/presto/spark/execution/TestNativeExecutionProcess.java @@ -27,6 +27,7 @@ import com.facebook.presto.spark.execution.property.NativeExecutionSystemConfig; import com.facebook.presto.spark.execution.property.PrestoSparkWorkerProperty; import com.facebook.presto.sql.analyzer.FeaturesConfig; +import com.google.common.collect.ImmutableMap; import org.testng.annotations.Test; import java.util.concurrent.ScheduledExecutorService; @@ -89,7 +90,7 @@ private NativeExecutionProcessFactory createNativeExecutionProcessFactory() PrestoSparkWorkerProperty workerProperty = new PrestoSparkWorkerProperty( new NativeExecutionConnectorConfig(), new NativeExecutionNodeConfig(), - new NativeExecutionSystemConfig()); + new NativeExecutionSystemConfig(ImmutableMap.of())); NativeExecutionProcessFactory factory = new NativeExecutionProcessFactory( new TestPrestoSparkHttpClient.TestingOkHttpClient( errorScheduler, diff --git a/presto-spark-base/src/test/java/com/facebook/presto/spark/execution/http/TestPrestoSparkHttpClient.java b/presto-spark-base/src/test/java/com/facebook/presto/spark/execution/http/TestPrestoSparkHttpClient.java index 2c0c3aa9ddb3e..6927065bb0842 100644 --- a/presto-spark-base/src/test/java/com/facebook/presto/spark/execution/http/TestPrestoSparkHttpClient.java +++ b/presto-spark-base/src/test/java/com/facebook/presto/spark/execution/http/TestPrestoSparkHttpClient.java @@ -899,7 +899,7 @@ private NativeExecutionProcess createNativeExecutionProcess( PrestoSparkWorkerProperty workerProperty = new PrestoSparkWorkerProperty( new NativeExecutionConnectorConfig(), new NativeExecutionNodeConfig(), - new NativeExecutionSystemConfig()); + new NativeExecutionSystemConfig(ImmutableMap.of())); NativeExecutionProcessFactory factory = new NativeExecutionProcessFactory( new TestingOkHttpClient(scheduledExecutorService, responseManager), scheduledExecutorService, diff --git a/presto-spark-base/src/test/java/com/facebook/presto/spark/execution/property/TestNativeExecutionSystemConfig.java b/presto-spark-base/src/test/java/com/facebook/presto/spark/execution/property/TestNativeExecutionSystemConfig.java index 1298fcc48d831..53200d798baba 100644 --- a/presto-spark-base/src/test/java/com/facebook/presto/spark/execution/property/TestNativeExecutionSystemConfig.java +++ b/presto-spark-base/src/test/java/com/facebook/presto/spark/execution/property/TestNativeExecutionSystemConfig.java @@ -15,8 +15,7 @@ import com.facebook.airlift.configuration.testing.ConfigAssertions; import com.facebook.airlift.units.DataSize; -import com.facebook.airlift.units.DataSize.Unit; -import com.facebook.airlift.units.Duration; +import com.google.common.collect.ImmutableMap; import org.testng.annotations.Test; import java.io.File; @@ -27,7 +26,6 @@ import java.nio.file.Paths; import java.util.Map; import java.util.Properties; -import java.util.concurrent.TimeUnit; import static com.facebook.airlift.configuration.testing.ConfigAssertions.assertFullMapping; import static com.facebook.airlift.configuration.testing.ConfigAssertions.assertRecordedDefaults; @@ -40,100 +38,145 @@ public class TestNativeExecutionSystemConfig public void testNativeExecutionSystemConfig() { // Test defaults - assertRecordedDefaults(ConfigAssertions.recordDefaults(NativeExecutionSystemConfig.class) - .setEnableSerializedPageChecksum(true) - .setEnableVeloxExpressionLogging(false) - .setEnableVeloxTaskLogging(true) - .setHttpServerReusePort(true) - .setHttpServerBindToNodeInternalAddressOnlyEnabled(true) - .setHttpServerPort(7777) - .setHttpServerNumIoThreadsHwMultiplier(1.0) - .setHttpsServerPort(7778) - .setEnableHttpsCommunication(false) - .setHttpsCiphers("AES128-SHA,AES128-SHA256,AES256-GCM-SHA384") - .setHttpsCertPath("") - .setHttpsKeyPath("") - .setExchangeHttpClientNumIoThreadsHwMultiplier(1.0) - .setAsyncDataCacheEnabled(false) - .setAsyncCacheSsdGb(0) - .setConnectorNumIoThreadsHwMultiplier(0.0) - .setShutdownOnsetSec(10) - .setSystemMemoryGb(10) - .setQueryMemoryGb(new DataSize(8, DataSize.Unit.GIGABYTE)) - .setUseMmapAllocator(true) - .setMemoryArbitratorKind("SHARED") - .setMemoryArbitratorCapacityGb(8) - .setSharedArbitratorReservedCapacity(new DataSize(0, Unit.GIGABYTE)) - .setSharedArbitratorMemoryPoolInitialCapacity(new DataSize(4, Unit.GIGABYTE)) - .setSharedArbitratorMemoryPoolReservedCapacity(new DataSize(32, Unit.MEGABYTE)) - .setSharedArbitratorMaxMemoryArbitrationTime(new Duration(5, TimeUnit.MINUTES)) - .setSpillerSpillPath("") - .setConcurrentLifespansPerTask(5) - .setMaxDriversPerTask(15) - .setOldTaskCleanupMs(false) - .setPrestoVersion("dummy.presto.version") - .setShuffleName("local") - .setEnableHttpServerAccessLog(true) - .setCoreOnAllocationFailureEnabled(false) - .setSpillEnabled(true) - .setAggregationSpillEnabled(true) - .setJoinSpillEnabled(true) - .setOrderBySpillEnabled(true) - .setMaxSpillBytes(600L << 30) - .setRegisterTestFunctions(false) - .setMemoryPoolInitCapacity(0) - .setMemoryPoolReservedCapacity(0) - .setMemoryPoolTransferCapacity(0) - .setMemoryReclaimWaitMs(0)); + NativeExecutionSystemConfig nativeExecutionSystemConfig = new NativeExecutionSystemConfig( + ImmutableMap.of()); + ImmutableMap.Builder builder = ImmutableMap.builder(); + ImmutableMap expectedConfigs = builder + .put("concurrent-lifespans-per-task", "5") + .put("enable-serialized-page-checksum", "true") + .put("enable_velox_expression_logging", "false") + .put("enable_velox_task_logging", "true") + .put("http-server.http.port", "7777") + .put("http-server.reuse-port", "true") + .put("http-server.bind-to-node-internal-address-only-enabled", "true") + .put("http-server.https.port", "7778") + .put("http-server.https.enabled", "false") + .put("https-supported-ciphers", "AES128-SHA,AES128-SHA256,AES256-GCM-SHA384") + .put("https-cert-path", "") + .put("https-key-path", "") + .put("http-server.num-io-threads-hw-multiplier", "1.0") + .put("exchange.http-client.num-io-threads-hw-multiplier", "1.0") + .put("async-data-cache-enabled", "false") + .put("async-cache-ssd-gb", "0") + .put("connector.num-io-threads-hw-multiplier", "0") + .put("presto.version", "dummy.presto.version") + .put("shutdown-onset-sec", "10") + .put("system-memory-gb", "10") + .put("query-memory-gb", "8") + .put("query.max-memory-per-node", "8GB") + .put("use-mmap-allocator", "true") + .put("memory-arbitrator-kind", "SHARED") + .put("shared-arbitrator.reserved-capacity", "0GB") + .put("shared-arbitrator.memory-pool-initial-capacity", "4GB") + .put("shared-arbitrator.max-memory-arbitration-time", "5m") + .put("experimental.spiller-spill-path", "") + .put("task.max-drivers-per-task", "15") + .put("enable-old-task-cleanup", "false") + .put("shuffle.name", "local") + .put("http-server.enable-access-log", "true") + .put("core-on-allocation-failure-enabled", "false") + .put("spill-enabled", "true") + .put("aggregation-spill-enabled", "true") + .put("join-spill-enabled", "true") + .put("order-by-spill-enabled", "true") + .put("max-spill-bytes", String.valueOf(600L << 30)) + .build(); + assertEquals(nativeExecutionSystemConfig.getAllProperties(), expectedConfigs); // Test explicit property mapping. Also makes sure properties returned by getAllProperties() covers full property list. - NativeExecutionSystemConfig expected = new NativeExecutionSystemConfig() - .setConcurrentLifespansPerTask(15) - .setEnableSerializedPageChecksum(false) - .setEnableVeloxExpressionLogging(true) - .setEnableVeloxTaskLogging(false) - .setHttpServerReusePort(false) - .setHttpServerBindToNodeInternalAddressOnlyEnabled(false) - .setHttpServerPort(8080) - .setHttpServerNumIoThreadsHwMultiplier(3.0) - .setHttpsServerPort(8081) - .setEnableHttpsCommunication(true) - .setHttpsCiphers("AES128-SHA") - .setHttpsCertPath("/tmp/non_existent.cert") - .setHttpsKeyPath("/tmp/non_existent.key") - .setExchangeHttpClientNumIoThreadsHwMultiplier(0.5) - .setAsyncDataCacheEnabled(true) - .setAsyncCacheSsdGb(1000) - .setConnectorNumIoThreadsHwMultiplier(1.0) - .setPrestoVersion("presto-version") - .setShutdownOnsetSec(30) - .setSystemMemoryGb(40) - .setQueryMemoryGb(new DataSize(20, Unit.GIGABYTE)) - .setUseMmapAllocator(false) - .setMemoryArbitratorKind("") - .setMemoryArbitratorCapacityGb(10) - .setSharedArbitratorReservedCapacity(new DataSize(8, Unit.GIGABYTE)) - .setSharedArbitratorMemoryPoolInitialCapacity(new DataSize(7, Unit.GIGABYTE)) - .setSharedArbitratorMemoryPoolReservedCapacity(new DataSize(6, Unit.GIGABYTE)) - .setSharedArbitratorMaxMemoryArbitrationTime(new Duration(123123123, TimeUnit.MILLISECONDS)) - .setSpillerSpillPath("dummy.spill.path") - .setMaxDriversPerTask(30) - .setOldTaskCleanupMs(true) - .setShuffleName("custom") - .setEnableHttpServerAccessLog(false) - .setCoreOnAllocationFailureEnabled(true) - .setSpillEnabled(false) - .setAggregationSpillEnabled(false) - .setJoinSpillEnabled(false) - .setOrderBySpillEnabled(false) - .setMaxSpillBytes(1L) - .setRegisterTestFunctions(true) - .setMemoryPoolInitCapacity(10) - .setMemoryPoolReservedCapacity(10) - .setMemoryPoolTransferCapacity(10) - .setMemoryReclaimWaitMs(10); - Map properties = expected.getAllProperties(); - assertFullMapping(properties, expected); + builder = ImmutableMap.builder(); + ImmutableMap systemConfig = builder + .put("concurrent-lifespans-per-task", "15") + .put("enable-serialized-page-checksum", "false") + .put("enable_velox_expression_logging", "true") + .put("enable_velox_task_logging", "false") + .put("http-server.http.port", "8777") + .put("http-server.reuse-port", "false") + .put("http-server.bind-to-node-internal-address-only-enabled", "false") + .put("http-server.https.port", "8778") + .put("http-server.https.enabled", "true") + .put("https-supported-ciphers", "override-cipher") + .put("https-cert-path", "/override/path/cert") + .put("https-key-path", "/override/path/key") + .put("http-server.num-io-threads-hw-multiplier", "1.5") + .put("exchange.http-client.num-io-threads-hw-multiplier", "1.5") + .put("async-data-cache-enabled", "true") + .put("async-cache-ssd-gb", "1") + .put("connector.num-io-threads-hw-multiplier", "0.1") + .put("presto.version", "override.presto.version") + .put("shutdown-onset-sec", "15") + .put("system-memory-gb", "5") + .put("query-memory-gb", "4") + .put("query.max-memory-per-node", "4GB") + .put("use-mmap-allocator", "false") + .put("memory-arbitrator-kind", "NOOP") + .put("shared-arbitrator.reserved-capacity", "1GB") + .put("shared-arbitrator.memory-pool-initial-capacity", "1GB") + .put("shared-arbitrator.max-memory-arbitration-time", "1s") + .put("experimental.spiller-spill-path", "/abc") + .put("task.max-drivers-per-task", "25") + .put("enable-old-task-cleanup", "true") + .put("shuffle.name", "remote") + .put("http-server.enable-access-log", "false") + .put("core-on-allocation-failure-enabled", "true") + .put("spill-enabled", "false") + .put("aggregation-spill-enabled", "false") + .put("join-spill-enabled", "false") + .put("order-by-spill-enabled", "false") + .put("non-defined-property-key-0", "non-defined-property-value-0") + .put("non-defined-property-key-1", "non-defined-property-value-1") + .put("non-defined-property-key-2", "non-defined-property-value-2") + .put("non-defined-property-key-3", "non-defined-property-value-3") + .build(); + nativeExecutionSystemConfig = new NativeExecutionSystemConfig(systemConfig); + + builder = ImmutableMap.builder(); + expectedConfigs = builder + .put("concurrent-lifespans-per-task", "15") + .put("enable-serialized-page-checksum", "false") + .put("enable_velox_expression_logging", "true") + .put("enable_velox_task_logging", "false") + .put("http-server.http.port", "8777") + .put("http-server.reuse-port", "false") + .put("http-server.bind-to-node-internal-address-only-enabled", "false") + .put("http-server.https.port", "8778") + .put("http-server.https.enabled", "true") + .put("https-supported-ciphers", "override-cipher") + .put("https-cert-path", "/override/path/cert") + .put("https-key-path", "/override/path/key") + .put("http-server.num-io-threads-hw-multiplier", "1.5") + .put("exchange.http-client.num-io-threads-hw-multiplier", "1.5") + .put("async-data-cache-enabled", "true") + .put("async-cache-ssd-gb", "1") + .put("connector.num-io-threads-hw-multiplier", "0.1") + .put("presto.version", "override.presto.version") + .put("shutdown-onset-sec", "15") + .put("system-memory-gb", "5") + .put("query-memory-gb", "4") + .put("query.max-memory-per-node", "4GB") + .put("use-mmap-allocator", "false") + .put("memory-arbitrator-kind", "NOOP") + .put("shared-arbitrator.reserved-capacity", "1GB") + .put("shared-arbitrator.memory-pool-initial-capacity", "1GB") + .put("shared-arbitrator.max-memory-arbitration-time", "1s") + .put("experimental.spiller-spill-path", "/abc") + .put("task.max-drivers-per-task", "25") + .put("enable-old-task-cleanup", "true") + .put("shuffle.name", "remote") + .put("http-server.enable-access-log", "false") + .put("core-on-allocation-failure-enabled", "true") + .put("spill-enabled", "false") + .put("aggregation-spill-enabled", "false") + .put("join-spill-enabled", "false") + .put("order-by-spill-enabled", "false") + .put("non-defined-property-key-0", "non-defined-property-value-0") + .put("non-defined-property-key-1", "non-defined-property-value-1") + .put("non-defined-property-key-2", "non-defined-property-value-2") + .put("non-defined-property-key-3", "non-defined-property-value-3") + .put("max-spill-bytes", String.valueOf(600L << 30)) // default spill bytes + .build(); + + assertEquals(nativeExecutionSystemConfig.getAllProperties(), expectedConfigs); } @Test @@ -182,7 +225,9 @@ public void testNativeExecutionConnectorConfig() @Test public void testFilePropertiesPopulator() { - PrestoSparkWorkerProperty workerProperty = new PrestoSparkWorkerProperty(new NativeExecutionConnectorConfig(), new NativeExecutionNodeConfig(), new NativeExecutionSystemConfig()); + PrestoSparkWorkerProperty workerProperty = new PrestoSparkWorkerProperty( + new NativeExecutionConnectorConfig(), new NativeExecutionNodeConfig(), + new NativeExecutionSystemConfig(ImmutableMap.of())); testPropertiesPopulate(workerProperty); } From 0c7260016f494d768b5d24da3934ae421ac7baf3 Mon Sep 17 00:00:00 2001 From: Matt Karrmann Date: Sun, 24 Aug 2025 20:03:24 -0700 Subject: [PATCH 002/113] feat(OSS Presto): Support tracking Page Sink Runtime Stats in TableWriterOperator (#25846) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Summary: Pull Request resolved: https://github.com/prestodb/presto/pull/25846 Pass the Operator Context's Runtime Stats down into the `TableWriteOperator`'s Page Sink. Specifically this diff makes the following changes: a) `TableWriteOperator` passes its `RuntimeStats` into the Page Sink it creates via `PageSinkManager.createPageSink` b) When the `PageSinkManager.createPageSink` is provided `RuntimeStats`, these `RuntimeStats` are passed into the `Session.toConnectorSession` call, which creates a `FullConnectorSession` instance c) When `Session.toConnectorSession` is provided `RuntimeStats`, it passes this into the `FullConnectorSession` instance it constructs d) Add a `Builder` to `FullConnectorSession`, which allows providing a `RuntimeStats` instance to `FullConnectorSession` at construction-time. `FullConnectorSession.getRuntimeStats()` now returns the `RuntimeStats` which was set at construction-time. If no `RuntimeStats` were provided at construction-time, then `FullConnectorSession.getRuntimeStats()` defaults to return the `Session` object's `RuntimeStats`—this preserves backwards compatibility. All changes preserve forward-compatibility. ## Context Without this change, the `FullConnectorSession`'s `RuntimeStats` points to the `Session`'s `RuntimeStat`s. All metrics added to the `Session`'s `RuntimeStats` within an Operator Worker-side are discarded. That is, all Runtime Metrics added to the Connector Session's RuntimeStats when executing `TableWriterOperator` were being completely discarded. Specifically, in Meta, the stats from our internal filesystem implementation were missing. Passing the Operator Context's `RuntimeStats` instance down into Connector Session is the simplest way to fix this. Additionally, since the previous `RuntimeStat`s for `TableWriteOperator`'s `FullConnectorSession` were always discarded, we can be confident that replacing them with the `OperatorContext` `RuntimeStat`s will not break anyone else's code. Differential Revision: D80675849 --- .../facebook/presto/FullConnectorSession.java | 137 +++++++++++++++--- .../java/com/facebook/presto/Session.java | 25 +++- .../presto/operator/TableWriterOperator.java | 13 +- .../presto/split/PageSinkManager.java | 23 ++- 4 files changed, 162 insertions(+), 36 deletions(-) diff --git a/presto-main-base/src/main/java/com/facebook/presto/FullConnectorSession.java b/presto-main-base/src/main/java/com/facebook/presto/FullConnectorSession.java index 1d8ee24735d8f..d085920337a36 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/FullConnectorSession.java +++ b/presto-main-base/src/main/java/com/facebook/presto/FullConnectorSession.java @@ -48,17 +48,11 @@ public class FullConnectorSession private final SessionPropertyManager sessionPropertyManager; private final SqlFunctionProperties sqlFunctionProperties; private final Map sessionFunctions; + private final RuntimeStats runtimeStats; public FullConnectorSession(Session session, ConnectorIdentity identity) { - this.session = requireNonNull(session, "session is null"); - this.identity = requireNonNull(identity, "identity is null"); - this.properties = null; - this.connectorId = null; - this.catalog = null; - this.sessionPropertyManager = null; - this.sqlFunctionProperties = session.getSqlFunctionProperties(); - this.sessionFunctions = ImmutableMap.copyOf(session.getSessionFunctions()); + this(builder(session, identity, null, null, null, null)); } public FullConnectorSession( @@ -69,14 +63,123 @@ public FullConnectorSession( String catalog, SessionPropertyManager sessionPropertyManager) { - this.session = requireNonNull(session, "session is null"); - this.identity = requireNonNull(identity, "identity is null"); - this.properties = ImmutableMap.copyOf(requireNonNull(properties, "properties is null")); - this.connectorId = requireNonNull(connectorId, "connectorId is null"); - this.catalog = requireNonNull(catalog, "catalog is null"); - this.sessionPropertyManager = requireNonNull(sessionPropertyManager, "sessionPropertyManager is null"); - this.sqlFunctionProperties = session.getSqlFunctionProperties(); - this.sessionFunctions = ImmutableMap.copyOf(session.getSessionFunctions()); + this(builder(session, identity, properties, connectorId, catalog, sessionPropertyManager)); + } + + private FullConnectorSession(Builder builder) + { + this.session = builder.getSession(); + this.identity = builder.getIdentity(); + this.properties = builder.getProperties(); + this.connectorId = builder.getConnectorId(); + this.catalog = builder.getCatalog(); + this.sessionPropertyManager = builder.getSessionPropertyManager(); + this.sqlFunctionProperties = builder.getSqlFunctionProperties() != null ? builder.getSqlFunctionProperties() : builder.getSession().getSqlFunctionProperties(); + this.sessionFunctions = builder.getSessionFunctions() != null ? builder.getSessionFunctions() : ImmutableMap.copyOf(builder.getSession().getSessionFunctions()); + this.runtimeStats = builder.getRuntimeStats() != null ? builder.getRuntimeStats() : builder.getSession().getRuntimeStats(); + } + + public static Builder builder( + Session session, + ConnectorIdentity identity, + Map properties, + ConnectorId connectorId, + String catalog, + SessionPropertyManager sessionPropertyManager) + { + return new Builder(session, identity, properties, connectorId, catalog, sessionPropertyManager); + } + + public static class Builder + { + private final Session session; + private final ConnectorIdentity identity; + private final Map properties; + private final ConnectorId connectorId; + private final String catalog; + private final SessionPropertyManager sessionPropertyManager; + + private SqlFunctionProperties sqlFunctionProperties; + private Map sessionFunctions; + private RuntimeStats runtimeStats; + + private Builder(Session session, ConnectorIdentity identity, Map properties, ConnectorId connectorId, String catalog, SessionPropertyManager sessionPropertyManager) + { + this.session = requireNonNull(session, "session is null"); + this.identity = requireNonNull(identity, "identity is null"); + this.properties = properties; + this.connectorId = connectorId; + this.catalog = catalog; + this.sessionPropertyManager = sessionPropertyManager; + } + + public Session getSession() + { + return session; + } + + public ConnectorIdentity getIdentity() + { + return identity; + } + + public Map getProperties() + { + return properties; + } + + public ConnectorId getConnectorId() + { + return connectorId; + } + + public String getCatalog() + { + return catalog; + } + + public SessionPropertyManager getSessionPropertyManager() + { + return sessionPropertyManager; + } + + public SqlFunctionProperties getSqlFunctionProperties() + { + return sqlFunctionProperties; + } + + public Builder setSqlFunctionProperties(SqlFunctionProperties sqlFunctionProperties) + { + this.sqlFunctionProperties = sqlFunctionProperties; + return this; + } + + public Map getSessionFunctions() + { + return sessionFunctions; + } + + public Builder setSessionFunctions(Map sessionFunctions) + { + this.sessionFunctions = sessionFunctions; + return this; + } + + public RuntimeStats getRuntimeStats() + { + return runtimeStats; + } + + public Builder setRuntimeStats(RuntimeStats runtimeStats) + { + this.runtimeStats = runtimeStats; + return this; + } + + public FullConnectorSession build() + { + return new FullConnectorSession(this); + } } public Session getSession() @@ -197,7 +300,7 @@ public WarningCollector getWarningCollector() @Override public RuntimeStats getRuntimeStats() { - return session.getRuntimeStats(); + return runtimeStats; } @Override diff --git a/presto-main-base/src/main/java/com/facebook/presto/Session.java b/presto-main-base/src/main/java/com/facebook/presto/Session.java index 14b7a05d142d7..327924fd9522f 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/Session.java +++ b/presto-main-base/src/main/java/com/facebook/presto/Session.java @@ -495,17 +495,26 @@ public SqlFunctionProperties getSqlFunctionProperties() .build(); } - public ConnectorSession toConnectorSession(ConnectorId connectorId) + public ConnectorSession toConnectorSession(ConnectorId connectorId, RuntimeStats runtimeStats) { requireNonNull(connectorId, "connectorId is null"); - return new FullConnectorSession( - this, - identity.toConnectorIdentity(connectorId.getCatalogName()), - connectorProperties.getOrDefault(connectorId, ImmutableMap.of()), - connectorId, - connectorId.getCatalogName(), - sessionPropertyManager); + FullConnectorSession.Builder connectorSessionBuilder = FullConnectorSession + .builder( + this, + identity.toConnectorIdentity(connectorId.getCatalogName()), + connectorProperties.getOrDefault(connectorId, ImmutableMap.of()), + connectorId, + connectorId.getCatalogName(), + sessionPropertyManager) + .setRuntimeStats(runtimeStats); + + return connectorSessionBuilder.build(); + } + + public ConnectorSession toConnectorSession(ConnectorId connectorId) + { + return toConnectorSession(connectorId, runtimeStats); } public SessionRepresentation toSessionRepresentation() diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/TableWriterOperator.java b/presto-main-base/src/main/java/com/facebook/presto/operator/TableWriterOperator.java index d46627202c6ec..5deda5f3ffb94 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/operator/TableWriterOperator.java +++ b/presto-main-base/src/main/java/com/facebook/presto/operator/TableWriterOperator.java @@ -20,6 +20,7 @@ import com.facebook.drift.annotations.ThriftStruct; import com.facebook.presto.Session; import com.facebook.presto.common.Page; +import com.facebook.presto.common.RuntimeStats; import com.facebook.presto.common.block.Block; import com.facebook.presto.common.block.BlockBuilder; import com.facebook.presto.common.block.RunLengthEncodedBlock; @@ -130,7 +131,7 @@ public Operator createOperator(DriverContext driverContext) boolean statisticsCpuTimerEnabled = !(statisticsAggregationOperator instanceof DevNullOperator) && isStatisticsCpuTimerEnabled(session); return new TableWriterOperator( context, - createPageSink(), + createPageSink(context), columnChannels, notNullChannelColumnNames, statisticsAggregationOperator, @@ -140,19 +141,21 @@ public Operator createOperator(DriverContext driverContext) pageSinkCommitStrategy); } - private ConnectorPageSink createPageSink() + private ConnectorPageSink createPageSink(OperatorContext operatorContext) { PageSinkContext.Builder pageSinkContextBuilder = PageSinkContext.builder() .setCommitRequired(pageSinkCommitStrategy.isCommitRequired()); + RuntimeStats runtimeStats = operatorContext.getRuntimeStats(); + if (target instanceof CreateHandle) { - return pageSinkManager.createPageSink(session, ((CreateHandle) target).getHandle(), pageSinkContextBuilder.build()); + return pageSinkManager.createPageSink(session, ((CreateHandle) target).getHandle(), pageSinkContextBuilder.build(), runtimeStats); } if (target instanceof InsertHandle) { - return pageSinkManager.createPageSink(session, ((InsertHandle) target).getHandle(), pageSinkContextBuilder.build()); + return pageSinkManager.createPageSink(session, ((InsertHandle) target).getHandle(), pageSinkContextBuilder.build(), runtimeStats); } if (target instanceof RefreshMaterializedViewHandle) { - return pageSinkManager.createPageSink(session, ((RefreshMaterializedViewHandle) target).getHandle(), pageSinkContextBuilder.build()); + return pageSinkManager.createPageSink(session, ((RefreshMaterializedViewHandle) target).getHandle(), pageSinkContextBuilder.build(), runtimeStats); } throw new UnsupportedOperationException("Unhandled target type: " + target.getClass().getName()); } diff --git a/presto-main-base/src/main/java/com/facebook/presto/split/PageSinkManager.java b/presto-main-base/src/main/java/com/facebook/presto/split/PageSinkManager.java index aee08454aefc9..39f6a860ac1d5 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/split/PageSinkManager.java +++ b/presto-main-base/src/main/java/com/facebook/presto/split/PageSinkManager.java @@ -14,6 +14,7 @@ package com.facebook.presto.split; import com.facebook.presto.Session; +import com.facebook.presto.common.RuntimeStats; import com.facebook.presto.metadata.InsertTableHandle; import com.facebook.presto.metadata.OutputTableHandle; import com.facebook.presto.spi.ConnectorId; @@ -46,22 +47,32 @@ public void removeConnectorPageSinkProvider(ConnectorId connectorId) pageSinkProviders.remove(connectorId); } - @Override - public ConnectorPageSink createPageSink(Session session, OutputTableHandle tableHandle, PageSinkContext pageSinkContext) + public ConnectorPageSink createPageSink(Session session, OutputTableHandle tableHandle, PageSinkContext pageSinkContext, RuntimeStats runtimeStats) { // assumes connectorId and catalog are the same - ConnectorSession connectorSession = session.toConnectorSession(tableHandle.getConnectorId()); + ConnectorSession connectorSession = session.toConnectorSession(tableHandle.getConnectorId(), runtimeStats); return providerFor(tableHandle.getConnectorId()).createPageSink(tableHandle.getTransactionHandle(), connectorSession, tableHandle.getConnectorHandle(), pageSinkContext); } - @Override - public ConnectorPageSink createPageSink(Session session, InsertTableHandle tableHandle, PageSinkContext pageSinkContext) + public ConnectorPageSink createPageSink(Session session, InsertTableHandle tableHandle, PageSinkContext pageSinkContext, RuntimeStats runtimeStats) { // assumes connectorId and catalog are the same - ConnectorSession connectorSession = session.toConnectorSession(tableHandle.getConnectorId()); + ConnectorSession connectorSession = session.toConnectorSession(tableHandle.getConnectorId(), runtimeStats); return providerFor(tableHandle.getConnectorId()).createPageSink(tableHandle.getTransactionHandle(), connectorSession, tableHandle.getConnectorHandle(), pageSinkContext); } + @Override + public ConnectorPageSink createPageSink(Session session, OutputTableHandle tableHandle, PageSinkContext pageSinkContext) + { + return createPageSink(session, tableHandle, pageSinkContext, null); + } + + @Override + public ConnectorPageSink createPageSink(Session session, InsertTableHandle tableHandle, PageSinkContext pageSinkContext) + { + return createPageSink(session, tableHandle, pageSinkContext, null); + } + private ConnectorPageSinkProvider providerFor(ConnectorId connectorId) { ConnectorPageSinkProvider provider = pageSinkProviders.get(connectorId); From 3b6bec0d0035f19bb846bef478b0c981daae25ba Mon Sep 17 00:00:00 2001 From: aspegren_david Date: Mon, 26 May 2025 16:29:11 +0200 Subject: [PATCH 003/113] Add session property to toggle Orc use column names feature There is an existing HiveClientConfig property hive.orc.use-column-names to access ORC file by column names, but no session property. This commit moves the existing HiveClientConfig property to HiveCommonClientConfig and introduces a session property in HiveCommonSessionProperties. It also implements changes accordingly in DwrfAggregatedPageSourceFactory, OrcAggregatedPageSourceFactory, OrcSelectivePageSourceFactory and OrcBatchPageSourceFactory. Constructors in those classes do not take boolean useOrcColumnNames anymore. Tests where those are used have also been changed. Hive connector documentation has been changed. An integration test has been added to TestHiveDistributedQueries.java. Helper function created in HiveTestUtils to replace function in TestHiveIntegrationSmokeTest. Remove superfluous constructors that have hiveClientConfig in parameter list from DwrfAggregatedPageSourceFactory.java and OrcAggregatedPageSourceFactory.java and change explicit calls in HiveTestUtils.java. Closes-Issue: #24134 Remove superfluous constructors that have hiveClientConfig in parameter list from DwrfAggregatedPageSourceFactory.java and OrcAggregatedPageSourceFactory.java and change explicit calls in HiveTestUtils.java. Add additional test with different column names to TestHiveDistributedQueries.java --- .../src/main/sphinx/connector/hive.rst | 4 + .../presto/hive/HiveCommonClientConfig.java | 14 +++ .../hive/HiveCommonSessionProperties.java | 12 ++ .../hive/TestHiveCommonClientConfig.java | 3 + .../presto/hive/HiveClientConfig.java | 14 --- .../orc/DwrfAggregatedPageSourceFactory.java | 26 +---- .../orc/OrcAggregatedPageSourceFactory.java | 26 +---- .../hive/orc/OrcBatchPageSourceFactory.java | 7 +- .../orc/OrcSelectivePageSourceFactory.java | 7 +- .../facebook/presto/hive/HiveTestUtils.java | 32 +++++- .../presto/hive/TestHiveClientConfig.java | 3 - .../hive/TestHiveDistributedQueries.java | 105 ++++++++++++++++++ .../presto/hive/TestHiveFileFormats.java | 15 +-- .../hive/TestHiveIntegrationSmokeTest.java | 25 +---- .../TestOrcBatchPageSourceMemoryTracking.java | 1 - .../presto/hive/benchmark/FileFormat.java | 1 - 16 files changed, 187 insertions(+), 108 deletions(-) diff --git a/presto-docs/src/main/sphinx/connector/hive.rst b/presto-docs/src/main/sphinx/connector/hive.rst index 7e427b9834b3b..7208faaff0b39 100644 --- a/presto-docs/src/main/sphinx/connector/hive.rst +++ b/presto-docs/src/main/sphinx/connector/hive.rst @@ -224,6 +224,10 @@ Property Name Description CopyOnFirstWriteConfiguration, it can result in silent failures where critical configuration properties are not correctly propagated. + + ``hive.orc.use-column-names`` Enable accessing ORC columns by name in the ORC file ``false`` + metadata, instead of their ordinal position. Also toggleable + through the ``hive.orc_use_column_names`` session property. ======================================================== ============================================================ ============ .. _constructor: https://github.com/apache/hadoop/blob/02a9190af5f8264e25966a80c8f9ea9bb6677899/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/conf/Configuration.java#L844-L875 diff --git a/presto-hive-common/src/main/java/com/facebook/presto/hive/HiveCommonClientConfig.java b/presto-hive-common/src/main/java/com/facebook/presto/hive/HiveCommonClientConfig.java index 0b5a34b2c8ed3..1e54401247e2f 100644 --- a/presto-hive-common/src/main/java/com/facebook/presto/hive/HiveCommonClientConfig.java +++ b/presto-hive-common/src/main/java/com/facebook/presto/hive/HiveCommonClientConfig.java @@ -37,6 +37,7 @@ public class HiveCommonClientConfig private DataSize orcStreamBufferSize = new DataSize(8, MEGABYTE); private OrcWriteValidationMode orcWriterValidationMode = OrcWriteValidationMode.BOTH; private double orcWriterValidationPercentage; + private boolean useOrcColumnNames; private DataSize orcTinyStripeThreshold = new DataSize(8, MEGABYTE); private boolean parquetBatchReadOptimizationEnabled; private boolean parquetEnableBatchReaderVerification; @@ -183,6 +184,19 @@ public HiveCommonClientConfig setOrcWriterValidationPercentage(double orcWriterV return this; } + public boolean isUseOrcColumnNames() + { + return useOrcColumnNames; + } + + @Config("hive.orc.use-column-names") + @ConfigDescription("Access ORC columns using names from the file first, and fallback to Hive schema column names if not found to ensure backward compatibility with old data") + public HiveCommonClientConfig setUseOrcColumnNames(boolean useOrcColumnNames) + { + this.useOrcColumnNames = useOrcColumnNames; + return this; + } + @NotNull public DataSize getOrcTinyStripeThreshold() { diff --git a/presto-hive-common/src/main/java/com/facebook/presto/hive/HiveCommonSessionProperties.java b/presto-hive-common/src/main/java/com/facebook/presto/hive/HiveCommonSessionProperties.java index de4bce0e2b661..b505cd37d8188 100644 --- a/presto-hive-common/src/main/java/com/facebook/presto/hive/HiveCommonSessionProperties.java +++ b/presto-hive-common/src/main/java/com/facebook/presto/hive/HiveCommonSessionProperties.java @@ -42,6 +42,8 @@ public class HiveCommonSessionProperties public static final String RANGE_FILTERS_ON_SUBSCRIPTS_ENABLED = "range_filters_on_subscripts_enabled"; @VisibleForTesting public static final String PARQUET_BATCH_READ_OPTIMIZATION_ENABLED = "parquet_batch_read_optimization_enabled"; + @VisibleForTesting + public static final String ORC_USE_COLUMN_NAMES = "orc_use_column_names"; public static final String NODE_SELECTION_STRATEGY = "node_selection_strategy"; private static final String ORC_BLOOM_FILTERS_ENABLED = "orc_bloom_filters_enabled"; @@ -153,6 +155,11 @@ public HiveCommonSessionProperties(HiveCommonClientConfig hiveCommonClientConfig "use JNI based zstd decompression for reading ORC files", hiveCommonClientConfig.isZstdJniDecompressionEnabled(), true), + booleanProperty( + ORC_USE_COLUMN_NAMES, + "Access ORC columns using names from the file first, and fallback to Hive schema column names if not found to ensure backward compatibility with old data", + hiveCommonClientConfig.isUseOrcColumnNames(), + false), booleanProperty( PARQUET_BATCH_READ_OPTIMIZATION_ENABLED, "Is Parquet batch read optimization enabled", @@ -262,6 +269,11 @@ public static boolean isOrcZstdJniDecompressionEnabled(ConnectorSession session) return session.getProperty(ORC_ZSTD_JNI_DECOMPRESSION_ENABLED, Boolean.class); } + public static boolean isUseOrcColumnNames(ConnectorSession session) + { + return session.getProperty(ORC_USE_COLUMN_NAMES, Boolean.class); + } + public static boolean isParquetBatchReadsEnabled(ConnectorSession session) { return session.getProperty(PARQUET_BATCH_READ_OPTIMIZATION_ENABLED, Boolean.class); diff --git a/presto-hive-common/src/test/java/com/facebook/presto/hive/TestHiveCommonClientConfig.java b/presto-hive-common/src/test/java/com/facebook/presto/hive/TestHiveCommonClientConfig.java index da24f4e64bc32..42579c3048b92 100644 --- a/presto-hive-common/src/test/java/com/facebook/presto/hive/TestHiveCommonClientConfig.java +++ b/presto-hive-common/src/test/java/com/facebook/presto/hive/TestHiveCommonClientConfig.java @@ -45,6 +45,7 @@ public void testDefaults() .setOrcOptimizedWriterEnabled(true) .setOrcWriterValidationPercentage(0.0) .setOrcWriterValidationMode(OrcWriteValidation.OrcWriteValidationMode.BOTH) + .setUseOrcColumnNames(false) .setZstdJniDecompressionEnabled(false) .setParquetBatchReaderVerificationEnabled(false) .setParquetBatchReadOptimizationEnabled(false) @@ -71,6 +72,7 @@ public void testExplicitPropertyMappings() .put("hive.orc.optimized-writer.enabled", "false") .put("hive.orc.writer.validation-percentage", "0.16") .put("hive.orc.writer.validation-mode", "DETAILED") + .put("hive.orc.use-column-names", "true") .put("hive.zstd-jni-decompression-enabled", "true") .put("hive.enable-parquet-batch-reader-verification", "true") .put("hive.parquet-batch-read-optimization-enabled", "true") @@ -94,6 +96,7 @@ public void testExplicitPropertyMappings() .setOrcOptimizedWriterEnabled(false) .setOrcWriterValidationPercentage(0.16) .setOrcWriterValidationMode(OrcWriteValidation.OrcWriteValidationMode.DETAILED) + .setUseOrcColumnNames(true) .setZstdJniDecompressionEnabled(true) .setParquetBatchReaderVerificationEnabled(true) .setParquetBatchReadOptimizationEnabled(true) diff --git a/presto-hive/src/main/java/com/facebook/presto/hive/HiveClientConfig.java b/presto-hive/src/main/java/com/facebook/presto/hive/HiveClientConfig.java index be0b6482d39f8..7f56998606165 100644 --- a/presto-hive/src/main/java/com/facebook/presto/hive/HiveClientConfig.java +++ b/presto-hive/src/main/java/com/facebook/presto/hive/HiveClientConfig.java @@ -107,7 +107,6 @@ public class HiveClientConfig private DataSize textMaxLineLength = new DataSize(100, MEGABYTE); private boolean assumeCanonicalPartitionKeys; - private boolean useOrcColumnNames; private double orcDefaultBloomFilterFpp = 0.05; private boolean rcfileOptimizedWriterEnabled = true; private boolean rcfileWriterValidate; @@ -720,19 +719,6 @@ public HiveClientConfig setS3FileSystemType(S3FileSystemType s3FileSystemType) return this; } - public boolean isUseOrcColumnNames() - { - return useOrcColumnNames; - } - - @Config("hive.orc.use-column-names") - @ConfigDescription("Access ORC columns using names from the file first, and fallback to Hive schema column names if not found to ensure backward compatibility with old data") - public HiveClientConfig setUseOrcColumnNames(boolean useOrcColumnNames) - { - this.useOrcColumnNames = useOrcColumnNames; - return this; - } - public double getOrcDefaultBloomFilterFpp() { return orcDefaultBloomFilterFpp; diff --git a/presto-hive/src/main/java/com/facebook/presto/hive/orc/DwrfAggregatedPageSourceFactory.java b/presto-hive/src/main/java/com/facebook/presto/hive/orc/DwrfAggregatedPageSourceFactory.java index 30fc1dcc49438..723759a0b162e 100644 --- a/presto-hive/src/main/java/com/facebook/presto/hive/orc/DwrfAggregatedPageSourceFactory.java +++ b/presto-hive/src/main/java/com/facebook/presto/hive/orc/DwrfAggregatedPageSourceFactory.java @@ -19,7 +19,6 @@ import com.facebook.presto.hive.FileFormatDataSourceStats; import com.facebook.presto.hive.HdfsEnvironment; import com.facebook.presto.hive.HiveAggregatedPageSourceFactory; -import com.facebook.presto.hive.HiveClientConfig; import com.facebook.presto.hive.HiveColumnHandle; import com.facebook.presto.hive.HiveFileContext; import com.facebook.presto.hive.HiveFileSplit; @@ -36,6 +35,7 @@ import java.util.List; import java.util.Optional; +import static com.facebook.presto.hive.HiveCommonSessionProperties.isUseOrcColumnNames; import static com.facebook.presto.hive.HiveErrorCode.HIVE_BAD_DATA; import static com.facebook.presto.hive.orc.OrcAggregatedPageSourceFactory.createOrcPageSource; import static com.facebook.presto.orc.DwrfEncryptionProvider.NO_ENCRYPTION; @@ -47,7 +47,6 @@ public class DwrfAggregatedPageSourceFactory { private final TypeManager typeManager; private final StandardFunctionResolution functionResolution; - private final boolean useOrcColumnNames; private final HdfsEnvironment hdfsEnvironment; private final FileFormatDataSourceStats stats; private final OrcFileTailSource orcFileTailSource; @@ -57,26 +56,6 @@ public class DwrfAggregatedPageSourceFactory public DwrfAggregatedPageSourceFactory( TypeManager typeManager, StandardFunctionResolution functionResolution, - HiveClientConfig config, - HdfsEnvironment hdfsEnvironment, - FileFormatDataSourceStats stats, - OrcFileTailSource orcFileTailSource, - StripeMetadataSourceFactory stripeMetadataSourceFactory) - { - this( - typeManager, - functionResolution, - requireNonNull(config, "hiveClientConfig is null").isUseOrcColumnNames(), - hdfsEnvironment, - stats, - orcFileTailSource, - stripeMetadataSourceFactory); - } - - public DwrfAggregatedPageSourceFactory( - TypeManager typeManager, - StandardFunctionResolution functionResolution, - boolean useOrcColumnNames, HdfsEnvironment hdfsEnvironment, FileFormatDataSourceStats stats, OrcFileTailSource orcFileTailSource, @@ -84,7 +63,6 @@ public DwrfAggregatedPageSourceFactory( { this.typeManager = requireNonNull(typeManager, "typeManager is null"); this.functionResolution = requireNonNull(functionResolution, "functionResolution is null"); - this.useOrcColumnNames = useOrcColumnNames; this.hdfsEnvironment = requireNonNull(hdfsEnvironment, "hdfsEnvironment is null"); this.stats = requireNonNull(stats, "stats is null"); this.orcFileTailSource = requireNonNull(orcFileTailSource, "orcFileTailCache is null"); @@ -116,7 +94,7 @@ public Optional createPageSource( configuration, fileSplit, columns, - useOrcColumnNames, + isUseOrcColumnNames(session), typeManager, functionResolution, stats, diff --git a/presto-hive/src/main/java/com/facebook/presto/hive/orc/OrcAggregatedPageSourceFactory.java b/presto-hive/src/main/java/com/facebook/presto/hive/orc/OrcAggregatedPageSourceFactory.java index c61507c7983bb..cd6313d484755 100644 --- a/presto-hive/src/main/java/com/facebook/presto/hive/orc/OrcAggregatedPageSourceFactory.java +++ b/presto-hive/src/main/java/com/facebook/presto/hive/orc/OrcAggregatedPageSourceFactory.java @@ -19,7 +19,6 @@ import com.facebook.presto.hive.FileFormatDataSourceStats; import com.facebook.presto.hive.HdfsEnvironment; import com.facebook.presto.hive.HiveAggregatedPageSourceFactory; -import com.facebook.presto.hive.HiveClientConfig; import com.facebook.presto.hive.HiveColumnHandle; import com.facebook.presto.hive.HiveFileContext; import com.facebook.presto.hive.HiveFileSplit; @@ -49,6 +48,7 @@ import static com.facebook.presto.hive.HiveCommonSessionProperties.getOrcMaxReadBlockSize; import static com.facebook.presto.hive.HiveCommonSessionProperties.getOrcTinyStripeThreshold; import static com.facebook.presto.hive.HiveCommonSessionProperties.isOrcZstdJniDecompressionEnabled; +import static com.facebook.presto.hive.HiveCommonSessionProperties.isUseOrcColumnNames; import static com.facebook.presto.hive.HiveUtil.getPhysicalHiveColumnHandles; import static com.facebook.presto.hive.orc.OrcPageSourceFactoryUtils.getOrcDataSource; import static com.facebook.presto.hive.orc.OrcPageSourceFactoryUtils.getOrcReader; @@ -62,7 +62,6 @@ public class OrcAggregatedPageSourceFactory { private final TypeManager typeManager; private final StandardFunctionResolution functionResolution; - private final boolean useOrcColumnNames; private final HdfsEnvironment hdfsEnvironment; private final FileFormatDataSourceStats stats; private final OrcFileTailSource orcFileTailSource; @@ -72,26 +71,6 @@ public class OrcAggregatedPageSourceFactory public OrcAggregatedPageSourceFactory( TypeManager typeManager, StandardFunctionResolution functionResolution, - HiveClientConfig config, - HdfsEnvironment hdfsEnvironment, - FileFormatDataSourceStats stats, - OrcFileTailSource orcFileTailSource, - StripeMetadataSourceFactory stripeMetadataSourceFactory) - { - this( - typeManager, - functionResolution, - requireNonNull(config, "hiveClientConfig is null").isUseOrcColumnNames(), - hdfsEnvironment, - stats, - orcFileTailSource, - stripeMetadataSourceFactory); - } - - public OrcAggregatedPageSourceFactory( - TypeManager typeManager, - StandardFunctionResolution functionResolution, - boolean useOrcColumnNames, HdfsEnvironment hdfsEnvironment, FileFormatDataSourceStats stats, OrcFileTailSource orcFileTailSource, @@ -99,7 +78,6 @@ public OrcAggregatedPageSourceFactory( { this.typeManager = requireNonNull(typeManager, "typeManager is null"); this.functionResolution = requireNonNull(functionResolution, "functionResolution is null"); - this.useOrcColumnNames = useOrcColumnNames; this.hdfsEnvironment = requireNonNull(hdfsEnvironment, "hdfsEnvironment is null"); this.stats = requireNonNull(stats, "stats is null"); this.orcFileTailSource = requireNonNull(orcFileTailSource, "orcFileTailCache is null"); @@ -132,7 +110,7 @@ public Optional createPageSource( configuration, fileSplit, columns, - useOrcColumnNames, + isUseOrcColumnNames(session), typeManager, functionResolution, stats, diff --git a/presto-hive/src/main/java/com/facebook/presto/hive/orc/OrcBatchPageSourceFactory.java b/presto-hive/src/main/java/com/facebook/presto/hive/orc/OrcBatchPageSourceFactory.java index 3c92bf8efc405..f1b892d4495b2 100644 --- a/presto-hive/src/main/java/com/facebook/presto/hive/orc/OrcBatchPageSourceFactory.java +++ b/presto-hive/src/main/java/com/facebook/presto/hive/orc/OrcBatchPageSourceFactory.java @@ -62,6 +62,7 @@ import static com.facebook.presto.hive.HiveCommonSessionProperties.getOrcTinyStripeThreshold; import static com.facebook.presto.hive.HiveCommonSessionProperties.isOrcBloomFiltersEnabled; import static com.facebook.presto.hive.HiveCommonSessionProperties.isOrcZstdJniDecompressionEnabled; +import static com.facebook.presto.hive.HiveCommonSessionProperties.isUseOrcColumnNames; import static com.facebook.presto.hive.HiveUtil.checkRowIDPartitionComponent; import static com.facebook.presto.hive.HiveUtil.getPhysicalHiveColumnHandles; import static com.facebook.presto.hive.orc.OrcPageSourceFactoryUtils.getOrcDataSource; @@ -78,7 +79,6 @@ public class OrcBatchPageSourceFactory implements HiveBatchPageSourceFactory { private final TypeManager typeManager; - private final boolean useOrcColumnNames; private final HdfsEnvironment hdfsEnvironment; private final FileFormatDataSourceStats stats; private final int domainCompactionThreshold; @@ -97,7 +97,6 @@ public OrcBatchPageSourceFactory( { this( typeManager, - requireNonNull(config, "hiveClientConfig is null").isUseOrcColumnNames(), hdfsEnvironment, stats, config.getDomainCompactionThreshold(), @@ -107,7 +106,6 @@ public OrcBatchPageSourceFactory( public OrcBatchPageSourceFactory( TypeManager typeManager, - boolean useOrcColumnNames, HdfsEnvironment hdfsEnvironment, FileFormatDataSourceStats stats, int domainCompactionThreshold, @@ -115,7 +113,6 @@ public OrcBatchPageSourceFactory( StripeMetadataSourceFactory stripeMetadataSourceFactory) { this.typeManager = requireNonNull(typeManager, "typeManager is null"); - this.useOrcColumnNames = useOrcColumnNames; this.hdfsEnvironment = requireNonNull(hdfsEnvironment, "hdfsEnvironment is null"); this.stats = requireNonNull(stats, "stats is null"); this.domainCompactionThreshold = domainCompactionThreshold; @@ -153,7 +150,7 @@ public Optional createPageSource( configuration, fileSplit, columns, - useOrcColumnNames, + isUseOrcColumnNames(session), effectivePredicate, hiveStorageTimeZone, typeManager, diff --git a/presto-hive/src/main/java/com/facebook/presto/hive/orc/OrcSelectivePageSourceFactory.java b/presto-hive/src/main/java/com/facebook/presto/hive/orc/OrcSelectivePageSourceFactory.java index 60354406ace2f..5fcdc9c89ad7c 100644 --- a/presto-hive/src/main/java/com/facebook/presto/hive/orc/OrcSelectivePageSourceFactory.java +++ b/presto-hive/src/main/java/com/facebook/presto/hive/orc/OrcSelectivePageSourceFactory.java @@ -104,6 +104,7 @@ import static com.facebook.presto.hive.HiveCommonSessionProperties.getOrcTinyStripeThreshold; import static com.facebook.presto.hive.HiveCommonSessionProperties.isOrcBloomFiltersEnabled; import static com.facebook.presto.hive.HiveCommonSessionProperties.isOrcZstdJniDecompressionEnabled; +import static com.facebook.presto.hive.HiveCommonSessionProperties.isUseOrcColumnNames; import static com.facebook.presto.hive.HiveErrorCode.HIVE_INVALID_BUCKET_FILES; import static com.facebook.presto.hive.HiveSessionProperties.isAdaptiveFilterReorderingEnabled; import static com.facebook.presto.hive.HiveSessionProperties.isLegacyTimestampBucketing; @@ -132,7 +133,6 @@ public class OrcSelectivePageSourceFactory private final TypeManager typeManager; private final StandardFunctionResolution functionResolution; private final RowExpressionService rowExpressionService; - private final boolean useOrcColumnNames; private final HdfsEnvironment hdfsEnvironment; private final FileFormatDataSourceStats stats; private final int domainCompactionThreshold; @@ -156,7 +156,6 @@ public OrcSelectivePageSourceFactory( typeManager, functionResolution, rowExpressionService, - requireNonNull(config, "hiveClientConfig is null").isUseOrcColumnNames(), hdfsEnvironment, stats, config.getDomainCompactionThreshold(), @@ -169,7 +168,6 @@ public OrcSelectivePageSourceFactory( TypeManager typeManager, StandardFunctionResolution functionResolution, RowExpressionService rowExpressionService, - boolean useOrcColumnNames, HdfsEnvironment hdfsEnvironment, FileFormatDataSourceStats stats, int domainCompactionThreshold, @@ -180,7 +178,6 @@ public OrcSelectivePageSourceFactory( this.typeManager = requireNonNull(typeManager, "typeManager is null"); this.functionResolution = requireNonNull(functionResolution, "functionResolution is null"); this.rowExpressionService = requireNonNull(rowExpressionService, "rowExpressionService is null"); - this.useOrcColumnNames = useOrcColumnNames; this.hdfsEnvironment = requireNonNull(hdfsEnvironment, "hdfsEnvironment is null"); this.stats = requireNonNull(stats, "stats is null"); this.domainCompactionThreshold = domainCompactionThreshold; @@ -230,7 +227,7 @@ public Optional createPageSource( outputColumns, domainPredicate, remainingPredicate, - useOrcColumnNames, + isUseOrcColumnNames(session), hiveStorageTimeZone, typeManager, functionResolution, diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/HiveTestUtils.java b/presto-hive/src/test/java/com/facebook/presto/hive/HiveTestUtils.java index 37b3b35139a27..14d00eee6dc81 100644 --- a/presto-hive/src/test/java/com/facebook/presto/hive/HiveTestUtils.java +++ b/presto-hive/src/test/java/com/facebook/presto/hive/HiveTestUtils.java @@ -17,7 +17,9 @@ import com.facebook.airlift.json.smile.SmileCodec; import com.facebook.airlift.log.Logger; import com.facebook.presto.PagesIndexPageSorter; +import com.facebook.presto.Session; import com.facebook.presto.cache.CacheConfig; +import com.facebook.presto.common.QualifiedObjectName; import com.facebook.presto.common.block.BlockEncodingManager; import com.facebook.presto.common.type.ArrayType; import com.facebook.presto.common.type.MapType; @@ -51,7 +53,9 @@ import com.facebook.presto.hive.s3.PrestoS3ConfigurationUpdater; import com.facebook.presto.hive.s3select.S3SelectRecordCursorProvider; import com.facebook.presto.metadata.FunctionAndTypeManager; +import com.facebook.presto.metadata.Metadata; import com.facebook.presto.metadata.MetadataManager; +import com.facebook.presto.metadata.TableLayout; import com.facebook.presto.operator.PagesIndex; import com.facebook.presto.orc.StorageStripeMetadataSource; import com.facebook.presto.orc.StripeMetadataSourceFactory; @@ -59,7 +63,9 @@ import com.facebook.presto.parquet.cache.MetadataReader; import com.facebook.presto.spi.ColumnHandle; import com.facebook.presto.spi.ConnectorSession; +import com.facebook.presto.spi.Constraint; import com.facebook.presto.spi.PageSorter; +import com.facebook.presto.spi.TableHandle; import com.facebook.presto.spi.function.StandardFunctionResolution; import com.facebook.presto.spi.plan.FilterStatsCalculatorService; import com.facebook.presto.spi.relation.DeterminismEvaluator; @@ -75,7 +81,9 @@ import com.facebook.presto.sql.relational.RowExpressionDeterminismEvaluator; import com.facebook.presto.sql.relational.RowExpressionDomainTranslator; import com.facebook.presto.sql.relational.RowExpressionOptimizer; +import com.facebook.presto.testing.QueryRunner; import com.facebook.presto.testing.TestingConnectorSession; +import com.facebook.presto.tests.DistributedQueryRunner; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; import io.airlift.slice.Slice; @@ -87,13 +95,17 @@ import java.util.List; import java.util.Optional; import java.util.Set; +import java.util.function.Function; import static com.facebook.airlift.json.JsonCodec.jsonCodec; import static com.facebook.airlift.json.smile.SmileCodec.smileCodec; import static com.facebook.presto.common.type.Decimals.encodeScaledValue; import static com.facebook.presto.hive.HiveDwrfEncryptionProvider.NO_ENCRYPTION; +import static com.facebook.presto.hive.HiveQueryRunner.TPCH_SCHEMA; +import static com.facebook.presto.transaction.TransactionBuilder.transaction; import static java.lang.String.format; import static java.util.stream.Collectors.toList; +import static org.testng.Assert.assertTrue; public final class HiveTestUtils { @@ -188,8 +200,8 @@ public static Set getDefaultHiveAggregatedPageS FileFormatDataSourceStats stats = new FileFormatDataSourceStats(); HdfsEnvironment testHdfsEnvironment = createTestHdfsEnvironment(hiveClientConfig, metastoreClientConfig); return ImmutableSet.builder() - .add(new OrcAggregatedPageSourceFactory(FUNCTION_AND_TYPE_MANAGER, FUNCTION_RESOLUTION, hiveClientConfig, testHdfsEnvironment, stats, new StorageOrcFileTailSource(), StripeMetadataSourceFactory.of(new StorageStripeMetadataSource()))) - .add(new DwrfAggregatedPageSourceFactory(FUNCTION_AND_TYPE_MANAGER, FUNCTION_RESOLUTION, hiveClientConfig, testHdfsEnvironment, stats, new StorageOrcFileTailSource(), StripeMetadataSourceFactory.of(new StorageStripeMetadataSource()))) + .add(new OrcAggregatedPageSourceFactory(FUNCTION_AND_TYPE_MANAGER, FUNCTION_RESOLUTION, testHdfsEnvironment, stats, new StorageOrcFileTailSource(), StripeMetadataSourceFactory.of(new StorageStripeMetadataSource()))) + .add(new DwrfAggregatedPageSourceFactory(FUNCTION_AND_TYPE_MANAGER, FUNCTION_RESOLUTION, testHdfsEnvironment, stats, new StorageOrcFileTailSource(), StripeMetadataSourceFactory.of(new StorageStripeMetadataSource()))) .add(new ParquetAggregatedPageSourceFactory(FUNCTION_AND_TYPE_MANAGER, FUNCTION_RESOLUTION, testHdfsEnvironment, stats, new MetadataReader())) .build(); } @@ -374,4 +386,20 @@ public static List> getAllSessionProperties(HiveClientConfig allSessionProperties.addAll(hiveCommonSessionProperties.getSessionProperties()); return allSessionProperties; } + + public static Object getHiveTableProperty(QueryRunner queryRunner, Session session, String tableName, Function propertyGetter) + { + Metadata metadata = ((DistributedQueryRunner) queryRunner).getCoordinator().getMetadata(); + + return transaction(queryRunner.getTransactionManager(), queryRunner.getAccessControl()) + .readOnly() + .execute(session, transactionSession -> { + Optional tableHandle = metadata.getMetadataResolver(transactionSession).getTableHandle(new QualifiedObjectName("hive", TPCH_SCHEMA, tableName)); + assertTrue(tableHandle.isPresent()); + + TableLayout layout = metadata.getLayout(transactionSession, tableHandle.get(), Constraint.alwaysTrue(), Optional.empty()) + .getLayout(); + return propertyGetter.apply((HiveTableLayoutHandle) layout.getNewTableHandle().getLayout().get()); + }); + } } diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveClientConfig.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveClientConfig.java index 878330c46764e..1d09a4993dff6 100644 --- a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveClientConfig.java +++ b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveClientConfig.java @@ -83,7 +83,6 @@ public void testDefaults() .setMaxPartitionsPerWriter(100) .setWriteValidationThreads(16) .setTextMaxLineLength(new DataSize(100, Unit.MEGABYTE)) - .setUseOrcColumnNames(false) .setAssumeCanonicalPartitionKeys(false) .setOrcDefaultBloomFilterFpp(0.05) .setRcfileOptimizedWriterEnabled(true) @@ -211,7 +210,6 @@ public void testExplicitPropertyMappings() .put("hive.max-concurrent-zero-row-file-creations", "100") .put("hive.assume-canonical-partition-keys", "true") .put("hive.text.max-line-length", "13MB") - .put("hive.orc.use-column-names", "true") .put("hive.orc.default-bloom-filter-fpp", "0.96") .put("hive.rcfile-optimized-writer.enabled", "false") .put("hive.rcfile.writer.validate", "true") @@ -334,7 +332,6 @@ public void testExplicitPropertyMappings() .setDomainSocketPath("/foo") .setS3FileSystemType(S3FileSystemType.EMRFS) .setTextMaxLineLength(new DataSize(13, Unit.MEGABYTE)) - .setUseOrcColumnNames(true) .setAssumeCanonicalPartitionKeys(true) .setOrcDefaultBloomFilterFpp(0.96) .setRcfileOptimizedWriterEnabled(false) diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueries.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueries.java index c47da01ea5281..7552535ae9a4f 100644 --- a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueries.java +++ b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueries.java @@ -21,10 +21,15 @@ import com.facebook.presto.testing.QueryRunner; import com.facebook.presto.tests.AbstractTestDistributedQueries; import com.google.common.collect.ImmutableSet; +import com.google.common.io.Files; import org.intellij.lang.annotations.Language; import org.testng.annotations.AfterClass; import org.testng.annotations.Test; +import java.io.File; +import java.io.IOException; +import java.net.URI; +import java.net.URISyntaxException; import java.util.Set; import static com.facebook.presto.SystemSessionProperties.CTE_MATERIALIZATION_STRATEGY; @@ -33,8 +38,12 @@ import static com.facebook.presto.SystemSessionProperties.PUSHDOWN_SUBFIELDS_FOR_MAP_FUNCTIONS; import static com.facebook.presto.SystemSessionProperties.PUSHDOWN_SUBFIELDS_FROM_LAMBDA_ENABLED; import static com.facebook.presto.SystemSessionProperties.VERBOSE_OPTIMIZER_INFO_ENABLED; +import static com.facebook.presto.hive.HiveCommonSessionProperties.ORC_USE_COLUMN_NAMES; +import static com.facebook.presto.hive.HiveTestUtils.getHiveTableProperty; import static com.facebook.presto.sql.tree.ExplainType.Type.LOGICAL; import static com.google.common.collect.Iterables.getOnlyElement; +import static com.google.common.io.MoreFiles.deleteRecursively; +import static com.google.common.io.RecursiveDeleteOption.ALLOW_INSECURE; import static io.airlift.tpch.TpchTable.getTables; import static java.lang.String.format; import static java.util.stream.Collectors.joining; @@ -290,5 +299,101 @@ public void testPushdownSubfieldForMapFunctions() } } + @Test + public void testOrcUseColumnNames() throws IOException, URISyntaxException + { + File externalTableDataDirectory = Files.createTempDir(); + String externalTableDataLocationUri = externalTableDataDirectory.toURI().toString(); + + try { + @Language("SQL") String createManagedTableSql = "" + + "create table test_orc_use_column_names (\n" + + " \"c1\" int,\n" + + " \"c2\" varchar\n" + + ")\n" + + "WITH (\n" + + " format = 'ORC'\n" + + ")"; + + assertUpdate(createManagedTableSql); + assertUpdate(format("insert into test_orc_use_column_names values (1, 'one')"), 1); + String tablePath = (String) getHiveTableProperty(getQueryRunner(), getSession(), "test_orc_use_column_names", (HiveTableLayoutHandle table) -> table.getTablePath()); + File managedTableDataDirectory = new File(new URI(tablePath).getRawPath()); + + assertTrue(managedTableDataDirectory.isDirectory(), "Source managed table data directory does not exist: " + managedTableDataDirectory); + File[] orcFiles = managedTableDataDirectory.listFiles(file -> file.isFile() + && !file.getName().contains(".")); + + if (orcFiles != null) { + for (File orcFile : orcFiles) { + File destinationFile = new File(externalTableDataDirectory, orcFile.getName()); + Files.copy(orcFile, destinationFile); + } + } + else { + throw new IllegalStateException("No ORC files found in managed table data directory: " + managedTableDataDirectory); + } + + @Language("SQL") String createMisMatchingExternalTableSql = format("" + + "CREATE TABLE test_orc_use_column_names_mismatching_ext (\n" + + " \"c2\" varchar,\n" + + " \"c1\" int\n" + + ")\n" + + "WITH (\n" + + " external_location = '%s',\n" + + " format = 'ORC'\n" + + ")", + externalTableDataLocationUri); + assertUpdate(createMisMatchingExternalTableSql); + + assertQueryFails(getSession(), + "select * from test_orc_use_column_names_mismatching_ext", + ".*java.io.IOException: Malformed ORC file. Can not read SQL type varchar from ORC stream .c1 of type INT.*"); + + Session useOrcColumnNamesSession = Session.builder(getQueryRunner().getDefaultSession()) + .setCatalogSessionProperty("hive", ORC_USE_COLUMN_NAMES, "true").build(); + assertQuerySucceeds(useOrcColumnNamesSession, "select * from test_orc_use_column_names_mismatching_ext"); + + @Language("SQL") String createDifferentColumnNameExternalTableSql = format("" + + "CREATE TABLE test_orc_use_column_names_different_column_name_ext (\n" + + " \"c1\" int,\n" + + " \"c3\" varchar\n" + + ")\n" + + "WITH (\n" + + " external_location = '%s',\n" + + " format = 'ORC'\n" + + ")", + externalTableDataLocationUri); + assertUpdate(createDifferentColumnNameExternalTableSql); + + assertQuery(useOrcColumnNamesSession, "select * from test_orc_use_column_names_different_column_name_ext", "VALUES (1, NULL)"); + + @Language("SQL") String createAdditionalColumnExternalTableSql = format("" + + "CREATE TABLE test_orc_use_column_names_additional_column_ext (\n" + + " \"c1\" int,\n" + + " \"c2\" varchar,\n" + + " \"c3\" varchar\n" + + ")\n" + + "WITH (\n" + + " external_location = '%s',\n" + + " format = 'ORC'\n" + + ")", + externalTableDataLocationUri); + assertUpdate(createAdditionalColumnExternalTableSql); + + assertQuery(useOrcColumnNamesSession, "select * from test_orc_use_column_names_additional_column_ext", "VALUES (1, 'one', NULL)"); + } + finally { + assertUpdate("DROP TABLE IF EXISTS test_orc_use_column_names"); + assertUpdate("DROP TABLE IF EXISTS test_orc_use_column_names_mismatching_ext"); + assertUpdate("DROP TABLE IF EXISTS test_orc_use_column_names_different_column_name_ext"); + assertUpdate("DROP TABLE IF EXISTS test_orc_use_column_names_additional_column_ext"); + + if (externalTableDataDirectory != null) { + deleteRecursively(externalTableDataDirectory.toPath(), ALLOW_INSECURE); + } + } + } + // Hive specific tests should normally go in TestHiveIntegrationSmokeTest } diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveFileFormats.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveFileFormats.java index 8c117033a686b..f9fed4fdf2024 100644 --- a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveFileFormats.java +++ b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveFileFormats.java @@ -339,7 +339,7 @@ public void testOrc(int rowCount) assertThatFileFormat(ORC) .withColumns(TEST_COLUMNS) .withRowsCount(rowCount) - .isReadableByPageSource(new OrcBatchPageSourceFactory(FUNCTION_AND_TYPE_MANAGER, false, HDFS_ENVIRONMENT, STATS, 100, new StorageOrcFileTailSource(), StripeMetadataSourceFactory.of(new StorageStripeMetadataSource()))); + .isReadableByPageSource(new OrcBatchPageSourceFactory(FUNCTION_AND_TYPE_MANAGER, HDFS_ENVIRONMENT, STATS, 100, new StorageOrcFileTailSource(), StripeMetadataSourceFactory.of(new StorageStripeMetadataSource()))); } @Test(dataProvider = "rowCount") @@ -362,7 +362,7 @@ public void testOrcOptimizedWriter(int rowCount) .withSession(session) .withFileWriterFactory(new OrcFileWriterFactory(HDFS_ENVIRONMENT, new OutputStreamDataSinkFactory(), FUNCTION_AND_TYPE_MANAGER, new NodeVersion("test"), HIVE_STORAGE_TIME_ZONE, STATS, new OrcFileWriterConfig(), NO_ENCRYPTION)) .isReadableByRecordCursor(new GenericHiveRecordCursorProvider(HDFS_ENVIRONMENT)) - .isReadableByPageSource(new OrcBatchPageSourceFactory(FUNCTION_AND_TYPE_MANAGER, false, HDFS_ENVIRONMENT, STATS, 100, new StorageOrcFileTailSource(), StripeMetadataSourceFactory.of(new StorageStripeMetadataSource()))); + .isReadableByPageSource(new OrcBatchPageSourceFactory(FUNCTION_AND_TYPE_MANAGER, HDFS_ENVIRONMENT, STATS, 100, new StorageOrcFileTailSource(), StripeMetadataSourceFactory.of(new StorageStripeMetadataSource()))); } @Test(dataProvider = "rowCount") @@ -394,14 +394,15 @@ public void testOrcUseColumnNames(int rowCount) { TestingConnectorSession session = new TestingConnectorSession(getAllSessionProperties( new HiveClientConfig(), - new HiveCommonClientConfig())); + new HiveCommonClientConfig() + .setUseOrcColumnNames(true))); assertThatFileFormat(ORC) .withWriteColumns(TEST_COLUMNS) .withRowsCount(rowCount) .withReadColumns(Lists.reverse(TEST_COLUMNS)) .withSession(session) - .isReadableByPageSource(new OrcBatchPageSourceFactory(FUNCTION_AND_TYPE_MANAGER, true, HDFS_ENVIRONMENT, STATS, 100, new StorageOrcFileTailSource(), StripeMetadataSourceFactory.of(new StorageStripeMetadataSource()))); + .isReadableByPageSource(new OrcBatchPageSourceFactory(FUNCTION_AND_TYPE_MANAGER, HDFS_ENVIRONMENT, STATS, 100, new StorageOrcFileTailSource(), StripeMetadataSourceFactory.of(new StorageStripeMetadataSource()))); } @Test(dataProvider = "rowCount") @@ -419,7 +420,7 @@ public void testOrcUseColumnNamesCompatibility(int rowCount) .withRowsCount(rowCount) .withReadColumns(TEST_COLUMNS) .withSession(session) - .isReadableByPageSource(new OrcBatchPageSourceFactory(FUNCTION_AND_TYPE_MANAGER, true, HDFS_ENVIRONMENT, STATS, 100, new StorageOrcFileTailSource(), StripeMetadataSourceFactory.of(new StorageStripeMetadataSource()))); + .isReadableByPageSource(new OrcBatchPageSourceFactory(FUNCTION_AND_TYPE_MANAGER, HDFS_ENVIRONMENT, STATS, 100, new StorageOrcFileTailSource(), StripeMetadataSourceFactory.of(new StorageStripeMetadataSource()))); } private static List getHiveColumnNameColumns() @@ -628,7 +629,7 @@ public void testTruncateVarcharColumn() assertThatFileFormat(ORC) .withWriteColumns(ImmutableList.of(writeColumn)) .withReadColumns(ImmutableList.of(readColumn)) - .isReadableByPageSource(new OrcBatchPageSourceFactory(FUNCTION_AND_TYPE_MANAGER, false, HDFS_ENVIRONMENT, STATS, 100, new StorageOrcFileTailSource(), StripeMetadataSourceFactory.of(new StorageStripeMetadataSource()))); + .isReadableByPageSource(new OrcBatchPageSourceFactory(FUNCTION_AND_TYPE_MANAGER, HDFS_ENVIRONMENT, STATS, 100, new StorageOrcFileTailSource(), StripeMetadataSourceFactory.of(new StorageStripeMetadataSource()))); assertThatFileFormat(PARQUET) .withWriteColumns(ImmutableList.of(writeColumn)) @@ -676,7 +677,7 @@ public void testFailForLongVarcharPartitionColumn() assertThatFileFormat(ORC) .withColumns(columns) - .isFailingForPageSource(new OrcBatchPageSourceFactory(FUNCTION_AND_TYPE_MANAGER, false, HDFS_ENVIRONMENT, STATS, 100, new StorageOrcFileTailSource(), StripeMetadataSourceFactory.of(new StorageStripeMetadataSource())), expectedErrorCode, expectedMessage); + .isFailingForPageSource(new OrcBatchPageSourceFactory(FUNCTION_AND_TYPE_MANAGER, HDFS_ENVIRONMENT, STATS, 100, new StorageOrcFileTailSource(), StripeMetadataSourceFactory.of(new StorageStripeMetadataSource())), expectedErrorCode, expectedMessage); assertThatFileFormat(PARQUET) .withColumns(columns) diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveIntegrationSmokeTest.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveIntegrationSmokeTest.java index 1f765eab22008..51a0c7807eb22 100644 --- a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveIntegrationSmokeTest.java +++ b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveIntegrationSmokeTest.java @@ -23,12 +23,10 @@ import com.facebook.presto.hive.HiveClientConfig.InsertExistingPartitionsBehavior; import com.facebook.presto.metadata.InsertTableHandle; import com.facebook.presto.metadata.Metadata; -import com.facebook.presto.metadata.TableLayout; import com.facebook.presto.spi.CatalogSchemaTableName; import com.facebook.presto.spi.ColumnMetadata; import com.facebook.presto.spi.ConnectorId; import com.facebook.presto.spi.ConnectorSession; -import com.facebook.presto.spi.Constraint; import com.facebook.presto.spi.TableHandle; import com.facebook.presto.spi.TableMetadata; import com.facebook.presto.spi.plan.MarkDistinctNode; @@ -75,7 +73,6 @@ import java.util.Optional; import java.util.function.BiConsumer; import java.util.function.Consumer; -import java.util.function.Function; import java.util.stream.LongStream; import static com.facebook.airlift.json.JsonCodec.jsonCodec; @@ -130,6 +127,7 @@ import static com.facebook.presto.hive.HiveTableProperties.PARTITIONED_BY_PROPERTY; import static com.facebook.presto.hive.HiveTableProperties.STORAGE_FORMAT_PROPERTY; import static com.facebook.presto.hive.HiveTestUtils.FUNCTION_AND_TYPE_MANAGER; +import static com.facebook.presto.hive.HiveTestUtils.getHiveTableProperty; import static com.facebook.presto.hive.HiveUtil.columnExtraInfo; import static com.facebook.presto.spi.security.SelectedRole.Type.ROLE; import static com.facebook.presto.sql.analyzer.FeaturesConfig.JoinDistributionType.BROADCAST; @@ -2083,31 +2081,14 @@ private TableMetadata getTableMetadata(String catalog, String schema, String tab }); } - private Object getHiveTableProperty(String tableName, Function propertyGetter) - { - Session session = getSession(); - Metadata metadata = ((DistributedQueryRunner) getQueryRunner()).getCoordinator().getMetadata(); - - return transaction(getQueryRunner().getTransactionManager(), getQueryRunner().getAccessControl()) - .readOnly() - .execute(session, transactionSession -> { - Optional tableHandle = metadata.getMetadataResolver(transactionSession).getTableHandle(new QualifiedObjectName(catalog, TPCH_SCHEMA, tableName)); - assertTrue(tableHandle.isPresent()); - - TableLayout layout = metadata.getLayout(transactionSession, tableHandle.get(), Constraint.alwaysTrue(), Optional.empty()) - .getLayout(); - return propertyGetter.apply((HiveTableLayoutHandle) layout.getNewTableHandle().getLayout().get()); - }); - } - private List getPartitions(String tableName) { - return (List) getHiveTableProperty(tableName, (HiveTableLayoutHandle table) -> getPartitions(table)); + return (List) getHiveTableProperty(getQueryRunner(), getSession(), tableName, (HiveTableLayoutHandle table) -> getPartitions(table)); } private int getBucketCount(String tableName) { - return (int) getHiveTableProperty(tableName, (HiveTableLayoutHandle table) -> table.getBucketHandle().get().getTableBucketCount()); + return (int) getHiveTableProperty(getQueryRunner(), getSession(), tableName, (HiveTableLayoutHandle table) -> table.getBucketHandle().get().getTableBucketCount()); } @Test diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestOrcBatchPageSourceMemoryTracking.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestOrcBatchPageSourceMemoryTracking.java index 42509b0089f7b..6144871f10040 100644 --- a/presto-hive/src/test/java/com/facebook/presto/hive/TestOrcBatchPageSourceMemoryTracking.java +++ b/presto-hive/src/test/java/com/facebook/presto/hive/TestOrcBatchPageSourceMemoryTracking.java @@ -456,7 +456,6 @@ public ConnectorPageSource newPageSource(FileFormatDataSourceStats stats, Connec OrcBatchPageSourceFactory orcPageSourceFactory = new OrcBatchPageSourceFactory( FUNCTION_AND_TYPE_MANAGER, - false, HDFS_ENVIRONMENT, stats, 100, diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/benchmark/FileFormat.java b/presto-hive/src/test/java/com/facebook/presto/hive/benchmark/FileFormat.java index 11b67fea8b68e..eea2d2a876fd3 100644 --- a/presto-hive/src/test/java/com/facebook/presto/hive/benchmark/FileFormat.java +++ b/presto-hive/src/test/java/com/facebook/presto/hive/benchmark/FileFormat.java @@ -161,7 +161,6 @@ public ConnectorPageSource createFileFormatReader(ConnectorSession session, Hdfs { HiveBatchPageSourceFactory pageSourceFactory = new OrcBatchPageSourceFactory( FUNCTION_AND_TYPE_MANAGER, - false, hdfsEnvironment, new FileFormatDataSourceStats(), 100, From b016faf6cd8c1d3e5a75729d9eddb1ccc08ba209 Mon Sep 17 00:00:00 2001 From: wangd Date: Tue, 26 Aug 2025 11:14:02 +0800 Subject: [PATCH 004/113] Refactor test framework to return transactionId control flags to client The test framework client now receives statement executing results with `clearTransactionId` and `startTransactionId` flags embedded. --- .../elasticsearch/ElasticsearchLoader.java | 2 +- .../TestIcebergDistributedQueries.java | 6 +++ .../presto/kafka/util/KafkaLoader.java | 2 +- .../presto/testing/MaterializedResult.java | 29 ++++++++++++-- .../AbstractTestNativeGeneralQueries.java | 2 +- .../sidecar/TestNativeSidecarPlugin.java | 2 +- .../presto/redis/util/RedisLoader.java | 2 +- .../TestSingleStoreDistributedQueries.java | 6 +++ .../presto/spark/PrestoSparkQueryRunner.java | 6 +++ .../tests/AbstractTestDistributedQueries.java | 38 ++++++++++++++++++- .../presto/tests/AbstractTestQueries.java | 3 +- .../tests/AbstractTestingPrestoClient.java | 2 +- .../facebook/presto/tests/ResultsSession.java | 3 +- .../presto/tests/TestingPrestoClient.java | 6 ++- ...nPropertiesExcludingInvalidProperties.java | 2 +- ...nPropertiesIncludingInvalidProperties.java | 4 +- 16 files changed, 99 insertions(+), 16 deletions(-) diff --git a/presto-elasticsearch/src/test/java/com/facebook/presto/elasticsearch/ElasticsearchLoader.java b/presto-elasticsearch/src/test/java/com/facebook/presto/elasticsearch/ElasticsearchLoader.java index 5a0e736c70c44..3c7d62d7a7567 100644 --- a/presto-elasticsearch/src/test/java/com/facebook/presto/elasticsearch/ElasticsearchLoader.java +++ b/presto-elasticsearch/src/test/java/com/facebook/presto/elasticsearch/ElasticsearchLoader.java @@ -119,7 +119,7 @@ public void addResults(QueryStatusInfo statusInfo, QueryData data) } @Override - public Void build(Map setSessionProperties, Set resetSessionProperties) + public Void build(Map setSessionProperties, Set resetSessionProperties, String startTransactionId, boolean clearTransactionId) { return null; } diff --git a/presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestIcebergDistributedQueries.java b/presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestIcebergDistributedQueries.java index c689545433980..e45c730bddd73 100644 --- a/presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestIcebergDistributedQueries.java +++ b/presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestIcebergDistributedQueries.java @@ -115,6 +115,12 @@ public void testDescribeOutputNamedAndUnnamed() assertEqualsIgnoreOrder(actual, expected); } + @Override + public void testClearTransactionId() + { + // Catalog iceberg only supports writes using autocommit + } + /** * Increased the optimizer timeout from 15000ms to 25000ms */ diff --git a/presto-kafka/src/test/java/com/facebook/presto/kafka/util/KafkaLoader.java b/presto-kafka/src/test/java/com/facebook/presto/kafka/util/KafkaLoader.java index 901baa5b23465..ddb31c6bd19cb 100644 --- a/presto-kafka/src/test/java/com/facebook/presto/kafka/util/KafkaLoader.java +++ b/presto-kafka/src/test/java/com/facebook/presto/kafka/util/KafkaLoader.java @@ -120,7 +120,7 @@ public void addResults(QueryStatusInfo statusInfo, QueryData data) } @Override - public Void build(Map setSessionProperties, Set resetSessionProperties) + public Void build(Map setSessionProperties, Set resetSessionProperties, String startTransactionId, boolean clearTransactionId) { return null; } diff --git a/presto-main-base/src/main/java/com/facebook/presto/testing/MaterializedResult.java b/presto-main-base/src/main/java/com/facebook/presto/testing/MaterializedResult.java index d46e595a537d4..848bc0de0ec0a 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/testing/MaterializedResult.java +++ b/presto-main-base/src/main/java/com/facebook/presto/testing/MaterializedResult.java @@ -18,6 +18,7 @@ import com.facebook.presto.common.PageBuilder; import com.facebook.presto.common.block.Block; import com.facebook.presto.common.block.BlockBuilder; +import com.facebook.presto.common.transaction.TransactionId; import com.facebook.presto.common.type.ArrayType; import com.facebook.presto.common.type.CharType; import com.facebook.presto.common.type.MapType; @@ -99,11 +100,13 @@ public class MaterializedResult private final Set resetSessionProperties; private final Optional updateType; private final OptionalLong updateCount; + private final Optional startedTransactionId; + private final boolean clearTransactionId; private final List warnings; public MaterializedResult(List rows, List types) { - this(rows, types, ImmutableMap.of(), ImmutableSet.of(), Optional.empty(), OptionalLong.empty(), ImmutableList.of()); + this(rows, types, ImmutableMap.of(), ImmutableSet.of(), Optional.empty(), OptionalLong.empty(), Optional.empty(), false, ImmutableList.of()); } public MaterializedResult( @@ -113,6 +116,8 @@ public MaterializedResult( Set resetSessionProperties, Optional updateType, OptionalLong updateCount, + Optional startedTransactionId, + boolean clearTransactionId, List warnings) { this.rows = ImmutableList.copyOf(requireNonNull(rows, "rows is null")); @@ -121,6 +126,8 @@ public MaterializedResult( this.resetSessionProperties = ImmutableSet.copyOf(requireNonNull(resetSessionProperties, "resetSessionProperties is null")); this.updateType = requireNonNull(updateType, "updateType is null"); this.updateCount = requireNonNull(updateCount, "updateCount is null"); + this.startedTransactionId = requireNonNull(startedTransactionId, "startedTransactionId is null"); + this.clearTransactionId = clearTransactionId; this.warnings = requireNonNull(warnings, "warnings is null"); } @@ -165,6 +172,16 @@ public OptionalLong getUpdateCount() return updateCount; } + public Optional getStartedTransactionId() + { + return startedTransactionId; + } + + public boolean isClearTransactionId() + { + return clearTransactionId; + } + public List getWarnings() { return warnings; @@ -185,13 +202,15 @@ public boolean equals(Object obj) Objects.equals(setSessionProperties, o.setSessionProperties) && Objects.equals(resetSessionProperties, o.resetSessionProperties) && Objects.equals(updateType, o.updateType) && - Objects.equals(updateCount, o.updateCount); + Objects.equals(updateCount, o.updateCount) && + Objects.equals(startedTransactionId, o.startedTransactionId) && + Objects.equals(clearTransactionId, o.clearTransactionId); } @Override public int hashCode() { - return Objects.hash(rows, types, setSessionProperties, resetSessionProperties, updateType, updateCount); + return Objects.hash(rows, types, setSessionProperties, resetSessionProperties, updateType, updateCount, startedTransactionId, clearTransactionId); } @Override @@ -204,6 +223,8 @@ public String toString() .add("resetSessionProperties", resetSessionProperties) .add("updateType", updateType.orElse(null)) .add("updateCount", updateCount.isPresent() ? updateCount.getAsLong() : null) + .add("startedTransactionId", startedTransactionId.orElse(null)) + .add("clearTransactionId", clearTransactionId) .omitNullValues() .toString(); } @@ -360,6 +381,8 @@ public MaterializedResult toTestTypes() resetSessionProperties, updateType, updateCount, + startedTransactionId, + clearTransactionId, warnings); } diff --git a/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/AbstractTestNativeGeneralQueries.java b/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/AbstractTestNativeGeneralQueries.java index e57585fd06eee..1549ccddbaaef 100644 --- a/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/AbstractTestNativeGeneralQueries.java +++ b/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/AbstractTestNativeGeneralQueries.java @@ -1211,7 +1211,7 @@ public void testSetSessionJavaWorkerSessionProperty() "MaterializedResult{rows=[[true]], " + "types=[boolean], " + "setSessionProperties={distinct_aggregation_spill_enabled=false}, " + - "resetSessionProperties=[], updateType=SET SESSION}"); + "resetSessionProperties=[], updateType=SET SESSION, clearTransactionId=false}"); } @Test diff --git a/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java b/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java index 05945f225456e..50bb5c252ffd2 100644 --- a/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java +++ b/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java @@ -155,7 +155,7 @@ public void testSetNativeWorkerSessionProperty() "MaterializedResult{rows=[[true]], " + "types=[boolean], " + "setSessionProperties={driver_cpu_time_slice_limit_ms=500}, " + - "resetSessionProperties=[], updateType=SET SESSION}"); + "resetSessionProperties=[], updateType=SET SESSION, clearTransactionId=false}"); } @Test diff --git a/presto-redis/src/test/java/com/facebook/presto/redis/util/RedisLoader.java b/presto-redis/src/test/java/com/facebook/presto/redis/util/RedisLoader.java index b33f6e14fcf7f..52b51368db4e3 100644 --- a/presto-redis/src/test/java/com/facebook/presto/redis/util/RedisLoader.java +++ b/presto-redis/src/test/java/com/facebook/presto/redis/util/RedisLoader.java @@ -143,7 +143,7 @@ public void addResults(QueryStatusInfo statusInfo, QueryData data) } @Override - public Void build(Map setSessionProperties, Set resetSessionProperties) + public Void build(Map setSessionProperties, Set resetSessionProperties, String startTransactionId, boolean clearTransactionId) { return null; } diff --git a/presto-singlestore/src/test/java/com/facebook/presto/plugin/singlestore/TestSingleStoreDistributedQueries.java b/presto-singlestore/src/test/java/com/facebook/presto/plugin/singlestore/TestSingleStoreDistributedQueries.java index 7a69686ebd915..93f453e13f3eb 100644 --- a/presto-singlestore/src/test/java/com/facebook/presto/plugin/singlestore/TestSingleStoreDistributedQueries.java +++ b/presto-singlestore/src/test/java/com/facebook/presto/plugin/singlestore/TestSingleStoreDistributedQueries.java @@ -183,6 +183,12 @@ public void testDescribeOutputNamedAndUnnamed() // this connector uses a non-canonical type for varchar columns in tpch } + @Override + public void testClearTransactionId() + { + // Catalog singlestore only supports writes using autocommit + } + @Override public void testInsert() { diff --git a/presto-spark-base/src/test/java/com/facebook/presto/spark/PrestoSparkQueryRunner.java b/presto-spark-base/src/test/java/com/facebook/presto/spark/PrestoSparkQueryRunner.java index b31117343954b..bd8f8c006c9d7 100644 --- a/presto-spark-base/src/test/java/com/facebook/presto/spark/PrestoSparkQueryRunner.java +++ b/presto-spark-base/src/test/java/com/facebook/presto/spark/PrestoSparkQueryRunner.java @@ -594,6 +594,8 @@ private MaterializedResult executeWithStrategies( ImmutableSet.of(), p.getUpdateType(), getOnlyElement(getOnlyElement(rows).getFields()) == null ? OptionalLong.empty() : OptionalLong.of((Long) getOnlyElement(getOnlyElement(rows).getFields())), + Optional.empty(), + false, ImmutableList.of()); } } @@ -606,6 +608,8 @@ else if (execution instanceof PrestoSparkAccessControlCheckerExecution) { ImmutableSet.of(), Optional.empty(), OptionalLong.empty(), + Optional.empty(), + false, ImmutableList.of()); } else { @@ -616,6 +620,8 @@ else if (execution instanceof PrestoSparkAccessControlCheckerExecution) { ImmutableSet.of(), Optional.empty(), OptionalLong.empty(), + Optional.empty(), + false, ImmutableList.of()); } } diff --git a/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestDistributedQueries.java b/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestDistributedQueries.java index aaeebab0b6bf1..4989a64ac92d0 100644 --- a/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestDistributedQueries.java +++ b/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestDistributedQueries.java @@ -249,6 +249,41 @@ public void testCreateTable() assertFalse(getQueryRunner().tableExists(getSession(), "test_create_like")); } + @Test + public void testClearTransactionId() + { + assertUpdate("create table test_clear_transaction_id_table(a int, b varchar)"); + assertUpdate("insert into test_clear_transaction_id_table values(1, '1001'), (2, '1002')", 2); + Session session = getQueryRunner().getDefaultSession(); + String defaultCatalog = session.getCatalog().get(); + transaction(getQueryRunner().getTransactionManager(), getQueryRunner().getAccessControl()) + .execute(Session.builder(session) + .setIdentity(new Identity("admin", + Optional.empty(), + ImmutableMap.of(defaultCatalog, new SelectedRole(ROLE, Optional.of("admin"))), + ImmutableMap.of(), + ImmutableMap.of(), + Optional.empty(), + Optional.empty())) + .build(), + txnSession -> { + MaterializedResult result = computeActual(txnSession, "select * from test_clear_transaction_id_table"); + assertEquals(result.getRowCount(), 2); + assertFalse(result.isClearTransactionId()); + + result = computeActual(txnSession, "insert into test_clear_transaction_id_table values(1, '1001'), (2, '1002')"); + assertEquals(result.getOnlyValue(), 2L); + assertFalse(result.isClearTransactionId()); + + // `Rollback` executes successfully, and the client gets a flag `clearTransactionId = true` + result = computeActual(txnSession, "rollback"); + assertTrue(result.isClearTransactionId()); + }); + + assertQuery("select * from test_clear_transaction_id_table", "values(1, '1001'), (2, '1002')"); + assertUpdate("drop table if exists test_clear_transaction_id_table"); + } + @Test public void testNonAutoCommitTransactionWithFailAndRollback() { @@ -276,7 +311,8 @@ public void testNonAutoCommitTransactionWithFailAndRollback() assertQueryFails(txnSession, "create table test_table(a int, b varchar)", "Current transaction is aborted, commands ignored until end of transaction block"); // execute `rollback` successfully - assertUpdate(txnSession, "rollback"); + MaterializedResult result = computeActual(txnSession, "rollback"); + assertTrue(result.isClearTransactionId()); }); assertQuery("select count(*) from test_non_autocommit_table", "values(0)"); diff --git a/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueries.java b/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueries.java index b35829e2760a8..52774715e6bdd 100644 --- a/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueries.java +++ b/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueries.java @@ -3218,7 +3218,8 @@ public void testSetSessionNativeWorkerSessionProperty() "MaterializedResult{rows=[[true]], " + "types=[boolean], " + "setSessionProperties={native_expression_max_array_size_in_reduce=50000}, " + - "resetSessionProperties=[], updateType=SET SESSION}"); + "resetSessionProperties=[], updateType=SET SESSION, " + + "clearTransactionId=false}"); } @Test diff --git a/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestingPrestoClient.java b/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestingPrestoClient.java index 9eb4d3eda142b..96d812c802e7f 100644 --- a/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestingPrestoClient.java +++ b/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestingPrestoClient.java @@ -117,7 +117,7 @@ public ResultWithQueryId execute(Session session, @Language("SQL") String sql resultsSession.setWarnings(results.getWarnings()); - T result = resultsSession.build(client.getSetSessionProperties(), client.getResetSessionProperties()); + T result = resultsSession.build(client.getSetSessionProperties(), client.getResetSessionProperties(), client.getStartedTransactionId(), client.isClearTransactionId()); return new ResultWithQueryId<>(new QueryId(results.getId()), result); } diff --git a/presto-tests/src/main/java/com/facebook/presto/tests/ResultsSession.java b/presto-tests/src/main/java/com/facebook/presto/tests/ResultsSession.java index e37f28ba2dc09..2be0ce2629c0e 100644 --- a/presto-tests/src/main/java/com/facebook/presto/tests/ResultsSession.java +++ b/presto-tests/src/main/java/com/facebook/presto/tests/ResultsSession.java @@ -16,6 +16,7 @@ import com.facebook.presto.client.QueryData; import com.facebook.presto.client.QueryStatusInfo; import com.facebook.presto.spi.PrestoWarning; +import jakarta.annotation.Nullable; import java.util.List; import java.util.Map; @@ -40,5 +41,5 @@ default void setWarnings(List warnings) void addResults(QueryStatusInfo statusInfo, QueryData data); - T build(Map setSessionProperties, Set resetSessionProperties); + T build(Map setSessionProperties, Set resetSessionProperties, @Nullable String startTransactionId, boolean clearTransactionId); } diff --git a/presto-tests/src/main/java/com/facebook/presto/tests/TestingPrestoClient.java b/presto-tests/src/main/java/com/facebook/presto/tests/TestingPrestoClient.java index 475143e1e543b..26ddac7e426e5 100644 --- a/presto-tests/src/main/java/com/facebook/presto/tests/TestingPrestoClient.java +++ b/presto-tests/src/main/java/com/facebook/presto/tests/TestingPrestoClient.java @@ -18,6 +18,7 @@ import com.facebook.presto.client.IntervalYearMonth; import com.facebook.presto.client.QueryData; import com.facebook.presto.client.QueryStatusInfo; +import com.facebook.presto.common.transaction.TransactionId; import com.facebook.presto.common.type.ArrayType; import com.facebook.presto.common.type.BigintEnumType; import com.facebook.presto.common.type.DecimalType; @@ -41,6 +42,7 @@ import com.facebook.presto.type.SqlIntervalYearMonth; import com.google.common.base.Function; import com.google.common.collect.ImmutableList; +import jakarta.annotation.Nullable; import java.math.BigDecimal; import java.time.LocalDate; @@ -147,7 +149,7 @@ public void addResults(QueryStatusInfo statusInfo, QueryData data) } @Override - public MaterializedResult build(Map setSessionProperties, Set resetSessionProperties) + public MaterializedResult build(Map setSessionProperties, Set resetSessionProperties, @Nullable String startTransactionId, boolean clearTransactionId) { checkState(types.get() != null, "never received types for the query"); return new MaterializedResult( @@ -157,6 +159,8 @@ public MaterializedResult build(Map setSessionProperties, Set Date: Mon, 25 Aug 2025 23:12:23 +0800 Subject: [PATCH 005/113] Enhance test framework to handle start/rollback/commit in a nature way --- .../TestAccumuloDistributedQueries.java | 6 + .../TestIcebergDistributedQueries.java | 8 +- .../java/com/facebook/presto/Session.java | 45 +++++++ .../TestSingleStoreDistributedQueries.java | 8 +- .../tests/AbstractTestDistributedQueries.java | 119 +++++++++--------- .../tests/AbstractTestQueryFramework.java | 10 ++ .../presto/tests/QueryAssertions.java | 34 +++++ 7 files changed, 167 insertions(+), 63 deletions(-) diff --git a/presto-accumulo/src/test/java/com/facebook/presto/accumulo/TestAccumuloDistributedQueries.java b/presto-accumulo/src/test/java/com/facebook/presto/accumulo/TestAccumuloDistributedQueries.java index a37560fae7a7a..957b10f793dad 100644 --- a/presto-accumulo/src/test/java/com/facebook/presto/accumulo/TestAccumuloDistributedQueries.java +++ b/presto-accumulo/src/test/java/com/facebook/presto/accumulo/TestAccumuloDistributedQueries.java @@ -121,6 +121,12 @@ public void testUpdate() // Updates are not supported by the connector } + @Override + public void testNonAutoCommitTransactionWithRollback() + { + // This connector do not support rollback for insert actions + } + @Override public void testInsert() { diff --git a/presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestIcebergDistributedQueries.java b/presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestIcebergDistributedQueries.java index e45c730bddd73..57dd62d142e3a 100644 --- a/presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestIcebergDistributedQueries.java +++ b/presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestIcebergDistributedQueries.java @@ -116,7 +116,13 @@ public void testDescribeOutputNamedAndUnnamed() } @Override - public void testClearTransactionId() + public void testNonAutoCommitTransactionWithRollback() + { + // Catalog iceberg only supports writes using autocommit + } + + @Override + public void testNonAutoCommitTransactionWithCommit() { // Catalog iceberg only supports writes using autocommit } diff --git a/presto-main-base/src/main/java/com/facebook/presto/Session.java b/presto-main-base/src/main/java/com/facebook/presto/Session.java index 327924fd9522f..05e01e17e5355 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/Session.java +++ b/presto-main-base/src/main/java/com/facebook/presto/Session.java @@ -42,6 +42,7 @@ import com.facebook.presto.sql.planner.optimizations.OptimizerInformationCollector; import com.facebook.presto.sql.planner.optimizations.OptimizerResultCollector; import com.facebook.presto.transaction.TransactionManager; +import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import com.google.common.collect.Maps; @@ -470,6 +471,50 @@ public Session beginTransactionId(TransactionId transactionId, boolean enableRol queryType); } + @VisibleForTesting + public Session clearTransaction(TransactionManager transactionManager, AccessControl accessControl) + { + checkArgument(this.transactionId.isPresent(), "Session does not have an active transaction"); + requireNonNull(transactionManager, "transactionManager is null"); + requireNonNull(accessControl, "accessControl is null"); + + for (Entry property : systemProperties.entrySet()) { + // verify permissions + accessControl.checkCanSetSystemSessionProperty(identity, context, property.getKey()); + + // validate session property value + sessionPropertyManager.validateSystemSessionProperty(property.getKey(), property.getValue()); + } + + return new Session( + queryId, + Optional.empty(), + clientTransactionSupport, + identity, + source, + catalog, + schema, + traceToken, + timeZoneKey, + locale, + remoteUserAddress, + userAgent, + clientInfo, + clientTags, + resourceEstimates, + startTime, + systemProperties, + connectorProperties, + unprocessedCatalogProperties, + sessionPropertyManager, + preparedStatements, + sessionFunctions, + tracer, + warningCollector, + runtimeStats, + queryType); + } + public ConnectorSession toConnectorSession() { return new FullConnectorSession(this, identity.toConnectorIdentity()); diff --git a/presto-singlestore/src/test/java/com/facebook/presto/plugin/singlestore/TestSingleStoreDistributedQueries.java b/presto-singlestore/src/test/java/com/facebook/presto/plugin/singlestore/TestSingleStoreDistributedQueries.java index 93f453e13f3eb..dbe8acb45a5bb 100644 --- a/presto-singlestore/src/test/java/com/facebook/presto/plugin/singlestore/TestSingleStoreDistributedQueries.java +++ b/presto-singlestore/src/test/java/com/facebook/presto/plugin/singlestore/TestSingleStoreDistributedQueries.java @@ -184,7 +184,13 @@ public void testDescribeOutputNamedAndUnnamed() } @Override - public void testClearTransactionId() + public void testNonAutoCommitTransactionWithRollback() + { + // Catalog singlestore only supports writes using autocommit + } + + @Override + public void testNonAutoCommitTransactionWithCommit() { // Catalog singlestore only supports writes using autocommit } diff --git a/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestDistributedQueries.java b/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestDistributedQueries.java index 4989a64ac92d0..48ba3ecfc11fa 100644 --- a/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestDistributedQueries.java +++ b/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestDistributedQueries.java @@ -77,7 +77,6 @@ import static com.facebook.presto.testing.TestingSession.TESTING_CATALOG; import static com.facebook.presto.testing.assertions.Assert.assertEquals; import static com.facebook.presto.tests.QueryAssertions.assertContains; -import static com.facebook.presto.transaction.TransactionBuilder.transaction; import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.collect.Iterables.getOnlyElement; import static com.google.common.util.concurrent.Uninterruptibles.sleepUninterruptibly; @@ -250,73 +249,71 @@ public void testCreateTable() } @Test - public void testClearTransactionId() + public void testNonAutoCommitTransactionWithRollback() { - assertUpdate("create table test_clear_transaction_id_table(a int, b varchar)"); - assertUpdate("insert into test_clear_transaction_id_table values(1, '1001'), (2, '1002')", 2); - Session session = getQueryRunner().getDefaultSession(); - String defaultCatalog = session.getCatalog().get(); - transaction(getQueryRunner().getTransactionManager(), getQueryRunner().getAccessControl()) - .execute(Session.builder(session) - .setIdentity(new Identity("admin", - Optional.empty(), - ImmutableMap.of(defaultCatalog, new SelectedRole(ROLE, Optional.of("admin"))), - ImmutableMap.of(), - ImmutableMap.of(), - Optional.empty(), - Optional.empty())) - .build(), - txnSession -> { - MaterializedResult result = computeActual(txnSession, "select * from test_clear_transaction_id_table"); - assertEquals(result.getRowCount(), 2); - assertFalse(result.isClearTransactionId()); - - result = computeActual(txnSession, "insert into test_clear_transaction_id_table values(1, '1001'), (2, '1002')"); - assertEquals(result.getOnlyValue(), 2L); - assertFalse(result.isClearTransactionId()); - - // `Rollback` executes successfully, and the client gets a flag `clearTransactionId = true` - result = computeActual(txnSession, "rollback"); - assertTrue(result.isClearTransactionId()); - }); - - assertQuery("select * from test_clear_transaction_id_table", "values(1, '1001'), (2, '1002')"); - assertUpdate("drop table if exists test_clear_transaction_id_table"); + assertUpdate("create table multi_statements_transaction_rollback(a int, b varchar)"); + assertUpdate("insert into multi_statements_transaction_rollback values(1, '1001'), (2, '1002')", 2); + + Session session = assertStartTransaction(getSession(), "start transaction"); + + assertQuery(session, "select * from multi_statements_transaction_rollback", "values(1, '1001'), (2, '1002')"); + assertUpdate(session, "insert into multi_statements_transaction_rollback values(3, '1003'), (4, '1004')", 2); + + // `Rollback` executes successfully, and the client gets a flag `clearTransactionId = true` + session = assertEndTransaction(session, "rollback"); + + assertQuery(session, "select * from multi_statements_transaction_rollback", "values(1, '1001'), (2, '1002')"); + assertUpdate(session, "drop table if exists multi_statements_transaction_rollback"); + } + + @Test + public void testNonAutoCommitTransactionWithCommit() + { + assertUpdate("create table multi_statements_transaction_commit(a int, b varchar)"); + assertUpdate("insert into multi_statements_transaction_commit values(1, '1001'), (2, '1002')", 2); + Session session = assertStartTransaction(getSession(), "start transaction"); + + assertQuery(session, "select * from multi_statements_transaction_commit", "values(1, '1001'), (2, '1002')"); + assertUpdate(session, "insert into multi_statements_transaction_commit values(3, '1003'), (4, '1004')", 2); + + // `Commit` executes successfully, and the client gets a flag `clearTransactionId = true` + session = assertEndTransaction(session, "commit"); + + assertQuery(session, "select * from multi_statements_transaction_commit", + "values(1, '1001'), (2, '1002'), (3, '1003'), (4, '1004')"); + assertUpdate(session, "drop table if exists multi_statements_transaction_commit"); } @Test public void testNonAutoCommitTransactionWithFailAndRollback() { assertUpdate("create table test_non_autocommit_table(a int, b varchar)"); - Session session = getQueryRunner().getDefaultSession(); - String defaultCatalog = session.getCatalog().get(); - transaction(getQueryRunner().getTransactionManager(), getQueryRunner().getAccessControl()) - .execute(Session.builder(session) - .setIdentity(new Identity("admin", - Optional.empty(), - ImmutableMap.of(defaultCatalog, new SelectedRole(ROLE, Optional.of("admin"))), - ImmutableMap.of(), - ImmutableMap.of(), - Optional.empty(), - Optional.empty())) - .build(), - txnSession -> { - // simulate failure of SQL statement execution - assertQueryFails(txnSession, "SELECT fail('forced failure')", "forced failure"); - - // cannot execute any SQLs except `rollback` in current session - assertQueryFails(txnSession, "select count(*) from test_non_autocommit_table", "Current transaction is aborted, commands ignored until end of transaction block"); - assertQueryFails(txnSession, "show tables", "Current transaction is aborted, commands ignored until end of transaction block"); - assertQueryFails(txnSession, "insert into test_non_autocommit_table values(1, '1001')", "Current transaction is aborted, commands ignored until end of transaction block"); - assertQueryFails(txnSession, "create table test_table(a int, b varchar)", "Current transaction is aborted, commands ignored until end of transaction block"); - - // execute `rollback` successfully - MaterializedResult result = computeActual(txnSession, "rollback"); - assertTrue(result.isClearTransactionId()); - }); - - assertQuery("select count(*) from test_non_autocommit_table", "values(0)"); - assertUpdate("drop table if exists test_non_autocommit_table"); + Session session = Session.builder(getSession()) + .setIdentity(new Identity("admin", + Optional.empty(), + ImmutableMap.of(getSession().getCatalog().get(), new SelectedRole(ROLE, Optional.of("admin"))), + ImmutableMap.of(), + ImmutableMap.of(), + Optional.empty(), + Optional.empty())) + .build(); + + session = assertStartTransaction(session, "start transaction"); + + // simulate failure of SQL statement execution + assertQueryFails(session, "SELECT fail('forced failure')", "forced failure"); + + // cannot execute any SQLs except `rollback` in current session + assertQueryFails(session, "select count(*) from test_non_autocommit_table", "Current transaction is aborted, commands ignored until end of transaction block"); + assertQueryFails(session, "show tables", "Current transaction is aborted, commands ignored until end of transaction block"); + assertQueryFails(session, "insert into test_non_autocommit_table values(1, '1001')", "Current transaction is aborted, commands ignored until end of transaction block"); + assertQueryFails(session, "create table test_table(a int, b varchar)", "Current transaction is aborted, commands ignored until end of transaction block"); + + // execute `rollback` successfully + session = assertEndTransaction(session, "rollback"); + + assertQuery(session, "select count(*) from test_non_autocommit_table", "values(0)"); + assertUpdate(session, "drop table if exists test_non_autocommit_table"); } @Test diff --git a/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueryFramework.java b/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueryFramework.java index 72999fbe97c16..4e6ac94dddc03 100644 --- a/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueryFramework.java +++ b/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueryFramework.java @@ -268,6 +268,16 @@ protected void assertUpdate(@Language("SQL") String sql) assertUpdate(getSession(), sql); } + protected Session assertStartTransaction(Session session, @Language("SQL") String sql) + { + return QueryAssertions.assertStartTransaction(queryRunner, session, sql); + } + + protected Session assertEndTransaction(Session session, @Language("SQL") String sql) + { + return QueryAssertions.assertEndTransaction(queryRunner, session, sql); + } + protected void assertUpdate(Session session, @Language("SQL") String sql) { QueryAssertions.assertUpdate(queryRunner, session, sql, OptionalLong.empty(), Optional.empty()); diff --git a/presto-tests/src/main/java/com/facebook/presto/tests/QueryAssertions.java b/presto-tests/src/main/java/com/facebook/presto/tests/QueryAssertions.java index 508b0694f0fb3..3da7b6e5b6320 100644 --- a/presto-tests/src/main/java/com/facebook/presto/tests/QueryAssertions.java +++ b/presto-tests/src/main/java/com/facebook/presto/tests/QueryAssertions.java @@ -17,6 +17,7 @@ import com.facebook.airlift.units.Duration; import com.facebook.presto.Session; import com.facebook.presto.common.QualifiedObjectName; +import com.facebook.presto.common.transaction.TransactionId; import com.facebook.presto.spi.WarningCollector; import com.facebook.presto.sql.planner.Plan; import com.facebook.presto.testing.ExpectedQueryRunner; @@ -46,7 +47,9 @@ import static java.util.concurrent.TimeUnit.MILLISECONDS; import static java.util.concurrent.TimeUnit.SECONDS; import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertTrue; import static org.testng.Assert.fail; public final class QueryAssertions @@ -57,6 +60,37 @@ private QueryAssertions() { } + public static Session assertStartTransaction(QueryRunner queryRunner, Session session, @Language("SQL") String sql) + { + MaterializedResult results = queryRunner.execute(session, sql); + if (!results.getUpdateType().isPresent()) { + fail("update type is not set"); + } + if (!results.getUpdateType().get().equals("START TRANSACTION")) { + fail("not a start transaction statement"); + } + assertTrue(results.getStartedTransactionId().isPresent()); + assertFalse(results.isClearTransactionId()); + + TransactionId transactionId = results.getStartedTransactionId().get(); + return session.beginTransactionId(transactionId, queryRunner.getTransactionManager(), queryRunner.getAccessControl()); + } + + public static Session assertEndTransaction(QueryRunner queryRunner, Session session, @Language("SQL") String sql) + { + MaterializedResult results = queryRunner.execute(session, sql); + if (!results.getUpdateType().isPresent()) { + fail("update type is not set"); + } + if (!results.getUpdateType().get().equals("ROLLBACK") && !results.getUpdateType().get().equals("COMMIT")) { + fail("not a end transaction statement"); + } + assertTrue(results.isClearTransactionId()); + assertFalse(results.getStartedTransactionId().isPresent()); + + return session.clearTransaction(queryRunner.getTransactionManager(), queryRunner.getAccessControl()); + } + public static void assertUpdate(QueryRunner queryRunner, Session session, @Language("SQL") String sql, OptionalLong count, Optional> planAssertion) { long start = System.nanoTime(); From fd534ddf3b37598c2ca20847fa9ea3b97bed7084 Mon Sep 17 00:00:00 2001 From: Nikhil Collooru Date: Tue, 26 Aug 2025 14:51:27 -0700 Subject: [PATCH 006/113] Fix bug in http message body extraction --- presto-native-execution/presto_cpp/main/TaskResource.cpp | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/presto-native-execution/presto_cpp/main/TaskResource.cpp b/presto-native-execution/presto_cpp/main/TaskResource.cpp index 863ad0494925b..ec08cf19bcf02 100644 --- a/presto-native-execution/presto_cpp/main/TaskResource.cpp +++ b/presto-native-execution/presto_cpp/main/TaskResource.cpp @@ -231,9 +231,13 @@ proxygen::RequestHandler* TaskResource::createOrUpdateTaskImpl( std::shared_ptr handlerState) { folly::via( httpSrvCpuExecutor_, - [this, &body, taskId, summarize, createOrUpdateFunc, receiveThrift]() { + [this, + requestBody = util::extractMessageBody(body), + taskId, + summarize, + createOrUpdateFunc, + receiveThrift]() { const auto startProcessCpuTimeNs = util::getProcessCpuTimeNs(); - std::string requestBody = util::extractMessageBody(body); std::unique_ptr taskInfo; try { From 27683bf37ab97a42433ca39270a137194fe64468 Mon Sep 17 00:00:00 2001 From: Christian Zentgraf Date: Fri, 13 Jun 2025 13:16:00 -0400 Subject: [PATCH 007/113] [native] Refactor arrow flight build in adapters script MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Velox provides a function to install the Arrow library. We don’t need to copy and paste the same code here and can re-use it. There is an EXTRA_ARROW_OPTIONS variable that allows custom Arrow library build options to be able to pass along that Arrow Flight should be built. --- .../tests/ArrowFlightConnectorTlsTest.cpp | 2 +- .../scripts/setup-adapters.sh | 70 ++++--------------- 2 files changed, 15 insertions(+), 57 deletions(-) diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/ArrowFlightConnectorTlsTest.cpp b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/ArrowFlightConnectorTlsTest.cpp index 4453183a39412..e0c611225bd65 100644 --- a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/ArrowFlightConnectorTlsTest.cpp +++ b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/ArrowFlightConnectorTlsTest.cpp @@ -116,7 +116,7 @@ class ArrowFlightTlsNoCertTest : public ArrowFlightConnectorTlsTestBase { }; TEST_F(ArrowFlightTlsNoCertTest, tlsNoCert) { - executeTest(false, "handshake failed"); + executeTest(false, "failed to connect"); } } // namespace facebook::presto::test diff --git a/presto-native-execution/scripts/setup-adapters.sh b/presto-native-execution/scripts/setup-adapters.sh index 3cb965fe71781..516b28c2302b7 100755 --- a/presto-native-execution/scripts/setup-adapters.sh +++ b/presto-native-execution/scripts/setup-adapters.sh @@ -15,14 +15,21 @@ set -eufx -o pipefail SCRIPT_DIR=$(readlink -f "$(dirname "${BASH_SOURCE[0]}")") -if [ -f "${SCRIPT_DIR}/setup-helper-functions.sh" ] +if [ -f "${SCRIPT_DIR}/setup-common.sh" ] then - source "${SCRIPT_DIR}/setup-helper-functions.sh" + source "${SCRIPT_DIR}/setup-common.sh" else - source "${SCRIPT_DIR}/../velox/scripts/setup-helper-functions.sh" + source "${SCRIPT_DIR}/../velox/scripts/setup-common.sh" fi DEPENDENCY_DIR=${DEPENDENCY_DIR:-$(pwd)} +OS=$(uname) +if [ "$OS" = "Darwin" ]; then + export INSTALL_PREFIX=${INSTALL_PREFIX:-"$(pwd)/deps-install"} +else + export INSTALL_PREFIX=${INSTALL_PREFIX:-"/usr/local"} +fi + function install_jwt_cpp { github_checkout Thalhammer/jwt-cpp v0.6.0 --depth 1 cmake_install -DBUILD_TESTS=OFF -DJWT_BUILD_EXAMPLES=OFF -DJWT_DISABLE_PICOJSON=ON -DJWT_CMAKE_FILES_INSTALL_DIR="${DEPENDENCY_DIR}/jwt-cpp" @@ -35,62 +42,13 @@ function install_prometheus_cpp { cmake_install -DBUILD_SHARED_LIBS=ON -DENABLE_PUSH=OFF -DENABLE_COMPRESSION=OFF } -function install_abseil { - # abseil-cpp - github_checkout abseil/abseil-cpp 20240116.2 --depth 1 - cmake_install \ - -DABSL_BUILD_TESTING=OFF \ - -DCMAKE_CXX_STANDARD=17 \ - -DABSL_PROPAGATE_CXX_STD=ON \ - -DABSL_ENABLE_INSTALL=ON -} - -function install_grpc { - # grpc - github_checkout grpc/grpc v1.48.1 --depth 1 - cmake_install \ - -DgRPC_BUILD_TESTS=OFF \ - -DgRPC_ABSL_PROVIDER=package \ - -DgRPC_ZLIB_PROVIDER=package \ - -DgRPC_CARES_PROVIDER=package \ - -DgRPC_RE2_PROVIDER=package \ - -DgRPC_SSL_PROVIDER=package \ - -DgRPC_PROTOBUF_PROVIDER=package \ - -DgRPC_INSTALL=ON -} - function install_arrow_flight { - ARROW_VERSION="${ARROW_VERSION:-15.0.0}" - if [[ "$OSTYPE" == "linux-gnu"* ]]; then - export INSTALL_PREFIX=${INSTALL_PREFIX:-"/usr/local"} - LINUX_DISTRIBUTION=$(. /etc/os-release && echo ${ID}) - if [[ "$LINUX_DISTRIBUTION" == "ubuntu" || "$LINUX_DISTRIBUTION" == "debian" ]]; then - SUDO="${SUDO:-"sudo --preserve-env"}" - ${SUDO} apt install -y libc-ares-dev - ${SUDO} ldconfig -v 2>/dev/null | grep "${INSTALL_PREFIX}/lib" || \ - echo "${INSTALL_PREFIX}/lib" | ${SUDO} tee /etc/ld.so.conf.d/local-libraries.conf > /dev/null \ - && ${SUDO} ldconfig - else - dnf -y install c-ares-devel - ldconfig -v 2>/dev/null | grep "${INSTALL_PREFIX}/lib" || \ - echo "${INSTALL_PREFIX}/lib" | tee /etc/ld.so.conf.d/local-libraries.conf > /dev/null \ - && ldconfig - fi - else - # The installation script for the Arrow Flight connector currently works only on Linux distributions. - return 0 - fi - - install_abseil - install_grpc - + # Velox provides an install for the Arrow library. Rebuild with the original Velox options and + # Arrow Flight enabled. The Velox version of Arrow is used. # NOTE: benchmarks are on due to a compilation error with v15.0.0, once updated that can be removed # see https://github.com/apache/arrow/issues/41617 - wget_and_untar https://github.com/apache/arrow/archive/apache-arrow-${ARROW_VERSION}.tar.gz arrow - cmake_install_dir arrow/cpp \ - -DARROW_FLIGHT=ON \ - -DARROW_BUILD_BENCHMARKS=ON \ - -DCMAKE_INSTALL_PREFIX=${INSTALL_PREFIX} + EXTRA_ARROW_OPTIONS=" -DARROW_FLIGHT=ON -DARROW_BUILD_BENCHMARKS=ON " + install_arrow } cd "${DEPENDENCY_DIR}" || exit From 95fc85e087a1abf34acb238a5046e7342a3a60d8 Mon Sep 17 00:00:00 2001 From: Anant Aneja <1797669+aaneja@users.noreply.github.com> Date: Mon, 25 Aug 2025 11:13:43 +0530 Subject: [PATCH 008/113] Fix incorrect session property names --- .../main/sphinx/optimizer/history-based-optimization.rst | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/presto-docs/src/main/sphinx/optimizer/history-based-optimization.rst b/presto-docs/src/main/sphinx/optimizer/history-based-optimization.rst index eeb5884e55be7..01e011889bbec 100644 --- a/presto-docs/src/main/sphinx/optimizer/history-based-optimization.rst +++ b/presto-docs/src/main/sphinx/optimizer/history-based-optimization.rst @@ -67,11 +67,11 @@ Session property Name Description ``restrict_history_based_optimization_to_complex_query`` Enable history based optimization only for complex queries, i.e. queries with join and aggregation. ``True`` ``history_input_table_statistics_matching_threshold`` Overrides the behavior of the configuration property ``hbo.history-matching-threshold`` ``hbo.history-matching-threshold`` in the current session. -``treat-low-confidence-zero-estimation-as-unknown`` Overrides the behavior of the configuration property +``treat_low_confidence_zero_estimation_unknown_enabled`` Overrides the behavior of the configuration property ``optimizer.treat-low-confidence-zero-estimation-as-unknown`` in the current session. ``optimizer.treat-low-confidence-zero-estimation-as-unknown`` -``confidence-based-broadcast`` Overrides the behavior of the configuration property +``confidence_based_broadcast_enabled`` Overrides the behavior of the configuration property ``optimizer.confidence-based-broadcast`` in the current session. ``optimizer.confidence-based-broadcast`` -``retry-query-with-history-based-optimization`` Overrides the behavior of the configuration property +``retry_query_with_history_based_optimization`` Overrides the behavior of the configuration property ``optimizer.retry-query-with-history-based-optimization`` in the current session. ``optimizer.retry-query-with-history-based-optimization`` =========================================================== ==================================================================================================== ============================================================== From 86d223bc83875b08df8065741f0efe3b6f539c57 Mon Sep 17 00:00:00 2001 From: wangd Date: Tue, 12 Aug 2025 01:00:35 +0800 Subject: [PATCH 009/113] Add missing content and fix typos in comments --- .../java/com/facebook/presto/sql/planner/iterative/Lookup.java | 2 +- .../rule/LeftJoinWithArrayContainsToEquiJoinCondition.java | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/iterative/Lookup.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/iterative/Lookup.java index 412e35666ddc8..42c860d36b2d8 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/iterative/Lookup.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/iterative/Lookup.java @@ -26,7 +26,7 @@ public interface Lookup /** * Resolves a node by materializing GroupReference nodes * representing symbolic references to other nodes. This method - * is deprecated since is assumes group contains only one node. + * is deprecated since it assumes group contains only one node. *

* If the node is not a GroupReference, it returns the * argument as is. diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/iterative/rule/LeftJoinWithArrayContainsToEquiJoinCondition.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/iterative/rule/LeftJoinWithArrayContainsToEquiJoinCondition.java index b5348879355cf..bd1e575f8e78b 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/iterative/rule/LeftJoinWithArrayContainsToEquiJoinCondition.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/iterative/rule/LeftJoinWithArrayContainsToEquiJoinCondition.java @@ -48,7 +48,7 @@ import static java.util.Objects.requireNonNull; /** - * When the join condition of a left join has pattern of contains(array, element) where array, we can rewrite it as a equi join condition. For example: + * When the join condition of a left join has pattern of contains(array, element) where array is from the right-side relation and element is from the left-side relation, we can rewrite it as an equi join condition. For example: *

  * - Left Join
  *      empty join clause

From f1f4ced45f44f3468c6e9f207afb0799722d4645 Mon Sep 17 00:00:00 2001
From: Ge Gao 
Date: Wed, 20 Aug 2025 14:58:01 -0700
Subject: [PATCH 010/113] Parse CharN in protocol to Varchar type

Reuse the existing Velox VarcharType to implement the type Char(n) in protocol.
Add a SystemConfig "char-n-type-enabled" to guard this feature.
Note this will make Char(n) type carry the behavior of VarcharType type. It is a different
behavior from Char(n) type in Presto today, where it has a fixed number
of characters. We suppose the user could call rpad() if today's behavior is needed.
---
 .../presto_cpp/main/common/Configs.cpp        |  5 +++
 .../presto_cpp/main/common/Configs.h          |  7 ++++
 .../main/common/tests/CMakeLists.txt          |  7 ++++
 .../{ => common}/tests/MutableConfigs.cpp     |  2 +-
 .../tests/MutableConfigs.h}                   |  0
 .../presto_cpp/main/tests/CMakeLists.txt      |  2 +-
 .../main/tests/CoordinatorDiscovererTest.cpp  |  2 +-
 .../main/tests/PrestoExchangeSourceTest.cpp   |  2 +-
 .../presto_cpp/main/tests/TaskManagerTest.cpp |  2 +-
 .../presto_cpp/main/types/TypeParser.cpp      |  8 +++++
 .../main/types/tests/CMakeLists.txt           |  1 +
 .../main/types/tests/RowExpressionTest.cpp    | 18 +++++++++-
 .../PrestoNativeQueryRunnerUtils.java         | 18 ++++++++--
 .../AbstractTestAggregationsNative.java       | 36 ++++++++++++++-----
 .../presto/nativetests/NativeTestsUtils.java  | 16 ++++++++-
 .../presto/nativetests/TestAggregations.java  |  9 +++--
 ...TestOptimizeMixedDistinctAggregations.java |  6 +++-
 17 files changed, 118 insertions(+), 23 deletions(-)
 rename presto-native-execution/presto_cpp/main/{ => common}/tests/MutableConfigs.cpp (96%)
 rename presto-native-execution/presto_cpp/main/{tests/MultableConfigs.h => common/tests/MutableConfigs.h} (100%)

diff --git a/presto-native-execution/presto_cpp/main/common/Configs.cpp b/presto-native-execution/presto_cpp/main/common/Configs.cpp
index 5d8a86af6b596..85e5f356ae815 100644
--- a/presto-native-execution/presto_cpp/main/common/Configs.cpp
+++ b/presto-native-execution/presto_cpp/main/common/Configs.cpp
@@ -261,6 +261,7 @@ SystemConfig::SystemConfig() {
           NUM_PROP(kHttpSrvIoEvbViolationThresholdMs, 1000),
           NUM_PROP(kMaxLocalExchangePartitionBufferSize, 65536),
           BOOL_PROP(kTextWriterEnabled, true),
+          BOOL_PROP(kCharNToVarcharImplicitCast, false),
       };
 }
 
@@ -926,6 +927,10 @@ bool SystemConfig::textWriterEnabled() const {
   return optionalProperty(kTextWriterEnabled).value();
 }
 
+bool SystemConfig::charNToVarcharImplicitCast() const {
+  return optionalProperty(kCharNToVarcharImplicitCast).value();
+}
+
 NodeConfig::NodeConfig() {
   registeredProps_ =
       std::unordered_map>{
diff --git a/presto-native-execution/presto_cpp/main/common/Configs.h b/presto-native-execution/presto_cpp/main/common/Configs.h
index baaa9dbf06d4c..fc49549ca1695 100644
--- a/presto-native-execution/presto_cpp/main/common/Configs.h
+++ b/presto-native-execution/presto_cpp/main/common/Configs.h
@@ -763,6 +763,11 @@ class SystemConfig : public ConfigBase {
   // TODO: remove once text writer is fully rolled out
   static constexpr std::string_view kTextWriterEnabled{"text-writer-enabled"};
 
+  /// Enable the type char(n) with the same behavior as unbounded varchar.
+  /// char(n) type is not supported by parser when set to false.
+  static constexpr std::string_view kCharNToVarcharImplicitCast{
+    "char-n-to-varchar-implicit-cast"};
+
   SystemConfig();
 
   virtual ~SystemConfig() = default;
@@ -1053,6 +1058,8 @@ class SystemConfig : public ConfigBase {
   uint64_t maxLocalExchangePartitionBufferSize() const;
 
   bool textWriterEnabled() const;
+
+  bool charNToVarcharImplicitCast() const;
 };
 
 /// Provides access to node properties defined in node.properties file.
diff --git a/presto-native-execution/presto_cpp/main/common/tests/CMakeLists.txt b/presto-native-execution/presto_cpp/main/common/tests/CMakeLists.txt
index ea638403e6640..267118dd1cd2f 100644
--- a/presto-native-execution/presto_cpp/main/common/tests/CMakeLists.txt
+++ b/presto-native-execution/presto_cpp/main/common/tests/CMakeLists.txt
@@ -9,6 +9,13 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+add_library(presto_mutable_configs MutableConfigs.cpp)
+
+target_link_libraries(
+  presto_mutable_configs
+  presto_common
+  velox_file)
+
 add_executable(presto_common_test CommonTest.cpp ConfigTest.cpp)
 
 add_test(presto_common_test presto_common_test)
diff --git a/presto-native-execution/presto_cpp/main/tests/MutableConfigs.cpp b/presto-native-execution/presto_cpp/main/common/tests/MutableConfigs.cpp
similarity index 96%
rename from presto-native-execution/presto_cpp/main/tests/MutableConfigs.cpp
rename to presto-native-execution/presto_cpp/main/common/tests/MutableConfigs.cpp
index cd59c17b2ca97..1232cc0e4c5cf 100644
--- a/presto-native-execution/presto_cpp/main/tests/MutableConfigs.cpp
+++ b/presto-native-execution/presto_cpp/main/common/tests/MutableConfigs.cpp
@@ -12,7 +12,7 @@
  * limitations under the License.
  */
 #include "presto_cpp/main/common/Configs.h"
-#include "presto_cpp/main/tests/MultableConfigs.h"
+#include "presto_cpp/main/common/tests/MutableConfigs.h"
 #include "velox/common/file/File.h"
 #include "velox/common/file/FileSystems.h"
 #include "velox/exec/tests/utils/TempDirectoryPath.h"
diff --git a/presto-native-execution/presto_cpp/main/tests/MultableConfigs.h b/presto-native-execution/presto_cpp/main/common/tests/MutableConfigs.h
similarity index 100%
rename from presto-native-execution/presto_cpp/main/tests/MultableConfigs.h
rename to presto-native-execution/presto_cpp/main/common/tests/MutableConfigs.h
diff --git a/presto-native-execution/presto_cpp/main/tests/CMakeLists.txt b/presto-native-execution/presto_cpp/main/tests/CMakeLists.txt
index f95483eb3ad25..2e24268aa1a07 100644
--- a/presto-native-execution/presto_cpp/main/tests/CMakeLists.txt
+++ b/presto-native-execution/presto_cpp/main/tests/CMakeLists.txt
@@ -14,7 +14,6 @@ add_executable(
   AnnouncerTest.cpp
   CoordinatorDiscovererTest.cpp
   HttpServerWrapper.cpp
-  MutableConfigs.cpp
   PeriodicMemoryCheckerTest.cpp
   PrestoExchangeSourceTest.cpp
   PrestoTaskTest.cpp
@@ -52,6 +51,7 @@ target_link_libraries(
   velox_functions_prestosql
   velox_aggregates
   velox_hive_partition_function
+  presto_mutable_configs
   ${RE2}
   GTest::gmock
   GTest::gtest
diff --git a/presto-native-execution/presto_cpp/main/tests/CoordinatorDiscovererTest.cpp b/presto-native-execution/presto_cpp/main/tests/CoordinatorDiscovererTest.cpp
index 76315aeb7bc4a..c942515aa535a 100644
--- a/presto-native-execution/presto_cpp/main/tests/CoordinatorDiscovererTest.cpp
+++ b/presto-native-execution/presto_cpp/main/tests/CoordinatorDiscovererTest.cpp
@@ -15,7 +15,7 @@
 #include "presto_cpp/main/CoordinatorDiscoverer.h"
 #include 
 #include "presto_cpp/main/common/Configs.h"
-#include "presto_cpp/main/tests/MultableConfigs.h"
+#include "presto_cpp/main/common/tests/MutableConfigs.h"
 #include "velox/common/file/FileSystems.h"
 
 using namespace facebook::velox;
diff --git a/presto-native-execution/presto_cpp/main/tests/PrestoExchangeSourceTest.cpp b/presto-native-execution/presto_cpp/main/tests/PrestoExchangeSourceTest.cpp
index 17fe5659484e4..a1e76624a6ebc 100644
--- a/presto-native-execution/presto_cpp/main/tests/PrestoExchangeSourceTest.cpp
+++ b/presto-native-execution/presto_cpp/main/tests/PrestoExchangeSourceTest.cpp
@@ -19,9 +19,9 @@
 #include 
 #include "folly/experimental/EventCount.h"
 #include "presto_cpp/main/PrestoExchangeSource.h"
+#include "presto_cpp/main/common/tests/MutableConfigs.h"
 #include "presto_cpp/main/common/Utils.h"
 #include "presto_cpp/main/tests/HttpServerWrapper.h"
-#include "presto_cpp/main/tests/MultableConfigs.h"
 #include "presto_cpp/presto_protocol/core/presto_protocol_core.h"
 #include "velox/common/base/tests/GTestUtils.h"
 #include "velox/common/file/FileSystems.h"
diff --git a/presto-native-execution/presto_cpp/main/tests/TaskManagerTest.cpp b/presto-native-execution/presto_cpp/main/tests/TaskManagerTest.cpp
index bab2407b56d8c..c5eba098be4e5 100644
--- a/presto-native-execution/presto_cpp/main/tests/TaskManagerTest.cpp
+++ b/presto-native-execution/presto_cpp/main/tests/TaskManagerTest.cpp
@@ -18,9 +18,9 @@
 #include "folly/experimental/EventCount.h"
 #include "presto_cpp/main/PrestoExchangeSource.h"
 #include "presto_cpp/main/TaskResource.h"
+#include "presto_cpp/main/common/tests/MutableConfigs.h"
 #include "presto_cpp/main/connectors/PrestoToVeloxConnector.h"
 #include "presto_cpp/main/tests/HttpServerWrapper.h"
-#include "presto_cpp/main/tests/MultableConfigs.h"
 #include "velox/common/base/Fs.h"
 #include "velox/common/base/tests/GTestUtils.h"
 #include "velox/common/file/FileSystems.h"
diff --git a/presto-native-execution/presto_cpp/main/types/TypeParser.cpp b/presto-native-execution/presto_cpp/main/types/TypeParser.cpp
index 47e8ce3592255..f8c757adf6e10 100644
--- a/presto-native-execution/presto_cpp/main/types/TypeParser.cpp
+++ b/presto-native-execution/presto_cpp/main/types/TypeParser.cpp
@@ -17,9 +17,17 @@
 #include "presto_cpp/main/types/TypeParser.h"
 #include "velox/type/parser/TypeParser.h"
 
+#include "presto_cpp/main/common/Configs.h"
+
 namespace facebook::presto {
 
 velox::TypePtr TypeParser::parse(const std::string& text) const {
+  if (SystemConfig::instance()->charNToVarcharImplicitCast()) {
+    if (text.find("char(") == 0 || text.find("CHAR(") == 0) {
+      return velox::VARCHAR();
+    }
+  }
+
   auto it = cache_.find(text);
   if (it != cache_.end()) {
     return it->second;
diff --git a/presto-native-execution/presto_cpp/main/types/tests/CMakeLists.txt b/presto-native-execution/presto_cpp/main/types/tests/CMakeLists.txt
index 27cbf06ffda08..9560036e501d4 100644
--- a/presto-native-execution/presto_cpp/main/types/tests/CMakeLists.txt
+++ b/presto-native-execution/presto_cpp/main/types/tests/CMakeLists.txt
@@ -54,6 +54,7 @@ target_link_libraries(
   $
   $
   presto_operators
+  presto_mutable_configs
   presto_type_test_utils
   velox_core
   velox_dwio_common_exception
diff --git a/presto-native-execution/presto_cpp/main/types/tests/RowExpressionTest.cpp b/presto-native-execution/presto_cpp/main/types/tests/RowExpressionTest.cpp
index 75a28f4db2d5f..6cb1d7d100815 100644
--- a/presto-native-execution/presto_cpp/main/types/tests/RowExpressionTest.cpp
+++ b/presto-native-execution/presto_cpp/main/types/tests/RowExpressionTest.cpp
@@ -14,10 +14,12 @@
 #include 
 #include 
 
+#include "presto_cpp/main/common/Configs.h"
+#include "presto_cpp/main/common/tests/MutableConfigs.h"
 #include "presto_cpp/main/types/PrestoToVeloxExpr.h"
 #include "presto_cpp/presto_protocol/core/presto_protocol_core.h"
+#include "velox/common/file/FileSystems.h"
 #include "velox/core/Expressions.h"
-#include "velox/type/Type.h"
 #include "velox/functions/prestosql/types/JsonRegistration.h"
 
 using namespace facebook::presto;
@@ -32,6 +34,8 @@ class RowExpressionTest : public ::testing::Test {
 
   void SetUp() override {
     registerJsonType();
+    filesystems::registerLocalFileSystem();
+    test::setupMutableSystemConfig();
     pool_ = memory::MemoryManager::getInstance()->addLeafPool();
     converter_ =
         std::make_unique(pool_.get(), &typeParser_);
@@ -414,6 +418,18 @@ TEST_F(RowExpressionTest, varbinary5) {
           '"');
 }
 
+TEST_F(RowExpressionTest, char) {
+  SystemConfig::instance()->setValue(std::string(SystemConfig::kCharNToVarcharImplicitCast), "true");
+  std::string str = R"##(
+        {
+            "@type": "constant",
+            "type": "char(3)",
+            "valueBlock": "DgAAAFZBUklBQkxFX1dJRFRIAQAAAAMAAAAAAwAAAGFiYw=="
+        }
+    )##";
+  testConstantExpression(str, "VARCHAR", "\"abc\"");
+}
+
 TEST_F(RowExpressionTest, timestamp) {
   std::string str = R"(
         {
diff --git a/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/PrestoNativeQueryRunnerUtils.java b/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/PrestoNativeQueryRunnerUtils.java
index 54980ba54782f..95095f796a72d 100644
--- a/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/PrestoNativeQueryRunnerUtils.java
+++ b/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/PrestoNativeQueryRunnerUtils.java
@@ -126,6 +126,7 @@ public static class HiveQueryRunnerBuilder
         private boolean enableRuntimeMetricsCollection;
         private boolean enableSsdCache;
         private boolean failOnNestedLoopJoin;
+        private boolean implicitCastCharNToVarchar;
         // External worker launcher is applicable only for the native hive query runner, since it depends on other
         // properties it should be created once all the other query runner configs are set. This variable indicates
         // whether the query runner returned by builder should use an external worker launcher, it will be true only
@@ -239,6 +240,12 @@ public HiveQueryRunnerBuilder setCacheMaxSize(Integer cacheMaxSize)
             return this;
         }
 
+        public HiveQueryRunnerBuilder setImplicitCastCharNToVarchar(boolean implicitCastCharNToVarchar)
+        {
+            this.implicitCastCharNToVarchar = implicitCastCharNToVarchar;
+            return this;
+        }
+
         public HiveQueryRunnerBuilder setExtraProperties(Map extraProperties)
         {
             this.extraProperties.putAll(extraProperties);
@@ -263,7 +270,7 @@ public QueryRunner build()
             Optional> externalWorkerLauncher = Optional.empty();
             if (this.useExternalWorkerLauncher) {
                 externalWorkerLauncher = getExternalWorkerLauncher("hive", serverBinary, cacheMaxSize, remoteFunctionServerUds,
-                        failOnNestedLoopJoin, coordinatorSidecarEnabled, enableRuntimeMetricsCollection, enableSsdCache);
+                        failOnNestedLoopJoin, coordinatorSidecarEnabled, enableRuntimeMetricsCollection, enableSsdCache, implicitCastCharNToVarchar);
             }
             return HiveQueryRunner.createQueryRunner(
                     ImmutableList.of(),
@@ -352,7 +359,7 @@ public QueryRunner build()
             Optional> externalWorkerLauncher = Optional.empty();
             if (this.useExternalWorkerLauncher) {
                 externalWorkerLauncher = getExternalWorkerLauncher("iceberg", serverBinary, cacheMaxSize, remoteFunctionServerUds,
-                        false, false, false, false);
+                        false, false, false, false, false);
             }
             return IcebergQueryRunner.builder()
                     .setExtraProperties(extraProperties)
@@ -449,7 +456,8 @@ public static Optional> getExternalWorkerLaunc
             Boolean failOnNestedLoopJoin,
             boolean isCoordinatorSidecarEnabled,
             boolean enableRuntimeMetricsCollection,
-            boolean enableSsdCache)
+            boolean enableSsdCache,
+            boolean implicitCastCharNToVarchar)
     {
         return
                 Optional.of((workerIndex, discoveryUri) -> {
@@ -497,6 +505,10 @@ public static Optional> getExternalWorkerLaunc
                             configProperties = format("%s%n" + "velox-plan-validator-fail-on-nested-loop-join=true%n", configProperties);
                         }
 
+                        if (implicitCastCharNToVarchar) {
+                            configProperties = format("%s%n" + "char-n-to-varchar-implicit-cast=true%n", configProperties);
+                        }
+
                         Files.write(tempDirectoryPath.resolve("config.properties"), configProperties.getBytes());
                         Files.write(tempDirectoryPath.resolve("node.properties"),
                                 format("node.id=%s%n" +
diff --git a/presto-native-tests/src/test/java/com/facebook/presto/nativetests/AbstractTestAggregationsNative.java b/presto-native-tests/src/test/java/com/facebook/presto/nativetests/AbstractTestAggregationsNative.java
index 39f787a1eb601..d668c65f098ea 100644
--- a/presto-native-tests/src/test/java/com/facebook/presto/nativetests/AbstractTestAggregationsNative.java
+++ b/presto-native-tests/src/test/java/com/facebook/presto/nativetests/AbstractTestAggregationsNative.java
@@ -24,13 +24,15 @@ public abstract class AbstractTestAggregationsNative
     private static final String QDIGEST_TYPE = "qdigest";
 
     private String storageFormat;
+    private boolean implicitCastCharNToVarchar;
     private String approxDistinctUnsupportedSignatureError;
     private String charTypeUnsupportedError;
     private String timeTypeUnsupportedError;
 
-    public void init(String storageFormat, boolean sidecarEnabled)
+    public void init(String storageFormat, boolean charNToVarcharImplicitCast, boolean sidecarEnabled)
     {
         this.storageFormat = storageFormat;
+        this.implicitCastCharNToVarchar = charNToVarcharImplicitCast;
         if (sidecarEnabled) {
             charTypeUnsupportedError = ".*Unknown type: char.*";
             timeTypeUnsupportedError = ".*Unknown type: time.*";
@@ -117,8 +119,16 @@ public void testApproximateCountDistinct()
         assertQuery("SELECT approx_distinct(CAST(custkey AS VARCHAR), 0.023) FROM orders", "SELECT 1036");
 
         // test char
-        assertQueryFails("SELECT approx_distinct(CAST(CAST(custkey AS VARCHAR) AS CHAR(20))) FROM orders", charTypeUnsupportedError, true);
-        assertQueryFails("SELECT approx_distinct(CAST(CAST(custkey AS VARCHAR) AS CHAR(20)), 0.023) FROM orders", charTypeUnsupportedError, true);
+        String charQuery = "SELECT approx_distinct(CAST(CAST(custkey AS VARCHAR) AS CHAR(20))) FROM orders";
+        String charWithErrorQuery = "SELECT approx_distinct(CAST(CAST(custkey AS VARCHAR) AS CHAR(20)), 0.023) FROM orders";
+        if (implicitCastCharNToVarchar) {
+            assertQuery(charQuery, "SELECT 1036");
+            assertQuery(charWithErrorQuery, "SELECT 1036");
+        }
+        else {
+            assertQueryFails(charQuery, charTypeUnsupportedError, true);
+            assertQueryFails(charWithErrorQuery, charTypeUnsupportedError, true);
+        }
 
         // test varbinary
         assertQuery("SELECT approx_distinct(to_utf8(CAST(custkey AS VARCHAR))) FROM orders", "SELECT 1036");
@@ -136,10 +146,13 @@ public void testSumDataSizeForStats()
         assertQuery("SELECT \"sum_data_size_for_stats\"(comment) FROM orders", "SELECT 787364");
 
         // char
-        // Presto removes trailing whitespaces when casting to CHAR.
-        // Hard code the expected data size since there is no easy to way to compute it in H2.
-        assertQueryFails("SELECT \"sum_data_size_for_stats\"(CAST(comment AS CHAR(1000))) FROM orders",
-                charTypeUnsupportedError, true);
+        String charQuery = "SELECT \"sum_data_size_for_stats\"(CAST(comment AS CHAR(1000))) FROM orders";
+        if (implicitCastCharNToVarchar) {
+            assertQuery(charQuery, "SELECT 787364");
+        }
+        else {
+            assertQueryFails(charQuery, charTypeUnsupportedError, true);
+        }
 
         // varbinary
         assertQuery("SELECT \"sum_data_size_for_stats\"(CAST(comment AS VARBINARY)) FROM orders", "SELECT 787364");
@@ -168,8 +181,13 @@ public void testMaxDataSizeForStats()
         assertQuery("SELECT \"max_data_size_for_stats\"(comment) FROM orders", "select 82");
 
         // char
-        assertQueryFails("SELECT \"max_data_size_for_stats\"(CAST(comment AS CHAR(1000))) FROM orders",
-                charTypeUnsupportedError, true);
+        String charQuery = "SELECT \"max_data_size_for_stats\"(CAST(comment AS CHAR(1000))) FROM orders";
+        if (implicitCastCharNToVarchar) {
+            assertQuery(charQuery, "SELECT 82");
+        }
+        else {
+            assertQueryFails(charQuery, charTypeUnsupportedError, true);
+        }
 
         // varbinary
         assertQuery("SELECT \"max_data_size_for_stats\"(CAST(comment AS VARBINARY)) FROM orders", "select 82");
diff --git a/presto-native-tests/src/test/java/com/facebook/presto/nativetests/NativeTestsUtils.java b/presto-native-tests/src/test/java/com/facebook/presto/nativetests/NativeTestsUtils.java
index b4dbaa197b36e..12d4cabed0220 100644
--- a/presto-native-tests/src/test/java/com/facebook/presto/nativetests/NativeTestsUtils.java
+++ b/presto-native-tests/src/test/java/com/facebook/presto/nativetests/NativeTestsUtils.java
@@ -24,13 +24,14 @@ public class NativeTestsUtils
 {
     private NativeTestsUtils() {}
 
-    public static QueryRunner createNativeQueryRunner(String storageFormat, boolean sidecarEnabled)
+    public static QueryRunner createNativeQueryRunner(String storageFormat, boolean charNToVarcharImplicitCast, boolean sidecarEnabled)
             throws Exception
     {
         QueryRunner queryRunner = nativeHiveQueryRunnerBuilder()
                 .setStorageFormat(storageFormat)
                 .setAddStorageFormatToPath(true)
                 .setUseThrift(true)
+                .setImplicitCastCharNToVarchar(charNToVarcharImplicitCast)
                 .setCoordinatorSidecarEnabled(sidecarEnabled)
                 .build();
         if (sidecarEnabled) {
@@ -39,6 +40,12 @@ public static QueryRunner createNativeQueryRunner(String storageFormat, boolean
         return queryRunner;
     }
 
+    public static QueryRunner createNativeQueryRunner(String storageFormat, boolean sidecarEnabled)
+            throws Exception
+    {
+        return createNativeQueryRunner(storageFormat, false, sidecarEnabled);
+    }
+
     public static void createTables(String storageFormat)
     {
         try {
@@ -58,4 +65,11 @@ public static void createTables(String storageFormat)
             throw new RuntimeException(e);
         }
     }
+
+    // TODO: remove and directly return charNToVarcharImplicitCast after addressing Issue #25894 and adding support for Char(n) type
+    // to class NativeTypeManager for Sidecar.
+    public static boolean getCharNToVarcharImplicitCastForTest(boolean sidecarEnabled, boolean charNToVarcharImplicitCast)
+    {
+        return sidecarEnabled ? false : charNToVarcharImplicitCast;
+    }
 }
diff --git a/presto-native-tests/src/test/java/com/facebook/presto/nativetests/TestAggregations.java b/presto-native-tests/src/test/java/com/facebook/presto/nativetests/TestAggregations.java
index 89d31a9257b95..220d14db2336a 100644
--- a/presto-native-tests/src/test/java/com/facebook/presto/nativetests/TestAggregations.java
+++ b/presto-native-tests/src/test/java/com/facebook/presto/nativetests/TestAggregations.java
@@ -22,6 +22,7 @@ public class TestAggregations
         extends AbstractTestAggregationsNative
 {
     private String storageFormat;
+    private boolean charNToVarcharImplicitCast;
     private boolean sidecarEnabled;
 
     @BeforeClass
@@ -30,15 +31,17 @@ public void init()
             throws Exception
     {
         storageFormat = System.getProperty("storageFormat", "PARQUET");
-        sidecarEnabled = parseBoolean(System.getProperty("sidecarEnabled", "true"));
-        super.init(storageFormat, sidecarEnabled);
+        sidecarEnabled = parseBoolean(System.getProperty("sidecarEnabled", "false"));
+        charNToVarcharImplicitCast = NativeTestsUtils.getCharNToVarcharImplicitCastForTest(
+                sidecarEnabled, parseBoolean(System.getProperty("charNToVarcharImplicitCast", "false")));
+        super.init(storageFormat, charNToVarcharImplicitCast, sidecarEnabled);
         super.init();
     }
 
     @Override
     protected QueryRunner createQueryRunner() throws Exception
     {
-        return NativeTestsUtils.createNativeQueryRunner(storageFormat, sidecarEnabled);
+        return NativeTestsUtils.createNativeQueryRunner(storageFormat, charNToVarcharImplicitCast, sidecarEnabled);
     }
 
     @Override
diff --git a/presto-native-tests/src/test/java/com/facebook/presto/nativetests/TestOptimizeMixedDistinctAggregations.java b/presto-native-tests/src/test/java/com/facebook/presto/nativetests/TestOptimizeMixedDistinctAggregations.java
index 9ff17c46fc9ef..8c775f51dd69a 100644
--- a/presto-native-tests/src/test/java/com/facebook/presto/nativetests/TestOptimizeMixedDistinctAggregations.java
+++ b/presto-native-tests/src/test/java/com/facebook/presto/nativetests/TestOptimizeMixedDistinctAggregations.java
@@ -26,6 +26,7 @@ public class TestOptimizeMixedDistinctAggregations
         extends AbstractTestAggregationsNative
 {
     private String storageFormat;
+    private boolean charNToVarcharImplicitCast;
     private boolean sidecarEnabled;
 
     @BeforeClass
@@ -35,7 +36,9 @@ public void init()
     {
         storageFormat = System.getProperty("storageFormat", "PARQUET");
         sidecarEnabled = parseBoolean(System.getProperty("sidecarEnabled", "true"));
-        super.init(storageFormat, sidecarEnabled);
+        charNToVarcharImplicitCast = NativeTestsUtils.getCharNToVarcharImplicitCastForTest(
+                sidecarEnabled, parseBoolean(System.getProperty("charNToVarcharImplicitCast", "false")));
+        super.init(storageFormat, charNToVarcharImplicitCast, sidecarEnabled);
         super.init();
     }
 
@@ -45,6 +48,7 @@ protected QueryRunner createQueryRunner() throws Exception
         QueryRunner queryRunner = nativeHiveQueryRunnerBuilder()
                 .setStorageFormat(storageFormat)
                 .setAddStorageFormatToPath(true)
+                .setImplicitCastCharNToVarchar(charNToVarcharImplicitCast)
                 .setUseThrift(true)
                 .setCoordinatorSidecarEnabled(sidecarEnabled)
                 .setExtraCoordinatorProperties(

From e52d33a4514b89ce6d9dfe45c2ea3170242c73b9 Mon Sep 17 00:00:00 2001
From: Li Zhou 
Date: Wed, 27 Aug 2025 19:23:45 +0100
Subject: [PATCH 011/113] Fix maven publish and add executable jars to github
 release (#25902)

## Description

This PR update the github action to publish maven artifacts with central
publishing method, since maven repo doesn't allow executable jar(with
shell script) to be published, so we will create a github release and
publish the jars

Need fix in release branch:
https://github.com/prestodb/presto/pull/25900

Sample release for executable jars
https://github.com/unix280/presto/releases

## Motivation and Context


## Impact
Release 0.294

## Test Plan
Tested the github release in myrepo:
https://github.com/unix280/presto/actions/runs/17272968441
Tested the maven publishing in local env

## Contributor checklist

- [ ] Please make sure your submission complies with our [contributing
guide](https://github.com/prestodb/presto/blob/master/CONTRIBUTING.md),
in particular [code
style](https://github.com/prestodb/presto/blob/master/CONTRIBUTING.md#code-style)
and [commit
standards](https://github.com/prestodb/presto/blob/master/CONTRIBUTING.md#commit-standards).
- [ ] PR description addresses the issue accurately and concisely. If
the change is non-trivial, a GitHub Issue is referenced.
- [ ] Documented new properties (with its default value), SQL syntax,
functions, or other functionality.
- [ ] If release notes are required, they follow the [release notes
guidelines](https://github.com/prestodb/presto/wiki/Release-Notes-Guidelines).
- [ ] Adequate tests were added if applicable.
- [ ] CI passed.

## Release Notes

```
== NO RELEASE NOTE ==
```
---
 .github/workflows/presto-release-publish.yml | 117 ++++++++++++-------
 1 file changed, 75 insertions(+), 42 deletions(-)

diff --git a/.github/workflows/presto-release-publish.yml b/.github/workflows/presto-release-publish.yml
index 25d7d3af2f7cf..bcd408405feb5 100644
--- a/.github/workflows/presto-release-publish.yml
+++ b/.github/workflows/presto-release-publish.yml
@@ -19,6 +19,11 @@ on:
         type: boolean
         default: true
         required: false
+      publish_github_release:
+        description: 'Publish executable jars to github release'
+        type: boolean
+        default: true
+        required: false
       publish_image:
         description: 'Publish presto docker image'
         type: boolean
@@ -33,11 +38,6 @@ on:
         description: 'prestissimo dependency image(e.g., prestodb/presto-native-dependency:latest)'
         required: false
         default: ''
-      publish_native_dep_image:
-        description: 'Publish prestissimo dependency docker image'
-        type: boolean
-        default: true
-        required: false
       tag_image_as_latest:
         description: 'Tag the published images as latest'
         type: boolean
@@ -102,8 +102,12 @@ jobs:
   publish-maven-artifacts:
     needs: publish-release-tag
     if: |
-      (!failure() &&!cancelled()) &&
-      (github.event.inputs.publish_maven == 'true' || github.event.inputs.publish_image == 'true' || github.event.inputs.publish_docs == 'true')
+      (!failure() &&!cancelled()) && (
+        github.event.inputs.publish_maven == 'true' ||
+        github.event.inputs.publish_image == 'true' ||
+        github.event.inputs.publish_docs == 'true' ||
+        github.event.inputs.publish_github_release == 'true'
+      )
     runs-on: ubuntu-latest
     environment: release
     timeout-minutes: 60
@@ -181,18 +185,8 @@ jobs:
           
           EOL
 
-      - name: Maven build
-        run: mvn clean install -DskipTests
-
-      - name: Upload artifacts
-        uses: actions/upload-artifact@v4
-        with:
-          name: presto-artifacts-${{ env.RELEASE_TAG }}
-          retention-days: 1
-          path: |
-            presto-server/target/presto-server-*.tar.gz
-            presto-cli/target/presto-cli-*-executable.jar
-            presto-docs/target/presto-docs-*.zip
+      - name: Build release artifacts
+        run: ./mvnw clean install -DskipTests
 
       - name: Set up GPG key
         env:
@@ -213,28 +207,67 @@ jobs:
           GPG_PASSPHRASE: "${{ secrets.GPG_PASSPHRASE }}"
         run: |
           unset MAVEN_CONFIG
-          modules=$(find . -maxdepth 1 -type d \( ! -name . \) -exec test -f '{}/pom.xml' \; -print | sed 's|^./||')
-          for module in $modules; do
-            if [ "$module" = "presto-test-coverage" ]; then
-              echo "Skipping $module"
-              continue
-            fi
-            MODULE_PATH=$(awk -F'[<>]' '/^    / {print $3; exit}' "$module/pom.xml" | sed 's|\.|/|g' | sed 's|^com/facebook/presto||')
-            echo "Checking https://repo1.maven.org/maven2/com/facebook/presto${MODULE_PATH}/${module}/${RELEASE_TAG}"
-            if ! curl --silent --fail "https://repo1.maven.org/maven2/com/facebook/presto${MODULE_PATH}/${module}/${RELEASE_TAG}" > /dev/null; then
-              echo "Publishing ${module}"
-              ./mvnw -s ${{ github.workspace }}/settings.xml -V -B -U -e -T1C deploy \
-                -Dgpg.passphrase="${{ secrets.GPG_PASSPHRASE }}" \
-                -Dmaven.wagon.http.retryHandler.count=8 \
-                -DskipTests \
-                -Drelease.autoPublish=${{ env.MAVEN_AUTO_PUBLISH }} \
-                -Poss-release \
-                -Pdeploy-to-ossrh \
-                -pl "$module"
-            else
-              echo "${module} ${RELEASE_TAG} already published, skipping."
-            fi
-          done
+          ./mvnw -s ${{ github.workspace }}/settings.xml -V -B -U -e -T1C deploy \
+            -Dgpg.passphrase="${{ secrets.GPG_PASSPHRASE }}" \
+            -Dmaven.wagon.http.retryHandler.count=8 \
+            -DskipTests \
+            -DstagingProfileId="${{ secrets.MAVEN_STAGING_PROFILE_ID }}" \
+            -DkeepStagingRepositoryOnFailure=true \
+            -DkeepStagingRepositoryOnCloseRuleFailure=true \
+            -DautoReleaseAfterClose=true \
+            -DstagingProgressTimeoutMinutes=60 \
+            -DdeploymentName=presto-${{ env.RELEASE_TAG }} \
+            -Drelease.autoPublish=${{ env.MAVEN_AUTO_PUBLISH }} \
+            -DskipExecutableJar \
+            -Poss-release \
+            -Pdeploy-to-ossrh \
+            -pl '!presto-test-coverage'
+
+      - name: Upload artifacts
+        uses: actions/upload-artifact@v4
+        with:
+          name: presto-artifacts-${{ env.RELEASE_TAG }}
+          retention-days: 7
+          path: |
+            presto-server/target/presto-server-*.tar.gz
+            presto-docs/target/presto-docs-*.zip
+            presto-cli/target/presto-cli-*-executable.jar
+            presto-benchmark-driver/target/presto-benchmark-driver-*-executable.jar
+            presto-testing-server-launcher/target/presto-testing-server-launcher-*-executable.jar
+
+  publish-github-release:
+    needs: publish-maven-artifacts
+    if: (!failure() && !cancelled()) && github.event.inputs.publish_github_release == 'true'
+    runs-on: ubuntu-latest
+    environment: release
+    timeout-minutes: 150
+    permissions:
+      packages: write
+      contents: read
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ env.RELEASE_TAG}}
+
+      - name: Download artifacts
+        uses: actions/download-artifact@v4
+        with:
+          name: presto-artifacts-${{ env.RELEASE_TAG }}
+          path: ./
+
+      - name: Create github release
+        uses: softprops/action-gh-release@v2
+        with:
+          tag_name: ${{ env.RELEASE_TAG }}
+          name: Presto ${{ env.RELEASE_TAG }}
+          token: ${{ secrets.PRESTODB_CI_TOKEN }}
+          body: |
+            See the release notes at https://prestodb.io/docs/current/release/release-${{ env.RELEASE_TAG }}.html
+          files: |
+            ./presto-cli/target/presto-cli-*-executable.jar
+            ./presto-benchmark-driver/target/presto-benchmark-driver-*-executable.jar
+            ./presto-testing-server-launcher/target/presto-testing-server-launcher-*-executable.jar
 
   publish-docker-image:
     needs: publish-maven-artifacts
@@ -349,7 +382,7 @@ jobs:
           else
             echo "Building new depedency image"
             docker compose build centos-native-dependency
-            if [[ "${{ github.event.inputs.publish_native_dep_image }}" == "true" ]]; then
+            if [[ "${{ github.event.inputs.publish_native_image }}" == "true" ]]; then
               docker tag presto/prestissimo-dependency:centos9 ${{ github.repository_owner }}/presto-native-dependency:${{ env.RELEASE_TAG }}-${{ env.COMMIT_SHA }}
               docker push ${{ github.repository_owner }}/presto-native-dependency:${{ env.RELEASE_TAG }}-${{ env.COMMIT_SHA }}
               

From 1cf0190f53ea393406d1215eaa9a6da608cf95ab Mon Sep 17 00:00:00 2001
From: Timothy Meehan 
Date: Wed, 27 Aug 2025 13:52:47 -0400
Subject: [PATCH 012/113] Add @pdabre12 as module committer

@pdabre12 has been voted as module committer for the Presto sidecar module.

Also, I fixed a bug that project committers could not approve some C++ code.  Per our contributing guide, project committers must be capable of approving all code (although C++ module committers are preferred for approving and merging C++ code).
---
 CODEOWNERS | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/CODEOWNERS b/CODEOWNERS
index 60f71c88b99ee..8796e20a27174 100644
--- a/CODEOWNERS
+++ b/CODEOWNERS
@@ -111,9 +111,9 @@ CODEOWNERS @prestodb/team-tsc
 
 #####################################################################
 # Prestissimo module
-/presto-native-execution @prestodb/team-velox
-/presto-native-sidecar-plugin @prestodb/team-velox
-/presto-native-tests @prestodb/team-velox
+/presto-native-execution @prestodb/team-velox @prestodb/committers
+/presto-native-sidecar-plugin @pdabre12 @prestodb/team-velox @prestodb/committers
+/presto-native-tests @prestodb/team-velox @prestodb/committers
 /.github/workflows/prestocpp-* @prestodb/team-velox @prestodb/committers
 
 #####################################################################

From 410bb8cbe3049ad5409c0d152a009876cfb92488 Mon Sep 17 00:00:00 2001
From: PRASHANT GOLASH <40184733+prashantgolash@users.noreply.github.com>
Date: Wed, 27 Aug 2025 15:01:59 -0700
Subject: [PATCH 013/113] [Coordinator throttling] Endpoint on Java worker
 reporting nodestats (#25687)

Summary:
Similar to cpp worker added the endpoint for java.
We won't be using the worker-load as going forward we will be focussing
on cpp worker only

Differential Revision: D79471792
---
 .../presto/server/ServerInfoResource.java     |  11 ++
 .../presto_cpp/main/PrestoServer.cpp          |  20 ++
 .../presto_cpp/main/PrestoServer.h            |   2 +
 .../core/presto_protocol_core.cpp             | 178 ++++++++++++++----
 .../core/presto_protocol_core.h               |  29 ++-
 .../core/presto_protocol_core.yml             |   5 +-
 .../facebook/presto/spi/NodeLoadMetrics.java  | 123 ++++++++++++
 .../com/facebook/presto/spi/NodeStats.java    |  90 +++++++++
 .../presto/server/TestServerInfoResource.java | 104 ++++++++--
 9 files changed, 503 insertions(+), 59 deletions(-)
 create mode 100644 presto-spi/src/main/java/com/facebook/presto/spi/NodeLoadMetrics.java
 create mode 100644 presto-spi/src/main/java/com/facebook/presto/spi/NodeStats.java

diff --git a/presto-main/src/main/java/com/facebook/presto/server/ServerInfoResource.java b/presto-main/src/main/java/com/facebook/presto/server/ServerInfoResource.java
index e24a08ab00559..8c686288caef8 100644
--- a/presto-main/src/main/java/com/facebook/presto/server/ServerInfoResource.java
+++ b/presto-main/src/main/java/com/facebook/presto/server/ServerInfoResource.java
@@ -19,6 +19,7 @@
 import com.facebook.presto.execution.resourceGroups.ResourceGroupManager;
 import com.facebook.presto.metadata.StaticCatalogStore;
 import com.facebook.presto.spi.NodeState;
+import com.facebook.presto.spi.NodeStats;
 import jakarta.annotation.security.RolesAllowed;
 import jakarta.inject.Inject;
 import jakarta.ws.rs.Consumes;
@@ -128,6 +129,16 @@ else if (!nodeResourceStatusProvider.hasResources() || !resourceGroupManager.isC
         }
     }
 
+    @GET
+    @Path("stats")
+    @Produces({APPLICATION_JSON, APPLICATION_THRIFT_BINARY, APPLICATION_THRIFT_COMPACT, APPLICATION_THRIFT_FB_COMPACT})
+    @RolesAllowed(ADMIN)
+    public NodeStats getServerStats()
+    {
+        NodeStats stats = new NodeStats(getServerState(), null);
+        return stats;
+    }
+
     @GET
     @Path("coordinator")
     @Produces(TEXT_PLAIN)
diff --git a/presto-native-execution/presto_cpp/main/PrestoServer.cpp b/presto-native-execution/presto_cpp/main/PrestoServer.cpp
index c89f91de77957..75a2bb1eedfe4 100644
--- a/presto-native-execution/presto_cpp/main/PrestoServer.cpp
+++ b/presto-native-execution/presto_cpp/main/PrestoServer.cpp
@@ -354,6 +354,14 @@ void PrestoServer::run() {
         json infoStateJson = convertNodeState(server->nodeState());
         http::sendOkResponse(downstream, infoStateJson);
       });
+  httpServer_->registerGet(
+      "/v1/info/stats",
+      [server = this](
+          proxygen::HTTPMessage* /*message*/,
+          const std::vector>& /*body*/,
+          proxygen::ResponseHandler* downstream) {
+        server->reportNodeStats(downstream);
+      });
   httpServer_->registerPut(
       "/v1/info/state",
       [server = this](
@@ -1743,4 +1751,16 @@ void PrestoServer::createTaskManager() {
       driverExecutor_.get(), httpSrvCpuExecutor_.get(), spillerExecutor_.get());
 }
 
+void PrestoServer::reportNodeStats(proxygen::ResponseHandler* downstream) {
+  protocol::NodeStats nodeStats;
+
+  auto loadMetrics = std::make_shared();
+  loadMetrics->cpuOverload = cpuOverloaded_;
+  loadMetrics->memoryOverload = memOverloaded_;
+
+  nodeStats.loadMetrics = loadMetrics;
+  nodeStats.nodeState = convertNodeState(this->nodeState());
+
+  http::sendOkResponse(downstream, json(nodeStats));
+}
 } // namespace facebook::presto
diff --git a/presto-native-execution/presto_cpp/main/PrestoServer.h b/presto-native-execution/presto_cpp/main/PrestoServer.h
index 00c7697218f4e..afeb45fc12092 100644
--- a/presto-native-execution/presto_cpp/main/PrestoServer.h
+++ b/presto-native-execution/presto_cpp/main/PrestoServer.h
@@ -208,6 +208,8 @@ class PrestoServer {
 
   void reportNodeStatus(proxygen::ResponseHandler* downstream);
 
+  void reportNodeStats(proxygen::ResponseHandler* downstream);
+
   void handleGracefulShutdown(
       const std::vector>& body,
       proxygen::ResponseHandler* downstream);
diff --git a/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.cpp b/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.cpp
index fd5029dfe20e1..e804c9a35925e 100644
--- a/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.cpp
+++ b/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.cpp
@@ -7034,6 +7034,148 @@ void from_json(const json& j, MergeJoinNode& p) {
 } // namespace facebook::presto::protocol
 namespace facebook::presto::protocol {
 
+void to_json(json& j, const NodeLoadMetrics& p) {
+  j = json::object();
+  to_json_key(
+      j,
+      "cpuUsedPercent",
+      p.cpuUsedPercent,
+      "NodeLoadMetrics",
+      "double",
+      "cpuUsedPercent");
+  to_json_key(
+      j,
+      "memoryUsedInBytes",
+      p.memoryUsedInBytes,
+      "NodeLoadMetrics",
+      "double",
+      "memoryUsedInBytes");
+  to_json_key(
+      j,
+      "numQueuedDrivers",
+      p.numQueuedDrivers,
+      "NodeLoadMetrics",
+      "int",
+      "numQueuedDrivers");
+  to_json_key(
+      j,
+      "cpuOverload",
+      p.cpuOverload,
+      "NodeLoadMetrics",
+      "bool",
+      "cpuOverload");
+  to_json_key(
+      j,
+      "memoryOverload",
+      p.memoryOverload,
+      "NodeLoadMetrics",
+      "bool",
+      "memoryOverload");
+}
+
+void from_json(const json& j, NodeLoadMetrics& p) {
+  from_json_key(
+      j,
+      "cpuUsedPercent",
+      p.cpuUsedPercent,
+      "NodeLoadMetrics",
+      "double",
+      "cpuUsedPercent");
+  from_json_key(
+      j,
+      "memoryUsedInBytes",
+      p.memoryUsedInBytes,
+      "NodeLoadMetrics",
+      "double",
+      "memoryUsedInBytes");
+  from_json_key(
+      j,
+      "numQueuedDrivers",
+      p.numQueuedDrivers,
+      "NodeLoadMetrics",
+      "int",
+      "numQueuedDrivers");
+  from_json_key(
+      j,
+      "cpuOverload",
+      p.cpuOverload,
+      "NodeLoadMetrics",
+      "bool",
+      "cpuOverload");
+  from_json_key(
+      j,
+      "memoryOverload",
+      p.memoryOverload,
+      "NodeLoadMetrics",
+      "bool",
+      "memoryOverload");
+}
+} // namespace facebook::presto::protocol
+namespace facebook::presto::protocol {
+// Loosly copied this here from NLOHMANN_JSON_SERIALIZE_ENUM()
+
+// NOLINTNEXTLINE: cppcoreguidelines-avoid-c-arrays
+static const std::pair NodeState_enum_table[] =
+    { // NOLINT: cert-err58-cpp
+        {NodeState::ACTIVE, "ACTIVE"},
+        {NodeState::INACTIVE, "INACTIVE"},
+        {NodeState::SHUTTING_DOWN, "SHUTTING_DOWN"}};
+void to_json(json& j, const NodeState& e) {
+  static_assert(std::is_enum::value, "NodeState must be an enum!");
+  const auto* it = std::find_if(
+      std::begin(NodeState_enum_table),
+      std::end(NodeState_enum_table),
+      [e](const std::pair& ej_pair) -> bool {
+        return ej_pair.first == e;
+      });
+  j = ((it != std::end(NodeState_enum_table))
+           ? it
+           : std::begin(NodeState_enum_table))
+          ->second;
+}
+void from_json(const json& j, NodeState& e) {
+  static_assert(std::is_enum::value, "NodeState must be an enum!");
+  const auto* it = std::find_if(
+      std::begin(NodeState_enum_table),
+      std::end(NodeState_enum_table),
+      [&j](const std::pair& ej_pair) -> bool {
+        return ej_pair.second == j;
+      });
+  e = ((it != std::end(NodeState_enum_table))
+           ? it
+           : std::begin(NodeState_enum_table))
+          ->first;
+}
+} // namespace facebook::presto::protocol
+namespace facebook::presto::protocol {
+
+void to_json(json& j, const NodeStats& p) {
+  j = json::object();
+  to_json_key(
+      j, "nodeState", p.nodeState, "NodeStats", "NodeState", "nodeState");
+  to_json_key(
+      j,
+      "loadMetrics",
+      p.loadMetrics,
+      "NodeStats",
+      "NodeLoadMetrics",
+      "loadMetrics");
+}
+
+void from_json(const json& j, NodeStats& p) {
+  from_json_key(
+      j, "nodeState", p.nodeState, "NodeStats", "NodeState", "nodeState");
+  from_json_key(
+      j,
+      "loadMetrics",
+      p.loadMetrics,
+      "NodeStats",
+      "NodeLoadMetrics",
+      "loadMetrics");
+}
+} // namespace facebook::presto::protocol
+namespace facebook::presto::protocol {
+
 void to_json(json& j, const NodeVersion& p) {
   j = json::object();
   to_json_key(j, "version", p.version, "NodeVersion", "String", "version");
@@ -11569,39 +11711,3 @@ void from_json(const json& j, WindowNode& p) {
       "preSortedOrderPrefix");
 }
 } // namespace facebook::presto::protocol
-namespace facebook::presto::protocol {
-// Loosly copied this here from NLOHMANN_JSON_SERIALIZE_ENUM()
-
-// NOLINTNEXTLINE: cppcoreguidelines-avoid-c-arrays
-static const std::pair NodeState_enum_table[] =
-    { // NOLINT: cert-err58-cpp
-        {NodeState::ACTIVE, "ACTIVE"},
-        {NodeState::INACTIVE, "INACTIVE"},
-        {NodeState::SHUTTING_DOWN, "SHUTTING_DOWN"}};
-void to_json(json& j, const NodeState& e) {
-  static_assert(std::is_enum::value, "NodeState must be an enum!");
-  const auto* it = std::find_if(
-      std::begin(NodeState_enum_table),
-      std::end(NodeState_enum_table),
-      [e](const std::pair& ej_pair) -> bool {
-        return ej_pair.first == e;
-      });
-  j = ((it != std::end(NodeState_enum_table))
-           ? it
-           : std::begin(NodeState_enum_table))
-          ->second;
-}
-void from_json(const json& j, NodeState& e) {
-  static_assert(std::is_enum::value, "NodeState must be an enum!");
-  const auto* it = std::find_if(
-      std::begin(NodeState_enum_table),
-      std::end(NodeState_enum_table),
-      [&j](const std::pair& ej_pair) -> bool {
-        return ej_pair.second == j;
-      });
-  e = ((it != std::end(NodeState_enum_table))
-           ? it
-           : std::begin(NodeState_enum_table))
-          ->first;
-}
-} // namespace facebook::presto::protocol
diff --git a/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.h b/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.h
index 946015a8a6f94..15656aa6e1780 100644
--- a/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.h
+++ b/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.h
@@ -1761,6 +1761,30 @@ void to_json(json& j, const MergeJoinNode& p);
 void from_json(const json& j, MergeJoinNode& p);
 } // namespace facebook::presto::protocol
 namespace facebook::presto::protocol {
+struct NodeLoadMetrics {
+  double cpuUsedPercent = {};
+  double memoryUsedInBytes = {};
+  int numQueuedDrivers = {};
+  bool cpuOverload = {};
+  bool memoryOverload = {};
+};
+void to_json(json& j, const NodeLoadMetrics& p);
+void from_json(const json& j, NodeLoadMetrics& p);
+} // namespace facebook::presto::protocol
+namespace facebook::presto::protocol {
+enum class NodeState { ACTIVE, INACTIVE, SHUTTING_DOWN };
+extern void to_json(json& j, const NodeState& e);
+extern void from_json(const json& j, NodeState& e);
+} // namespace facebook::presto::protocol
+namespace facebook::presto::protocol {
+struct NodeStats {
+  NodeState nodeState = {};
+  std::shared_ptr loadMetrics = {};
+};
+void to_json(json& j, const NodeStats& p);
+void from_json(const json& j, NodeStats& p);
+} // namespace facebook::presto::protocol
+namespace facebook::presto::protocol {
 struct NodeVersion {
   String version = {};
 };
@@ -2527,8 +2551,3 @@ struct WindowNode : public PlanNode {
 void to_json(json& j, const WindowNode& p);
 void from_json(const json& j, WindowNode& p);
 } // namespace facebook::presto::protocol
-namespace facebook::presto::protocol {
-enum class NodeState { ACTIVE, INACTIVE, SHUTTING_DOWN };
-extern void to_json(json& j, const NodeState& e);
-extern void from_json(const json& j, NodeState& e);
-} // namespace facebook::presto::protocol
diff --git a/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.yml b/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.yml
index 7f90f5d510c9b..bd710882c51d1 100644
--- a/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.yml
+++ b/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.yml
@@ -42,9 +42,6 @@ ExtraFields:
   RemoteTransactionHandle:
     Optional: dummy
 
-AddToOutput:
-  - NodeState
-
 AbstractClasses:
   ColumnHandle:
     super: JsonEncodedSubclass
@@ -345,3 +342,5 @@ JavaClasses:
   - presto-spi/src/main/java/com/facebook/presto/spi/plan/DeleteNode.java
   - presto-spi/src/main/java/com/facebook/presto/spi/plan/BaseInputDistribution.java
   - presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInFunctionKind.java
+  - presto-spi/src/main/java/com/facebook/presto/spi/NodeStats.java
+  - presto-spi/src/main/java/com/facebook/presto/spi/NodeLoadMetrics.java
diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/NodeLoadMetrics.java b/presto-spi/src/main/java/com/facebook/presto/spi/NodeLoadMetrics.java
new file mode 100644
index 0000000000000..493018b344587
--- /dev/null
+++ b/presto-spi/src/main/java/com/facebook/presto/spi/NodeLoadMetrics.java
@@ -0,0 +1,123 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.spi;
+
+import com.facebook.drift.annotations.ThriftConstructor;
+import com.facebook.drift.annotations.ThriftField;
+import com.facebook.drift.annotations.ThriftStruct;
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+import java.util.Objects;
+
+/**
+ * Load Statistics for a worker node in the Presto cluster along with the state.
+ */
+@ThriftStruct
+public class NodeLoadMetrics
+{
+    private final double cpuUsedPercent;
+    private final double memoryUsedInBytes;
+    private final int numQueuedDrivers;
+    private final boolean cpuOverload;
+    private final boolean memoryOverload;
+
+    @JsonCreator
+    @ThriftConstructor
+    public NodeLoadMetrics(
+            @JsonProperty("cpuUsedPercent") double cpuUsedPercent,
+            @JsonProperty("memoryUsedInBytes") double memoryUsedInBytes,
+            @JsonProperty("numQueuedDrivers") int numQueuedDrivers,
+            @JsonProperty("cpuOverload") boolean cpuOverload,
+            @JsonProperty("memoryOverload") boolean memoryOverload)
+    {
+        this.cpuUsedPercent = cpuUsedPercent;
+        this.memoryUsedInBytes = memoryUsedInBytes;
+        this.numQueuedDrivers = numQueuedDrivers;
+        this.cpuOverload = cpuOverload;
+        this.memoryOverload = memoryOverload;
+    }
+
+    @JsonProperty
+    @ThriftField(1)
+    public double getCpuUsedPercent()
+    {
+        return cpuUsedPercent;
+    }
+
+    @JsonProperty
+    @ThriftField(2)
+    public double getMemoryUsedInBytes()
+    {
+        return memoryUsedInBytes;
+    }
+
+    @JsonProperty
+    @ThriftField(3)
+    public int getNumQueuedDrivers()
+    {
+        return numQueuedDrivers;
+    }
+
+    @JsonProperty
+    @ThriftField(4)
+    public boolean getCpuOverload()
+    {
+        return cpuOverload;
+    }
+
+    @JsonProperty
+    @ThriftField(5)
+    public boolean getMemoryOverload()
+    {
+        return memoryOverload;
+    }
+
+    @Override
+    public boolean equals(Object o)
+    {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        NodeLoadMetrics that = (NodeLoadMetrics) o;
+        return Double.compare(cpuUsedPercent, that.cpuUsedPercent) == 0 &&
+                Double.compare(memoryUsedInBytes, that.memoryUsedInBytes) == 0 &&
+                numQueuedDrivers == that.numQueuedDrivers &&
+                cpuOverload == that.cpuOverload &&
+                memoryOverload == that.memoryOverload;
+    }
+
+    @Override
+    public int hashCode()
+    {
+        return Objects.hash(cpuUsedPercent, memoryUsedInBytes, numQueuedDrivers, cpuOverload, memoryOverload);
+    }
+
+    @Override
+    public String toString()
+    {
+        StringBuilder sb = new StringBuilder();
+        sb.append("NodeLoadMetrics{");
+        sb.append("cpuUsedPercent=").append(cpuUsedPercent);
+        sb.append(", memoryUsedInBytes=").append(memoryUsedInBytes);
+        sb.append(", numQueuedDrivers=").append(numQueuedDrivers);
+        sb.append(", cpuOverload=").append(cpuOverload);
+        sb.append(", memoryOverload=").append(memoryOverload);
+        sb.append('}');
+        return sb.toString();
+    }
+}
diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/NodeStats.java b/presto-spi/src/main/java/com/facebook/presto/spi/NodeStats.java
new file mode 100644
index 0000000000000..6eba1949d0199
--- /dev/null
+++ b/presto-spi/src/main/java/com/facebook/presto/spi/NodeStats.java
@@ -0,0 +1,90 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.spi;
+
+import com.facebook.drift.annotations.ThriftConstructor;
+import com.facebook.drift.annotations.ThriftField;
+import com.facebook.drift.annotations.ThriftStruct;
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+import java.util.Objects;
+import java.util.Optional;
+
+import static java.util.Objects.requireNonNull;
+
+/**
+ * Load Statistics for a node in the Presto cluster along with the state.
+ */
+@ThriftStruct
+public class NodeStats
+{
+    private final NodeState nodeState;
+    private final Optional loadMetrics;
+
+    @JsonCreator
+    @ThriftConstructor
+    public NodeStats(
+            @JsonProperty("nodeState") NodeState nodeState,
+            @JsonProperty("loadMetrics") Optional loadMetrics)
+    {
+        this.nodeState = requireNonNull(nodeState, "nodeState is null");
+        this.loadMetrics = loadMetrics;
+    }
+
+    @JsonProperty
+    @ThriftField(1)
+    public NodeState getNodeState()
+    {
+        return nodeState;
+    }
+
+    @JsonProperty
+    @ThriftField(2)
+    public Optional getLoadMetrics()
+    {
+        return loadMetrics;
+    }
+
+    @Override
+    public boolean equals(Object o)
+    {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        NodeStats nodeStats = (NodeStats) o;
+        return Objects.equals(nodeState, nodeStats.nodeState) &&
+                Objects.equals(loadMetrics, nodeStats.loadMetrics);
+    }
+
+    @Override
+    public int hashCode()
+    {
+        return Objects.hash(nodeState, loadMetrics);
+    }
+
+    @Override
+    public String toString()
+    {
+        StringBuilder sb = new StringBuilder();
+        sb.append("NodeStats{");
+        sb.append("nodeState=").append(nodeState);
+        sb.append(", loadMetrics=").append(loadMetrics);
+        sb.append('}');
+        return sb.toString();
+    }
+}
diff --git a/presto-tests/src/test/java/com/facebook/presto/server/TestServerInfoResource.java b/presto-tests/src/test/java/com/facebook/presto/server/TestServerInfoResource.java
index 75f743d5e710c..f4130226f75f7 100644
--- a/presto-tests/src/test/java/com/facebook/presto/server/TestServerInfoResource.java
+++ b/presto-tests/src/test/java/com/facebook/presto/server/TestServerInfoResource.java
@@ -22,10 +22,10 @@
 import com.facebook.drift.transport.netty.codec.Protocol;
 import com.facebook.presto.server.testing.TestingPrestoServer;
 import com.facebook.presto.spi.NodeState;
+import com.facebook.presto.spi.NodeStats;
 import com.facebook.presto.tests.DistributedQueryRunner;
 import com.google.common.collect.ImmutableMap;
 import org.testng.annotations.AfterClass;
-import org.testng.annotations.AfterGroups;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.BeforeGroups;
 import org.testng.annotations.DataProvider;
@@ -44,12 +44,14 @@
 import static com.facebook.presto.tests.tpch.TpchQueryRunner.createQueryRunnerWithNoClusterReadyCheck;
 import static org.testng.Assert.assertEquals;
 
-@Test(singleThreaded = true)
+@Test(singleThreaded = true, groups = "TestServerInfoResource")
 public class TestServerInfoResource
 {
     private HttpClient client;
-    private DistributedQueryRunner queryRunner;
-    private DistributedQueryRunner queryRunnerWithNoClusterReadyCheck;
+    private DistributedQueryRunner queryRunnerActive;
+    private DistributedQueryRunner queryRunnerInactiveResourceManagers;
+    private DistributedQueryRunner queryRunnerInactiveCoordinators;
+    private DistributedQueryRunner queryRunnerInactiveResourceGroup;
     private ThriftCodecManager thriftCodeManager;
 
     @BeforeClass
@@ -73,11 +75,28 @@ public void teardown()
         this.client = null;
     }
 
-    @AfterGroups(groups = {"createQueryRunner", "getServerStateWithoutRequiredResourceManagers", "getServerStateWithoutRequiredCoordinators", "getServerStateWithoutRequiredCoordinators", "createQueryRunnerWithNoClusterReadyCheckSkipLoadingResourceGroupConfigurationManager"})
+    @AfterClass(alwaysRun = true)
     public void serverTearDown()
     {
-        for (TestingPrestoServer server : queryRunner.getServers()) {
-            closeQuietly(server);
+        if (queryRunnerActive != null) {
+            for (TestingPrestoServer server : queryRunnerActive.getServers()) {
+                closeQuietly(server);
+            }
+        }
+        if (queryRunnerInactiveResourceManagers != null) {
+            for (TestingPrestoServer server : queryRunnerInactiveResourceManagers.getServers()) {
+                closeQuietly(server);
+            }
+        }
+        if (queryRunnerInactiveCoordinators != null) {
+            for (TestingPrestoServer server : queryRunnerInactiveCoordinators.getServers()) {
+                closeQuietly(server);
+            }
+        }
+        if (queryRunnerInactiveResourceGroup != null) {
+            for (TestingPrestoServer server : queryRunnerInactiveResourceGroup.getServers()) {
+                closeQuietly(server);
+            }
         }
     }
 
@@ -85,7 +104,7 @@ public void serverTearDown()
     public void createQueryRunnerSetup()
             throws Exception
     {
-        queryRunner = createQueryRunner(
+        queryRunnerActive = createQueryRunner(
                 ImmutableMap.of(),
                 ImmutableMap.of(),
                 ImmutableMap.of(),
@@ -96,17 +115,27 @@ public void createQueryRunnerSetup()
     @Test(timeOut = 30_000, groups = {"createQueryRunner"}, dataProvider = "thriftEncodingToggle")
     public void testGetServerStateWithRequiredResourceManagerCoordinators(boolean useThriftEncoding, Protocol thriftProtocol)
     {
-        TestingPrestoServer server = queryRunner.getCoordinator(0);
+        TestingPrestoServer server = queryRunnerActive.getCoordinator(0);
         URI uri = uriBuilderFrom(server.getBaseUrl().resolve("/v1/info/state")).build();
         NodeState state = getNodeState(uri, useThriftEncoding, thriftProtocol);
         assertEquals(state, NodeState.ACTIVE);
     }
 
+    @Test(timeOut = 30_000, groups = {"createQueryRunner"}, dataProvider = "thriftEncodingToggle", dependsOnMethods = "testGetServerStateWithRequiredResourceManagerCoordinators")
+    public void testGetServerStatsWithRequiredResourceManagerCoordinators(boolean useThriftEncoding, Protocol thriftProtocol)
+    {
+        TestingPrestoServer server = queryRunnerActive.getCoordinator(0);
+        URI uri = uriBuilderFrom(server.getBaseUrl().resolve("/v1/info/stats")).build();
+        NodeStats stats = getNodeStats(uri, useThriftEncoding, thriftProtocol);
+        assertEquals(stats.getNodeState(), NodeState.ACTIVE);
+        assertEquals(stats.getLoadMetrics().isPresent(), false);
+    }
+
     @BeforeGroups("getServerStateWithoutRequiredResourceManagers")
     public void createQueryRunnerWithNoClusterReadyCheckSetup()
             throws Exception
     {
-        queryRunner = createQueryRunnerWithNoClusterReadyCheck(
+        queryRunnerInactiveResourceManagers = createQueryRunnerWithNoClusterReadyCheck(
                         ImmutableMap.of(),
                         ImmutableMap.of(),
                         ImmutableMap.of(),
@@ -117,17 +146,27 @@ public void createQueryRunnerWithNoClusterReadyCheckSetup()
     @Test(timeOut = 30_000, groups = {"getServerStateWithoutRequiredResourceManagers"}, dataProvider = "thriftEncodingToggle")
     public void testGetServerStateWithoutRequiredResourceManagers(boolean useThriftEncoding, Protocol thriftProtocol)
     {
-        TestingPrestoServer server = queryRunner.getCoordinator(0);
+        TestingPrestoServer server = queryRunnerInactiveResourceManagers.getCoordinator(0);
         URI uri = uriBuilderFrom(server.getBaseUrl().resolve("/v1/info/state")).build();
         NodeState state = getNodeState(uri, useThriftEncoding, thriftProtocol);
         assertEquals(state, NodeState.INACTIVE);
     }
 
+    @Test(timeOut = 30_000, groups = {"getServerStateWithoutRequiredResourceManagers"}, dataProvider = "thriftEncodingToggle", dependsOnMethods = "testGetServerStateWithoutRequiredResourceManagers")
+    public void testGetServerStatsWithoutRequiredResourceManagers(boolean useThriftEncoding, Protocol thriftProtocol)
+    {
+        TestingPrestoServer server = queryRunnerInactiveResourceManagers.getCoordinator(0);
+        URI uri = uriBuilderFrom(server.getBaseUrl().resolve("/v1/info/stats")).build();
+        NodeStats stats = getNodeStats(uri, useThriftEncoding, thriftProtocol);
+        assertEquals(stats.getNodeState(), NodeState.INACTIVE);
+        assertEquals(stats.getLoadMetrics().isPresent(), false);
+    }
+
     @BeforeGroups("getServerStateWithoutRequiredCoordinators")
     public void getServerStateWithoutRequiredCoordinatorsSetup()
             throws Exception
     {
-        queryRunner = createQueryRunnerWithNoClusterReadyCheck(
+        queryRunnerInactiveCoordinators = createQueryRunnerWithNoClusterReadyCheck(
                 ImmutableMap.of(),
                 ImmutableMap.of(),
                 ImmutableMap.of(),
@@ -138,18 +177,28 @@ public void getServerStateWithoutRequiredCoordinatorsSetup()
     @Test(timeOut = 30_000, groups = {"getServerStateWithoutRequiredCoordinators"}, dataProvider = "thriftEncodingToggle")
     public void testGetServerStateWithoutRequiredCoordinators(boolean useThriftEncoding, Protocol thriftProtocol)
     {
-        TestingPrestoServer server = queryRunner.getCoordinator(0);
+        TestingPrestoServer server = queryRunnerInactiveCoordinators.getCoordinator(0);
         URI uri = uriBuilderFrom(server.getBaseUrl().resolve("/v1/info/state")).build();
         NodeState state = getNodeState(uri, useThriftEncoding, thriftProtocol);
 
         assertEquals(state, NodeState.INACTIVE);
     }
 
+    @Test(timeOut = 30_000, groups = {"getServerStateWithoutRequiredCoordinators"}, dataProvider = "thriftEncodingToggle", dependsOnMethods = "testGetServerStateWithoutRequiredCoordinators")
+    public void testGetServerStatsWithoutRequiredCoordinators(boolean useThriftEncoding, Protocol thriftProtocol)
+    {
+        TestingPrestoServer server = queryRunnerInactiveCoordinators.getCoordinator(0);
+        URI uri = uriBuilderFrom(server.getBaseUrl().resolve("/v1/info/stats")).build();
+        NodeStats stats = getNodeStats(uri, useThriftEncoding, thriftProtocol);
+        assertEquals(stats.getNodeState(), NodeState.INACTIVE);
+        assertEquals(stats.getLoadMetrics().isPresent(), false);
+    }
+
     @BeforeGroups("createQueryRunnerWithNoClusterReadyCheckSkipLoadingResourceGroupConfigurationManager")
     public void createQueryRunnerWithNoClusterReadyCheckSkipLoadingResourceGroupConfigurationManager()
             throws Exception
     {
-        queryRunner = createQueryRunnerWithNoClusterReadyCheck(
+        queryRunnerInactiveResourceGroup = createQueryRunnerWithNoClusterReadyCheck(
                 ImmutableMap.of(),
                 ImmutableMap.of(),
                 ImmutableMap.of(),
@@ -161,12 +210,22 @@ public void createQueryRunnerWithNoClusterReadyCheckSkipLoadingResourceGroupConf
     public void testGetServerStateWhenResourceGroupConfigurationManagerNotLoaded()
             throws Exception
     {
-        TestingPrestoServer server = queryRunner.getCoordinator(0);
+        TestingPrestoServer server = queryRunnerInactiveResourceGroup.getCoordinator(0);
         URI uri = uriBuilderFrom(server.getBaseUrl().resolve("/v1/info/state")).build();
         NodeState state = getNodeState(uri, false, null);
         assertEquals(state, NodeState.INACTIVE);
     }
 
+    @Test(groups = {"createQueryRunnerWithNoClusterReadyCheckSkipLoadingResourceGroupConfigurationManager"}, dependsOnMethods = "testGetServerStateWhenResourceGroupConfigurationManagerNotLoaded")
+    public void testGetServerStatsWhenResourceGroupConfigurationManagerNotLoaded()
+    {
+        TestingPrestoServer server = queryRunnerInactiveResourceGroup.getCoordinator(0);
+        URI uri = uriBuilderFrom(server.getBaseUrl().resolve("/v1/info/stats")).build();
+        NodeStats stats = getNodeStats(uri, false, null);
+        assertEquals(stats.getNodeState(), NodeState.INACTIVE);
+        assertEquals(stats.getLoadMetrics().isPresent(), false);
+    }
+
     private NodeState getNodeState(URI uri, boolean useThriftEncoding, Protocol thriftProtocol)
     {
         Request.Builder requestBuilder = useThriftEncoding ? ThriftRequestUtils.prepareThriftGet(thriftProtocol) : getJsonTransportBuilder(prepareGet());
@@ -181,4 +240,19 @@ private NodeState getNodeState(URI uri, boolean useThriftEncoding, Protocol thri
             return client.execute(request, createJsonResponseHandler(jsonCodec(NodeState.class)));
         }
     }
+
+    private NodeStats getNodeStats(URI uri, boolean useThriftEncoding, Protocol thriftProtocol)
+    {
+        Request.Builder requestBuilder = useThriftEncoding ? ThriftRequestUtils.prepareThriftGet(thriftProtocol) : getJsonTransportBuilder(prepareGet());
+        Request request = requestBuilder
+                .setHeader(PRESTO_USER, "user")
+                .setUri(uri)
+                .build();
+        if (useThriftEncoding) {
+            return client.execute(request, new ThriftResponseHandler<>(thriftCodeManager.getCodec(NodeStats.class))).getValue();
+        }
+        else {
+            return client.execute(request, createJsonResponseHandler(jsonCodec(NodeStats.class)));
+        }
+    }
 }

From b553f71021f32ef5b266725e331a9e4eacedff4c Mon Sep 17 00:00:00 2001
From: Pratik Joseph Dabre 
Date: Wed, 27 Aug 2025 15:51:35 -0700
Subject: [PATCH 014/113] [native] Return false for queries with spherical
 geometry in sidecar enabled clusters

---
 .../facebook/presto/common/type/TypeManager.java   |  5 +++++
 .../presto/common/type/TestingTypeManager.java     |  6 ++++++
 .../facebook/presto/hudi/TestingTypeManager.java   |  6 ++++++
 .../presto/iceberg/TestIcebergFileWriter.java      |  6 ++++++
 .../facebook/presto/execution/AddColumnTask.java   |  2 +-
 .../facebook/presto/execution/CreateTableTask.java |  2 +-
 .../BuiltInTypeAndFunctionNamespaceManager.java    | 14 +++++++++++++-
 .../presto/metadata/FunctionAndTypeManager.java    | 12 ++++++++++++
 .../facebook/presto/metadata/SignatureBinder.java  |  2 +-
 .../presto/sql/analyzer/ExpressionAnalyzer.java    |  2 +-
 .../presto/sql/analyzer/ExpressionTreeUtils.java   |  2 +-
 .../iterative/rule/ExtractSpatialJoins.java        |  5 +++++
 .../sql/rewrite/NativeExecutionTypeRewrite.java    |  2 +-
 .../presto/type/TestBuiltInTypeRegistry.java       |  2 +-
 .../sidecar/typemanager/NativeTypeManager.java     | 13 +++++++++++++
 .../presto/sidecar/TestNativeSidecarPlugin.java    | 10 ++++++++++
 .../presto/spi/type}/UnknownTypeException.java     |  2 +-
 17 files changed, 84 insertions(+), 9 deletions(-)
 rename {presto-main-base/src/main/java/com/facebook/presto => presto-spi/src/main/java/com/facebook/presto/spi/type}/UnknownTypeException.java (96%)

diff --git a/presto-common/src/main/java/com/facebook/presto/common/type/TypeManager.java b/presto-common/src/main/java/com/facebook/presto/common/type/TypeManager.java
index 277ef86ffb881..2dc2aa18cfaa2 100644
--- a/presto-common/src/main/java/com/facebook/presto/common/type/TypeManager.java
+++ b/presto-common/src/main/java/com/facebook/presto/common/type/TypeManager.java
@@ -44,4 +44,9 @@ default Collection getParametricTypes()
     {
         throw new UnsupportedOperationException();
     }
+
+    /**
+     * Checks for the existence of this type.
+     */
+    boolean hasType(TypeSignature signature);
 }
diff --git a/presto-common/src/test/java/com/facebook/presto/common/type/TestingTypeManager.java b/presto-common/src/test/java/com/facebook/presto/common/type/TestingTypeManager.java
index f2e6f847a48a5..e94a0a9465cf0 100644
--- a/presto-common/src/test/java/com/facebook/presto/common/type/TestingTypeManager.java
+++ b/presto-common/src/test/java/com/facebook/presto/common/type/TestingTypeManager.java
@@ -59,4 +59,10 @@ public List getTypes()
     {
         return ImmutableList.of(BOOLEAN, INTEGER, BIGINT, DOUBLE, VARCHAR, VARBINARY, TIMESTAMP, DATE, ID, HYPER_LOG_LOG);
     }
+
+    @Override
+    public boolean hasType(TypeSignature signature)
+    {
+        return getType(signature) != null;
+    }
 }
diff --git a/presto-hudi/src/test/java/com/facebook/presto/hudi/TestingTypeManager.java b/presto-hudi/src/test/java/com/facebook/presto/hudi/TestingTypeManager.java
index 8b9f9fb7a747a..c4a0658622fc6 100644
--- a/presto-hudi/src/test/java/com/facebook/presto/hudi/TestingTypeManager.java
+++ b/presto-hudi/src/test/java/com/facebook/presto/hudi/TestingTypeManager.java
@@ -63,4 +63,10 @@ public List getTypes()
     {
         return ImmutableList.of(BOOLEAN, INTEGER, BIGINT, DOUBLE, VARCHAR, VARBINARY, TIMESTAMP, DATE, HYPER_LOG_LOG);
     }
+
+    @Override
+    public boolean hasType(TypeSignature signature)
+    {
+        return getType(signature) != null;
+    }
 }
diff --git a/presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestIcebergFileWriter.java b/presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestIcebergFileWriter.java
index 4e14cf3a1c639..31f0729907da2 100644
--- a/presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestIcebergFileWriter.java
+++ b/presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestIcebergFileWriter.java
@@ -189,5 +189,11 @@ public List getTypes()
         {
             return ImmutableList.of(BooleanType.BOOLEAN, INTEGER, BIGINT, DoubleType.DOUBLE, VARCHAR, VARBINARY, TIMESTAMP, DATE, HYPER_LOG_LOG);
         }
+
+        @Override
+        public boolean hasType(TypeSignature signature)
+        {
+            return getType(signature) != null;
+        }
     }
 }
diff --git a/presto-main-base/src/main/java/com/facebook/presto/execution/AddColumnTask.java b/presto-main-base/src/main/java/com/facebook/presto/execution/AddColumnTask.java
index 3fbe29ff989a2..db9f05c402538 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/execution/AddColumnTask.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/execution/AddColumnTask.java
@@ -14,7 +14,6 @@
 package com.facebook.presto.execution;
 
 import com.facebook.presto.Session;
-import com.facebook.presto.UnknownTypeException;
 import com.facebook.presto.common.QualifiedObjectName;
 import com.facebook.presto.common.type.Type;
 import com.facebook.presto.metadata.Metadata;
@@ -25,6 +24,7 @@
 import com.facebook.presto.spi.TableHandle;
 import com.facebook.presto.spi.WarningCollector;
 import com.facebook.presto.spi.security.AccessControl;
+import com.facebook.presto.spi.type.UnknownTypeException;
 import com.facebook.presto.sql.analyzer.SemanticException;
 import com.facebook.presto.sql.tree.AddColumn;
 import com.facebook.presto.sql.tree.ColumnDefinition;
diff --git a/presto-main-base/src/main/java/com/facebook/presto/execution/CreateTableTask.java b/presto-main-base/src/main/java/com/facebook/presto/execution/CreateTableTask.java
index f2199283cf43b..306bbde275f6e 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/execution/CreateTableTask.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/execution/CreateTableTask.java
@@ -14,7 +14,6 @@
 package com.facebook.presto.execution;
 
 import com.facebook.presto.Session;
-import com.facebook.presto.UnknownTypeException;
 import com.facebook.presto.common.QualifiedObjectName;
 import com.facebook.presto.common.type.Type;
 import com.facebook.presto.metadata.Metadata;
@@ -28,6 +27,7 @@
 import com.facebook.presto.spi.constraints.PrimaryKeyConstraint;
 import com.facebook.presto.spi.constraints.TableConstraint;
 import com.facebook.presto.spi.security.AccessControl;
+import com.facebook.presto.spi.type.UnknownTypeException;
 import com.facebook.presto.sql.analyzer.SemanticException;
 import com.facebook.presto.sql.tree.ColumnDefinition;
 import com.facebook.presto.sql.tree.ConstraintSpecification;
diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInTypeAndFunctionNamespaceManager.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInTypeAndFunctionNamespaceManager.java
index f7565a43b47d2..1c3afcf3fa115 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInTypeAndFunctionNamespaceManager.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInTypeAndFunctionNamespaceManager.java
@@ -13,7 +13,6 @@
  */
 package com.facebook.presto.metadata;
 
-import com.facebook.presto.UnknownTypeException;
 import com.facebook.presto.common.CatalogSchemaName;
 import com.facebook.presto.common.Page;
 import com.facebook.presto.common.QualifiedObjectName;
@@ -236,6 +235,7 @@
 import com.facebook.presto.spi.function.SqlFunctionVisibility;
 import com.facebook.presto.spi.function.SqlInvokedFunction;
 import com.facebook.presto.spi.function.SqlInvokedScalarFunctionImplementation;
+import com.facebook.presto.spi.type.UnknownTypeException;
 import com.facebook.presto.sql.analyzer.FunctionsConfig;
 import com.facebook.presto.type.BigintOperators;
 import com.facebook.presto.type.BooleanOperators;
@@ -1285,6 +1285,18 @@ public Type getType(TypeSignature typeSignature)
         }
     }
 
+    @Override
+    public boolean hasType(TypeSignature typeSignature)
+    {
+        try {
+            getType(typeSignature);
+            return true;
+        }
+        catch (UnknownTypeException e) {
+            return false;
+        }
+    }
+
     @Override
     public Type getParameterizedType(String baseTypeName, List typeParameters)
     {
diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionAndTypeManager.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionAndTypeManager.java
index a409fdddb4afc..88036bbbe55bc 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionAndTypeManager.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionAndTypeManager.java
@@ -278,6 +278,12 @@ public Collection getParametricTypes()
                 return FunctionAndTypeManager.this.getParametricTypes();
             }
 
+            @Override
+            public boolean hasType(TypeSignature signature)
+            {
+                return FunctionAndTypeManager.this.hasType(signature);
+            }
+
             @Override
             public Collection listBuiltInFunctions()
             {
@@ -395,6 +401,12 @@ public Type getType(TypeSignature signature)
         return getUserDefinedType(signature);
     }
 
+    @Override
+    public boolean hasType(TypeSignature signature)
+    {
+        return servingTypeManager.get().hasType(signature);
+    }
+
     @Override
     public Type getParameterizedType(String baseTypeName, List typeParameters)
     {
diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/SignatureBinder.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/SignatureBinder.java
index a71844c54b7d4..bd383aa6ef87b 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/metadata/SignatureBinder.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/SignatureBinder.java
@@ -13,7 +13,6 @@
  */
 package com.facebook.presto.metadata;
 
-import com.facebook.presto.UnknownTypeException;
 import com.facebook.presto.common.type.FunctionType;
 import com.facebook.presto.common.type.NamedTypeSignature;
 import com.facebook.presto.common.type.ParameterKind;
@@ -25,6 +24,7 @@
 import com.facebook.presto.spi.function.LongVariableConstraint;
 import com.facebook.presto.spi.function.Signature;
 import com.facebook.presto.spi.function.TypeVariableConstraint;
+import com.facebook.presto.spi.type.UnknownTypeException;
 import com.facebook.presto.sql.analyzer.TypeSignatureProvider;
 import com.google.common.base.VerifyException;
 import com.google.common.collect.ImmutableList;
diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/ExpressionAnalyzer.java b/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/ExpressionAnalyzer.java
index b7446727feb8b..5a6f4c787a0ff 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/ExpressionAnalyzer.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/ExpressionAnalyzer.java
@@ -14,7 +14,6 @@
 package com.facebook.presto.sql.analyzer;
 
 import com.facebook.presto.Session;
-import com.facebook.presto.UnknownTypeException;
 import com.facebook.presto.common.ErrorCode;
 import com.facebook.presto.common.QualifiedObjectName;
 import com.facebook.presto.common.Subfield;
@@ -47,6 +46,7 @@
 import com.facebook.presto.spi.function.SqlInvokedFunction;
 import com.facebook.presto.spi.security.AccessControl;
 import com.facebook.presto.spi.security.DenyAllAccessControl;
+import com.facebook.presto.spi.type.UnknownTypeException;
 import com.facebook.presto.sql.parser.SqlParser;
 import com.facebook.presto.sql.planner.TypeProvider;
 import com.facebook.presto.sql.relational.FunctionResolution;
diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/ExpressionTreeUtils.java b/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/ExpressionTreeUtils.java
index 7ad6060e5a90d..048e57ee7138d 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/ExpressionTreeUtils.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/ExpressionTreeUtils.java
@@ -13,13 +13,13 @@
  */
 package com.facebook.presto.sql.analyzer;
 
-import com.facebook.presto.UnknownTypeException;
 import com.facebook.presto.common.type.EnumType;
 import com.facebook.presto.common.type.Type;
 import com.facebook.presto.common.type.TypeWithName;
 import com.facebook.presto.spi.SourceLocation;
 import com.facebook.presto.spi.function.FunctionHandle;
 import com.facebook.presto.spi.relation.VariableReferenceExpression;
+import com.facebook.presto.spi.type.UnknownTypeException;
 import com.facebook.presto.sql.tree.ArrayConstructor;
 import com.facebook.presto.sql.tree.Cast;
 import com.facebook.presto.sql.tree.ComparisonExpression;
diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/iterative/rule/ExtractSpatialJoins.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/iterative/rule/ExtractSpatialJoins.java
index 851e40b05825f..4b30f9de7f9aa 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/iterative/rule/ExtractSpatialJoins.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/iterative/rule/ExtractSpatialJoins.java
@@ -469,6 +469,11 @@ else if (alignment < 0) {
 
     private static boolean isSphericalJoin(Metadata metadata, RowExpression firstArgument, RowExpression secondArgument)
     {
+        // In sidecar-enabled clusters, SphericalGeography isn't a supported type.
+        // If SphericalGeography is not supported, it can be assumed that this join isn't a spherical join, hence returning False.
+        if (!metadata.getFunctionAndTypeManager().hasType(SPHERICAL_GEOGRAPHY_TYPE_SIGNATURE)) {
+            return false;
+        }
         Type sphericalGeographyType = metadata.getType(SPHERICAL_GEOGRAPHY_TYPE_SIGNATURE);
         return firstArgument.getType().equals(sphericalGeographyType) || secondArgument.getType().equals(sphericalGeographyType);
     }
diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/rewrite/NativeExecutionTypeRewrite.java b/presto-main-base/src/main/java/com/facebook/presto/sql/rewrite/NativeExecutionTypeRewrite.java
index c0abeb230ae87..d458bccba14e8 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/sql/rewrite/NativeExecutionTypeRewrite.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/sql/rewrite/NativeExecutionTypeRewrite.java
@@ -16,7 +16,6 @@
 import com.facebook.airlift.log.Logger;
 import com.facebook.presto.Session;
 import com.facebook.presto.SystemSessionProperties;
-import com.facebook.presto.UnknownTypeException;
 import com.facebook.presto.common.type.BigintEnumType;
 import com.facebook.presto.common.type.EnumType;
 import com.facebook.presto.common.type.Type;
@@ -25,6 +24,7 @@
 import com.facebook.presto.metadata.Metadata;
 import com.facebook.presto.spi.WarningCollector;
 import com.facebook.presto.spi.security.AccessControl;
+import com.facebook.presto.spi.type.UnknownTypeException;
 import com.facebook.presto.sql.analyzer.FunctionAndTypeResolver;
 import com.facebook.presto.sql.analyzer.QueryExplainer;
 import com.facebook.presto.sql.analyzer.SemanticException;
diff --git a/presto-main-base/src/test/java/com/facebook/presto/type/TestBuiltInTypeRegistry.java b/presto-main-base/src/test/java/com/facebook/presto/type/TestBuiltInTypeRegistry.java
index b5a4b8f49f3d5..0c626e5f0fc3d 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/type/TestBuiltInTypeRegistry.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/type/TestBuiltInTypeRegistry.java
@@ -13,11 +13,11 @@
  */
 package com.facebook.presto.type;
 
-import com.facebook.presto.UnknownTypeException;
 import com.facebook.presto.common.type.Type;
 import com.facebook.presto.common.type.TypeSignature;
 import com.facebook.presto.metadata.FunctionAndTypeManager;
 import com.facebook.presto.metadata.OperatorNotFoundException;
+import com.facebook.presto.spi.type.UnknownTypeException;
 import com.google.common.collect.ImmutableSet;
 import org.testng.annotations.Test;
 
diff --git a/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/typemanager/NativeTypeManager.java b/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/typemanager/NativeTypeManager.java
index a1cd8cb0f00af..e86ce277c0f31 100644
--- a/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/typemanager/NativeTypeManager.java
+++ b/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/typemanager/NativeTypeManager.java
@@ -21,6 +21,7 @@
 import com.facebook.presto.common.type.TypeSignature;
 import com.facebook.presto.common.type.TypeSignatureParameter;
 import com.facebook.presto.common.type.VarcharType;
+import com.facebook.presto.spi.type.UnknownTypeException;
 import com.google.common.cache.CacheBuilder;
 import com.google.common.cache.CacheLoader;
 import com.google.common.cache.LoadingCache;
@@ -150,6 +151,18 @@ public Type getType(TypeSignature typeSignature)
         }
     }
 
+    @Override
+    public boolean hasType(TypeSignature typeSignature)
+    {
+        try {
+            getType(typeSignature);
+            return true;
+        }
+        catch (UnknownTypeException e) {
+            return false;
+        }
+    }
+
     @Override
     public Collection getParametricTypes()
     {
diff --git a/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java b/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java
index 50bb5c252ffd2..1c5d05d55b274 100644
--- a/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java
+++ b/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java
@@ -393,6 +393,16 @@ public void testGeometryQueries()
         assertQuery("SELECT " +
                 "ST_DISTANCE(ST_POINT(a.nationkey, a.regionkey), ST_POINT(b.nationkey, b.regionkey)) " +
                 "FROM nation a JOIN nation b ON a.nationkey < b.nationkey");
+        assertQueryFails(
+                "WITH regions(name, geom) AS (VALUES" +
+                        "        ('A', ST_GeometryFromText('POLYGON ((0 0, 0 5, 5 5, 5 0, 0 0))'))," +
+                        "        ('B', ST_GeometryFromText('POLYGON ((5 0, 5 5, 10 5, 10 0, 5 0))')))," +
+                        "points(id, geom) AS (VALUES" +
+                        "        ('P1', ST_Point(1, 1))," +
+                        "        ('P2', ST_Point(6, 1))," +
+                        "        ('P3', ST_Point(8, 4)))" +
+                        "SELECT p.id, r.name FROM points p LEFT JOIN regions r ON ST_Within(p.geom, r.geom)",
+                "Error from native plan checker: .SpatialJoinNode no abstract type PlanNode ");
     }
 
     private String generateRandomTableName()
diff --git a/presto-main-base/src/main/java/com/facebook/presto/UnknownTypeException.java b/presto-spi/src/main/java/com/facebook/presto/spi/type/UnknownTypeException.java
similarity index 96%
rename from presto-main-base/src/main/java/com/facebook/presto/UnknownTypeException.java
rename to presto-spi/src/main/java/com/facebook/presto/spi/type/UnknownTypeException.java
index 984f953cbc0ee..9138efcf5d6ae 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/UnknownTypeException.java
+++ b/presto-spi/src/main/java/com/facebook/presto/spi/type/UnknownTypeException.java
@@ -11,7 +11,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.facebook.presto;
+package com.facebook.presto.spi.type;
 
 import com.facebook.presto.common.type.TypeSignature;
 import com.facebook.presto.spi.PrestoException;

From 44cdb14f683183f25de005f7e7486f7b51db5335 Mon Sep 17 00:00:00 2001
From: Hazmi 
Date: Wed, 6 Aug 2025 17:13:56 +0300
Subject: [PATCH 015/113] Added DecimalType WriteMapping to QueryBuilder

WriteMapping support for decimal type is already present for writing values but is missing from the query builder.
This PR adds the write function to the query builder buildSql function
---
 .../jdbc/mapping/StandardColumnMappings.java   |  3 +++
 .../mysql/TestMySqlIntegrationSmokeTest.java   | 18 ++++++++++++++++++
 2 files changed, 21 insertions(+)

diff --git a/presto-base-jdbc/src/main/java/com/facebook/presto/plugin/jdbc/mapping/StandardColumnMappings.java b/presto-base-jdbc/src/main/java/com/facebook/presto/plugin/jdbc/mapping/StandardColumnMappings.java
index 3d705e4a2e94e..78b60ec7bb5b4 100644
--- a/presto-base-jdbc/src/main/java/com/facebook/presto/plugin/jdbc/mapping/StandardColumnMappings.java
+++ b/presto-base-jdbc/src/main/java/com/facebook/presto/plugin/jdbc/mapping/StandardColumnMappings.java
@@ -428,6 +428,9 @@ else if (type.equals(DOUBLE)) {
         else if (type instanceof CharType || type instanceof VarcharType) {
             return Optional.of(charWriteMapping());
         }
+        else if (type instanceof DecimalType) {
+            return Optional.of(decimalWriteMapping((DecimalType) type));
+        }
         else if (type.equals(DateType.DATE)) {
             return Optional.of(dateWriteMapping());
         }
diff --git a/presto-mysql/src/test/java/com/facebook/presto/plugin/mysql/TestMySqlIntegrationSmokeTest.java b/presto-mysql/src/test/java/com/facebook/presto/plugin/mysql/TestMySqlIntegrationSmokeTest.java
index adbbd0b928a35..f60a227fd3861 100644
--- a/presto-mysql/src/test/java/com/facebook/presto/plugin/mysql/TestMySqlIntegrationSmokeTest.java
+++ b/presto-mysql/src/test/java/com/facebook/presto/plugin/mysql/TestMySqlIntegrationSmokeTest.java
@@ -216,6 +216,24 @@ public void testMysqlGeometry()
         assertUpdate("DROP TABLE tpch.test_geometry");
     }
 
+    @Test
+    public void testMysqlDecimal()
+    {
+        assertUpdate("CREATE TABLE test_decimal (d DECIMAL(10, 2))");
+
+        assertUpdate("INSERT INTO test_decimal VALUES (123.45)", 1);
+        assertUpdate("INSERT INTO test_decimal VALUES (67890.12)", 1);
+        assertUpdate("INSERT INTO test_decimal VALUES (0.99)", 1);
+
+        assertQuery(
+                "SELECT d FROM test_decimal WHERE d<200.00 AND d>0.00",
+                "VALUES " +
+                        "CAST('123.45' AS DECIMAL), " +
+                        "CAST('0.99' AS DECIMAL)");
+
+        assertUpdate("DROP TABLE test_decimal");
+    }
+
     @Test
     public void testCharTrailingSpace()
             throws Exception

From 1d61f3eedcc69ed213b0fda297d44341560acbf4 Mon Sep 17 00:00:00 2001
From: Hazmi 
Date: Mon, 18 Aug 2025 12:26:46 +0300
Subject: [PATCH 016/113] Add TIME, TIME_WITH_TIMEZONE,
 TIMESTAMP_WITH_TIME_ZONE to standard jdbc write mappings

These types are missing in the new write mapping interface. If implemented, this will add them back.
---
 .../jdbc/mapping/StandardColumnMappings.java  |  9 ++++++++
 .../mysql/TestMySqlIntegrationSmokeTest.java  | 22 +++++++++++++++++++
 2 files changed, 31 insertions(+)

diff --git a/presto-base-jdbc/src/main/java/com/facebook/presto/plugin/jdbc/mapping/StandardColumnMappings.java b/presto-base-jdbc/src/main/java/com/facebook/presto/plugin/jdbc/mapping/StandardColumnMappings.java
index 78b60ec7bb5b4..ed4d18eab0cfa 100644
--- a/presto-base-jdbc/src/main/java/com/facebook/presto/plugin/jdbc/mapping/StandardColumnMappings.java
+++ b/presto-base-jdbc/src/main/java/com/facebook/presto/plugin/jdbc/mapping/StandardColumnMappings.java
@@ -396,6 +396,15 @@ else if (type instanceof DateType) {
         else if (type instanceof TimestampType) {
             return Optional.of(timestampWriteMapping((TimestampType) type));
         }
+        else if (type.equals(TIME)) {
+            return Optional.of(timeWriteMapping());
+        }
+        else if (type.equals(TIME_WITH_TIME_ZONE)) {
+            return Optional.of(timeWithTimeZoneWriteMapping());
+        }
+        else if (type.equals(TIMESTAMP_WITH_TIME_ZONE)) {
+            return Optional.of(timestampWithTimeZoneWriteMapping());
+        }
         else if (type.equals(UUID)) {
             return Optional.of(uuidWriteMapping());
         }
diff --git a/presto-mysql/src/test/java/com/facebook/presto/plugin/mysql/TestMySqlIntegrationSmokeTest.java b/presto-mysql/src/test/java/com/facebook/presto/plugin/mysql/TestMySqlIntegrationSmokeTest.java
index f60a227fd3861..65b97f1788d7d 100644
--- a/presto-mysql/src/test/java/com/facebook/presto/plugin/mysql/TestMySqlIntegrationSmokeTest.java
+++ b/presto-mysql/src/test/java/com/facebook/presto/plugin/mysql/TestMySqlIntegrationSmokeTest.java
@@ -234,6 +234,28 @@ public void testMysqlDecimal()
         assertUpdate("DROP TABLE test_decimal");
     }
 
+    @Test
+    public void testMysqlTime()
+    {
+        assertUpdate("CREATE TABLE test_time (datatype_time time)");
+
+        assertUpdate("INSERT INTO test_time VALUES (time '01:02:03.456')", 1);
+
+        assertQuery(
+                "SELECT datatype_time FROM test_time",
+                "VALUES " +
+                        "CAST('01:02:03.456' AS time)");
+
+        assertUpdate("DROP TABLE test_time");
+    }
+
+    @Test
+    public void testMysqlUnsupportedTimeTypes()
+    {
+        assertQueryFails("CREATE TABLE test_timestamp_with_timezone (timestamp_with_time_zone timestamp with time zone)", "Unsupported column type: timestamp with time zone");
+        assertQueryFails("CREATE TABLE test_time_with_timezone (time_with_with_time_zone time with time zone)", "Unsupported column type: time with time zone");
+    }
+
     @Test
     public void testCharTrailingSpace()
             throws Exception

From 024d57a248da3e91bd69c4462d573b296a1e610d Mon Sep 17 00:00:00 2001
From: Hazmi 
Date: Thu, 24 Jul 2025 17:14:52 +0300
Subject: [PATCH 017/113] Added `iceberg.engine.hive.lock-enabled`
 configuration

Added the `iceberg.engine.hive.lock-enabled` to enable or disable table locks
when iceberg accesses a hive table. This can be overridden with the table property
`engine.hive.lock-enabled`
---
 .../src/main/sphinx/connector/iceberg.rst     |  4 +
 .../hive/metastore/ExtendedHiveMetastore.java |  3 +
 .../InMemoryCachingHiveMetastore.java         | 13 ++++
 .../metastore/RecordingHiveMetastore.java     |  7 ++
 .../metastore/file/FileHiveMetastore.java     | 41 +++++++++-
 .../metastore/glue/GlueHiveMetastore.java     | 25 +++++++
 .../thrift/BridgingHiveMetastore.java         | 19 +++++
 .../hive/metastore/thrift/HiveMetastore.java  |  3 +
 .../metastore/thrift/HiveMetastoreClient.java |  4 +
 .../metastore/thrift/ThriftHiveMetastore.java | 30 ++++++++
 .../thrift/ThriftHiveMetastoreClient.java     |  8 ++
 .../thrift/ThriftHiveMetastoreStats.java      |  8 ++
 .../metastore/UnimplementedHiveMetastore.java |  7 ++
 .../thrift/InMemoryHiveMetastore.java         |  7 ++
 .../thrift/MockHiveMetastoreClient.java       |  8 ++
 .../presto/iceberg/HiveTableOperations.java   | 35 +++++++--
 .../IcebergHiveTableOperationsConfig.java     | 14 ++++
 .../hive/TestHiveTableOperationsConfig.java   |  7 +-
 .../hive/TestIcebergDistributedHive.java      | 74 +++++++++++++++++++
 19 files changed, 307 insertions(+), 10 deletions(-)

diff --git a/presto-docs/src/main/sphinx/connector/iceberg.rst b/presto-docs/src/main/sphinx/connector/iceberg.rst
index 2f8cec12f1c35..9267f5a8d8df1 100644
--- a/presto-docs/src/main/sphinx/connector/iceberg.rst
+++ b/presto-docs/src/main/sphinx/connector/iceberg.rst
@@ -93,6 +93,10 @@ Property Name                                            Description
 
 ``iceberg.hive.table-refresh.backoff-scale-factor``      The multiple used to scale subsequent wait time between       4.0
                                                          retries.
+
+``iceberg.engine.hive.lock-enabled``                     Whether to use locks to ensure atomicity of commits.          true
+                                                         This will turn off locks but is overridden at a table level
+                                                         with the table configuration ``engine.hive.lock-enabled``.
 ======================================================== ============================================================= ============
 
 Nessie catalog
diff --git a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/ExtendedHiveMetastore.java b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/ExtendedHiveMetastore.java
index dc972904a0ac4..4cfdc0d799a86 100644
--- a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/ExtendedHiveMetastore.java
+++ b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/ExtendedHiveMetastore.java
@@ -31,6 +31,7 @@
 import java.util.Optional;
 import java.util.Set;
 import java.util.function.Function;
+import java.util.function.Supplier;
 
 public interface ExtendedHiveMetastore
 {
@@ -89,6 +90,8 @@ default void dropTableFromMetastore(MetastoreContext metastoreContext, String da
      */
     MetastoreOperationResult replaceTable(MetastoreContext metastoreContext, String databaseName, String tableName, Table newTable, PrincipalPrivileges principalPrivileges);
 
+    MetastoreOperationResult persistTable(MetastoreContext metastoreContext, String databaseName, String tableName, Table newTable, PrincipalPrivileges principalPrivileges, Supplier update, Map additionalParameters);
+
     MetastoreOperationResult renameTable(MetastoreContext metastoreContext, String databaseName, String tableName, String newDatabaseName, String newTableName);
 
     MetastoreOperationResult addColumn(MetastoreContext metastoreContext, String databaseName, String tableName, String columnName, HiveType columnType, String columnComment);
diff --git a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/InMemoryCachingHiveMetastore.java b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/InMemoryCachingHiveMetastore.java
index f281c0f600c3b..fdf142b4a25ef 100644
--- a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/InMemoryCachingHiveMetastore.java
+++ b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/InMemoryCachingHiveMetastore.java
@@ -49,6 +49,7 @@
 import java.util.concurrent.ThreadLocalRandom;
 import java.util.function.Function;
 import java.util.function.Predicate;
+import java.util.function.Supplier;
 
 import static com.facebook.presto.hive.HiveErrorCode.HIVE_CORRUPTED_PARTITION_CACHE;
 import static com.facebook.presto.hive.HiveErrorCode.HIVE_PARTITION_DROPPED_DURING_QUERY;
@@ -450,6 +451,18 @@ private Map, PartitionStatistics> loadPartition
         return result.build();
     }
 
+    @Override
+    public MetastoreOperationResult persistTable(MetastoreContext metastoreContext, String databaseName, String tableName, Table newTable, PrincipalPrivileges principalPrivileges, Supplier update, Map additionalParameters)
+    {
+        try {
+            return getDelegate().persistTable(metastoreContext, databaseName, tableName, newTable, principalPrivileges, update, additionalParameters);
+        }
+        finally {
+            invalidateTableCache(databaseName, tableName);
+            invalidateTableCache(newTable.getDatabaseName(), newTable.getTableName());
+        }
+    }
+
     @Override
     public void updateTableStatistics(MetastoreContext metastoreContext, String databaseName, String tableName, Function update)
     {
diff --git a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/RecordingHiveMetastore.java b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/RecordingHiveMetastore.java
index 6163ae58b8d55..4c70b9005690c 100644
--- a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/RecordingHiveMetastore.java
+++ b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/RecordingHiveMetastore.java
@@ -315,6 +315,13 @@ public MetastoreOperationResult replaceTable(MetastoreContext metastoreContext,
         return delegate.replaceTable(metastoreContext, databaseName, tableName, newTable, principalPrivileges);
     }
 
+    @Override
+    public MetastoreOperationResult persistTable(MetastoreContext metastoreContext, String databaseName, String tableName, Table newTable, PrincipalPrivileges principalPrivileges, Supplier update, Map additionalParameters)
+    {
+        verifyRecordingMode();
+        return delegate.persistTable(metastoreContext, databaseName, tableName, newTable, principalPrivileges, update, additionalParameters);
+    }
+
     @Override
     public MetastoreOperationResult renameTable(MetastoreContext metastoreContext, String databaseName, String tableName, String newDatabaseName, String newTableName)
     {
diff --git a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/file/FileHiveMetastore.java b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/file/FileHiveMetastore.java
index 0d2953f65046d..ee1879b73e698 100644
--- a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/file/FileHiveMetastore.java
+++ b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/file/FileHiveMetastore.java
@@ -86,6 +86,7 @@
 import java.util.Optional;
 import java.util.Set;
 import java.util.function.Function;
+import java.util.function.Supplier;
 
 import static com.facebook.presto.hive.HiveErrorCode.HIVE_METASTORE_ERROR;
 import static com.facebook.presto.hive.HiveErrorCode.HIVE_PARTITION_DROPPED_DURING_QUERY;
@@ -134,6 +135,7 @@ public class FileHiveMetastore
 
     protected final HdfsEnvironment hdfsEnvironment;
     protected final HdfsContext hdfsContext;
+
     protected final FileSystem metadataFileSystem;
 
     private final Path catalogDirectory;
@@ -477,6 +479,44 @@ public synchronized MetastoreOperationResult replaceTable(MetastoreContext metas
         return EMPTY_RESULT;
     }
 
+    @Override
+    public synchronized MetastoreOperationResult persistTable(
+            MetastoreContext metastoreContext,
+            String databaseName,
+            String tableName,
+            Table newTable,
+            PrincipalPrivileges principalPrivileges,
+            Supplier update,
+            Map additionalParameters)
+    {
+        checkArgument(!newTable.getTableType().equals(TEMPORARY_TABLE), "temporary tables must never be stored in the metastore");
+
+        Table oldTable = getRequiredTable(metastoreContext, databaseName, tableName);
+        validateReplaceTableType(oldTable, newTable);
+        if (!oldTable.getDatabaseName().equals(databaseName) || !oldTable.getTableName().equals(tableName)) {
+            throw new PrestoException(HIVE_METASTORE_ERROR, "Replacement table must have same name");
+        }
+
+        Path tableMetadataDirectory = getTableMetadataDirectory(oldTable);
+
+        deleteTablePrivileges(oldTable);
+        for (Entry> entry : principalPrivileges.getUserPrivileges().asMap().entrySet()) {
+            setTablePrivileges(metastoreContext, new PrestoPrincipal(USER, entry.getKey()), databaseName, tableName, entry.getValue());
+        }
+        for (Entry> entry : principalPrivileges.getRolePrivileges().asMap().entrySet()) {
+            setTablePrivileges(metastoreContext, new PrestoPrincipal(ROLE, entry.getKey()), databaseName, tableName, entry.getValue());
+        }
+        PartitionStatistics updatedStatistics = update.get();
+
+        TableMetadata updatedMetadata = new TableMetadata(newTable)
+                .withParameters(updateStatisticsParameters(newTable.getParameters(), updatedStatistics.getBasicStatistics()))
+                .withColumnStatistics(updatedStatistics.getColumnStatistics());
+
+        writeSchemaFile("table", tableMetadataDirectory, tableCodec, updatedMetadata, true);
+
+        return EMPTY_RESULT;
+    }
+
     @Override
     public synchronized MetastoreOperationResult renameTable(MetastoreContext metastoreContext, String databaseName, String tableName, String newDatabaseName, String newTableName)
     {
@@ -1292,7 +1332,6 @@ private List getChildSchemaDirectories(Path metadataDirectory)
             if (!metadataFileSystem.isDirectory(metadataDirectory)) {
                 return ImmutableList.of();
             }
-
             ImmutableList.Builder childSchemaDirectories = ImmutableList.builder();
             for (FileStatus child : metadataFileSystem.listStatus(metadataDirectory)) {
                 if (!child.isDirectory()) {
diff --git a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/glue/GlueHiveMetastore.java b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/glue/GlueHiveMetastore.java
index 20876e451db2a..ef782cd889fac 100644
--- a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/glue/GlueHiveMetastore.java
+++ b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/glue/GlueHiveMetastore.java
@@ -116,6 +116,7 @@
 import java.util.concurrent.ExecutorCompletionService;
 import java.util.concurrent.Future;
 import java.util.function.Function;
+import java.util.function.Supplier;
 
 import static com.facebook.presto.hive.HiveErrorCode.HIVE_METASTORE_ERROR;
 import static com.facebook.presto.hive.HiveErrorCode.HIVE_PARTITION_DROPPED_DURING_QUERY;
@@ -589,6 +590,30 @@ public MetastoreOperationResult replaceTable(MetastoreContext metastoreContext,
         }
     }
 
+    public MetastoreOperationResult persistTable(MetastoreContext metastoreContext, String databaseName, String tableName, Table newTable, PrincipalPrivileges principalPrivileges, Supplier update, Map additionalParameters)
+    {
+        PartitionStatistics updatedStatistics = update.get();
+        if (!updatedStatistics.getColumnStatistics().isEmpty()) {
+            throw new PrestoException(NOT_SUPPORTED, "Glue metastore does not support column level statistics");
+        }
+        try {
+            TableInput newTableInput = GlueInputConverter.convertTable(newTable);
+            newTableInput.setParameters(updateStatisticsParameters(newTableInput.getParameters(), updatedStatistics.getBasicStatistics()));
+            stats.getUpdateTable().record(() -> glueClient.updateTable(new UpdateTableRequest()
+                    .withCatalogId(catalogId)
+                    .withDatabaseName(databaseName)
+                    .withTableInput(newTableInput)));
+
+            return EMPTY_RESULT;
+        }
+        catch (EntityNotFoundException e) {
+            throw new TableNotFoundException(new SchemaTableName(databaseName, tableName));
+        }
+        catch (AmazonServiceException e) {
+            throw new PrestoException(HIVE_METASTORE_ERROR, e);
+        }
+    }
+
     @Override
     public MetastoreOperationResult renameTable(MetastoreContext metastoreContext, String databaseName, String tableName, String newDatabaseName, String newTableName)
     {
diff --git a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/BridgingHiveMetastore.java b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/BridgingHiveMetastore.java
index d80d5f16570ef..e3787b7ebe91a 100644
--- a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/BridgingHiveMetastore.java
+++ b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/BridgingHiveMetastore.java
@@ -45,7 +45,10 @@
 import com.facebook.presto.spi.statistics.ColumnStatisticType;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Maps;
 import jakarta.inject.Inject;
+import org.apache.hadoop.hive.common.StatsSetupConst;
+import org.apache.hadoop.hive.metastore.api.EnvironmentContext;
 import org.apache.hadoop.hive.metastore.api.FieldSchema;
 
 import java.util.List;
@@ -53,6 +56,7 @@
 import java.util.Optional;
 import java.util.Set;
 import java.util.function.Function;
+import java.util.function.Supplier;
 import java.util.stream.Collectors;
 
 import static com.facebook.presto.hive.metastore.MetastoreUtil.getPartitionNamesWithEmptyVersion;
@@ -276,6 +280,21 @@ private MetastoreOperationResult alterTable(MetastoreContext metastoreContext, S
         return delegate.alterTable(metastoreContext, databaseName, tableName, table);
     }
 
+    private MetastoreOperationResult alterTableWithEnvironmentContext(MetastoreContext metastoreContext, String databaseName, String tableName, org.apache.hadoop.hive.metastore.api.Table table, EnvironmentContext environmentContext)
+    {
+        return delegate.alterTableWithEnvironmentContext(metastoreContext, databaseName, tableName, table, environmentContext);
+    }
+
+    @Override
+    public MetastoreOperationResult persistTable(MetastoreContext metastoreContext, String databaseName, String tableName, Table newTable, PrincipalPrivileges principalPrivileges, Supplier update, Map additionalParameters)
+    {
+        checkArgument(!newTable.getTableType().equals(TEMPORARY_TABLE), "temporary tables must never be stored in the metastore");
+        Map env = Maps.newHashMapWithExpectedSize(additionalParameters.size() + 1);
+        env.putAll(additionalParameters);
+        env.put(StatsSetupConst.DO_NOT_UPDATE_STATS, StatsSetupConst.TRUE);
+        return alterTableWithEnvironmentContext(metastoreContext, databaseName, tableName, toMetastoreApiTable(newTable, principalPrivileges, metastoreContext.getColumnConverter()), new EnvironmentContext(env));
+    }
+
     @Override
     public Optional getPartition(MetastoreContext metastoreContext, String databaseName, String tableName, List partitionValues)
     {
diff --git a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/HiveMetastore.java b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/HiveMetastore.java
index 0d499b18eca8d..f66484a7bdcba 100644
--- a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/HiveMetastore.java
+++ b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/HiveMetastore.java
@@ -36,6 +36,7 @@
 import com.facebook.presto.spi.statistics.ColumnStatisticType;
 import com.google.common.collect.ImmutableList;
 import org.apache.hadoop.hive.metastore.api.Database;
+import org.apache.hadoop.hive.metastore.api.EnvironmentContext;
 import org.apache.hadoop.hive.metastore.api.FieldSchema;
 import org.apache.hadoop.hive.metastore.api.Partition;
 import org.apache.hadoop.hive.metastore.api.Table;
@@ -62,6 +63,8 @@ public interface HiveMetastore
 
     MetastoreOperationResult alterTable(MetastoreContext metastoreContext, String databaseName, String tableName, Table table);
 
+    MetastoreOperationResult alterTableWithEnvironmentContext(MetastoreContext metastoreContext, String databaseName, String tableName, Table table, EnvironmentContext environmentContext);
+
     default List getDatabases(MetastoreContext metastoreContext, String pattern)
     {
         return getAllDatabases(metastoreContext);
diff --git a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/HiveMetastoreClient.java b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/HiveMetastoreClient.java
index 4880eedcdd6ee..caf27aeed8619 100644
--- a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/HiveMetastoreClient.java
+++ b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/HiveMetastoreClient.java
@@ -16,6 +16,7 @@
 import org.apache.hadoop.hive.metastore.api.CheckLockRequest;
 import org.apache.hadoop.hive.metastore.api.ColumnStatisticsObj;
 import org.apache.hadoop.hive.metastore.api.Database;
+import org.apache.hadoop.hive.metastore.api.EnvironmentContext;
 import org.apache.hadoop.hive.metastore.api.FieldSchema;
 import org.apache.hadoop.hive.metastore.api.HiveObjectPrivilege;
 import org.apache.hadoop.hive.metastore.api.HiveObjectRef;
@@ -86,6 +87,9 @@ void dropTable(String databaseName, String name, boolean deleteData)
     void alterTable(String databaseName, String tableName, Table newTable)
             throws TException;
 
+    void alterTableWithEnvironmentContext(String databaseName, String tableName, Table newTable, EnvironmentContext context)
+            throws TException;
+
     Table getTable(String databaseName, String tableName)
             throws TException;
 
diff --git a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/ThriftHiveMetastore.java b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/ThriftHiveMetastore.java
index 7df3f397e4cc2..c3bb696979e17 100644
--- a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/ThriftHiveMetastore.java
+++ b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/ThriftHiveMetastore.java
@@ -64,6 +64,7 @@
 import org.apache.hadoop.hive.metastore.api.CheckLockRequest;
 import org.apache.hadoop.hive.metastore.api.ColumnStatisticsObj;
 import org.apache.hadoop.hive.metastore.api.Database;
+import org.apache.hadoop.hive.metastore.api.EnvironmentContext;
 import org.apache.hadoop.hive.metastore.api.FieldSchema;
 import org.apache.hadoop.hive.metastore.api.HiveObjectPrivilege;
 import org.apache.hadoop.hive.metastore.api.HiveObjectRef;
@@ -1160,6 +1161,35 @@ public MetastoreOperationResult alterTable(MetastoreContext metastoreContext, St
         }
     }
 
+    @Override
+    public MetastoreOperationResult alterTableWithEnvironmentContext(MetastoreContext metastoreContext, String databaseName, String tableName, Table table, EnvironmentContext environmentContext)
+    {
+        try {
+            retry()
+                    .stopOn(InvalidOperationException.class, MetaException.class)
+                    .stopOnIllegalExceptions()
+                    .run("alterTableWithEnvironmentContext", stats.getAlterTableWithEnvironmentContext().wrap(() ->
+                            getMetastoreClientThenCall(metastoreContext, client -> {
+                                Optional source = getTable(metastoreContext, databaseName, tableName);
+                                if (!source.isPresent()) {
+                                    throw new TableNotFoundException(new SchemaTableName(databaseName, tableName));
+                                }
+                                client.alterTableWithEnvironmentContext(databaseName, tableName, table, environmentContext);
+                                return null;
+                            })));
+            return EMPTY_RESULT;
+        }
+        catch (NoSuchObjectException e) {
+            throw new TableNotFoundException(new SchemaTableName(databaseName, tableName));
+        }
+        catch (TException e) {
+            throw new PrestoException(HIVE_METASTORE_ERROR, e);
+        }
+        catch (Exception e) {
+            throw propagate(e);
+        }
+    }
+
     @Override
     public Optional> getPartitionNames(MetastoreContext metastoreContext, String databaseName, String tableName)
     {
diff --git a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/ThriftHiveMetastoreClient.java b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/ThriftHiveMetastoreClient.java
index f34955d09432e..93ed690cc87a6 100644
--- a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/ThriftHiveMetastoreClient.java
+++ b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/ThriftHiveMetastoreClient.java
@@ -23,6 +23,7 @@
 import org.apache.hadoop.hive.metastore.api.ColumnStatisticsObj;
 import org.apache.hadoop.hive.metastore.api.Database;
 import org.apache.hadoop.hive.metastore.api.DropConstraintRequest;
+import org.apache.hadoop.hive.metastore.api.EnvironmentContext;
 import org.apache.hadoop.hive.metastore.api.FieldSchema;
 import org.apache.hadoop.hive.metastore.api.GetRoleGrantsForPrincipalRequest;
 import org.apache.hadoop.hive.metastore.api.GetRoleGrantsForPrincipalResponse;
@@ -199,6 +200,13 @@ public void alterTable(String databaseName, String tableName, Table newTable)
         client.alter_table(constructSchemaName(catalogName, databaseName), tableName, newTable);
     }
 
+    @Override
+    public void alterTableWithEnvironmentContext(String databaseName, String tableName, Table newTable, EnvironmentContext context)
+            throws TException
+    {
+        client.alter_table_with_environment_context(constructSchemaName(catalogName, databaseName), tableName, newTable, context);
+    }
+
     @Override
     public Table getTable(String databaseName, String tableName)
             throws TException
diff --git a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/ThriftHiveMetastoreStats.java b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/ThriftHiveMetastoreStats.java
index 543c1a2ba27da..aadc7f225b097 100644
--- a/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/ThriftHiveMetastoreStats.java
+++ b/presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/thrift/ThriftHiveMetastoreStats.java
@@ -40,6 +40,7 @@ public class ThriftHiveMetastoreStats
     private final HiveMetastoreApiStats createTableWithConstraints = new HiveMetastoreApiStats();
     private final HiveMetastoreApiStats dropTable = new HiveMetastoreApiStats();
     private final HiveMetastoreApiStats alterTable = new HiveMetastoreApiStats();
+    private final HiveMetastoreApiStats alterTableWithEnvironmentContext = new HiveMetastoreApiStats();
     private final HiveMetastoreApiStats addPartitions = new HiveMetastoreApiStats();
     private final HiveMetastoreApiStats dropPartition = new HiveMetastoreApiStats();
     private final HiveMetastoreApiStats alterPartition = new HiveMetastoreApiStats();
@@ -216,6 +217,13 @@ public HiveMetastoreApiStats getAlterTable()
         return alterTable;
     }
 
+    @Managed
+    @Nested
+    public HiveMetastoreApiStats getAlterTableWithEnvironmentContext()
+    {
+        return alterTableWithEnvironmentContext;
+    }
+
     @Managed
     @Nested
     public HiveMetastoreApiStats getAddPartitions()
diff --git a/presto-hive-metastore/src/test/java/com/facebook/presto/hive/metastore/UnimplementedHiveMetastore.java b/presto-hive-metastore/src/test/java/com/facebook/presto/hive/metastore/UnimplementedHiveMetastore.java
index 018ed253bdb0f..2e3bbe18952f5 100644
--- a/presto-hive-metastore/src/test/java/com/facebook/presto/hive/metastore/UnimplementedHiveMetastore.java
+++ b/presto-hive-metastore/src/test/java/com/facebook/presto/hive/metastore/UnimplementedHiveMetastore.java
@@ -28,6 +28,7 @@
 import java.util.Optional;
 import java.util.Set;
 import java.util.function.Function;
+import java.util.function.Supplier;
 
 public class UnimplementedHiveMetastore
         implements ExtendedHiveMetastore
@@ -128,6 +129,12 @@ public MetastoreOperationResult replaceTable(MetastoreContext metastoreContext,
         throw new UnsupportedOperationException();
     }
 
+    @Override
+    public MetastoreOperationResult persistTable(MetastoreContext metastoreContext, String databaseName, String tableName, Table newTable, PrincipalPrivileges principalPrivileges, Supplier update, Map additionalParameters)
+    {
+        throw new UnsupportedOperationException();
+    }
+
     @Override
     public MetastoreOperationResult renameTable(MetastoreContext metastoreContext, String databaseName, String tableName, String newDatabaseName, String newTableName)
     {
diff --git a/presto-hive-metastore/src/test/java/com/facebook/presto/hive/metastore/thrift/InMemoryHiveMetastore.java b/presto-hive-metastore/src/test/java/com/facebook/presto/hive/metastore/thrift/InMemoryHiveMetastore.java
index 636e99b38196e..866697dd088df 100644
--- a/presto-hive-metastore/src/test/java/com/facebook/presto/hive/metastore/thrift/InMemoryHiveMetastore.java
+++ b/presto-hive-metastore/src/test/java/com/facebook/presto/hive/metastore/thrift/InMemoryHiveMetastore.java
@@ -40,6 +40,7 @@
 import org.apache.hadoop.fs.Path;
 import org.apache.hadoop.hive.metastore.TableType;
 import org.apache.hadoop.hive.metastore.api.Database;
+import org.apache.hadoop.hive.metastore.api.EnvironmentContext;
 import org.apache.hadoop.hive.metastore.api.Partition;
 import org.apache.hadoop.hive.metastore.api.PrincipalPrivilegeSet;
 import org.apache.hadoop.hive.metastore.api.PrincipalType;
@@ -283,6 +284,12 @@ public synchronized MetastoreOperationResult alterTable(MetastoreContext metasto
         return EMPTY_RESULT;
     }
 
+    @Override
+    public synchronized MetastoreOperationResult alterTableWithEnvironmentContext(MetastoreContext metastoreContext, String databaseName, String tableName, Table newTable, EnvironmentContext environmentContext)
+    {
+        return alterTable(metastoreContext, databaseName, tableName, newTable);
+    }
+
     @Override
     public synchronized Optional> getAllTables(MetastoreContext metastoreContext, String databaseName)
     {
diff --git a/presto-hive-metastore/src/test/java/com/facebook/presto/hive/metastore/thrift/MockHiveMetastoreClient.java b/presto-hive-metastore/src/test/java/com/facebook/presto/hive/metastore/thrift/MockHiveMetastoreClient.java
index 4ee92ab0e2b47..ffb7a77670d65 100644
--- a/presto-hive-metastore/src/test/java/com/facebook/presto/hive/metastore/thrift/MockHiveMetastoreClient.java
+++ b/presto-hive-metastore/src/test/java/com/facebook/presto/hive/metastore/thrift/MockHiveMetastoreClient.java
@@ -29,6 +29,7 @@
 import org.apache.hadoop.hive.metastore.api.CheckLockRequest;
 import org.apache.hadoop.hive.metastore.api.ColumnStatisticsObj;
 import org.apache.hadoop.hive.metastore.api.Database;
+import org.apache.hadoop.hive.metastore.api.EnvironmentContext;
 import org.apache.hadoop.hive.metastore.api.FieldSchema;
 import org.apache.hadoop.hive.metastore.api.HiveObjectPrivilege;
 import org.apache.hadoop.hive.metastore.api.HiveObjectRef;
@@ -365,6 +366,13 @@ public void alterTable(String databaseName, String tableName, Table newTable)
         throw new UnsupportedOperationException();
     }
 
+    @Override
+    public void alterTableWithEnvironmentContext(String databaseName, String tableName, Table newTable, EnvironmentContext context)
+            throws TException
+    {
+        throw new UnsupportedOperationException();
+    }
+
     @Override
     public int addPartitions(List newPartitions)
     {
diff --git a/presto-iceberg/src/main/java/com/facebook/presto/iceberg/HiveTableOperations.java b/presto-iceberg/src/main/java/com/facebook/presto/iceberg/HiveTableOperations.java
index b608c5e93d332..29b85a11e4fc1 100644
--- a/presto-iceberg/src/main/java/com/facebook/presto/iceberg/HiveTableOperations.java
+++ b/presto-iceberg/src/main/java/com/facebook/presto/iceberg/HiveTableOperations.java
@@ -30,15 +30,18 @@
 import com.facebook.presto.spi.SchemaTableName;
 import com.facebook.presto.spi.TableNotFoundException;
 import com.facebook.presto.spi.security.PrestoPrincipal;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.cache.CacheBuilder;
 import com.google.common.cache.CacheLoader;
 import com.google.common.cache.LoadingCache;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableMultimap;
 import com.google.common.collect.Sets;
 import jakarta.annotation.Nullable;
 import org.apache.hadoop.hive.serde2.lazy.LazySimpleSerDe;
 import org.apache.hadoop.mapred.FileInputFormat;
 import org.apache.hadoop.mapred.FileOutputFormat;
+import org.apache.iceberg.BaseMetastoreTableOperations;
 import org.apache.iceberg.LocationProviders;
 import org.apache.iceberg.TableMetadata;
 import org.apache.iceberg.TableMetadata.MetadataLogEntry;
@@ -106,6 +109,7 @@ public class HiveTableOperations
     private final Optional owner;
     private final Optional location;
     private final HdfsFileIO fileIO;
+
     private final IcebergHiveTableOperationsConfig config;
 
     private TableMetadata currentMetadata;
@@ -251,14 +255,19 @@ public void commit(@Nullable TableMetadata base, TableMetadata metadata)
         String newMetadataLocation = writeNewMetadata(metadata, version + 1);
 
         Table table;
-        // getting a process-level lock per table to avoid concurrent commit attempts to the same table from the same
-        // JVM process, which would result in unnecessary and costly HMS lock acquisition requests
         Optional lockId = Optional.empty();
+        boolean useHMSLock = Optional.ofNullable(metadata.property(TableProperties.HIVE_LOCK_ENABLED, null))
+                .map(Boolean::parseBoolean)
+                .orElse(config.getLockingEnabled());
         ReentrantLock tableLevelMutex = commitLockCache.getUnchecked(database + "." + tableName);
+        // getting a process-level lock per table to avoid concurrent commit attempts to the same table from the same
+        // JVM process, which would result in unnecessary and costly HMS lock acquisition requests
         tableLevelMutex.lock();
         try {
             try {
-                lockId = metastore.lock(metastoreContext, database, tableName);
+                if (useHMSLock) {
+                    lockId = metastore.lock(metastoreContext, database, tableName);
+                }
                 if (base == null) {
                     String tableComment = metadata.properties().get(TABLE_COMMENT);
                     Map parameters = new HashMap<>();
@@ -318,10 +327,7 @@ public void commit(@Nullable TableMetadata base, TableMetadata metadata)
             }
             else {
                 PartitionStatistics tableStats = metastore.getTableStatistics(metastoreContext, database, tableName);
-                metastore.replaceTable(metastoreContext, database, tableName, table, privileges);
-
-                // attempt to put back previous table statistics
-                metastore.updateTableStatistics(metastoreContext, database, tableName, oldStats -> tableStats);
+                metastore.persistTable(metastoreContext, database, tableName, table, privileges, () -> tableStats, useHMSLock ? ImmutableMap.of() : hmsEnvContext(base.metadataFileLocation()));
             }
             deleteRemovedMetadataFiles(base, metadata);
         }
@@ -500,4 +506,19 @@ private void deleteRemovedMetadataFiles(TableMetadata base, TableMetadata metada
                     .run(previousMetadataFile -> io().deleteFile(previousMetadataFile.file()));
         }
     }
+
+    private Map hmsEnvContext(String metadataLocation)
+    {
+        return ImmutableMap.of(
+                org.apache.iceberg.hive.HiveTableOperations.NO_LOCK_EXPECTED_KEY,
+                BaseMetastoreTableOperations.METADATA_LOCATION_PROP,
+                org.apache.iceberg.hive.HiveTableOperations.NO_LOCK_EXPECTED_VALUE,
+                metadataLocation);
+    }
+
+    @VisibleForTesting
+    public IcebergHiveTableOperationsConfig getConfig()
+    {
+        return config;
+    }
 }
diff --git a/presto-iceberg/src/main/java/com/facebook/presto/iceberg/IcebergHiveTableOperationsConfig.java b/presto-iceberg/src/main/java/com/facebook/presto/iceberg/IcebergHiveTableOperationsConfig.java
index a0cb74928b530..9d9d358044030 100644
--- a/presto-iceberg/src/main/java/com/facebook/presto/iceberg/IcebergHiveTableOperationsConfig.java
+++ b/presto-iceberg/src/main/java/com/facebook/presto/iceberg/IcebergHiveTableOperationsConfig.java
@@ -31,6 +31,7 @@ public class IcebergHiveTableOperationsConfig
     private Duration tableRefreshMaxRetryTime = succinctDuration(1, MINUTES);
     private double tableRefreshBackoffScaleFactor = 4.0;
     private int tableRefreshRetries = 20;
+    private boolean lockingEnabled = true;
 
     @MinDuration("1ms")
     public Duration getTableRefreshBackoffMinSleepTime()
@@ -101,4 +102,17 @@ public int getTableRefreshRetries()
     {
         return tableRefreshRetries;
     }
+
+    @Config("iceberg.engine.hive.lock-enabled")
+    @ConfigDescription("Whether to use HMS locks to ensure atomicity of commits")
+    public IcebergHiveTableOperationsConfig setLockingEnabled(boolean lockingEnabled)
+    {
+        this.lockingEnabled = lockingEnabled;
+        return this;
+    }
+
+    public boolean getLockingEnabled()
+    {
+        return lockingEnabled;
+    }
 }
diff --git a/presto-iceberg/src/test/java/com/facebook/presto/iceberg/hive/TestHiveTableOperationsConfig.java b/presto-iceberg/src/test/java/com/facebook/presto/iceberg/hive/TestHiveTableOperationsConfig.java
index 5ac51d56b975b..42265f3fa5470 100644
--- a/presto-iceberg/src/test/java/com/facebook/presto/iceberg/hive/TestHiveTableOperationsConfig.java
+++ b/presto-iceberg/src/test/java/com/facebook/presto/iceberg/hive/TestHiveTableOperationsConfig.java
@@ -37,7 +37,8 @@ public void testDefaults()
                 .setTableRefreshBackoffMaxSleepTime(succinctDuration(5, SECONDS))
                 .setTableRefreshMaxRetryTime(succinctDuration(1, MINUTES))
                 .setTableRefreshBackoffScaleFactor(4.0)
-                .setTableRefreshRetries(20));
+                .setTableRefreshRetries(20)
+                .setLockingEnabled(true));
     }
 
     @Test
@@ -49,6 +50,7 @@ public void testExplicitPropertyMappings()
                 .put("iceberg.hive.table-refresh.max-retry-time", "30s")
                 .put("iceberg.hive.table-refresh.retries", "42")
                 .put("iceberg.hive.table-refresh.backoff-scale-factor", "2.0")
+                .put("iceberg.engine.hive.lock-enabled", "false")
                 .build();
 
         IcebergHiveTableOperationsConfig expected = new IcebergHiveTableOperationsConfig()
@@ -56,7 +58,8 @@ public void testExplicitPropertyMappings()
                 .setTableRefreshBackoffMaxSleepTime(succinctDuration(20, SECONDS))
                 .setTableRefreshMaxRetryTime(succinctDuration(30, SECONDS))
                 .setTableRefreshBackoffScaleFactor(2.0)
-                .setTableRefreshRetries(42);
+                .setTableRefreshRetries(42)
+                .setLockingEnabled(false);
 
         assertFullMapping(properties, expected);
     }
diff --git a/presto-iceberg/src/test/java/com/facebook/presto/iceberg/hive/TestIcebergDistributedHive.java b/presto-iceberg/src/test/java/com/facebook/presto/iceberg/hive/TestIcebergDistributedHive.java
index f0ff6cf69b0c0..6e0697f4ebb65 100644
--- a/presto-iceberg/src/test/java/com/facebook/presto/iceberg/hive/TestIcebergDistributedHive.java
+++ b/presto-iceberg/src/test/java/com/facebook/presto/iceberg/hive/TestIcebergDistributedHive.java
@@ -16,7 +16,11 @@
 import com.facebook.presto.Session;
 import com.facebook.presto.common.QualifiedObjectName;
 import com.facebook.presto.common.transaction.TransactionId;
+import com.facebook.presto.hive.HdfsContext;
+import com.facebook.presto.hive.HiveColumnConverterProvider;
 import com.facebook.presto.hive.metastore.ExtendedHiveMetastore;
+import com.facebook.presto.hive.metastore.MetastoreContext;
+import com.facebook.presto.iceberg.HiveTableOperations;
 import com.facebook.presto.iceberg.IcebergCatalogName;
 import com.facebook.presto.iceberg.IcebergDistributedTestBase;
 import com.facebook.presto.iceberg.IcebergHiveMetadata;
@@ -26,7 +30,9 @@
 import com.facebook.presto.metadata.CatalogManager;
 import com.facebook.presto.metadata.CatalogMetadata;
 import com.facebook.presto.metadata.MetadataUtil;
+import com.facebook.presto.spi.ColumnMetadata;
 import com.facebook.presto.spi.ConnectorId;
+import com.facebook.presto.spi.ConnectorSession;
 import com.facebook.presto.spi.SchemaTableName;
 import com.facebook.presto.spi.TableHandle;
 import com.facebook.presto.spi.connector.classloader.ClassLoaderSafeConnectorMetadata;
@@ -34,20 +40,33 @@
 import com.google.common.cache.CacheBuilder;
 import com.google.common.cache.CacheStats;
 import com.google.common.collect.ImmutableMap;
+import org.apache.iceberg.BaseTable;
+import org.apache.iceberg.PartitionSpec;
 import org.apache.iceberg.Table;
+import org.apache.iceberg.TableMetadata;
+import org.apache.iceberg.Transaction;
 import org.testng.annotations.Test;
 
 import java.lang.reflect.Field;
+import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 
+import static com.facebook.presto.common.type.IntegerType.INTEGER;
 import static com.facebook.presto.hive.metastore.InMemoryCachingHiveMetastore.memoizeMetastore;
+import static com.facebook.presto.hive.metastore.MetastoreUtil.getMetastoreHeaders;
+import static com.facebook.presto.hive.metastore.MetastoreUtil.isUserDefinedTypeEncodingEnabled;
 import static com.facebook.presto.iceberg.CatalogType.HIVE;
+import static com.facebook.presto.iceberg.IcebergAbstractMetadata.toIcebergSchema;
 import static com.facebook.presto.iceberg.IcebergQueryRunner.ICEBERG_CATALOG;
 import static com.facebook.presto.spi.statistics.ColumnStatisticType.NUMBER_OF_DISTINCT_VALUES;
 import static com.facebook.presto.spi.statistics.ColumnStatisticType.TOTAL_SIZE_IN_BYTES;
+import static com.google.common.io.Files.createTempDir;
 import static java.lang.String.format;
+import static org.apache.iceberg.TableMetadata.newTableMetadata;
+import static org.apache.iceberg.Transactions.createTableTransaction;
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertFalse;
 import static org.testng.Assert.assertTrue;
@@ -56,6 +75,15 @@
 public class TestIcebergDistributedHive
         extends IcebergDistributedTestBase
 {
+    public TestIcebergDistributedHive(Map extraConnectorProperties)
+    {
+        super(HIVE, ImmutableMap.builder()
+                .put("iceberg.hive-statistics-merge-strategy", Joiner.on(",").join(
+                        NUMBER_OF_DISTINCT_VALUES.name(),
+                        TOTAL_SIZE_IN_BYTES.name()))
+                .putAll(extraConnectorProperties)
+                .build());
+    }
     public TestIcebergDistributedHive()
     {
         super(HIVE, ImmutableMap.of("iceberg.hive-statistics-merge-strategy", Joiner.on(",").join(NUMBER_OF_DISTINCT_VALUES.name(), TOTAL_SIZE_IN_BYTES.name())));
@@ -187,6 +215,21 @@ public void testManifestFileCachingDisabled()
         assertQuerySucceeds(session, "DROP SCHEMA default");
     }
 
+    @Test
+    public void testCommitTableMetadataForNoLock()
+    {
+        createTable("iceberg-test-table", createTempDir().toURI().toString(), ImmutableMap.of("engine.hive.lock-enabled", "false"), 2);
+        BaseTable table = (BaseTable) loadTable("iceberg-test-table");
+        HiveTableOperations operations = (HiveTableOperations) table.operations();
+        TableMetadata currentMetadata = operations.current();
+
+        ImmutableMap.Builder builder = ImmutableMap.builder();
+        builder.putAll(currentMetadata.properties());
+        builder.put("test_property_new", "test_value_new");
+        operations.commit(currentMetadata, TableMetadata.buildFrom(currentMetadata).setProperties(builder.build()).build());
+        assertEquals(operations.current().properties(), builder.build());
+    }
+
     @Override
     protected Table loadTable(String tableName)
     {
@@ -210,4 +253,35 @@ protected ExtendedHiveMetastore getFileHiveMetastore()
                 "test");
         return memoizeMetastore(fileHiveMetastore, false, 1000, 0);
     }
+
+    protected Table createTable(String tableName, String targetPath, Map tableProperties, int columns)
+    {
+        CatalogManager catalogManager = getDistributedQueryRunner().getCoordinator().getCatalogManager();
+        ConnectorId connectorId = catalogManager.getCatalog(getDistributedQueryRunner().getDefaultSession().getCatalog().get()).get().getConnectorId();
+        ConnectorSession session = getQueryRunner().getDefaultSession().toConnectorSession(connectorId);
+        MetastoreContext context = new MetastoreContext(session.getIdentity(), session.getQueryId(), session.getClientInfo(), session.getClientTags(), session.getSource(), getMetastoreHeaders(session), isUserDefinedTypeEncodingEnabled(session), HiveColumnConverterProvider.DEFAULT_COLUMN_CONVERTER_PROVIDER, session.getWarningCollector(), session.getRuntimeStats());
+        HdfsContext hdfsContext = new HdfsContext(session, "tpch", tableName);
+        HiveTableOperations operations = new HiveTableOperations(
+                getFileHiveMetastore(),
+                context,
+                getHdfsEnvironment(),
+                hdfsContext,
+                new IcebergHiveTableOperationsConfig().setLockingEnabled(false),
+                new ManifestFileCache(CacheBuilder.newBuilder().build(), false, 0, 1024),
+                "tpch",
+                tableName,
+                session.getUser(),
+                targetPath);
+        List columnMetadataList = new ArrayList<>();
+        for (int i = 0; i < columns; i++) {
+            columnMetadataList.add(ColumnMetadata.builder().setName("column" + i).setType(INTEGER).build());
+        }
+        TableMetadata metadata = newTableMetadata(
+                toIcebergSchema(columnMetadataList),
+                PartitionSpec.unpartitioned(), targetPath,
+                tableProperties);
+        Transaction transaction = createTableTransaction(tableName, operations, metadata);
+        transaction.commitTransaction();
+        return transaction.table();
+    }
 }

From f3907c8b1a34b0d223f74624676ed9b4c6ae1a3c Mon Sep 17 00:00:00 2001
From: Hazmi 
Date: Tue, 10 Dec 2024 16:50:16 +0300
Subject: [PATCH 018/113] Resolve map(varchar, json) canonicalization bug

The map function will not sort a json object by its keys, despite the
json_parse function sorting the same input.
If implemented, this will sort json objects.

Resolves #24207
---
 .../com/facebook/presto/util/JsonUtil.java     | 18 ++++++++++++++++--
 .../facebook/presto/type/TestMapOperators.java |  7 ++-----
 2 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/presto-main-base/src/main/java/com/facebook/presto/util/JsonUtil.java b/presto-main-base/src/main/java/com/facebook/presto/util/JsonUtil.java
index 4b48dcbd8fff8..e8e7cb4802de2 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/util/JsonUtil.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/util/JsonUtil.java
@@ -13,6 +13,7 @@
  */
 package com.facebook.presto.util;
 
+import com.facebook.airlift.json.JsonObjectMapperProvider;
 import com.facebook.presto.common.block.Block;
 import com.facebook.presto.common.block.BlockBuilder;
 import com.facebook.presto.common.block.SingleRowBlockWriter;
@@ -41,6 +42,7 @@
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.primitives.Shorts;
 import com.google.common.primitives.SignedBytes;
+import io.airlift.slice.DynamicSliceOutput;
 import io.airlift.slice.Slice;
 import io.airlift.slice.SliceOutput;
 import io.airlift.slice.Slices;
@@ -86,6 +88,7 @@
 import static com.fasterxml.jackson.core.JsonToken.FIELD_NAME;
 import static com.fasterxml.jackson.core.JsonToken.START_ARRAY;
 import static com.fasterxml.jackson.core.JsonToken.START_OBJECT;
+import static com.fasterxml.jackson.databind.SerializationFeature.ORDER_MAP_ENTRIES_BY_KEYS;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Verify.verify;
 import static it.unimi.dsi.fastutil.HashCommon.arraySize;
@@ -106,6 +109,7 @@ public final class JsonUtil
     // `OBJECT_MAPPER.writeValueAsString(parser.readValueAsTree());` preserves input order.
     // Be aware. Using it arbitrarily can produce invalid json (ordered by key is required in Presto).
     private static final ObjectMapper OBJECT_MAPPED_UNORDERED = new ObjectMapper(JSON_FACTORY);
+    private static final ObjectMapper OBJECT_MAPPED_SORTED = new JsonObjectMapperProvider().get().configure(ORDER_MAP_ENTRIES_BY_KEYS, true);
 
     private static final int MAX_JSON_LENGTH_IN_ERROR_MESSAGE = 10_000;
 
@@ -956,8 +960,18 @@ static BlockBuilderAppender createBlockBuilderAppender(Type type)
                     return new VarcharBlockBuilderAppender(type);
                 case StandardTypes.JSON:
                     return (parser, blockBuilder, sqlFunctionProperties) -> {
-                        String json = OBJECT_MAPPED_UNORDERED.writeValueAsString(parser.readValueAsTree());
-                        JSON.writeSlice(blockBuilder, Slices.utf8Slice(json));
+                        Slice slice = Slices.utf8Slice(OBJECT_MAPPED_UNORDERED.writeValueAsString(parser.readValueAsTree()));
+                        try (JsonParser jsonParser = createJsonParser(JSON_FACTORY, slice)) {
+                            SliceOutput dynamicSliceOutput = new DynamicSliceOutput(slice.length());
+                            OBJECT_MAPPED_SORTED.writeValue((OutputStream) dynamicSliceOutput, OBJECT_MAPPED_SORTED.readValue(jsonParser, Object.class));
+                            // nextToken() returns null if the input is parsed correctly,
+                            // but will throw an exception if there are trailing characters.
+                            jsonParser.nextToken();
+                            JSON.writeSlice(blockBuilder, dynamicSliceOutput.slice());
+                        }
+                        catch (Exception e) {
+                            throw new PrestoException(INVALID_FUNCTION_ARGUMENT, format("Cannot convert '%s' to JSON", slice.toStringUtf8()));
+                        }
                     };
                 case StandardTypes.ARRAY:
                     return new ArrayBlockBuilderAppender(createBlockBuilderAppender(((ArrayType) type).getElementType()));
diff --git a/presto-main-base/src/test/java/com/facebook/presto/type/TestMapOperators.java b/presto-main-base/src/test/java/com/facebook/presto/type/TestMapOperators.java
index 0690252130ce8..c00ad1b5baf86 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/type/TestMapOperators.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/type/TestMapOperators.java
@@ -527,16 +527,13 @@ public void testJsonToMap()
                         .put("k8", "[null]")
                         .build());
 
-        // These two tests verifies that partial json cast preserves input order
-        // The second test should never happen in real life because valid json in presto requires natural key ordering.
-        // However, it is added to make sure that the order in the first test is not a coincidence.
         assertFunction("CAST(JSON '{\"k1\": {\"1klmnopq\":1, \"2klmnopq\":2, \"3klmnopq\":3, \"4klmnopq\":4, \"5klmnopq\":5, \"6klmnopq\":6, \"7klmnopq\":7}}' AS MAP)",
                 mapType(VARCHAR, JSON),
                 ImmutableMap.of("k1", "{\"1klmnopq\":1,\"2klmnopq\":2,\"3klmnopq\":3,\"4klmnopq\":4,\"5klmnopq\":5,\"6klmnopq\":6,\"7klmnopq\":7}"));
+
         assertFunction("CAST(unchecked_to_json('{\"k1\": {\"7klmnopq\":7, \"6klmnopq\":6, \"5klmnopq\":5, \"4klmnopq\":4, \"3klmnopq\":3, \"2klmnopq\":2, \"1klmnopq\":1}}') AS MAP)",
                 mapType(VARCHAR, JSON),
-                ImmutableMap.of("k1", "{\"7klmnopq\":7,\"6klmnopq\":6,\"5klmnopq\":5,\"4klmnopq\":4,\"3klmnopq\":3,\"2klmnopq\":2,\"1klmnopq\":1}"));
-
+                ImmutableMap.of("k1", "{\"1klmnopq\":1,\"2klmnopq\":2,\"3klmnopq\":3,\"4klmnopq\":4,\"5klmnopq\":5,\"6klmnopq\":6,\"7klmnopq\":7}"));
         // nested array/map
         assertFunction("CAST(JSON '{\"1\": [1, 2], \"2\": [3, null], \"3\": [], \"5\": [null, null], \"8\": null}' AS MAP>)",
                 mapType(BIGINT, new ArrayType(BIGINT)),

From 688d4013a28a2eb824ac02ad091ffd0f25b95604 Mon Sep 17 00:00:00 2001
From: Kevin Tang <90079564+kevintang2022@users.noreply.github.com>
Date: Thu, 28 Aug 2025 10:11:30 -0700
Subject: [PATCH 019/113] Native built in namespace manager (#25826)

Summary:
- Add abstract class BuiltInSpecialFunctionNamespaceManager
  - Add BuiltInNativeFunctionNamespaceManager
- Refactor BuiltInPluginFunctionNamespaceManager to extend the abstract
class
- Deduplicate sidecar function registry logic by moving some of it to
presto-main-base module from presto-native-sidecar-plugin module
- Add function name conflict logic to FunctionAndTypeManager that
overrides SQL built in functions but does not override Java built in
functions.
- Add retry logic in to fetch function registry from worker: retry
interval is every 1 minute

Note: `show functions` will show both built in functions in the same
namespace. This is already similar behavior to regular Native sidecar
namespace enabled with default presto.default prefix. The `show
functions` logic is not addressed in this change. Can add some unit
tests for show functions as well

Tests:
Added unit tests that enable to flag for this feature, and it is
overriding the SQL function implementation properly.

## Release Notes
Please follow [release notes
guidelines](https://github.com/prestodb/presto/wiki/Release-Notes-Guidelines)
and fill in the release notes below.

```
== NO RELEASE NOTE ==
```
---
 pom.xml                                       |   7 +
 presto-built-in-worker-function-tools/pom.xml |  62 +++++
 .../tools/ForNativeFunctionRegistryInfo.java  |  26 ++
 .../NativeSidecarFunctionRegistryTool.java    | 120 +++++++++
 .../NativeSidecarRegistryToolConfig.java      |  51 ++++
 .../tools/WorkerFunctionRegistryTool.java     |  23 ++
 .../builtin/tools/WorkerFunctionUtil.java     | 175 ++++++++++++
 .../TestNativeSidecarRegistryToolConfig.java  |  50 ++++
 presto-main-base/pom.xml                      |   5 +
 .../presto/metadata/BuiltInFunctionKind.java  |   3 +-
 ...BuiltInPluginFunctionNamespaceManager.java | 221 +--------------
 ...uiltInSpecialFunctionNamespaceManager.java | 253 ++++++++++++++++++
 ...BuiltInWorkerFunctionNamespaceManager.java |  59 ++++
 .../metadata/FunctionAndTypeManager.java      |  63 ++++-
 .../presto/sql/analyzer/FeaturesConfig.java   |  15 ++
 .../TestConvertApplicableTypeToVariable.java  |   4 +-
 .../sql/analyzer/TestFeaturesConfig.java      |   3 +
 presto-main/pom.xml                           |  10 +
 .../facebook/presto/server/PrestoServer.java  |   8 +
 .../presto/server/ServerMainModule.java       |  32 +++
 .../server/testing/TestingPrestoServer.java   |  12 +
 presto-native-execution/pom.xml               |  28 ++
 .../PrestoNativeQueryRunnerUtils.java         |  16 +-
 .../TestBuiltInNativeFunctions.java           | 178 ++++++++++++
 presto-native-sidecar-plugin/pom.xml          |   4 +
 .../NativeFunctionNamespaceManager.java       | 150 +----------
 presto-spark-base/pom.xml                     |   4 +
 .../tests/AbstractTestQueryFramework.java     |   8 +-
 .../presto/tests/DistributedQueryRunner.java  |   7 +
 29 files changed, 1230 insertions(+), 367 deletions(-)
 create mode 100644 presto-built-in-worker-function-tools/pom.xml
 create mode 100644 presto-built-in-worker-function-tools/src/main/java/com/facebook/presto/builtin/tools/ForNativeFunctionRegistryInfo.java
 create mode 100644 presto-built-in-worker-function-tools/src/main/java/com/facebook/presto/builtin/tools/NativeSidecarFunctionRegistryTool.java
 create mode 100644 presto-built-in-worker-function-tools/src/main/java/com/facebook/presto/builtin/tools/NativeSidecarRegistryToolConfig.java
 create mode 100644 presto-built-in-worker-function-tools/src/main/java/com/facebook/presto/builtin/tools/WorkerFunctionRegistryTool.java
 create mode 100644 presto-built-in-worker-function-tools/src/main/java/com/facebook/presto/builtin/tools/WorkerFunctionUtil.java
 create mode 100644 presto-built-in-worker-function-tools/src/test/java/com/facebook/presto/builtin/tools/TestNativeSidecarRegistryToolConfig.java
 create mode 100644 presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInSpecialFunctionNamespaceManager.java
 create mode 100644 presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInWorkerFunctionNamespaceManager.java
 rename {presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar => presto-main-base/src/test/java/com/facebook/presto/metadata}/TestConvertApplicableTypeToVariable.java (98%)
 create mode 100644 presto-native-execution/src/test/java/com/facebook/presto/nativeworker/TestBuiltInNativeFunctions.java

diff --git a/pom.xml b/pom.xml
index 52ed78bd8ad51..d3f5e1fb01c06 100644
--- a/pom.xml
+++ b/pom.xml
@@ -153,6 +153,7 @@
         presto-parser
         presto-main-base
         presto-main
+        presto-built-in-worker-function-tools
         presto-ml
         presto-geospatial-toolkit
         presto-benchmark
@@ -815,6 +816,12 @@
                 ${project.version}
             
 
+            
+                com.facebook.presto
+                presto-built-in-worker-function-tools
+                ${project.version}
+            
+
             
                 com.facebook.presto
                 presto-main-base
diff --git a/presto-built-in-worker-function-tools/pom.xml b/presto-built-in-worker-function-tools/pom.xml
new file mode 100644
index 0000000000000..8159c62c45b5d
--- /dev/null
+++ b/presto-built-in-worker-function-tools/pom.xml
@@ -0,0 +1,62 @@
+
+
+    
+        presto-root
+        com.facebook.presto
+        0.295-SNAPSHOT
+    
+    4.0.0
+
+    presto-built-in-worker-function-tools
+    presto-built-in-worker-function-tools
+
+    
+        ${project.parent.basedir}
+        17
+        true
+    
+
+    
+        
+            com.facebook.presto
+            presto-spi
+        
+        
+            com.facebook.presto
+            presto-function-namespace-managers-common
+        
+        
+            com.facebook.airlift
+            http-client
+        
+        
+            com.google.inject
+            guice
+        
+        
+            com.google.guava
+            guava
+        
+        
+            com.facebook.airlift
+            json
+        
+        
+            com.facebook.presto
+            presto-common
+        
+        
+            com.facebook.airlift
+            log
+        
+        
+            com.facebook.airlift
+            configuration
+        
+        
+            org.testng
+            testng
+            test
+        
+    
+
diff --git a/presto-built-in-worker-function-tools/src/main/java/com/facebook/presto/builtin/tools/ForNativeFunctionRegistryInfo.java b/presto-built-in-worker-function-tools/src/main/java/com/facebook/presto/builtin/tools/ForNativeFunctionRegistryInfo.java
new file mode 100644
index 0000000000000..000a6b2e6d0e1
--- /dev/null
+++ b/presto-built-in-worker-function-tools/src/main/java/com/facebook/presto/builtin/tools/ForNativeFunctionRegistryInfo.java
@@ -0,0 +1,26 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.builtin.tools;
+
+import com.google.inject.BindingAnnotation;
+
+import java.lang.annotation.Retention;
+
+import static java.lang.annotation.RetentionPolicy.RUNTIME;
+
+@Retention(RUNTIME)
+@BindingAnnotation
+public @interface ForNativeFunctionRegistryInfo
+{
+}
diff --git a/presto-built-in-worker-function-tools/src/main/java/com/facebook/presto/builtin/tools/NativeSidecarFunctionRegistryTool.java b/presto-built-in-worker-function-tools/src/main/java/com/facebook/presto/builtin/tools/NativeSidecarFunctionRegistryTool.java
new file mode 100644
index 0000000000000..1990918df633f
--- /dev/null
+++ b/presto-built-in-worker-function-tools/src/main/java/com/facebook/presto/builtin/tools/NativeSidecarFunctionRegistryTool.java
@@ -0,0 +1,120 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.builtin.tools;
+
+import com.facebook.airlift.http.client.HttpClient;
+import com.facebook.airlift.http.client.HttpUriBuilder;
+import com.facebook.airlift.http.client.JsonResponseHandler;
+import com.facebook.airlift.http.client.Request;
+import com.facebook.airlift.json.JsonCodec;
+import com.facebook.airlift.log.Logger;
+import com.facebook.presto.functionNamespace.JsonBasedUdfFunctionMetadata;
+import com.facebook.presto.functionNamespace.UdfFunctionSignatureMap;
+import com.facebook.presto.spi.Node;
+import com.facebook.presto.spi.NodeManager;
+import com.facebook.presto.spi.PrestoException;
+import com.facebook.presto.spi.StandardErrorCode;
+import com.facebook.presto.spi.function.SqlFunction;
+import com.google.common.collect.ImmutableMap;
+
+import java.net.URI;
+import java.util.List;
+import java.util.Map;
+
+import static com.google.common.collect.ImmutableList.toImmutableList;
+import static java.util.Objects.requireNonNull;
+
+public class NativeSidecarFunctionRegistryTool
+        implements WorkerFunctionRegistryTool
+{
+    private final int maxRetries;
+    private final long retryDelayMs;
+    private static final Logger log = Logger.get(NativeSidecarFunctionRegistryTool.class);
+    private final JsonCodec>> nativeFunctionSignatureMapJsonCodec;
+    private final NodeManager nodeManager;
+    private final HttpClient httpClient;
+    private static final String FUNCTION_SIGNATURES_ENDPOINT = "/v1/functions";
+
+    public NativeSidecarFunctionRegistryTool(
+            HttpClient httpClient,
+            JsonCodec>> nativeFunctionSignatureMapJsonCodec,
+            NodeManager nodeManager,
+            int nativeSidecarRegistryToolNumRetries,
+            long nativeSidecarRegistryToolRetryDelayMs)
+    {
+        this.nativeFunctionSignatureMapJsonCodec =
+                requireNonNull(nativeFunctionSignatureMapJsonCodec, "nativeFunctionSignatureMapJsonCodec is null");
+        this.nodeManager = requireNonNull(nodeManager, "nodeManager is null");
+        this.httpClient = requireNonNull(httpClient, "typeManager is null");
+        this.maxRetries = nativeSidecarRegistryToolNumRetries;
+        this.retryDelayMs = nativeSidecarRegistryToolRetryDelayMs;
+    }
+
+    @Override
+    public List getWorkerFunctions()
+    {
+        return getNativeFunctionSignatureMap()
+                .getUDFSignatureMap()
+                .entrySet()
+                .stream()
+                .flatMap(entry -> entry.getValue().stream()
+                        .map(metaInfo -> WorkerFunctionUtil.createSqlInvokedFunction(entry.getKey(), metaInfo, "presto")))
+                .collect(toImmutableList());
+    }
+
+    private UdfFunctionSignatureMap getNativeFunctionSignatureMap()
+    {
+        try {
+            Request request = Request.Builder.prepareGet().setUri(getSidecarLocationOnStartup(nodeManager, maxRetries, retryDelayMs)).build();
+            Map> nativeFunctionSignatureMap = httpClient.execute(request, JsonResponseHandler.createJsonResponseHandler(nativeFunctionSignatureMapJsonCodec));
+            return new UdfFunctionSignatureMap(ImmutableMap.copyOf(nativeFunctionSignatureMap));
+        }
+        catch (Exception e) {
+            throw new PrestoException(StandardErrorCode.INVALID_ARGUMENTS, "Failed to get functions from sidecar.", e);
+        }
+    }
+
+    public static URI getSidecarLocationOnStartup(NodeManager nodeManager, int maxRetries, long retryDelayMs)
+    {
+        Node sidecarNode = null;
+        for (int attempt = 1; attempt <= maxRetries; attempt++) {
+            try {
+                sidecarNode = nodeManager.getSidecarNode();
+                if (sidecarNode != null) {
+                    break;
+                }
+            }
+            catch (Exception e) {
+                log.error("Error getting sidecar node (attempt " + attempt + "): " + e.getMessage());
+                if (attempt == maxRetries) {
+                    throw new RuntimeException("Failed to get sidecar node", e);
+                }
+                else {
+                    try {
+                        Thread.sleep(retryDelayMs);
+                    }
+                    catch (InterruptedException ie) {
+                        Thread.currentThread().interrupt();
+                        throw new RuntimeException("Retry fetching sidecar function registry interrupted", ie);
+                    }
+                }
+            }
+        }
+
+        return HttpUriBuilder
+                .uriBuilderFrom(sidecarNode.getHttpUri())
+                .appendPath(FUNCTION_SIGNATURES_ENDPOINT)
+                .build();
+    }
+}
diff --git a/presto-built-in-worker-function-tools/src/main/java/com/facebook/presto/builtin/tools/NativeSidecarRegistryToolConfig.java b/presto-built-in-worker-function-tools/src/main/java/com/facebook/presto/builtin/tools/NativeSidecarRegistryToolConfig.java
new file mode 100644
index 0000000000000..4ce944572406d
--- /dev/null
+++ b/presto-built-in-worker-function-tools/src/main/java/com/facebook/presto/builtin/tools/NativeSidecarRegistryToolConfig.java
@@ -0,0 +1,51 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.builtin.tools;
+
+import com.facebook.airlift.configuration.Config;
+import com.facebook.airlift.configuration.ConfigDescription;
+
+import java.time.Duration;
+
+public class NativeSidecarRegistryToolConfig
+{
+    private int nativeSidecarRegistryToolNumRetries = 8;
+    private long nativeSidecarRegistryToolRetryDelayMs = Duration.ofMinutes(1).toMillis();
+
+    public int getNativeSidecarRegistryToolNumRetries()
+    {
+        return nativeSidecarRegistryToolNumRetries;
+    }
+
+    @Config("native-sidecar-registry-tool.num-retries")
+    @ConfigDescription("Max times to retry fetching sidecar node")
+    public NativeSidecarRegistryToolConfig setNativeSidecarRegistryToolNumRetries(int nativeSidecarRegistryToolNumRetries)
+    {
+        this.nativeSidecarRegistryToolNumRetries = nativeSidecarRegistryToolNumRetries;
+        return this;
+    }
+
+    public long getNativeSidecarRegistryToolRetryDelayMs()
+    {
+        return nativeSidecarRegistryToolRetryDelayMs;
+    }
+
+    @Config("native-sidecar-registry-tool.retry-delay-ms")
+    @ConfigDescription("Cooldown period to retry when fetching sidecar node fails")
+    public NativeSidecarRegistryToolConfig setNativeSidecarRegistryToolRetryDelayMs(long nativeSidecarRegistryToolRetryDelayMs)
+    {
+        this.nativeSidecarRegistryToolRetryDelayMs = nativeSidecarRegistryToolRetryDelayMs;
+        return this;
+    }
+}
diff --git a/presto-built-in-worker-function-tools/src/main/java/com/facebook/presto/builtin/tools/WorkerFunctionRegistryTool.java b/presto-built-in-worker-function-tools/src/main/java/com/facebook/presto/builtin/tools/WorkerFunctionRegistryTool.java
new file mode 100644
index 0000000000000..99c7cfe13bc78
--- /dev/null
+++ b/presto-built-in-worker-function-tools/src/main/java/com/facebook/presto/builtin/tools/WorkerFunctionRegistryTool.java
@@ -0,0 +1,23 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.builtin.tools;
+
+import com.facebook.presto.spi.function.SqlFunction;
+
+import java.util.List;
+
+public interface WorkerFunctionRegistryTool
+{
+    List getWorkerFunctions();
+}
diff --git a/presto-built-in-worker-function-tools/src/main/java/com/facebook/presto/builtin/tools/WorkerFunctionUtil.java b/presto-built-in-worker-function-tools/src/main/java/com/facebook/presto/builtin/tools/WorkerFunctionUtil.java
new file mode 100644
index 0000000000000..d880429f7489f
--- /dev/null
+++ b/presto-built-in-worker-function-tools/src/main/java/com/facebook/presto/builtin/tools/WorkerFunctionUtil.java
@@ -0,0 +1,175 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.facebook.presto.builtin.tools;
+
+import com.facebook.presto.common.CatalogSchemaName;
+import com.facebook.presto.common.QualifiedObjectName;
+import com.facebook.presto.common.type.NamedTypeSignature;
+import com.facebook.presto.common.type.StandardTypes;
+import com.facebook.presto.common.type.TypeSignature;
+import com.facebook.presto.common.type.TypeSignatureParameter;
+import com.facebook.presto.functionNamespace.JsonBasedUdfFunctionMetadata;
+import com.facebook.presto.spi.function.AggregationFunctionMetadata;
+import com.facebook.presto.spi.function.LongVariableConstraint;
+import com.facebook.presto.spi.function.Parameter;
+import com.facebook.presto.spi.function.RoutineCharacteristics;
+import com.facebook.presto.spi.function.SqlInvokedFunction;
+import com.facebook.presto.spi.function.TypeVariableConstraint;
+import com.google.common.collect.ImmutableList;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Optional;
+
+import static com.facebook.presto.spi.function.FunctionVersion.notVersioned;
+import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Preconditions.checkState;
+
+public class WorkerFunctionUtil
+{
+    private WorkerFunctionUtil() {}
+
+    public static synchronized SqlInvokedFunction createSqlInvokedFunction(String functionName, JsonBasedUdfFunctionMetadata jsonBasedUdfFunctionMetaData, String catalogName)
+    {
+        checkState(jsonBasedUdfFunctionMetaData.getRoutineCharacteristics().getLanguage().equals(RoutineCharacteristics.Language.CPP), "WorkerFunctionUtil only supports CPP UDF");
+        QualifiedObjectName qualifiedFunctionName = QualifiedObjectName.valueOf(new CatalogSchemaName(catalogName, jsonBasedUdfFunctionMetaData.getSchema()), functionName);
+        List parameterNameList = jsonBasedUdfFunctionMetaData.getParamNames();
+        List parameterTypeList = convertApplicableTypeToVariable(jsonBasedUdfFunctionMetaData.getParamTypes());
+        List typeVariableConstraintsList = jsonBasedUdfFunctionMetaData.getTypeVariableConstraints().isPresent() ?
+                jsonBasedUdfFunctionMetaData.getTypeVariableConstraints().get() : ImmutableList.of();
+        List longVariableConstraintList = jsonBasedUdfFunctionMetaData.getLongVariableConstraints().isPresent() ?
+                jsonBasedUdfFunctionMetaData.getLongVariableConstraints().get() : ImmutableList.of();
+
+        TypeSignature outputType = convertApplicableTypeToVariable(jsonBasedUdfFunctionMetaData.getOutputType());
+        ImmutableList.Builder parameterBuilder = ImmutableList.builder();
+        for (int i = 0; i < parameterNameList.size(); i++) {
+            parameterBuilder.add(new Parameter(parameterNameList.get(i), parameterTypeList.get(i)));
+        }
+
+        Optional aggregationFunctionMetadata =
+                jsonBasedUdfFunctionMetaData.getAggregateMetadata()
+                        .map(metadata -> new AggregationFunctionMetadata(
+                                convertApplicableTypeToVariable(metadata.getIntermediateType()),
+                                metadata.isOrderSensitive()));
+
+        return new SqlInvokedFunction(
+                qualifiedFunctionName,
+                parameterBuilder.build(),
+                typeVariableConstraintsList,
+                longVariableConstraintList,
+                outputType,
+                jsonBasedUdfFunctionMetaData.getDocString(),
+                jsonBasedUdfFunctionMetaData.getRoutineCharacteristics(),
+                "",
+                jsonBasedUdfFunctionMetaData.getVariableArity(),
+                notVersioned(),
+                jsonBasedUdfFunctionMetaData.getFunctionKind(),
+                aggregationFunctionMetadata);
+    }
+
+    // Todo: Improve the handling of parameter type differentiation in native execution.
+    // HACK: Currently, we lack support for correctly identifying the parameterKind, specifically between TYPE and VARIABLE,
+    // in native execution. The following utility functions help bridge this gap by parsing the type signature and verifying whether its base
+    // and parameters are of a supported type. The valid types list are non - parametric types that Presto supports.
+    public static List convertApplicableTypeToVariable(List typeSignatures)
+    {
+        List newTypeSignaturesList = new ArrayList<>();
+        for (TypeSignature typeSignature : typeSignatures) {
+            if (!typeSignature.getParameters().isEmpty()) {
+                TypeSignature newTypeSignature =
+                        new TypeSignature(
+                                typeSignature.getBase(),
+                                getTypeSignatureParameters(
+                                        typeSignature,
+                                        typeSignature.getParameters()));
+                newTypeSignaturesList.add(newTypeSignature);
+            }
+            else {
+                newTypeSignaturesList.add(typeSignature);
+            }
+        }
+        return newTypeSignaturesList;
+    }
+
+    public static TypeSignature convertApplicableTypeToVariable(TypeSignature typeSignature)
+    {
+        List typeSignaturesList = convertApplicableTypeToVariable(ImmutableList.of(typeSignature));
+        checkArgument(!typeSignaturesList.isEmpty(), "Type signature list is empty for : " + typeSignature);
+        return typeSignaturesList.get(0);
+    }
+
+    private static List getTypeSignatureParameters(
+            TypeSignature typeSignature,
+            List typeSignatureParameterList)
+    {
+        List newParameterTypeList = new ArrayList<>();
+        for (TypeSignatureParameter parameter : typeSignatureParameterList) {
+            if (parameter.isLongLiteral()) {
+                newParameterTypeList.add(parameter);
+                continue;
+            }
+
+            boolean isNamedTypeSignature = parameter.isNamedTypeSignature();
+            TypeSignature parameterTypeSignature;
+            // If it's a named type signatures only in the case of row signature types.
+            if (isNamedTypeSignature) {
+                parameterTypeSignature = parameter.getNamedTypeSignature().getTypeSignature();
+            }
+            else {
+                parameterTypeSignature = parameter.getTypeSignature();
+            }
+
+            if (parameterTypeSignature.getParameters().isEmpty()) {
+                boolean changeTypeToVariable = isDecimalTypeBase(typeSignature.getBase());
+                if (changeTypeToVariable) {
+                    newParameterTypeList.add(
+                            TypeSignatureParameter.of(parameterTypeSignature.getBase()));
+                }
+                else {
+                    if (isNamedTypeSignature) {
+                        newParameterTypeList.add(TypeSignatureParameter.of(parameter.getNamedTypeSignature()));
+                    }
+                    else {
+                        newParameterTypeList.add(TypeSignatureParameter.of(parameterTypeSignature));
+                    }
+                }
+            }
+            else {
+                TypeSignature newTypeSignature =
+                        new TypeSignature(
+                                parameterTypeSignature.getBase(),
+                                getTypeSignatureParameters(
+                                        parameterTypeSignature.getStandardTypeSignature(),
+                                        parameterTypeSignature.getParameters()));
+                if (isNamedTypeSignature) {
+                    newParameterTypeList.add(
+                            TypeSignatureParameter.of(
+                                    new NamedTypeSignature(
+                                            Optional.empty(),
+                                            newTypeSignature)));
+                }
+                else {
+                    newParameterTypeList.add(TypeSignatureParameter.of(newTypeSignature));
+                }
+            }
+        }
+        return newParameterTypeList;
+    }
+
+    private static boolean isDecimalTypeBase(String typeBase)
+    {
+        return typeBase.equals(StandardTypes.DECIMAL);
+    }
+}
diff --git a/presto-built-in-worker-function-tools/src/test/java/com/facebook/presto/builtin/tools/TestNativeSidecarRegistryToolConfig.java b/presto-built-in-worker-function-tools/src/test/java/com/facebook/presto/builtin/tools/TestNativeSidecarRegistryToolConfig.java
new file mode 100644
index 0000000000000..34df3f89b88f3
--- /dev/null
+++ b/presto-built-in-worker-function-tools/src/test/java/com/facebook/presto/builtin/tools/TestNativeSidecarRegistryToolConfig.java
@@ -0,0 +1,50 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.builtin.tools;
+
+import com.google.common.collect.ImmutableMap;
+import org.testng.annotations.Test;
+
+import java.util.Map;
+
+import static com.facebook.airlift.configuration.testing.ConfigAssertions.assertFullMapping;
+import static com.facebook.airlift.configuration.testing.ConfigAssertions.assertRecordedDefaults;
+import static com.facebook.airlift.configuration.testing.ConfigAssertions.recordDefaults;
+
+public class TestNativeSidecarRegistryToolConfig
+{
+    @Test
+    public void testDefaults()
+    {
+        assertRecordedDefaults(recordDefaults(NativeSidecarRegistryToolConfig.class)
+                .setNativeSidecarRegistryToolNumRetries(8)
+                .setNativeSidecarRegistryToolRetryDelayMs(60_000L));
+    }
+
+    @Test
+    public void testExplicitPropertyMappings()
+            throws Exception
+    {
+        Map properties = new ImmutableMap.Builder()
+                .put("native-sidecar-registry-tool.num-retries", "15")
+                .put("native-sidecar-registry-tool.retry-delay-ms", "11115")
+                .build();
+
+        NativeSidecarRegistryToolConfig expected = new NativeSidecarRegistryToolConfig()
+                .setNativeSidecarRegistryToolNumRetries(15)
+                .setNativeSidecarRegistryToolRetryDelayMs(11_115L);
+
+        assertFullMapping(properties, expected);
+    }
+}
diff --git a/presto-main-base/pom.xml b/presto-main-base/pom.xml
index f724e286320a5..c51c38e73f2f7 100644
--- a/presto-main-base/pom.xml
+++ b/presto-main-base/pom.xml
@@ -461,6 +461,11 @@
             io.netty
             netty-transport
         
+        
+            com.facebook.presto
+            presto-built-in-worker-function-tools
+            test
+        
     
 
     
diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInFunctionKind.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInFunctionKind.java
index 4d12bb7f97d16..71ca25dd468cc 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInFunctionKind.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInFunctionKind.java
@@ -20,7 +20,8 @@
 public enum BuiltInFunctionKind
 {
     ENGINE(0),
-    PLUGIN(1);
+    PLUGIN(1),
+    WORKER(2);
 
     private final int value;
 
diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInPluginFunctionNamespaceManager.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInPluginFunctionNamespaceManager.java
index 6a694f436f17e..75e0d3415d935 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInPluginFunctionNamespaceManager.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInPluginFunctionNamespaceManager.java
@@ -13,226 +13,33 @@
  */
 package com.facebook.presto.metadata;
 
-import com.facebook.presto.common.Page;
-import com.facebook.presto.common.QualifiedObjectName;
-import com.facebook.presto.common.block.BlockEncodingSerde;
-import com.facebook.presto.common.function.SqlFunctionResult;
-import com.facebook.presto.common.type.TypeManager;
-import com.facebook.presto.common.type.UserDefinedType;
-import com.facebook.presto.spi.PrestoException;
-import com.facebook.presto.spi.function.AlterRoutineCharacteristics;
-import com.facebook.presto.spi.function.FunctionHandle;
-import com.facebook.presto.spi.function.FunctionMetadata;
-import com.facebook.presto.spi.function.FunctionNamespaceManager;
-import com.facebook.presto.spi.function.FunctionNamespaceTransactionHandle;
-import com.facebook.presto.spi.function.Parameter;
-import com.facebook.presto.spi.function.ScalarFunctionImplementation;
-import com.facebook.presto.spi.function.Signature;
+import com.facebook.presto.spi.function.FunctionImplementationType;
 import com.facebook.presto.spi.function.SqlFunction;
-import com.facebook.presto.spi.function.SqlInvokedFunction;
-import com.facebook.presto.spi.function.SqlInvokedScalarFunctionImplementation;
-import com.google.common.base.Supplier;
-import com.google.common.base.Suppliers;
-import com.google.common.cache.CacheBuilder;
-import com.google.common.cache.CacheLoader;
-import com.google.common.cache.LoadingCache;
-import com.google.common.util.concurrent.UncheckedExecutionException;
 
 import java.util.Collection;
 import java.util.List;
-import java.util.Optional;
-import java.util.concurrent.CompletableFuture;
 
 import static com.facebook.presto.metadata.BuiltInFunctionKind.PLUGIN;
 import static com.facebook.presto.spi.function.FunctionImplementationType.SQL;
-import static com.facebook.presto.spi.function.FunctionKind.SCALAR;
 import static com.google.common.base.Preconditions.checkArgument;
-import static com.google.common.base.Throwables.throwIfInstanceOf;
-import static com.google.common.collect.ImmutableList.toImmutableList;
-import static java.util.Collections.emptyList;
-import static java.util.Objects.requireNonNull;
-import static java.util.concurrent.TimeUnit.HOURS;
 
 public class BuiltInPluginFunctionNamespaceManager
-        implements FunctionNamespaceManager
+        extends BuiltInSpecialFunctionNamespaceManager
 {
-    private volatile FunctionMap functions = new FunctionMap();
-    private final FunctionAndTypeManager functionAndTypeManager;
-    private final Supplier cachedFunctions =
-            Suppliers.memoize(this::checkForNamingConflicts);
-    private final LoadingCache specializedFunctionKeyCache;
-    private final LoadingCache specializedScalarCache;
-
     public BuiltInPluginFunctionNamespaceManager(FunctionAndTypeManager functionAndTypeManager)
     {
-        this.functionAndTypeManager = requireNonNull(functionAndTypeManager, "functionAndTypeManager is null");
-        specializedFunctionKeyCache = CacheBuilder.newBuilder()
-                .maximumSize(1000)
-                .expireAfterWrite(1, HOURS)
-                .build(CacheLoader.from(this::doGetSpecializedFunctionKey));
-        specializedScalarCache = CacheBuilder.newBuilder()
-                .maximumSize(1000)
-                .expireAfterWrite(1, HOURS)
-                .build(CacheLoader.from(key -> {
-                    checkArgument(
-                            key.getFunction() instanceof SqlInvokedFunction,
-                            "Unsupported scalar function class: %s",
-                            key.getFunction().getClass());
-                    return new SqlInvokedScalarFunctionImplementation(((SqlInvokedFunction) key.getFunction()).getBody());
-                }));
+        super(functionAndTypeManager);
     }
 
-    public synchronized void registerPluginFunctions(List functions)
+    @Override
+    public synchronized void registerBuiltInSpecialFunctions(List functions)
     {
         checkForNamingConflicts(functions);
         this.functions = new FunctionMap(this.functions, functions);
     }
 
     @Override
-    public FunctionHandle getFunctionHandle(Optional transactionHandle, Signature signature)
-    {
-        return new BuiltInFunctionHandle(signature, PLUGIN);
-    }
-
-    @Override
-    public Collection getFunctions(Optional transactionHandle, QualifiedObjectName functionName)
-    {
-        if (functions.list().isEmpty() ||
-                (!functionName.getCatalogSchemaName().equals(functionAndTypeManager.getDefaultNamespace()))) {
-            return emptyList();
-        }
-        return cachedFunctions.get().get(functionName);
-    }
-
-    /**
-     * likePattern / escape is not used for optimization, returning all functions.
-     */
-    @Override
-    public Collection listFunctions(Optional likePattern, Optional escape)
-    {
-        return cachedFunctions.get().list();
-    }
-
-    public FunctionMetadata getFunctionMetadata(FunctionHandle functionHandle)
-    {
-        checkArgument(functionHandle instanceof BuiltInFunctionHandle, "Expect BuiltInFunctionHandle");
-        Signature signature = ((BuiltInFunctionHandle) functionHandle).getSignature();
-        SpecializedFunctionKey functionKey;
-        try {
-            functionKey = specializedFunctionKeyCache.getUnchecked(signature);
-        }
-        catch (UncheckedExecutionException e) {
-            throwIfInstanceOf(e.getCause(), PrestoException.class);
-            throw e;
-        }
-        SqlFunction function = functionKey.getFunction();
-        checkArgument(function instanceof SqlInvokedFunction, "BuiltInPluginFunctionNamespaceManager only support SqlInvokedFunctions");
-        SqlInvokedFunction sqlFunction = (SqlInvokedFunction) function;
-        List argumentNames = sqlFunction.getParameters().stream().map(Parameter::getName).collect(toImmutableList());
-        return new FunctionMetadata(
-                signature.getName(),
-                signature.getArgumentTypes(),
-                argumentNames,
-                signature.getReturnType(),
-                signature.getKind(),
-                sqlFunction.getRoutineCharacteristics().getLanguage(),
-                SQL,
-                function.isDeterministic(),
-                function.isCalledOnNullInput(),
-                sqlFunction.getVersion(),
-                sqlFunction.getComplexTypeFunctionDescriptor());
-    }
-
-    public ScalarFunctionImplementation getScalarFunctionImplementation(FunctionHandle functionHandle)
-    {
-        checkArgument(functionHandle instanceof BuiltInFunctionHandle, "Expect BuiltInFunctionHandle");
-        return getScalarFunctionImplementation(((BuiltInFunctionHandle) functionHandle).getSignature());
-    }
-
-    @Override
-    public void setBlockEncodingSerde(BlockEncodingSerde blockEncodingSerde)
-    {
-        throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support setting block encoding");
-    }
-
-    @Override
-    public FunctionNamespaceTransactionHandle beginTransaction()
-    {
-        throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support setting block encoding");
-    }
-
-    @Override
-    public void commit(FunctionNamespaceTransactionHandle transactionHandle)
-    {
-        throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support setting block encoding");
-    }
-
-    @Override
-    public void abort(FunctionNamespaceTransactionHandle transactionHandle)
-    {
-        throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support setting block encoding");
-    }
-
-    @Override
-    public void createFunction(SqlInvokedFunction function, boolean replace)
-    {
-        throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support setting block encoding");
-    }
-
-    @Override
-    public void dropFunction(QualifiedObjectName functionName, Optional parameterTypes, boolean exists)
-    {
-        throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support drop function");
-    }
-
-    @Override
-    public void alterFunction(QualifiedObjectName functionName, Optional parameterTypes, AlterRoutineCharacteristics alterRoutineCharacteristics)
-    {
-        throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not alter function");
-    }
-
-    @Override
-    public void addUserDefinedType(UserDefinedType userDefinedType)
-    {
-        throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support adding user defined types");
-    }
-
-    @Override
-    public Optional getUserDefinedType(QualifiedObjectName typeName)
-    {
-        throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support getting user defined types");
-    }
-
-    @Override
-    public CompletableFuture executeFunction(String source, FunctionHandle functionHandle, Page input, List channels, TypeManager typeManager)
-    {
-        throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not execute function");
-    }
-
-    private ScalarFunctionImplementation getScalarFunctionImplementation(Signature signature)
-    {
-        checkArgument(signature.getKind() == SCALAR, "%s is not a scalar function", signature);
-        checkArgument(signature.getTypeVariableConstraints().isEmpty(), "%s has unbound type parameters", signature);
-
-        try {
-            return specializedScalarCache.getUnchecked(getSpecializedFunctionKey(signature));
-        }
-        catch (UncheckedExecutionException e) {
-            throwIfInstanceOf(e.getCause(), PrestoException.class);
-            throw e;
-        }
-    }
-
-    private synchronized FunctionMap checkForNamingConflicts()
-    {
-        Optional> functionNamespaceManager =
-                functionAndTypeManager.getServingFunctionNamespaceManager(functionAndTypeManager.getDefaultNamespace());
-        checkArgument(functionNamespaceManager.isPresent(), "Cannot find function namespace for catalog '%s'", functionAndTypeManager.getDefaultNamespace().getCatalogName());
-        checkForNamingConflicts(functionNamespaceManager.get().listFunctions(Optional.empty(), Optional.empty()));
-        return functions;
-    }
-
-    private synchronized void checkForNamingConflicts(Collection functions)
+    protected synchronized void checkForNamingConflicts(Collection functions)
     {
         for (SqlFunction function : functions) {
             for (SqlFunction existingFunction : this.functions.list()) {
@@ -241,19 +48,15 @@ private synchronized void checkForNamingConflicts(Collection
+{
+    protected volatile FunctionMap functions = new FunctionMap();
+    private final FunctionAndTypeManager functionAndTypeManager;
+    private final Supplier cachedFunctions =
+            Suppliers.memoize(this::createFunctionMap);
+    private final LoadingCache specializedFunctionKeyCache;
+    private final LoadingCache specializedScalarCache;
+
+    public BuiltInSpecialFunctionNamespaceManager(FunctionAndTypeManager functionAndTypeManager)
+    {
+        this.functionAndTypeManager = requireNonNull(functionAndTypeManager, "functionAndTypeManager is null");
+        specializedFunctionKeyCache = CacheBuilder.newBuilder()
+                .maximumSize(1000)
+                .expireAfterWrite(1, HOURS)
+                .build(CacheLoader.from(this::doGetSpecializedFunctionKey));
+        specializedScalarCache = CacheBuilder.newBuilder()
+                .maximumSize(1000)
+                .expireAfterWrite(1, HOURS)
+                .build(CacheLoader.from(key -> {
+                    checkArgument(
+                            key.getFunction() instanceof SqlInvokedFunction,
+                            "Unsupported scalar function class: %s",
+                            key.getFunction().getClass());
+                    return new SqlInvokedScalarFunctionImplementation(((SqlInvokedFunction) key.getFunction()).getBody());
+                }));
+    }
+
+    @Override
+    public Collection getFunctions(Optional transactionHandle, QualifiedObjectName functionName)
+    {
+        if (functions.list().isEmpty() ||
+                (!functionName.getCatalogSchemaName().equals(functionAndTypeManager.getDefaultNamespace()))) {
+            return emptyList();
+        }
+        return cachedFunctions.get().get(functionName);
+    }
+
+    /**
+     * likePattern / escape is not used for optimization, returning all functions.
+     */
+    @Override
+    public Collection listFunctions(Optional likePattern, Optional escape)
+    {
+        return cachedFunctions.get().list();
+    }
+
+    @Override
+    public FunctionHandle getFunctionHandle(Optional transactionHandle, Signature signature)
+    {
+        return new BuiltInFunctionHandle(signature, getBuiltInFunctionKind());
+    }
+
+    public FunctionMetadata getFunctionMetadata(FunctionHandle functionHandle)
+    {
+        checkArgument(functionHandle instanceof BuiltInFunctionHandle, "Expect BuiltInFunctionHandle");
+        Signature signature = ((BuiltInFunctionHandle) functionHandle).getSignature();
+        SpecializedFunctionKey functionKey;
+        try {
+            functionKey = specializedFunctionKeyCache.getUnchecked(signature);
+        }
+        catch (UncheckedExecutionException e) {
+            throwIfInstanceOf(e.getCause(), PrestoException.class);
+            throw e;
+        }
+        SqlFunction function = functionKey.getFunction();
+        checkArgument(function instanceof SqlInvokedFunction, "BuiltInPluginFunctionNamespaceManager only support SqlInvokedFunctions");
+        SqlInvokedFunction sqlFunction = (SqlInvokedFunction) function;
+        List argumentNames = sqlFunction.getParameters().stream().map(Parameter::getName).collect(toImmutableList());
+        return new FunctionMetadata(
+                signature.getName(),
+                signature.getArgumentTypes(),
+                argumentNames,
+                signature.getReturnType(),
+                signature.getKind(),
+                sqlFunction.getRoutineCharacteristics().getLanguage(),
+                getDefaultFunctionMetadataImplementationType(),
+                function.isDeterministic(),
+                function.isCalledOnNullInput(),
+                sqlFunction.getVersion(),
+                sqlFunction.getComplexTypeFunctionDescriptor());
+    }
+
+    @Override
+    public void setBlockEncodingSerde(BlockEncodingSerde blockEncodingSerde)
+    {
+        throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support setting block encoding");
+    }
+
+    @Override
+    public FunctionNamespaceTransactionHandle beginTransaction()
+    {
+        throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support setting block encoding");
+    }
+
+    @Override
+    public void commit(FunctionNamespaceTransactionHandle transactionHandle)
+    {
+        throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support setting block encoding");
+    }
+
+    @Override
+    public void abort(FunctionNamespaceTransactionHandle transactionHandle)
+    {
+        throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support setting block encoding");
+    }
+
+    @Override
+    public void createFunction(SqlInvokedFunction function, boolean replace)
+    {
+        throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support setting block encoding");
+    }
+
+    @Override
+    public void dropFunction(QualifiedObjectName functionName, Optional parameterTypes, boolean exists)
+    {
+        throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support drop function");
+    }
+
+    @Override
+    public void alterFunction(QualifiedObjectName functionName, Optional parameterTypes, AlterRoutineCharacteristics alterRoutineCharacteristics)
+    {
+        throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not alter function");
+    }
+
+    @Override
+    public void addUserDefinedType(UserDefinedType userDefinedType)
+    {
+        throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support adding user defined types");
+    }
+
+    @Override
+    public Optional getUserDefinedType(QualifiedObjectName typeName)
+    {
+        throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not support getting user defined types");
+    }
+
+    @Override
+    public CompletableFuture executeFunction(String source, FunctionHandle functionHandle, Page input, List channels, TypeManager typeManager)
+    {
+        throw new UnsupportedOperationException("BuiltInPluginFunctionNamespaceManager does not execute function");
+    }
+
+    protected abstract void checkForNamingConflicts(Collection functions);
+
+    protected abstract BuiltInFunctionKind getBuiltInFunctionKind();
+
+    protected abstract FunctionImplementationType getDefaultFunctionMetadataImplementationType();
+
+    public abstract void registerBuiltInSpecialFunctions(List functions);
+
+    @Override
+    public ScalarFunctionImplementation getScalarFunctionImplementation(FunctionHandle functionHandle)
+    {
+        checkArgument(functionHandle instanceof BuiltInFunctionHandle, "Expect BuiltInFunctionHandle");
+        return getScalarFunctionImplementation(((BuiltInFunctionHandle) functionHandle).getSignature());
+    }
+
+    private synchronized FunctionMap createFunctionMap()
+    {
+        Optional> functionNamespaceManager =
+                functionAndTypeManager.getServingFunctionNamespaceManager(functionAndTypeManager.getDefaultNamespace());
+        checkArgument(functionNamespaceManager.isPresent(), "Cannot find function namespace for catalog '%s'", functionAndTypeManager.getDefaultNamespace().getCatalogName());
+        checkForNamingConflicts(functionNamespaceManager.get().listFunctions(Optional.empty(), Optional.empty()));
+        return functions;
+    }
+
+    private ScalarFunctionImplementation getScalarFunctionImplementation(Signature signature)
+    {
+        checkArgument(signature.getKind() == SCALAR, "%s is not a scalar function", signature);
+        checkArgument(signature.getTypeVariableConstraints().isEmpty(), "%s has unbound type parameters", signature);
+
+        try {
+            return specializedScalarCache.getUnchecked(getSpecializedFunctionKey(signature));
+        }
+        catch (UncheckedExecutionException e) {
+            throwIfInstanceOf(e.getCause(), PrestoException.class);
+            throw e;
+        }
+    }
+
+    private SpecializedFunctionKey doGetSpecializedFunctionKey(Signature signature)
+    {
+        return functionAndTypeManager.getSpecializedFunctionKey(signature, getFunctions(Optional.empty(), signature.getName()));
+    }
+
+    private SpecializedFunctionKey getSpecializedFunctionKey(Signature signature)
+    {
+        try {
+            return specializedFunctionKeyCache.getUnchecked(signature);
+        }
+        catch (UncheckedExecutionException e) {
+            throwIfInstanceOf(e.getCause(), PrestoException.class);
+            throw e;
+        }
+    }
+}
diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInWorkerFunctionNamespaceManager.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInWorkerFunctionNamespaceManager.java
new file mode 100644
index 0000000000000..41ef943001cb3
--- /dev/null
+++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInWorkerFunctionNamespaceManager.java
@@ -0,0 +1,59 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.metadata;
+
+import com.facebook.presto.spi.function.FunctionImplementationType;
+import com.facebook.presto.spi.function.SqlFunction;
+
+import java.util.Collection;
+import java.util.List;
+
+import static com.facebook.presto.metadata.BuiltInFunctionKind.WORKER;
+import static com.facebook.presto.spi.function.FunctionImplementationType.CPP;
+
+public class BuiltInWorkerFunctionNamespaceManager
+        extends BuiltInSpecialFunctionNamespaceManager
+{
+    public BuiltInWorkerFunctionNamespaceManager(FunctionAndTypeManager functionAndTypeManager)
+    {
+        super(functionAndTypeManager);
+    }
+
+    @Override
+    public synchronized void registerBuiltInSpecialFunctions(List functions)
+    {
+        // only register functions once
+        if (!this.functions.list().isEmpty()) {
+            return;
+        }
+        this.functions = new FunctionMap(this.functions, functions);
+    }
+
+    @Override
+    protected synchronized void checkForNamingConflicts(Collection functions)
+    {
+    }
+
+    @Override
+    protected BuiltInFunctionKind getBuiltInFunctionKind()
+    {
+        return WORKER;
+    }
+
+    @Override
+    protected FunctionImplementationType getDefaultFunctionMetadataImplementationType()
+    {
+        return CPP;
+    }
+}
diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionAndTypeManager.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionAndTypeManager.java
index 88036bbbe55bc..85dffe5487b7c 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionAndTypeManager.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionAndTypeManager.java
@@ -13,6 +13,7 @@
  */
 package com.facebook.presto.metadata;
 
+import com.facebook.airlift.log.Logger;
 import com.facebook.presto.Session;
 import com.facebook.presto.common.CatalogSchemaName;
 import com.facebook.presto.common.Page;
@@ -39,6 +40,7 @@
 import com.facebook.presto.spi.function.AggregationFunctionImplementation;
 import com.facebook.presto.spi.function.AlterRoutineCharacteristics;
 import com.facebook.presto.spi.function.FunctionHandle;
+import com.facebook.presto.spi.function.FunctionImplementationType;
 import com.facebook.presto.spi.function.FunctionMetadata;
 import com.facebook.presto.spi.function.FunctionMetadataManager;
 import com.facebook.presto.spi.function.FunctionNamespaceManager;
@@ -94,6 +96,7 @@
 import static com.facebook.presto.SystemSessionProperties.isListBuiltInFunctionsOnly;
 import static com.facebook.presto.common.type.TypeSignature.parseTypeSignature;
 import static com.facebook.presto.metadata.BuiltInFunctionKind.PLUGIN;
+import static com.facebook.presto.metadata.BuiltInFunctionKind.WORKER;
 import static com.facebook.presto.metadata.BuiltInTypeAndFunctionNamespaceManager.JAVA_BUILTIN_NAMESPACE;
 import static com.facebook.presto.metadata.CastType.toOperatorType;
 import static com.facebook.presto.metadata.FunctionSignatureMatcher.constructFunctionNotFoundErrorMessage;
@@ -132,6 +135,7 @@ public class FunctionAndTypeManager
         implements FunctionMetadataManager, TypeManager
 {
     private static final Pattern DEFAULT_NAMESPACE_PREFIX_PATTERN = Pattern.compile("[a-z]+\\.[a-z]+");
+    private static final Logger log = Logger.get(FunctionAndTypeManager.class);
     private final TransactionManager transactionManager;
     private final TableFunctionRegistry tableFunctionRegistry;
     private final BlockEncodingSerde blockEncodingSerde;
@@ -147,9 +151,11 @@ public class FunctionAndTypeManager
     private final LoadingCache functionCache;
     private final CacheStatsMBean cacheStatsMBean;
     private final boolean nativeExecution;
+    private final boolean isBuiltInSidecarFunctionsEnabled;
     private final CatalogSchemaName defaultNamespace;
     private final AtomicReference servingTypeManager;
     private final AtomicReference>> servingTypeManagerParametricTypesSupplier;
+    private final BuiltInWorkerFunctionNamespaceManager builtInWorkerFunctionNamespaceManager;
     private final BuiltInPluginFunctionNamespaceManager builtInPluginFunctionNamespaceManager;
     private final FunctionsConfig functionsConfig;
     private final Set types;
@@ -185,9 +191,11 @@ public FunctionAndTypeManager(
         this.functionSignatureMatcher = new FunctionSignatureMatcher(this);
         this.typeCoercer = new TypeCoercer(functionsConfig, this);
         this.nativeExecution = featuresConfig.isNativeExecutionEnabled();
+        this.isBuiltInSidecarFunctionsEnabled = featuresConfig.isBuiltInSidecarFunctionsEnabled();
         this.defaultNamespace = configureDefaultNamespace(functionsConfig.getDefaultNamespacePrefix());
         this.servingTypeManager = new AtomicReference<>(builtInTypeAndFunctionNamespaceManager);
         this.servingTypeManagerParametricTypesSupplier = new AtomicReference<>(this::getServingTypeManagerParametricTypes);
+        this.builtInWorkerFunctionNamespaceManager = new BuiltInWorkerFunctionNamespaceManager(this);
         this.builtInPluginFunctionNamespaceManager = new BuiltInPluginFunctionNamespaceManager(this);
     }
 
@@ -367,8 +375,12 @@ public FunctionMetadata getFunctionMetadata(FunctionHandle functionHandle)
         if (isBuiltInPluginFunctionHandle(functionHandle)) {
             return builtInPluginFunctionNamespaceManager.getFunctionMetadata(functionHandle);
         }
+        if (isBuiltInWorkerFunctionHandle(functionHandle)) {
+            return builtInWorkerFunctionNamespaceManager.getFunctionMetadata(functionHandle);
+        }
         Optional> functionNamespaceManager = getServingFunctionNamespaceManager(functionHandle.getCatalogSchemaName());
         checkArgument(functionNamespaceManager.isPresent(), "Cannot find function namespace for '%s'", functionHandle.getCatalogSchemaName());
+
         return functionNamespaceManager.get().getFunctionMetadata(functionHandle);
     }
 
@@ -469,9 +481,16 @@ public void registerBuiltInFunctions(List functions)
         builtInTypeAndFunctionNamespaceManager.registerBuiltInFunctions(functions);
     }
 
+    public void registerWorkerFunctions(List functions)
+    {
+        if (isBuiltInSidecarFunctionsEnabled) {
+            builtInWorkerFunctionNamespaceManager.registerBuiltInSpecialFunctions(functions);
+        }
+    }
+
     public void registerPluginFunctions(List functions)
     {
-        builtInPluginFunctionNamespaceManager.registerPluginFunctions(functions);
+        builtInPluginFunctionNamespaceManager.registerBuiltInSpecialFunctions(functions);
     }
 
     public void registerConnectorFunctions(String catalogName, List functions)
@@ -502,6 +521,7 @@ public List listFunctions(Session session, Optional likePat
                             defaultNamespace.getCatalogName()).listFunctions(likePattern, escape).stream()
                     .collect(toImmutableList()));
             functions.addAll(builtInPluginFunctionNamespaceManager.listFunctions(likePattern, escape).stream().collect(toImmutableList()));
+            functions.addAll(builtInWorkerFunctionNamespaceManager.listFunctions(likePattern, escape).stream().collect(toImmutableList()));
         }
         else {
             functions.addAll(SessionFunctionUtils.listFunctions(session.getSessionFunctions()));
@@ -509,6 +529,7 @@ public List listFunctions(Session session, Optional likePat
                     .flatMap(manager -> manager.listFunctions(likePattern, escape).stream())
                     .collect(toImmutableList()));
             functions.addAll(builtInPluginFunctionNamespaceManager.listFunctions(likePattern, escape).stream().collect(toImmutableList()));
+            functions.addAll(builtInWorkerFunctionNamespaceManager.listFunctions(likePattern, escape).stream().collect(toImmutableList()));
         }
 
         return functions.build().stream()
@@ -654,6 +675,10 @@ public ScalarFunctionImplementation getScalarFunctionImplementation(FunctionHand
         if (isBuiltInPluginFunctionHandle(functionHandle)) {
             return builtInPluginFunctionNamespaceManager.getScalarFunctionImplementation(functionHandle);
         }
+        if (isBuiltInWorkerFunctionHandle(functionHandle)) {
+            return builtInWorkerFunctionNamespaceManager.getScalarFunctionImplementation(functionHandle);
+        }
+
         Optional> functionNamespaceManager = getServingFunctionNamespaceManager(functionHandle.getCatalogSchemaName());
         checkArgument(functionNamespaceManager.isPresent(), "Cannot find function namespace for '%s'", functionHandle.getCatalogSchemaName());
         return functionNamespaceManager.get().getScalarFunctionImplementation(functionHandle);
@@ -977,13 +1002,15 @@ private Collection getFunctions(
         return ImmutableList.builder()
                 .addAll(functionNamespaceManager.getFunctions(transactionHandle, functionName))
                 .addAll(builtInPluginFunctionNamespaceManager.getFunctions(transactionHandle, functionName))
+                .addAll(builtInWorkerFunctionNamespaceManager.getFunctions(transactionHandle, functionName))
                 .build();
     }
 
     /**
      * Gets the function handle of the function if there is a match. We enforce explicit naming for dynamic function namespaces.
      * All unqualified function names will only be resolved against the built-in default function namespace. We get all the candidates
-     * from the current default namespace and additionally all the candidates from builtInPluginFunctionNamespaceManager.
+     * from the current default namespace and additionally all the candidates from builtInPluginFunctionNamespaceManager and
+     * builtInWorkerFunctionNamespaceManager.
      *
      * @throws PrestoException if there are no matches or multiple matches
      */
@@ -998,20 +1025,43 @@ private FunctionHandle getMatchingFunctionHandle(
                 getMatchingFunction(functionNamespaceManager.getFunctions(transactionHandle, functionName), parameterTypes, coercionAllowed);
         Optional matchingPluginFunctionSignature =
                 getMatchingFunction(builtInPluginFunctionNamespaceManager.getFunctions(transactionHandle, functionName), parameterTypes, coercionAllowed);
+        Optional matchingWorkerFunctionSignature =
+                getMatchingFunction(builtInWorkerFunctionNamespaceManager.getFunctions(transactionHandle, functionName), parameterTypes, coercionAllowed);
 
         if (matchingDefaultFunctionSignature.isPresent() && matchingPluginFunctionSignature.isPresent()) {
             throw new PrestoException(AMBIGUOUS_FUNCTION_CALL, format("Function '%s' has two matching signatures. Please specify parameter types. \n" +
                     "First match : '%s', Second match: '%s'", functionName, matchingDefaultFunctionSignature.get(), matchingPluginFunctionSignature.get()));
         }
 
-        if (matchingDefaultFunctionSignature.isPresent()) {
-            return functionNamespaceManager.getFunctionHandle(transactionHandle, matchingDefaultFunctionSignature.get());
+        if (matchingDefaultFunctionSignature.isPresent() && matchingWorkerFunctionSignature.isPresent()) {
+            FunctionHandle defaultFunctionHandle = functionNamespaceManager.getFunctionHandle(transactionHandle, matchingDefaultFunctionSignature.get());
+            FunctionHandle workerFunctionHandle = builtInWorkerFunctionNamespaceManager.getFunctionHandle(transactionHandle, matchingWorkerFunctionSignature.get());
+
+            if (functionNamespaceManager.getFunctionMetadata(defaultFunctionHandle).getImplementationType().equals(FunctionImplementationType.JAVA)) {
+                return defaultFunctionHandle;
+            }
+            if (functionNamespaceManager.getFunctionMetadata(defaultFunctionHandle).getImplementationType().equals(FunctionImplementationType.SQL)) {
+                return workerFunctionHandle;
+            }
+        }
+
+        if (matchingPluginFunctionSignature.isPresent() && matchingWorkerFunctionSignature.isPresent()) {
+            // built in plugin function namespace manager always has SQL as implementation type
+            return builtInWorkerFunctionNamespaceManager.getFunctionHandle(transactionHandle, matchingWorkerFunctionSignature.get());
+        }
+
+        if (matchingWorkerFunctionSignature.isPresent()) {
+            return builtInWorkerFunctionNamespaceManager.getFunctionHandle(transactionHandle, matchingWorkerFunctionSignature.get());
         }
 
         if (matchingPluginFunctionSignature.isPresent()) {
             return builtInPluginFunctionNamespaceManager.getFunctionHandle(transactionHandle, matchingPluginFunctionSignature.get());
         }
 
+        if (matchingDefaultFunctionSignature.isPresent()) {
+            return functionNamespaceManager.getFunctionHandle(transactionHandle, matchingDefaultFunctionSignature.get());
+        }
+
         throw new PrestoException(FUNCTION_NOT_FOUND, constructFunctionNotFoundErrorMessage(functionName, parameterTypes,
                 getFunctions(functionName, transactionHandle, functionNamespaceManager)));
     }
@@ -1029,6 +1079,11 @@ private boolean isBuiltInPluginFunctionHandle(FunctionHandle functionHandle)
         return (functionHandle instanceof BuiltInFunctionHandle) && ((BuiltInFunctionHandle) functionHandle).getBuiltInFunctionKind().equals(PLUGIN);
     }
 
+    private boolean isBuiltInWorkerFunctionHandle(FunctionHandle functionHandle)
+    {
+        return (functionHandle instanceof BuiltInFunctionHandle) && ((BuiltInFunctionHandle) functionHandle).getBuiltInFunctionKind().equals(WORKER);
+    }
+
     private static class FunctionResolutionCacheKey
     {
         private final QualifiedObjectName functionName;
diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/FeaturesConfig.java b/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/FeaturesConfig.java
index 2b0b89d99b472..06b1fee1fdce7 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/FeaturesConfig.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/FeaturesConfig.java
@@ -310,6 +310,8 @@ public class FeaturesConfig
     private boolean pushdownSubfieldForMapFunctions = true;
     private long maxSerializableObjectSize = 1000;
 
+    private boolean builtInSidecarFunctionsEnabled;
+
     public enum PartitioningPrecisionStrategy
     {
         // Let Presto decide when to repartition
@@ -3111,4 +3113,17 @@ public long getMaxSerializableObjectSize()
     {
         return maxSerializableObjectSize;
     }
+
+    @Config("built-in-sidecar-functions-enabled")
+    @ConfigDescription("Enable using CPP functions from sidecar over coordinator SQL implementations.")
+    public FeaturesConfig setBuiltInSidecarFunctionsEnabled(boolean builtInSidecarFunctionsEnabled)
+    {
+        this.builtInSidecarFunctionsEnabled = builtInSidecarFunctionsEnabled;
+        return this;
+    }
+
+    public boolean isBuiltInSidecarFunctionsEnabled()
+    {
+        return this.builtInSidecarFunctionsEnabled;
+    }
 }
diff --git a/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestConvertApplicableTypeToVariable.java b/presto-main-base/src/test/java/com/facebook/presto/metadata/TestConvertApplicableTypeToVariable.java
similarity index 98%
rename from presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestConvertApplicableTypeToVariable.java
rename to presto-main-base/src/test/java/com/facebook/presto/metadata/TestConvertApplicableTypeToVariable.java
index 93f9ee75f2b82..9e9a8b629124a 100644
--- a/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestConvertApplicableTypeToVariable.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/metadata/TestConvertApplicableTypeToVariable.java
@@ -11,7 +11,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.facebook.presto.sidecar;
+package com.facebook.presto.metadata;
 
 import com.facebook.presto.common.type.NamedTypeSignature;
 import com.facebook.presto.common.type.TypeSignature;
@@ -21,8 +21,8 @@
 import java.util.List;
 import java.util.Optional;
 
+import static com.facebook.presto.builtin.tools.WorkerFunctionUtil.convertApplicableTypeToVariable;
 import static com.facebook.presto.common.type.TypeSignature.parseTypeSignature;
-import static com.facebook.presto.sidecar.functionNamespace.NativeFunctionNamespaceManager.convertApplicableTypeToVariable;
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertTrue;
 
diff --git a/presto-main-base/src/test/java/com/facebook/presto/sql/analyzer/TestFeaturesConfig.java b/presto-main-base/src/test/java/com/facebook/presto/sql/analyzer/TestFeaturesConfig.java
index 25a802b1e06ea..cc29fe7160751 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/sql/analyzer/TestFeaturesConfig.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/sql/analyzer/TestFeaturesConfig.java
@@ -201,6 +201,7 @@ public void testDefaults()
                 .setPushRemoteExchangeThroughGroupId(false)
                 .setOptimizeMultipleApproxPercentileOnSameFieldEnabled(true)
                 .setNativeExecutionEnabled(false)
+                .setBuiltInSidecarFunctionsEnabled(false)
                 .setDisableTimeStampWithTimeZoneForNative(false)
                 .setDisableIPAddressForNative(false)
                 .setNativeExecutionExecutablePath("./presto_server")
@@ -416,6 +417,7 @@ public void testExplicitPropertyMappings()
                 .put("optimizer.push-remote-exchange-through-group-id", "true")
                 .put("optimizer.optimize-multiple-approx-percentile-on-same-field", "false")
                 .put("native-execution-enabled", "true")
+                .put("built-in-sidecar-functions-enabled", "true")
                 .put("disable-timestamp-with-timezone-for-native-execution", "true")
                 .put("disable-ipaddress-for-native-execution", "true")
                 .put("native-execution-executable-path", "/bin/echo")
@@ -628,6 +630,7 @@ public void testExplicitPropertyMappings()
                 .setPushRemoteExchangeThroughGroupId(true)
                 .setOptimizeMultipleApproxPercentileOnSameFieldEnabled(false)
                 .setNativeExecutionEnabled(true)
+                .setBuiltInSidecarFunctionsEnabled(true)
                 .setDisableTimeStampWithTimeZoneForNative(true)
                 .setDisableIPAddressForNative(true)
                 .setNativeExecutionExecutablePath("/bin/echo")
diff --git a/presto-main/pom.xml b/presto-main/pom.xml
index 07e0f4828ceb0..194d70a0610c8 100644
--- a/presto-main/pom.xml
+++ b/presto-main/pom.xml
@@ -23,6 +23,11 @@
             presto-main-base
         
 
+        
+            com.facebook.presto
+            presto-built-in-worker-function-tools
+        
+
         
             com.facebook.airlift
             jmx-http
@@ -53,6 +58,11 @@
             presto-common
         
 
+        
+            com.facebook.presto
+            presto-function-namespace-managers-common
+        
+
         
             io.jsonwebtoken
             jjwt-api
diff --git a/presto-main/src/main/java/com/facebook/presto/server/PrestoServer.java b/presto-main/src/main/java/com/facebook/presto/server/PrestoServer.java
index 552b6da5d812e..6462e9b8fb57b 100644
--- a/presto-main/src/main/java/com/facebook/presto/server/PrestoServer.java
+++ b/presto-main/src/main/java/com/facebook/presto/server/PrestoServer.java
@@ -33,6 +33,7 @@
 import com.facebook.drift.transport.netty.server.DriftNettyServerTransport;
 import com.facebook.presto.ClientRequestFilterManager;
 import com.facebook.presto.ClientRequestFilterModule;
+import com.facebook.presto.builtin.tools.WorkerFunctionRegistryTool;
 import com.facebook.presto.dispatcher.QueryPrerequisitesManager;
 import com.facebook.presto.dispatcher.QueryPrerequisitesManagerModule;
 import com.facebook.presto.eventlistener.EventListenerManager;
@@ -43,6 +44,7 @@
 import com.facebook.presto.metadata.Catalog;
 import com.facebook.presto.metadata.CatalogManager;
 import com.facebook.presto.metadata.DiscoveryNodeManager;
+import com.facebook.presto.metadata.FunctionAndTypeManager;
 import com.facebook.presto.metadata.InternalNodeManager;
 import com.facebook.presto.metadata.SessionPropertyManager;
 import com.facebook.presto.metadata.StaticCatalogStore;
@@ -54,6 +56,7 @@
 import com.facebook.presto.server.security.PasswordAuthenticatorManager;
 import com.facebook.presto.server.security.PrestoAuthenticatorManager;
 import com.facebook.presto.server.security.ServerSecurityModule;
+import com.facebook.presto.spi.function.SqlFunction;
 import com.facebook.presto.sql.analyzer.FeaturesConfig;
 import com.facebook.presto.sql.expressions.ExpressionOptimizerManager;
 import com.facebook.presto.sql.parser.SqlParserOptions;
@@ -198,6 +201,11 @@ public void run()
             PluginNodeManager pluginNodeManager = new PluginNodeManager(nodeManager, nodeInfo.getEnvironment());
             planCheckerProviderManager.loadPlanCheckerProviders(pluginNodeManager);
 
+            if (injector.getInstance(FeaturesConfig.class).isBuiltInSidecarFunctionsEnabled()) {
+                List functions = injector.getInstance(WorkerFunctionRegistryTool.class).getWorkerFunctions();
+                injector.getInstance(FunctionAndTypeManager.class).registerWorkerFunctions(functions);
+            }
+
             injector.getInstance(ClientRequestFilterManager.class).loadClientRequestFilters();
             injector.getInstance(ExpressionOptimizerManager.class).loadExpressionOptimizerFactories();
 
diff --git a/presto-main/src/main/java/com/facebook/presto/server/ServerMainModule.java b/presto-main/src/main/java/com/facebook/presto/server/ServerMainModule.java
index 634b8654634a6..c564764143ce7 100644
--- a/presto-main/src/main/java/com/facebook/presto/server/ServerMainModule.java
+++ b/presto-main/src/main/java/com/facebook/presto/server/ServerMainModule.java
@@ -16,7 +16,10 @@
 import com.facebook.airlift.concurrent.BoundedExecutor;
 import com.facebook.airlift.configuration.AbstractConfigurationAwareModule;
 import com.facebook.airlift.discovery.client.ServiceAnnouncement;
+import com.facebook.airlift.http.client.HttpClient;
 import com.facebook.airlift.http.server.TheServlet;
+import com.facebook.airlift.json.JsonCodec;
+import com.facebook.airlift.json.JsonCodecFactory;
 import com.facebook.airlift.json.JsonObjectMapperProvider;
 import com.facebook.airlift.stats.GcMonitor;
 import com.facebook.airlift.stats.JmxGcMonitor;
@@ -32,6 +35,10 @@
 import com.facebook.presto.PagesIndexPageSorter;
 import com.facebook.presto.SystemSessionProperties;
 import com.facebook.presto.block.BlockJsonSerde;
+import com.facebook.presto.builtin.tools.ForNativeFunctionRegistryInfo;
+import com.facebook.presto.builtin.tools.NativeSidecarFunctionRegistryTool;
+import com.facebook.presto.builtin.tools.NativeSidecarRegistryToolConfig;
+import com.facebook.presto.builtin.tools.WorkerFunctionRegistryTool;
 import com.facebook.presto.catalogserver.CatalogServerClient;
 import com.facebook.presto.catalogserver.RandomCatalogServerAddressSelector;
 import com.facebook.presto.catalogserver.RemoteMetadataManager;
@@ -79,6 +86,7 @@
 import com.facebook.presto.execution.scheduler.TableWriteInfo;
 import com.facebook.presto.execution.scheduler.nodeSelection.NodeSelectionStats;
 import com.facebook.presto.execution.scheduler.nodeSelection.SimpleTtlNodeSelectorConfig;
+import com.facebook.presto.functionNamespace.JsonBasedUdfFunctionMetadata;
 import com.facebook.presto.index.IndexManager;
 import com.facebook.presto.memory.LocalMemoryManager;
 import com.facebook.presto.memory.LocalMemoryManagerExporter;
@@ -243,6 +251,7 @@
 import com.google.inject.Module;
 import com.google.inject.Provides;
 import com.google.inject.Scopes;
+import com.google.inject.TypeLiteral;
 import com.google.inject.multibindings.MapBinder;
 import io.airlift.slice.Slice;
 import jakarta.annotation.PreDestroy;
@@ -251,6 +260,7 @@
 import jakarta.servlet.Servlet;
 
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.LinkedBlockingQueue;
@@ -266,6 +276,7 @@
 import static com.facebook.airlift.http.client.HttpClientBinder.httpClientBinder;
 import static com.facebook.airlift.jaxrs.JaxrsBinder.jaxrsBinder;
 import static com.facebook.airlift.json.JsonBinder.jsonBinder;
+import static com.facebook.airlift.json.JsonCodec.listJsonCodec;
 import static com.facebook.airlift.json.JsonCodecBinder.jsonCodecBinder;
 import static com.facebook.airlift.json.smile.SmileCodecBinder.smileCodecBinder;
 import static com.facebook.airlift.units.DataSize.Unit.MEGABYTE;
@@ -394,6 +405,11 @@ else if (serverConfig.isCoordinator()) {
                 .withAddressSelector(((addressSelectorBinder, annotation, prefix) ->
                         addressSelectorBinder.bind(AddressSelector.class).annotatedWith(annotation).to(FixedAddressSelector.class)));
 
+        binder.bind(new TypeLiteral>>>() {})
+                .toInstance(new JsonCodecFactory().mapJsonCodec(String.class, listJsonCodec(JsonBasedUdfFunctionMetadata.class)));
+        httpClientBinder(binder).bindHttpClient("native-function-registry", ForNativeFunctionRegistryInfo.class);
+        configBinder(binder).bindConfig(NativeSidecarRegistryToolConfig.class);
+
         // node scheduler
         // TODO: remove from NodePartitioningManager and move to CoordinatorModule
         configBinder(binder).bindConfig(NodeSchedulerConfig.class);
@@ -891,6 +907,22 @@ public static FragmentResultCacheManager createFragmentResultCacheManager(FileFr
         return new NoOpFragmentResultCacheManager();
     }
 
+    @Provides
+    @Singleton
+    public WorkerFunctionRegistryTool provideWorkerFunctionRegistryTool(
+            NativeSidecarRegistryToolConfig config,
+            @ForNativeFunctionRegistryInfo HttpClient httpClient,
+            JsonCodec>> nativeFunctionSignatureMapJsonCodec,
+            NodeManager nodeManager)
+    {
+        return new NativeSidecarFunctionRegistryTool(
+                httpClient,
+                nativeFunctionSignatureMapJsonCodec,
+                nodeManager,
+                config.getNativeSidecarRegistryToolNumRetries(),
+                config.getNativeSidecarRegistryToolRetryDelayMs());
+    }
+
     public static class ExecutorCleanup
     {
         private final List executors;
diff --git a/presto-main/src/main/java/com/facebook/presto/server/testing/TestingPrestoServer.java b/presto-main/src/main/java/com/facebook/presto/server/testing/TestingPrestoServer.java
index 1cb04591b15da..43c574058f2bb 100644
--- a/presto-main/src/main/java/com/facebook/presto/server/testing/TestingPrestoServer.java
+++ b/presto-main/src/main/java/com/facebook/presto/server/testing/TestingPrestoServer.java
@@ -34,6 +34,7 @@
 import com.facebook.drift.transport.netty.server.DriftNettyServerTransport;
 import com.facebook.presto.ClientRequestFilterManager;
 import com.facebook.presto.ClientRequestFilterModule;
+import com.facebook.presto.builtin.tools.WorkerFunctionRegistryTool;
 import com.facebook.presto.connector.ConnectorManager;
 import com.facebook.presto.cost.StatsCalculator;
 import com.facebook.presto.dispatcher.DispatchManager;
@@ -51,6 +52,7 @@
 import com.facebook.presto.memory.LocalMemoryManager;
 import com.facebook.presto.metadata.AllNodes;
 import com.facebook.presto.metadata.CatalogManager;
+import com.facebook.presto.metadata.FunctionAndTypeManager;
 import com.facebook.presto.metadata.InternalNode;
 import com.facebook.presto.metadata.InternalNodeManager;
 import com.facebook.presto.metadata.Metadata;
@@ -147,6 +149,8 @@ public class TestingPrestoServer
     private final boolean preserveData;
     private final LifeCycleManager lifeCycleManager;
     private final PluginManager pluginManager;
+    private final FunctionAndTypeManager functionAndTypeManager;
+    private final WorkerFunctionRegistryTool workerFunctionRegistryTool;
     private final ConnectorManager connectorManager;
     private final TestingHttpServer server;
     private final CatalogManager catalogManager;
@@ -367,6 +371,9 @@ public TestingPrestoServer(
 
         connectorManager = injector.getInstance(ConnectorManager.class);
 
+        functionAndTypeManager = injector.getInstance(FunctionAndTypeManager.class);
+        workerFunctionRegistryTool = injector.getInstance(WorkerFunctionRegistryTool.class);
+
         server = injector.getInstance(TestingHttpServer.class);
         catalogManager = injector.getInstance(CatalogManager.class);
         transactionManager = injector.getInstance(TransactionManager.class);
@@ -501,6 +508,11 @@ private Map getServerProperties(
         return ImmutableMap.copyOf(serverProperties);
     }
 
+    public void registerWorkerFunctions()
+    {
+        functionAndTypeManager.registerWorkerFunctions(workerFunctionRegistryTool.getWorkerFunctions());
+    }
+
     @Override
     public void close()
             throws IOException
diff --git a/presto-native-execution/pom.xml b/presto-native-execution/pom.xml
index ab6d81aba382b..8f849f866335a 100644
--- a/presto-native-execution/pom.xml
+++ b/presto-native-execution/pom.xml
@@ -36,6 +36,18 @@
             test
         
 
+        
+            org.jetbrains
+            annotations
+            test
+        
+
+        
+            org.weakref
+            jmxutils
+            test
+        
+
         
             io.airlift.tpch
             tpch
@@ -116,6 +128,22 @@
             test
         
 
+        
+            com.facebook.presto
+            presto-parser
+            test
+            
+                
+                    com.google.guava
+                    guava
+                
+                
+                    com.facebook.presto
+                    presto-spi
+                
+            
+        
+
         
             com.facebook.presto
             presto-iceberg
diff --git a/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/PrestoNativeQueryRunnerUtils.java b/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/PrestoNativeQueryRunnerUtils.java
index 95095f796a72d..e287ef7e0cb20 100644
--- a/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/PrestoNativeQueryRunnerUtils.java
+++ b/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/PrestoNativeQueryRunnerUtils.java
@@ -123,6 +123,7 @@ public static class HiveQueryRunnerBuilder
         private String security;
         private boolean addStorageFormatToPath;
         private boolean coordinatorSidecarEnabled;
+        private boolean builtInWorkerFunctionsEnabled;
         private boolean enableRuntimeMetricsCollection;
         private boolean enableSsdCache;
         private boolean failOnNestedLoopJoin;
@@ -221,6 +222,12 @@ public HiveQueryRunnerBuilder setCoordinatorSidecarEnabled(boolean coordinatorSi
             return this;
         }
 
+        public HiveQueryRunnerBuilder setBuiltInWorkerFunctionsEnabled(boolean builtInWorkerFunctionsEnabled)
+        {
+            this.builtInWorkerFunctionsEnabled = builtInWorkerFunctionsEnabled;
+            return this;
+        }
+
         public HiveQueryRunnerBuilder setStorageFormat(String storageFormat)
         {
             this.storageFormat = storageFormat;
@@ -270,7 +277,7 @@ public QueryRunner build()
             Optional> externalWorkerLauncher = Optional.empty();
             if (this.useExternalWorkerLauncher) {
                 externalWorkerLauncher = getExternalWorkerLauncher("hive", serverBinary, cacheMaxSize, remoteFunctionServerUds,
-                        failOnNestedLoopJoin, coordinatorSidecarEnabled, enableRuntimeMetricsCollection, enableSsdCache, implicitCastCharNToVarchar);
+                        failOnNestedLoopJoin, coordinatorSidecarEnabled, builtInWorkerFunctionsEnabled, enableRuntimeMetricsCollection, enableSsdCache, implicitCastCharNToVarchar);
             }
             return HiveQueryRunner.createQueryRunner(
                     ImmutableList.of(),
@@ -359,7 +366,7 @@ public QueryRunner build()
             Optional> externalWorkerLauncher = Optional.empty();
             if (this.useExternalWorkerLauncher) {
                 externalWorkerLauncher = getExternalWorkerLauncher("iceberg", serverBinary, cacheMaxSize, remoteFunctionServerUds,
-                        false, false, false, false, false);
+                        false, false, false, false, false, false);
             }
             return IcebergQueryRunner.builder()
                     .setExtraProperties(extraProperties)
@@ -455,6 +462,7 @@ public static Optional> getExternalWorkerLaunc
             Optional remoteFunctionServerUds,
             Boolean failOnNestedLoopJoin,
             boolean isCoordinatorSidecarEnabled,
+            boolean isBuiltInWorkerFunctionsEnabled,
             boolean enableRuntimeMetricsCollection,
             boolean enableSsdCache,
             boolean implicitCastCharNToVarchar)
@@ -478,6 +486,10 @@ public static Optional> getExternalWorkerLaunc
                                     "native-sidecar=true%n" +
                                     "presto.default-namespace=native.default%n", configProperties);
                         }
+                        else if (isBuiltInWorkerFunctionsEnabled) {
+                            configProperties = format("%s%n" +
+                                    "native-sidecar=true%n", configProperties);
+                        }
 
                         if (enableRuntimeMetricsCollection) {
                             configProperties = format("%s%n" +
diff --git a/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/TestBuiltInNativeFunctions.java b/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/TestBuiltInNativeFunctions.java
new file mode 100644
index 0000000000000..f22707cf0bfad
--- /dev/null
+++ b/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/TestBuiltInNativeFunctions.java
@@ -0,0 +1,178 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.facebook.presto.nativeworker;
+
+import com.facebook.presto.cost.CostCalculator;
+import com.facebook.presto.cost.CostCalculatorUsingExchanges;
+import com.facebook.presto.cost.CostCalculatorWithEstimatedExchanges;
+import com.facebook.presto.cost.CostComparator;
+import com.facebook.presto.cost.TaskCountEstimator;
+import com.facebook.presto.execution.QueryManagerConfig;
+import com.facebook.presto.execution.TaskManagerConfig;
+import com.facebook.presto.metadata.InMemoryNodeManager;
+import com.facebook.presto.metadata.Metadata;
+import com.facebook.presto.nodeManager.PluginNodeManager;
+import com.facebook.presto.spi.WarningCollector;
+import com.facebook.presto.sql.analyzer.FeaturesConfig;
+import com.facebook.presto.sql.analyzer.QueryExplainer;
+import com.facebook.presto.sql.expressions.ExpressionOptimizerManager;
+import com.facebook.presto.sql.planner.PartitioningProviderManager;
+import com.facebook.presto.sql.planner.PlanFragmenter;
+import com.facebook.presto.sql.planner.PlanOptimizers;
+import com.facebook.presto.sql.planner.optimizations.PlanOptimizer;
+import com.facebook.presto.sql.planner.sanity.PlanChecker;
+import com.facebook.presto.sql.tree.ExplainType;
+import com.facebook.presto.testing.QueryRunner;
+import com.facebook.presto.tests.AbstractTestQueryFramework;
+import com.facebook.presto.tests.DistributedQueryRunner;
+import com.google.common.collect.ImmutableMap;
+import org.intellij.lang.annotations.Language;
+import org.testng.annotations.Test;
+import org.weakref.jmx.MBeanExporter;
+import org.weakref.jmx.testing.TestingMBeanServer;
+
+import java.util.List;
+import java.util.regex.Pattern;
+
+import static com.facebook.presto.nativeworker.NativeQueryRunnerUtils.createLineitem;
+import static com.facebook.presto.nativeworker.NativeQueryRunnerUtils.createNation;
+import static com.facebook.presto.nativeworker.NativeQueryRunnerUtils.createOrders;
+import static com.facebook.presto.nativeworker.NativeQueryRunnerUtils.createOrdersEx;
+import static com.facebook.presto.nativeworker.NativeQueryRunnerUtils.createRegion;
+import static com.facebook.presto.transaction.TransactionBuilder.transaction;
+import static com.facebook.presto.util.AnalyzerUtil.createParsingOptions;
+import static java.util.Collections.emptyList;
+import static org.testng.Assert.fail;
+
+public class TestBuiltInNativeFunctions
+        extends AbstractTestQueryFramework
+{
+    @Override
+    protected void createTables()
+    {
+        QueryRunner queryRunner = (QueryRunner) getExpectedQueryRunner();
+        createLineitem(queryRunner);
+        createNation(queryRunner);
+        createOrders(queryRunner);
+        createOrdersEx(queryRunner);
+        createRegion(queryRunner);
+    }
+
+    @Override
+    protected QueryRunner createQueryRunner()
+            throws Exception
+    {
+        DistributedQueryRunner queryRunner = (DistributedQueryRunner) PrestoNativeQueryRunnerUtils.nativeHiveQueryRunnerBuilder()
+                .setExtraProperties(ImmutableMap.of("built-in-sidecar-functions-enabled", "true"))
+                .setAddStorageFormatToPath(true)
+                .setBuiltInWorkerFunctionsEnabled(true)
+                .build();
+
+        queryRunner.registerNativeFunctions();
+
+        return queryRunner;
+    }
+
+    @Override
+    protected QueryRunner createExpectedQueryRunner()
+            throws Exception
+    {
+        return PrestoNativeQueryRunnerUtils.javaHiveQueryRunnerBuilder()
+                .setAddStorageFormatToPath(true)
+                .build();
+    }
+
+    private void assertJsonPlan(@Language("SQL") String query, boolean withBuiltInSidecarEnabled, @Language("RegExp") String jsonPlanRegex, boolean shouldContainRegex)
+    {
+        QueryRunner queryRunner;
+        if (withBuiltInSidecarEnabled) {
+            queryRunner = getQueryRunner();
+        }
+        else {
+            queryRunner = (QueryRunner) getExpectedQueryRunner();
+        }
+
+        QueryExplainer explainer = getQueryExplainerFromProvidedQueryRunner(queryRunner);
+        transaction(queryRunner.getTransactionManager(), queryRunner.getAccessControl())
+                .singleStatement()
+                .execute(queryRunner.getDefaultSession(), transactionSession -> {
+                    String actualPlan = explainer.getJsonPlan(transactionSession, getSqlParser().createStatement(query, createParsingOptions(transactionSession)), ExplainType.Type.LOGICAL, emptyList(), WarningCollector.NOOP, query);
+                    Pattern p = Pattern.compile(jsonPlanRegex, Pattern.MULTILINE);
+                    if (shouldContainRegex) {
+                        if (!p.matcher(actualPlan).find()) {
+                            fail("Query plan text does not contain regex");
+                        }
+                    }
+                    else {
+                        if (p.matcher(actualPlan).find()) {
+                            fail("Query plan text contains bad pattern");
+                        }
+                    }
+
+                    return null;
+                });
+    }
+
+    private QueryExplainer getQueryExplainerFromProvidedQueryRunner(QueryRunner queryRunner)
+    {
+        Metadata metadata = queryRunner.getMetadata();
+        FeaturesConfig featuresConfig = createFeaturesConfig();
+        boolean noExchange = queryRunner.getNodeCount() == 1;
+        TaskCountEstimator taskCountEstimator = new TaskCountEstimator(queryRunner::getNodeCount);
+        CostCalculator costCalculator = new CostCalculatorUsingExchanges(taskCountEstimator);
+        List optimizers = new PlanOptimizers(
+                metadata,
+                getSqlParser(),
+                noExchange,
+                new MBeanExporter(new TestingMBeanServer()),
+                queryRunner.getSplitManager(),
+                queryRunner.getPlanOptimizerManager(),
+                queryRunner.getPageSourceManager(),
+                queryRunner.getStatsCalculator(),
+                costCalculator,
+                new CostCalculatorWithEstimatedExchanges(costCalculator, taskCountEstimator),
+                new CostComparator(featuresConfig),
+                taskCountEstimator,
+                new PartitioningProviderManager(),
+                featuresConfig,
+                new ExpressionOptimizerManager(
+                        new PluginNodeManager(new InMemoryNodeManager()),
+                        queryRunner.getMetadata().getFunctionAndTypeManager()),
+                new TaskManagerConfig())
+                .getPlanningTimeOptimizers();
+        return new QueryExplainer(
+                optimizers,
+                new PlanFragmenter(metadata, queryRunner.getNodePartitioningManager(), new QueryManagerConfig(), featuresConfig, queryRunner.getPlanCheckerProviderManager()),
+                metadata,
+                queryRunner.getAccessControl(),
+                getSqlParser(),
+                queryRunner.getStatsCalculator(),
+                costCalculator,
+                ImmutableMap.of(),
+                new PlanChecker(featuresConfig, false, queryRunner.getPlanCheckerProviderManager()));
+    }
+
+    @Test
+    public void testUdfQueries()
+    {
+        assertQuery("SELECT ARRAY['abc']");
+        assertQuery("SELECT ARRAY[1, 2, 3]");
+        assertQuery("SELECT map_remove_null_values( MAP( ARRAY['a', 'b', 'c'], ARRAY[1, NULL, 3] ) )");
+        assertQuery("SELECT presto.default.map_remove_null_values( MAP( ARRAY['a', 'b', 'c'], ARRAY[1, NULL, 3] ) )");
+        assertQueryFails("SELECT native.default.map_remove_null_values( MAP( ARRAY['a', 'b', 'c'], ARRAY[1, NULL, 3] ) )", ".*Function native.default.map_remove_null_values not registered.*");
+        assertJsonPlan("SELECT map_remove_null_values( MAP( ARRAY['a', 'b', 'c'], ARRAY[1, NULL, 3] ) )", true, "lambda", false);
+        assertJsonPlan("SELECT map_remove_null_values( MAP( ARRAY['a', 'b', 'c'], ARRAY[1, NULL, 3] ) )", false, "lambda", true);
+    }
+}
diff --git a/presto-native-sidecar-plugin/pom.xml b/presto-native-sidecar-plugin/pom.xml
index 153ad18295f8c..b2bc40e8f2bd4 100644
--- a/presto-native-sidecar-plugin/pom.xml
+++ b/presto-native-sidecar-plugin/pom.xml
@@ -260,6 +260,10 @@
                 
             
         
+        
+            com.facebook.presto
+            presto-built-in-worker-function-tools
+        
     
 
     
diff --git a/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/functionNamespace/NativeFunctionNamespaceManager.java b/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/functionNamespace/NativeFunctionNamespaceManager.java
index 11ef917f5522d..f1b7142a188ee 100644
--- a/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/functionNamespace/NativeFunctionNamespaceManager.java
+++ b/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/functionNamespace/NativeFunctionNamespaceManager.java
@@ -14,14 +14,10 @@
 package com.facebook.presto.sidecar.functionNamespace;
 
 import com.facebook.airlift.log.Logger;
-import com.facebook.presto.common.CatalogSchemaName;
 import com.facebook.presto.common.QualifiedObjectName;
-import com.facebook.presto.common.type.NamedTypeSignature;
-import com.facebook.presto.common.type.StandardTypes;
 import com.facebook.presto.common.type.Type;
 import com.facebook.presto.common.type.TypeManager;
 import com.facebook.presto.common.type.TypeSignature;
-import com.facebook.presto.common.type.TypeSignatureParameter;
 import com.facebook.presto.common.type.UserDefinedType;
 import com.facebook.presto.functionNamespace.AbstractSqlInvokedFunctionNamespaceManager;
 import com.facebook.presto.functionNamespace.JsonBasedUdfFunctionMetadata;
@@ -38,7 +34,6 @@
 import com.facebook.presto.spi.function.FunctionMetadata;
 import com.facebook.presto.spi.function.FunctionMetadataManager;
 import com.facebook.presto.spi.function.FunctionNamespaceTransactionHandle;
-import com.facebook.presto.spi.function.LongVariableConstraint;
 import com.facebook.presto.spi.function.Parameter;
 import com.facebook.presto.spi.function.ScalarFunctionImplementation;
 import com.facebook.presto.spi.function.Signature;
@@ -48,35 +43,29 @@
 import com.facebook.presto.spi.function.SqlFunctionSupplier;
 import com.facebook.presto.spi.function.SqlInvokedAggregationFunctionImplementation;
 import com.facebook.presto.spi.function.SqlInvokedFunction;
-import com.facebook.presto.spi.function.TypeVariableConstraint;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Suppliers;
 import com.google.common.cache.CacheBuilder;
 import com.google.common.cache.CacheLoader;
 import com.google.common.cache.LoadingCache;
-import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.util.concurrent.UncheckedExecutionException;
 import jakarta.inject.Inject;
 
-import java.util.ArrayList;
 import java.util.Collection;
-import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Supplier;
 
+import static com.facebook.presto.builtin.tools.WorkerFunctionUtil.createSqlInvokedFunction;
 import static com.facebook.presto.common.type.TypeSignatureUtils.resolveIntermediateType;
 import static com.facebook.presto.spi.StandardErrorCode.DUPLICATE_FUNCTION_ERROR;
 import static com.facebook.presto.spi.StandardErrorCode.GENERIC_INTERNAL_ERROR;
 import static com.facebook.presto.spi.StandardErrorCode.GENERIC_USER_ERROR;
 import static com.facebook.presto.spi.StandardErrorCode.NOT_SUPPORTED;
-import static com.facebook.presto.spi.function.FunctionVersion.notVersioned;
-import static com.facebook.presto.spi.function.RoutineCharacteristics.Language.CPP;
 import static com.google.common.base.Preconditions.checkArgument;
-import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.MoreCollectors.onlyElement;
 import static java.lang.String.format;
@@ -138,7 +127,7 @@ private synchronized void populateNamespaceManager(UdfFunctionSignatureMap udfFu
     {
         Map> udfSignatureMap = udfFunctionSignatureMap.getUDFSignatureMap();
         udfSignatureMap.forEach((name, metaInfoList) -> {
-            List functions = metaInfoList.stream().map(metaInfo -> createSqlInvokedFunction(name, metaInfo)).collect(toImmutableList());
+            List functions = metaInfoList.stream().map(metaInfo -> createSqlInvokedFunction(name, metaInfo, getCatalogName())).collect(toImmutableList());
             functions.forEach(this::createFunction);
         });
     }
@@ -200,44 +189,6 @@ private AggregationFunctionImplementation processSqlFunctionHandle(SqlFunctionHa
         return aggregationImplementationByHandle.get(sqlFunctionHandle);
     }
 
-    protected synchronized SqlInvokedFunction createSqlInvokedFunction(String functionName, JsonBasedUdfFunctionMetadata jsonBasedUdfFunctionMetaData)
-    {
-        checkState(jsonBasedUdfFunctionMetaData.getRoutineCharacteristics().getLanguage().equals(CPP), "NativeFunctionNamespaceManager only supports CPP UDF");
-        QualifiedObjectName qualifiedFunctionName = QualifiedObjectName.valueOf(new CatalogSchemaName(getCatalogName(), jsonBasedUdfFunctionMetaData.getSchema()), functionName);
-        List parameterNameList = jsonBasedUdfFunctionMetaData.getParamNames();
-        List parameterTypeList = convertApplicableTypeToVariable(jsonBasedUdfFunctionMetaData.getParamTypes());
-        List typeVariableConstraintsList = jsonBasedUdfFunctionMetaData.getTypeVariableConstraints().isPresent() ?
-                jsonBasedUdfFunctionMetaData.getTypeVariableConstraints().get() : Collections.emptyList();
-        List longVariableConstraintList = jsonBasedUdfFunctionMetaData.getLongVariableConstraints().isPresent() ?
-                jsonBasedUdfFunctionMetaData.getLongVariableConstraints().get() : Collections.emptyList();
-
-        TypeSignature outputType = convertApplicableTypeToVariable(jsonBasedUdfFunctionMetaData.getOutputType());
-        ImmutableList.Builder parameterBuilder = ImmutableList.builder();
-        for (int i = 0; i < parameterNameList.size(); i++) {
-            parameterBuilder.add(new Parameter(parameterNameList.get(i), parameterTypeList.get(i)));
-        }
-
-        Optional aggregationFunctionMetadata =
-                jsonBasedUdfFunctionMetaData.getAggregateMetadata()
-                        .map(metadata -> new AggregationFunctionMetadata(
-                                convertApplicableTypeToVariable(metadata.getIntermediateType()),
-                                metadata.isOrderSensitive()));
-
-        return new SqlInvokedFunction(
-                qualifiedFunctionName,
-                parameterBuilder.build(),
-                typeVariableConstraintsList,
-                longVariableConstraintList,
-                outputType,
-                jsonBasedUdfFunctionMetaData.getDocString(),
-                jsonBasedUdfFunctionMetaData.getRoutineCharacteristics(),
-                "",
-                jsonBasedUdfFunctionMetaData.getVariableArity(),
-                notVersioned(),
-                jsonBasedUdfFunctionMetaData.getFunctionKind(),
-                aggregationFunctionMetadata);
-    }
-
     @Override
     protected Collection fetchFunctionsDirect(QualifiedObjectName functionName)
     {
@@ -320,109 +271,12 @@ public final FunctionHandle getFunctionHandle(Optional typeSignaturesList = convertApplicableTypeToVariable(ImmutableList.of(typeSignature));
-        checkArgument(!typeSignaturesList.isEmpty(), "Type signature list is empty for : " + typeSignature);
-        return typeSignaturesList.get(0);
-    }
-
-    public static List convertApplicableTypeToVariable(List typeSignatures)
-    {
-        List newTypeSignaturesList = new ArrayList<>();
-        for (TypeSignature typeSignature : typeSignatures) {
-            if (!typeSignature.getParameters().isEmpty()) {
-                TypeSignature newTypeSignature =
-                        new TypeSignature(
-                                typeSignature.getBase(),
-                                getTypeSignatureParameters(
-                                        typeSignature,
-                                        typeSignature.getParameters()));
-                newTypeSignaturesList.add(newTypeSignature);
-            }
-            else {
-                newTypeSignaturesList.add(typeSignature);
-            }
-        }
-        return newTypeSignaturesList;
-    }
-
     @VisibleForTesting
     public FunctionDefinitionProvider getFunctionDefinitionProvider()
     {
         return functionDefinitionProvider;
     }
 
-    private static List getTypeSignatureParameters(
-            TypeSignature typeSignature,
-            List typeSignatureParameterList)
-    {
-        List newParameterTypeList = new ArrayList<>();
-        for (TypeSignatureParameter parameter : typeSignatureParameterList) {
-            if (parameter.isLongLiteral()) {
-                newParameterTypeList.add(parameter);
-                continue;
-            }
-
-            boolean isNamedTypeSignature = parameter.isNamedTypeSignature();
-            TypeSignature parameterTypeSignature;
-            // If it's a named type signatures only in the case of row signature types.
-            if (isNamedTypeSignature) {
-                parameterTypeSignature = parameter.getNamedTypeSignature().getTypeSignature();
-            }
-            else {
-                parameterTypeSignature = parameter.getTypeSignature();
-            }
-
-            if (parameterTypeSignature.getParameters().isEmpty()) {
-                boolean changeTypeToVariable = isDecimalTypeBase(typeSignature.getBase());
-                if (changeTypeToVariable) {
-                    newParameterTypeList.add(
-                            TypeSignatureParameter.of(parameterTypeSignature.getBase()));
-                }
-                else {
-                    if (isNamedTypeSignature) {
-                        newParameterTypeList.add(TypeSignatureParameter.of(parameter.getNamedTypeSignature()));
-                    }
-                    else {
-                        newParameterTypeList.add(TypeSignatureParameter.of(parameterTypeSignature));
-                    }
-                }
-            }
-            else {
-                TypeSignature newTypeSignature =
-                        new TypeSignature(
-                                parameterTypeSignature.getBase(),
-                                getTypeSignatureParameters(
-                                        parameterTypeSignature.getStandardTypeSignature(),
-                                        parameterTypeSignature.getParameters()));
-                if (isNamedTypeSignature) {
-                    newParameterTypeList.add(
-                            TypeSignatureParameter.of(
-                                    new NamedTypeSignature(
-                                            Optional.empty(),
-                                            newTypeSignature)));
-                }
-                else {
-                    newParameterTypeList.add(TypeSignatureParameter.of(newTypeSignature));
-                }
-            }
-        }
-        return newParameterTypeList;
-    }
-
-    private static boolean isDecimalTypeBase(String typeBase)
-    {
-        return typeBase.equals(StandardTypes.DECIMAL);
-    }
-
-    // Hack ends here
-
     private synchronized void createFunction(SqlInvokedFunction function)
     {
         checkFunctionLanguageSupported(function);
diff --git a/presto-spark-base/pom.xml b/presto-spark-base/pom.xml
index 8f3dd0f91a66d..37dc1d4fb541b 100644
--- a/presto-spark-base/pom.xml
+++ b/presto-spark-base/pom.xml
@@ -117,6 +117,10 @@
                     com.facebook.airlift.drift
                     *
                 
+                
+                    com.facebook.presto
+                    presto-function-namespace-managers-common
+                
             
         
 
diff --git a/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueryFramework.java b/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueryFramework.java
index 4e6ac94dddc03..9115059ca9e6d 100644
--- a/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueryFramework.java
+++ b/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueryFramework.java
@@ -575,7 +575,7 @@ protected SubPlan subplan(String sql, Session session)
         }
     }
 
-    private QueryExplainer getQueryExplainer()
+    protected QueryExplainer getQueryExplainer()
     {
         Metadata metadata = queryRunner.getMetadata();
         FeaturesConfig featuresConfig = createFeaturesConfig();
@@ -640,6 +640,12 @@ protected ExpectedQueryRunner getExpectedQueryRunner()
         return expectedQueryRunner;
     }
 
+    protected SqlParser getSqlParser()
+    {
+        checkState(sqlParser != null, "sqlParser not set");
+        return sqlParser;
+    }
+
     public interface QueryRunnerSupplier
     {
         QueryRunner get()
diff --git a/presto-tests/src/main/java/com/facebook/presto/tests/DistributedQueryRunner.java b/presto-tests/src/main/java/com/facebook/presto/tests/DistributedQueryRunner.java
index d9ce3439011ab..67bf65f67eb55 100644
--- a/presto-tests/src/main/java/com/facebook/presto/tests/DistributedQueryRunner.java
+++ b/presto-tests/src/main/java/com/facebook/presto/tests/DistributedQueryRunner.java
@@ -1040,6 +1040,13 @@ public void loadPlanCheckerProviderManager(String planCheckerProviderName, Map
Date: Tue, 26 Aug 2025 15:10:48 -0700
Subject: [PATCH 020/113]  [native] Build for Gcc14 (#25861)

Summary:
Fix to support GCC14 build

- Replace `{}` with explicit empty container to avoid the following error within optionals.

          error: converting to 'std::in_place_t' from list would use explicit contructor
     `{}` leads to copy initialization which is not allowed since in_place_t is marked explicit

- Add Import `chrono` in `Duration.h` as gcc14 mandates having it

- Correct include directory path for proxygen

- Ignore errors associated with template-id-cdtor as gcc14 fails build for constructors having template support

Rollback Plan:


```
== NO RELEASE NOTE ==
```


Differential Revision: D80784416

Pulled By: pratikpugalia
---
 presto-native-execution/CMakeLists.txt                    | 8 +++++++-
 .../main/operators/tests/BinarySortableSerializerTest.cpp | 2 +-
 .../presto_cpp/main/operators/tests/BroadcastTest.cpp     | 2 +-
 .../presto_cpp/presto_protocol/core/Duration.h            | 1 +
 4 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/presto-native-execution/CMakeLists.txt b/presto-native-execution/CMakeLists.txt
index b85d27d8c691e..d09c9fad6c98b 100644
--- a/presto-native-execution/CMakeLists.txt
+++ b/presto-native-execution/CMakeLists.txt
@@ -33,6 +33,12 @@ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${SCRIPT_CXX_FLAGS}")
 set(DISABLED_WARNINGS
     "-Wno-nullability-completeness -Wno-deprecated-declarations")
 
+if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU")
+  if(CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL "14.0.0")
+    string(APPEND DISABLED_WARNINGS " -Wno-error=template-id-cdtor")
+  endif()
+endif()
+
 # Important warnings that must be explicitly enabled.
 set(ENABLE_WARNINGS "-Wreorder")
 
@@ -221,7 +227,7 @@ include_directories(SYSTEM ${FBTHRIFT_INCLUDE_DIR})
 set(PROXYGEN_LIBRARIES ${PROXYGEN_HTTP_SERVER} ${PROXYGEN} ${WANGLE} ${FIZZ}
                        ${MVFST_EXCEPTION})
 find_path(PROXYGEN_DIR NAMES include/proxygen)
-set(PROXYGEN_INCLUDE_DIR "${PROXYGEN_DIR}/include/proxygen")
+set(PROXYGEN_INCLUDE_DIR "${PROXYGEN_DIR}/include/")
 
 include_directories(SYSTEM ${OPENSSL_INCLUDE_DIR} ${PROXYGEN_INCLUDE_DIR})
 include_directories(.)
diff --git a/presto-native-execution/presto_cpp/main/operators/tests/BinarySortableSerializerTest.cpp b/presto-native-execution/presto_cpp/main/operators/tests/BinarySortableSerializerTest.cpp
index 4db5793e31708..77481b697befb 100644
--- a/presto-native-execution/presto_cpp/main/operators/tests/BinarySortableSerializerTest.cpp
+++ b/presto-native-execution/presto_cpp/main/operators/tests/BinarySortableSerializerTest.cpp
@@ -1045,7 +1045,7 @@ TEST_F(BinarySortableSerializerTest, ArrayTypeSingleFieldTests) {
   // null < []
   EXPECT_TRUE(
       singleArrayFieldCompare(
-          {std::nullopt, {{}}}, velox::core::kAscNullsFirst) < 0);
+          {std::nullopt, {std::vector>{}}}, velox::core::kAscNullsFirst) < 0);
 }
 
 TEST_F(BinarySortableSerializerTest, RowTypeSingleFieldTests) {
diff --git a/presto-native-execution/presto_cpp/main/operators/tests/BroadcastTest.cpp b/presto-native-execution/presto_cpp/main/operators/tests/BroadcastTest.cpp
index c685ede460fee..95e63ca080251 100644
--- a/presto-native-execution/presto_cpp/main/operators/tests/BroadcastTest.cpp
+++ b/presto-native-execution/presto_cpp/main/operators/tests/BroadcastTest.cpp
@@ -272,7 +272,7 @@ TEST_F(BroadcastTest, endToEndSerdeLayout) {
   runBroadcastTest({data}, {{"c1", "c1", "c2"}});
 
   // Skip all.
-  runBroadcastTest({data}, {{}});
+  runBroadcastTest({data}, {std::vector{}});
 }
 
 TEST_F(BroadcastTest, endToEndWithNoRows) {
diff --git a/presto-native-execution/presto_cpp/presto_protocol/core/Duration.h b/presto-native-execution/presto_cpp/presto_protocol/core/Duration.h
index e0aae14020320..0a5bf10f0589f 100644
--- a/presto-native-execution/presto_cpp/presto_protocol/core/Duration.h
+++ b/presto-native-execution/presto_cpp/presto_protocol/core/Duration.h
@@ -13,6 +13,7 @@
  */
 #pragma once
 #include 
+#include 
 
 namespace facebook::presto::protocol {
 

From daff2fe625434eecd1f8abbb8cd12263123290cf Mon Sep 17 00:00:00 2001
From: Rebecca Schlussel 
Date: Thu, 28 Aug 2025 10:59:29 -0400
Subject: [PATCH 021/113] Add Prestissimo committers as codeowners for native
 session properties

---
 CODEOWNERS | 1 +
 1 file changed, 1 insertion(+)

diff --git a/CODEOWNERS b/CODEOWNERS
index 8796e20a27174..ecdd3c163c347 100644
--- a/CODEOWNERS
+++ b/CODEOWNERS
@@ -114,6 +114,7 @@ CODEOWNERS @prestodb/team-tsc
 /presto-native-execution @prestodb/team-velox @prestodb/committers
 /presto-native-sidecar-plugin @pdabre12 @prestodb/team-velox @prestodb/committers
 /presto-native-tests @prestodb/team-velox @prestodb/committers
+/presto-main-base/src/main/java/com/facebook/presto/sessionpropertyproviders/NativeWorkerSessionPropertyProvider.java  @prestodb/team-velox @prestodb/committers
 /.github/workflows/prestocpp-* @prestodb/team-velox @prestodb/committers
 
 #####################################################################

From bcfdd43a3eb86c4a366cc47a14b6059ab3dd169a Mon Sep 17 00:00:00 2001
From: Rebecca Schlussel 
Date: Thu, 28 Aug 2025 11:01:21 -0400
Subject: [PATCH 022/113] Fix codeowners file for presto-main/presto-main-base

Presto-main was split into presto-main and presto-main-base. Update
paths in codeowners file to reflect the change.
---
 CODEOWNERS | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/CODEOWNERS b/CODEOWNERS
index ecdd3c163c347..faf25f724f484 100644
--- a/CODEOWNERS
+++ b/CODEOWNERS
@@ -50,6 +50,7 @@
 /presto-lark-sheets @prestodb/committers
 /presto-local-file @prestodb/committers
 /presto-main @prestodb/committers
+/presto-main-base @prestodb/committers
 /presto-matching @prestodb/committers
 /presto-memory @prestodb/committers
 /presto-memory-context @prestodb/committers
@@ -100,10 +101,10 @@ CODEOWNERS @prestodb/team-tsc
 # Presto core
 
 # Presto analyzer and optimizer
-/presto-main/src/*/java/com/facebook/presto/sql @jaystarshot @feilong-liu @ClarenceThreepwood @prestodb/committers
+/presto-main-base/src/*/java/com/facebook/presto/sql @jaystarshot @feilong-liu @ClarenceThreepwood @prestodb/committers
 
 # Presto cost based optimizer framework
-/presto-main/src/*/java/com/facebook/presto/cost @jaystarshot @feilong-liu @ClarenceThreepwood @prestodb/committers
+/presto-main-base/src/*/java/com/facebook/presto/cost @jaystarshot @feilong-liu @ClarenceThreepwood @prestodb/committers
 
 # Testing module
 # Note: all code owners in Presto core should be included here as well

From ac1f4f7b57e4761e328064bd9d21f79e15a7155d Mon Sep 17 00:00:00 2001
From: Artem Selishchev 
Date: Thu, 28 Aug 2025 15:27:16 -0700
Subject: [PATCH 023/113] feat: Add session property for
 debugMemoryPoolWarnThresholdBytes (#25750)

Summary:
I added threshold for logging memory pool allocations":
https://github.com/facebookincubator/velox/pull/14437
In this adding I'm adding corresponding session property to configure
the threshold.

Differential Revision: D80066283
---
 .../sphinx/presto_cpp/properties-session.rst  | 23 +++++++++++++++++++
 .../NativeWorkerSessionPropertyProvider.java  | 10 ++++++++
 .../presto_cpp/main/QueryContextManager.cpp   |  9 +++++---
 .../presto_cpp/main/SessionProperties.cpp     | 13 +++++++++++
 .../presto_cpp/main/SessionProperties.h       |  6 +++++
 .../main/tests/SessionPropertiesTest.cpp      |  2 ++
 6 files changed, 60 insertions(+), 3 deletions(-)

diff --git a/presto-docs/src/main/sphinx/presto_cpp/properties-session.rst b/presto-docs/src/main/sphinx/presto_cpp/properties-session.rst
index 629040783be03..687dffbcd2f0a 100644
--- a/presto-docs/src/main/sphinx/presto_cpp/properties-session.rst
+++ b/presto-docs/src/main/sphinx/presto_cpp/properties-session.rst
@@ -97,6 +97,29 @@ If set to ``true``, disables the optimization in expression evaluation to delay
 
 This should only be used for debugging purposes.
 
+``native_debug_memory_pool_name_regex``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+* **Type:** ``varchar``
+* **Default value:** ``""``
+
+Native Execution only. Regular expression pattern to match memory pool names for allocation callsite tracking.
+Matched pools will also perform leak checks at destruction. Empty string disables tracking.
+
+This should only be used for debugging purposes.
+
+``native_debug_memory_pool_warn_threshold_bytes``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+* **Type:** ``bigint``
+* **Default value:** ``0``
+
+Native Execution only. Warning threshold for memory pool allocations. Logs callsites when exceeded.
+Requires allocation tracking to be enabled with ``native_debug_memory_pool_name_regex``.
+Accepts B/KB/MB/GB units. Set to 0B to disable.
+
+This should only be used for debugging purposes.
+
 ``native_execution_type_rewrite_enabled``
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
diff --git a/presto-main-base/src/main/java/com/facebook/presto/sessionpropertyproviders/NativeWorkerSessionPropertyProvider.java b/presto-main-base/src/main/java/com/facebook/presto/sessionpropertyproviders/NativeWorkerSessionPropertyProvider.java
index a4d6c81eb62ba..0e54809e87b80 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/sessionpropertyproviders/NativeWorkerSessionPropertyProvider.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/sessionpropertyproviders/NativeWorkerSessionPropertyProvider.java
@@ -53,6 +53,7 @@ public class NativeWorkerSessionPropertyProvider
     public static final String NATIVE_DEBUG_DISABLE_EXPRESSION_WITH_MEMOIZATION = "native_debug_disable_expression_with_memoization";
     public static final String NATIVE_DEBUG_DISABLE_EXPRESSION_WITH_LAZY_INPUTS = "native_debug_disable_expression_with_lazy_inputs";
     public static final String NATIVE_DEBUG_MEMORY_POOL_NAME_REGEX = "native_debug_memory_pool_name_regex";
+    public static final String NATIVE_DEBUG_MEMORY_POOL_WARN_THRESHOLD_BYTES = "native_debug_memory_pool_warn_threshold_bytes";
     public static final String NATIVE_SELECTIVE_NIMBLE_READER_ENABLED = "native_selective_nimble_reader_enabled";
     public static final String NATIVE_MAX_PARTIAL_AGGREGATION_MEMORY = "native_max_partial_aggregation_memory";
     public static final String NATIVE_MAX_EXTENDED_PARTIAL_AGGREGATION_MEMORY = "native_max_extended_partial_aggregation_memory";
@@ -213,6 +214,15 @@ public NativeWorkerSessionPropertyProvider(FeaturesConfig featuresConfig)
                                 " string means no match for all.",
                         "",
                         true),
+                stringProperty(
+                        NATIVE_DEBUG_MEMORY_POOL_WARN_THRESHOLD_BYTES,
+                        "Warning threshold in bytes for debug memory pools. When set to a " +
+                                "non-zero value, a warning will be logged once per memory pool when " +
+                                "allocations cause the pool to exceed this threshold. This is useful for " +
+                                "identifying memory usage patterns during debugging. A value of " +
+                                "0 means no warning threshold is enforced.",
+                        "0B",
+                        true),
                 booleanProperty(
                         NATIVE_SELECTIVE_NIMBLE_READER_ENABLED,
                         "Temporary flag to control whether selective Nimble reader should be " +
diff --git a/presto-native-execution/presto_cpp/main/QueryContextManager.cpp b/presto-native-execution/presto_cpp/main/QueryContextManager.cpp
index 6387be0fd87fe..b96f03dce2909 100644
--- a/presto-native-execution/presto_cpp/main/QueryContextManager.cpp
+++ b/presto-native-execution/presto_cpp/main/QueryContextManager.cpp
@@ -63,7 +63,8 @@ std::shared_ptr QueryContextManager::createAndCacheQueryCtx(
     QueryContextCache& cache,
     const QueryId& queryId,
     velox::core::QueryConfig&& queryConfig,
-    std::unordered_map>&& connectorConfigs,
+    std::unordered_map>&&
+        connectorConfigs,
     std::shared_ptr&& pool) {
   auto queryCtx = core::QueryCtx::create(
       driverExecutor_,
@@ -96,10 +97,12 @@ std::shared_ptr QueryContextManager::findOrCreateQueryCtx(
   // is still indexed by the query id.
   static std::atomic_uint64_t poolId{0};
   std::optional poolDbgOpts;
-  const auto debugMemoryPoolNameRegex = queryConfig.debugMemoryPoolNameRegex();
+  auto debugMemoryPoolNameRegex = queryConfig.debugMemoryPoolNameRegex();
   if (!debugMemoryPoolNameRegex.empty()) {
     poolDbgOpts = memory::MemoryPool::DebugOptions{
-        .debugPoolNameRegex = debugMemoryPoolNameRegex};
+        .debugPoolNameRegex = std::move(debugMemoryPoolNameRegex),
+        .debugPoolWarnThresholdBytes =
+            queryConfig.debugMemoryPoolWarnThresholdBytes()};
   }
   auto pool = memory::MemoryManager::getInstance()->addRootPool(
       fmt::format("{}_{}", queryId, poolId++),
diff --git a/presto-native-execution/presto_cpp/main/SessionProperties.cpp b/presto-native-execution/presto_cpp/main/SessionProperties.cpp
index e3a03a152b121..d87ff9167cafc 100644
--- a/presto-native-execution/presto_cpp/main/SessionProperties.cpp
+++ b/presto-native-execution/presto_cpp/main/SessionProperties.cpp
@@ -280,6 +280,19 @@ SessionProperties::SessionProperties() {
       QueryConfig::kDebugMemoryPoolNameRegex,
       c.debugMemoryPoolNameRegex());
 
+  addSessionProperty(
+      kDebugMemoryPoolWarnThresholdBytes,
+      "Warning threshold in bytes for debug memory pools. When set to a "
+      "non-zero value, a warning will be logged once per memory pool when "
+      "allocations cause the pool to exceed this threshold. This is useful for "
+      "identifying memory usage patterns during debugging. Requires allocation "
+      "tracking to be enabled with `native_debug_memory_pool_name_regex` "
+      "for the pool. A value of 0 means no warning threshold is enforced.",
+      BIGINT(),
+      false,
+      QueryConfig::kDebugMemoryPoolWarnThresholdBytes,
+      std::to_string(c.debugMemoryPoolWarnThresholdBytes()));
+
   addSessionProperty(
       kSelectiveNimbleReaderEnabled,
       "Temporary flag to control whether selective Nimble reader should be "
diff --git a/presto-native-execution/presto_cpp/main/SessionProperties.h b/presto-native-execution/presto_cpp/main/SessionProperties.h
index 10758ee4c95d8..ab888852d5e00 100644
--- a/presto-native-execution/presto_cpp/main/SessionProperties.h
+++ b/presto-native-execution/presto_cpp/main/SessionProperties.h
@@ -184,6 +184,12 @@ class SessionProperties {
   static constexpr const char* kDebugMemoryPoolNameRegex =
       "native_debug_memory_pool_name_regex";
 
+  /// Warning threshold in bytes for memory pool allocations. Logs callsites 
+  /// when exceeded. Requires allocation tracking to be enabled with 
+  /// `native_debug_memory_pool_name_regex` property for the pool.
+  static constexpr const char* kDebugMemoryPoolWarnThresholdBytes =
+      "native_debug_memory_pool_warn_threshold_bytes";
+
   /// Temporary flag to control whether selective Nimble reader should be used
   /// in this query or not.  Will be removed after the selective Nimble reader
   /// is fully rolled out.
diff --git a/presto-native-execution/presto_cpp/main/tests/SessionPropertiesTest.cpp b/presto-native-execution/presto_cpp/main/tests/SessionPropertiesTest.cpp
index 20cc3644270f2..c88abee0e3719 100644
--- a/presto-native-execution/presto_cpp/main/tests/SessionPropertiesTest.cpp
+++ b/presto-native-execution/presto_cpp/main/tests/SessionPropertiesTest.cpp
@@ -69,6 +69,8 @@ TEST_F(SessionPropertiesTest, validateMapping) {
        core::QueryConfig::kDebugDisableExpressionWithLazyInputs},
       {SessionProperties::kDebugMemoryPoolNameRegex,
        core::QueryConfig::kDebugMemoryPoolNameRegex},
+      {SessionProperties::kDebugMemoryPoolWarnThresholdBytes,
+       core::QueryConfig::kDebugMemoryPoolWarnThresholdBytes},
       {SessionProperties::kSelectiveNimbleReaderEnabled,
        core::QueryConfig::kSelectiveNimbleReaderEnabled},
       {SessionProperties::kQueryTraceEnabled,

From acb54223da11722fdd9eeec08d18f78deab791dc Mon Sep 17 00:00:00 2001
From: Pramod Satya 
Date: Thu, 28 Aug 2025 09:28:24 -0700
Subject: [PATCH 024/113] [native] Advance velox

Co-authored-by: Christian Zentgraf 
---
 .../runtime-metrics/PrometheusStatsReporter.h | 21 +++++++
 .../main/types/PrestoToVeloxQueryPlan.cpp     | 61 +++++++++++--------
 .../main/types/PrestoToVeloxQueryPlan.h       |  2 +-
 presto-native-execution/velox                 |  2 +-
 4 files changed, 60 insertions(+), 26 deletions(-)

diff --git a/presto-native-execution/presto_cpp/main/runtime-metrics/PrometheusStatsReporter.h b/presto-native-execution/presto_cpp/main/runtime-metrics/PrometheusStatsReporter.h
index bf8574a624f4e..e5e8031c1e36e 100644
--- a/presto-native-execution/presto_cpp/main/runtime-metrics/PrometheusStatsReporter.h
+++ b/presto-native-execution/presto_cpp/main/runtime-metrics/PrometheusStatsReporter.h
@@ -65,6 +65,18 @@ class PrometheusStatsReporter : public facebook::velox::BaseStatsReporter {
       int64_t max,
       const std::vector& pcts) const override;
 
+  void registerQuantileMetricExportType(
+      const char* /* key */,
+      const std::vector& /* statTypes */,
+      const std::vector& /* pcts */,
+      const std::vector& /* slidingWindowsSeconds */) const override {};
+
+  void registerQuantileMetricExportType(
+      folly::StringPiece /* key */,
+      const std::vector& /* statTypes */,
+      const std::vector& /* pcts */,
+      const std::vector& /* slidingWindowsSeconds */) const override {};
+
   void addMetricValue(const std::string& key, size_t value = 1) const override;
 
   void addMetricValue(const char* key, size_t value = 1) const override;
@@ -79,6 +91,15 @@ class PrometheusStatsReporter : public facebook::velox::BaseStatsReporter {
   void addHistogramMetricValue(folly::StringPiece key, size_t value)
       const override;
 
+  void addQuantileMetricValue(const std::string& /* key */, size_t /* value */)
+      const override {};
+
+  void addQuantileMetricValue(const char* /* key */, size_t /* value */)
+      const override {};
+
+  void addQuantileMetricValue(folly::StringPiece /* key */, size_t /* value */)
+      const override{};
+
   std::string fetchMetrics() override;
 
   /**
diff --git a/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.cpp b/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.cpp
index e484dab310148..caecfd18f7deb 100644
--- a/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.cpp
+++ b/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.cpp
@@ -818,8 +818,8 @@ VectorPtr VeloxQueryPlanConverterBase::evaluateConstantExpression(
   return result[0];
 }
 
-std::shared_ptr
-VeloxQueryPlanConverterBase::generateAggregationNode(
+std::optional
+VeloxQueryPlanConverterBase::toColumnStatsSpec(
     const std::shared_ptr&
         statisticsAggregation,
     core::AggregationNode::Step step,
@@ -828,7 +828,7 @@ VeloxQueryPlanConverterBase::generateAggregationNode(
     const std::shared_ptr& /*tableWriteInfo*/,
     const protocol::TaskId& /*taskId*/) {
   if (statisticsAggregation == nullptr) {
-    return nullptr;
+    return std::nullopt;
   }
   const auto outputVariables = statisticsAggregation->outputVariables;
   const auto aggregationMap = statisticsAggregation->aggregations;
@@ -844,8 +844,7 @@ VeloxQueryPlanConverterBase::generateAggregationNode(
   std::vector aggregateNames;
   std::vector aggregates;
   toAggregations(outputVariables, aggregationMap, aggregates, aggregateNames);
-
-  return std::make_shared(
+  const auto aggregationNode = std::make_shared(
       id,
       step,
       toVeloxExprs(statisticsAggregation->groupingVariables),
@@ -854,6 +853,22 @@ VeloxQueryPlanConverterBase::generateAggregationNode(
       aggregates,
       /*ignoreNullKeys=*/false,
       sourceVeloxPlan);
+
+  // Sanity checks on aggregation node.
+  if (aggregationNode->ignoreNullKeys() ||
+      aggregationNode->groupId().has_value() ||
+      aggregationNode->isPreGrouped() ||
+      !aggregationNode->globalGroupingSets().empty() ||
+      aggregationNode->aggregateNames().empty() ||
+      (aggregationNode->aggregateNames().size() !=
+       aggregationNode->aggregates().size())) {
+    return std::nullopt;
+  }
+  return std::make_optional(core::ColumnStatsSpec{
+      aggregationNode->groupingKeys(),
+      aggregationNode->step(),
+      aggregationNode->aggregateNames(),
+      aggregationNode->aggregates()});
 }
 
 std::vector
@@ -1451,19 +1466,18 @@ VeloxQueryPlanConverterBase::toVeloxQueryPlan(
       typeParser_);
   const auto sourceVeloxPlan =
       toVeloxQueryPlan(node->source, tableWriteInfo, taskId);
-  std::shared_ptr aggregationNode =
-      generateAggregationNode(
-          node->statisticsAggregation,
-          core::AggregationNode::Step::kPartial,
-          node->id,
-          sourceVeloxPlan,
-          tableWriteInfo,
-          taskId);
+  std::optional columnStatsSpec = toColumnStatsSpec(
+      node->statisticsAggregation,
+      core::AggregationNode::Step::kPartial,
+      node->id,
+      sourceVeloxPlan,
+      tableWriteInfo,
+      taskId);
   return std::make_shared(
       node->id,
       toRowType(node->columns, typeParser_),
       node->columnNames,
-      std::move(aggregationNode),
+      columnStatsSpec,
       std::move(insertTableHandle),
       node->partitioningScheme != nullptr,
       outputType,
@@ -1521,7 +1535,7 @@ VeloxQueryPlanConverterBase::toVeloxQueryPlan(
       node->id,
       inputColumns,
       inputColumns->names(),
-      /*aggregationNode=*/nullptr,
+      /*columnStatsSpec=*/std::nullopt,
       std::move(insertTableHandle),
       true, // delete only supported on partitioned tables
       outputType,
@@ -1543,17 +1557,16 @@ VeloxQueryPlanConverterBase::toVeloxQueryPlan(
       typeParser_);
   const auto sourceVeloxPlan =
       toVeloxQueryPlan(node->source, tableWriteInfo, taskId);
-  std::shared_ptr aggregationNode =
-      generateAggregationNode(
-          node->statisticsAggregation,
-          core::AggregationNode::Step::kIntermediate,
-          node->id,
-          sourceVeloxPlan,
-          tableWriteInfo,
-          taskId);
+  std::optional columnStatsSpec = toColumnStatsSpec(
+      node->statisticsAggregation,
+      core::AggregationNode::Step::kIntermediate,
+      node->id,
+      sourceVeloxPlan,
+      tableWriteInfo,
+      taskId);
 
   return std::make_shared(
-      node->id, outputType, aggregationNode, sourceVeloxPlan);
+      node->id, outputType, columnStatsSpec, sourceVeloxPlan);
 }
 
 std::shared_ptr
diff --git a/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.h b/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.h
index 6a728a0b4c44c..3d13090b3801c 100644
--- a/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.h
+++ b/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.h
@@ -204,7 +204,7 @@ class VeloxQueryPlanConverterBase {
   velox::VectorPtr evaluateConstantExpression(
       const velox::core::TypedExprPtr& expression);
 
-  std::shared_ptr generateAggregationNode(
+  std::optional toColumnStatsSpec(
       const std::shared_ptr&
           statisticsAggregation,
       velox::core::AggregationNode::Step step,
diff --git a/presto-native-execution/velox b/presto-native-execution/velox
index 8d01456cf77a5..9aea41b8a8c5f 160000
--- a/presto-native-execution/velox
+++ b/presto-native-execution/velox
@@ -1 +1 @@
-Subproject commit 8d01456cf77a56d56c371ecc9509c5ae111157d8
+Subproject commit 9aea41b8a8c5fc28afbbedc25f4488165fae6d21

From 366daeeb4ac89a5b5639f072b6118f7533a645d2 Mon Sep 17 00:00:00 2001
From: adkharat 
Date: Wed, 13 Aug 2025 16:56:07 +0530
Subject: [PATCH 025/113] Enable case-sensitive identifier support for BigQuery
 connector

---
 .../presto/plugin/bigquery/BigQueryClient.java  |  7 +++++--
 .../presto/plugin/bigquery/BigQueryConfig.java  | 17 +++++++++++++++++
 .../plugin/bigquery/BigQueryMetadata.java       |  9 +++++++++
 .../plugin/bigquery/TestBigQueryConfig.java     |  6 +++++-
 .../src/main/sphinx/connector/bigquery.rst      |  3 +++
 5 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/presto-bigquery/src/main/java/com/facebook/presto/plugin/bigquery/BigQueryClient.java b/presto-bigquery/src/main/java/com/facebook/presto/plugin/bigquery/BigQueryClient.java
index 59215dc81fe79..221ece8ae3b74 100644
--- a/presto-bigquery/src/main/java/com/facebook/presto/plugin/bigquery/BigQueryClient.java
+++ b/presto-bigquery/src/main/java/com/facebook/presto/plugin/bigquery/BigQueryClient.java
@@ -50,6 +50,7 @@ public class BigQueryClient
     private final Optional viewMaterializationProject;
     private final Optional viewMaterializationDataset;
     private final String tablePrefix = "_pbc_";
+    private final boolean caseSensitiveNameMatching;
 
     // presto converts the dataset and table names to lower case, while BigQuery is case sensitive
     private final ConcurrentMap tableIds = new ConcurrentHashMap<>();
@@ -60,6 +61,7 @@ public class BigQueryClient
         this.bigQuery = requireNonNull(bigQuery, "bigQuery is null");
         this.viewMaterializationProject = requireNonNull(config.getViewMaterializationProject(), "viewMaterializationProject is null");
         this.viewMaterializationDataset = requireNonNull(config.getViewMaterializationDataset(), "viewMaterializationDataset is null");
+        this.caseSensitiveNameMatching = config.isCaseSensitiveNameMatching();
     }
 
     public TableInfo getTable(TableId tableId)
@@ -108,7 +110,7 @@ private void addTableMappingIfNeeded(DatasetId datasetID, Table table)
     private Dataset addDataSetMappingIfNeeded(Dataset dataset)
     {
         DatasetId bigQueryDatasetId = dataset.getDatasetId();
-        DatasetId prestoDatasetId = DatasetId.of(bigQueryDatasetId.getProject(), bigQueryDatasetId.getDataset().toLowerCase(ENGLISH));
+        DatasetId prestoDatasetId = DatasetId.of(bigQueryDatasetId.getProject(), bigQueryDatasetId.getDataset());
         datasetIds.putIfAbsent(prestoDatasetId, bigQueryDatasetId);
         return dataset;
     }
@@ -123,7 +125,8 @@ protected TableId createDestinationTable(TableId tableId)
 
     private String createTableName()
     {
-        return format(tablePrefix + "%s", randomUUID().toString().toLowerCase(ENGLISH).replace("-", ""));
+        String uuid = randomUUID().toString().replace("-", "");
+        return caseSensitiveNameMatching ? format("%s%s", tablePrefix, uuid) : format("%s%s", tablePrefix, uuid).toLowerCase(ENGLISH);
     }
 
     private DatasetId mapIfNeeded(String project, String dataset)
diff --git a/presto-bigquery/src/main/java/com/facebook/presto/plugin/bigquery/BigQueryConfig.java b/presto-bigquery/src/main/java/com/facebook/presto/plugin/bigquery/BigQueryConfig.java
index 86e09dba9077e..1247f385e1514 100644
--- a/presto-bigquery/src/main/java/com/facebook/presto/plugin/bigquery/BigQueryConfig.java
+++ b/presto-bigquery/src/main/java/com/facebook/presto/plugin/bigquery/BigQueryConfig.java
@@ -35,6 +35,7 @@ public class BigQueryConfig
     private Optional parentProjectId = Optional.empty();
     private OptionalInt parallelism = OptionalInt.empty();
     private boolean viewsEnabled;
+    private boolean caseSensitiveNameMatching;
     private Optional viewMaterializationProject = Optional.empty();
     private Optional viewMaterializationDataset = Optional.empty();
     private int maxReadRowsRetries = DEFAULT_MAX_READ_ROWS_RETRIES;
@@ -181,6 +182,22 @@ public BigQueryConfig setMaxReadRowsRetries(int maxReadRowsRetries)
         return this;
     }
 
+    public boolean isCaseSensitiveNameMatching()
+    {
+        return caseSensitiveNameMatching;
+    }
+
+    @Config("case-sensitive-name-matching")
+    @ConfigDescription(
+            "Case sensitivity for schema and table name matching. " +
+             "true = preserve case and require exact matches; " +
+             "false (default) = normalize to lower case and match case-insensitively.")
+    public BigQueryConfig setCaseSensitiveNameMatching(boolean caseSensitiveNameMatching)
+    {
+        this.caseSensitiveNameMatching = caseSensitiveNameMatching;
+        return this;
+    }
+
     ReadSessionCreatorConfig createReadSessionCreatorConfig()
     {
         return new ReadSessionCreatorConfig(
diff --git a/presto-bigquery/src/main/java/com/facebook/presto/plugin/bigquery/BigQueryMetadata.java b/presto-bigquery/src/main/java/com/facebook/presto/plugin/bigquery/BigQueryMetadata.java
index 8c028fbf59b7f..ad0bf4775a42b 100644
--- a/presto-bigquery/src/main/java/com/facebook/presto/plugin/bigquery/BigQueryMetadata.java
+++ b/presto-bigquery/src/main/java/com/facebook/presto/plugin/bigquery/BigQueryMetadata.java
@@ -51,6 +51,7 @@
 import static com.google.cloud.bigquery.TableDefinition.Type.TABLE;
 import static com.google.cloud.bigquery.TableDefinition.Type.VIEW;
 import static com.google.common.collect.ImmutableList.toImmutableList;
+import static java.util.Locale.ENGLISH;
 import static java.util.Objects.requireNonNull;
 import static java.util.stream.Collectors.toMap;
 
@@ -63,12 +64,14 @@ public class BigQueryMetadata
     private static final Logger log = Logger.get(BigQueryMetadata.class);
     private final BigQueryClient bigQueryClient;
     private final String projectId;
+    private final boolean caseSensitiveNameMatching;
 
     @Inject
     public BigQueryMetadata(BigQueryClient bigQueryClient, BigQueryConfig config)
     {
         this.bigQueryClient = bigQueryClient;
         this.projectId = config.getProjectId().orElse(bigQueryClient.getProjectId());
+        this.caseSensitiveNameMatching = config.isCaseSensitiveNameMatching();
     }
 
     @Override
@@ -233,4 +236,10 @@ private List listTables(ConnectorSession session, SchemaTablePr
                 ImmutableList.of(tableName) :
                 ImmutableList.of(); // table does not exist
     }
+
+    @Override
+    public String normalizeIdentifier(ConnectorSession session, String identifier)
+    {
+        return caseSensitiveNameMatching ? identifier : identifier.toLowerCase(ENGLISH);
+    }
 }
diff --git a/presto-bigquery/src/test/java/com/facebook/presto/plugin/bigquery/TestBigQueryConfig.java b/presto-bigquery/src/test/java/com/facebook/presto/plugin/bigquery/TestBigQueryConfig.java
index 053f83600eb85..8829d86cd6e23 100644
--- a/presto-bigquery/src/test/java/com/facebook/presto/plugin/bigquery/TestBigQueryConfig.java
+++ b/presto-bigquery/src/test/java/com/facebook/presto/plugin/bigquery/TestBigQueryConfig.java
@@ -36,7 +36,8 @@ public void testDefaults()
                 .setParallelism(20)
                 .setViewMaterializationProject("vmproject")
                 .setViewMaterializationDataset("vmdataset")
-                .setMaxReadRowsRetries(10);
+                .setMaxReadRowsRetries(10)
+                .setCaseSensitiveNameMatching(false);
 
         assertEquals(config.getCredentialsKey(), Optional.of("ckey"));
         assertEquals(config.getCredentialsFile(), Optional.of("cfile"));
@@ -46,6 +47,7 @@ public void testDefaults()
         assertEquals(config.getViewMaterializationProject(), Optional.of("vmproject"));
         assertEquals(config.getViewMaterializationDataset(), Optional.of("vmdataset"));
         assertEquals(config.getMaxReadRowsRetries(), 10);
+        assertEquals(config.isCaseSensitiveNameMatching(), false);
     }
 
     @Test
@@ -59,6 +61,7 @@ public void testExplicitPropertyMappingsWithCredentialsKey()
                 .put("bigquery.view-materialization-project", "vmproject")
                 .put("bigquery.view-materialization-dataset", "vmdataset")
                 .put("bigquery.max-read-rows-retries", "10")
+                .put("case-sensitive-name-matching", "true")
                 .build();
 
         ConfigurationFactory configurationFactory = new ConfigurationFactory(properties);
@@ -71,6 +74,7 @@ public void testExplicitPropertyMappingsWithCredentialsKey()
         assertEquals(config.getViewMaterializationProject(), Optional.of("vmproject"));
         assertEquals(config.getViewMaterializationDataset(), Optional.of("vmdataset"));
         assertEquals(config.getMaxReadRowsRetries(), 10);
+        assertEquals(config.isCaseSensitiveNameMatching(), true);
     }
 
     @Test
diff --git a/presto-docs/src/main/sphinx/connector/bigquery.rst b/presto-docs/src/main/sphinx/connector/bigquery.rst
index d71c600b758c8..4e90e92555da0 100644
--- a/presto-docs/src/main/sphinx/connector/bigquery.rst
+++ b/presto-docs/src/main/sphinx/connector/bigquery.rst
@@ -137,6 +137,9 @@ Property                                  Description
 ``bigquery.max-read-rows-retries``        The number of retries in case of retryable server issues       ``3``
 ``bigquery.credentials-key``              credentials key (base64 encoded)                               None. See `authentication <#authentication>`_
 ``bigquery.credentials-file``             JSON credentials file path                                     None. See `authentication <#authentication>`_
+``case-sensitive-name-matching``          Enable case sensitive identifier support for schema and table  ``false``
+                                          names for the connector. When disabled, names are matched
+                                          case-insensitively using lowercase normalization.          
 ========================================= ============================================================== ==============================================
 
 Data Types

From b323a87deb81b37f5f8c394de6f01ad8e8d11266 Mon Sep 17 00:00:00 2001
From: Li Zhou 
Date: Fri, 29 Aug 2025 17:04:04 +0100
Subject: [PATCH 026/113] Exclude executable jars from maven publishing
 (#25909)

## Description

This PR is [the fix from branch release-0.294](
https://github.com/prestodb/presto/pull/25900), to fix maven release
issues

## Motivation and Context
Merge the fix from release branch into master branch

## Impact
Newer releases

## Test Plan
Tested with release 0.294

## Contributor checklist

- [ ] Please make sure your submission complies with our [contributing
guide](https://github.com/prestodb/presto/blob/master/CONTRIBUTING.md),
in particular [code
style](https://github.com/prestodb/presto/blob/master/CONTRIBUTING.md#code-style)
and [commit
standards](https://github.com/prestodb/presto/blob/master/CONTRIBUTING.md#commit-standards).
- [ ] PR description addresses the issue accurately and concisely. If
the change is non-trivial, a GitHub Issue is referenced.
- [ ] Documented new properties (with its default value), SQL syntax,
functions, or other functionality.
- [ ] If release notes are required, they follow the [release notes
guidelines](https://github.com/prestodb/presto/wiki/Release-Notes-Guidelines).
- [ ] Adequate tests were added if applicable.
- [ ] CI passed.

## Release Notes

```
== NO RELEASE NOTE ==
```
---
 pom.xml                                       | 16 +--
 presto-benchmark-driver/pom.xml               | 94 ++++++++++--------
 .../presto/benchto/benchmarks/Dummy.java      | 21 ++++
 presto-cli/pom.xml                            | 98 +++++++++++--------
 presto-testing-server-launcher/pom.xml        | 95 ++++++++++--------
 5 files changed, 183 insertions(+), 141 deletions(-)
 create mode 100644 presto-benchto-benchmarks/src/main/java/com/facebook/presto/benchto/benchmarks/Dummy.java

diff --git a/pom.xml b/pom.xml
index d3f5e1fb01c06..f979056f76891 100644
--- a/pom.xml
+++ b/pom.xml
@@ -61,7 +61,6 @@
         8.11.3
         3.8.0
         1.13.1
-        1.6.8
         9.7.1
         1.9.17
         313
@@ -108,6 +107,7 @@
         -missing
         1.17.1
         2.32.9
+        true
     
 
     
@@ -2798,18 +2798,6 @@
                     
                 
 
-                
-                    org.sonatype.plugins
-                    nexus-staging-maven-plugin
-                    ${dep.nexus-staging-plugin.version}
-                    
-                    
-                        ossrh
-                        https://oss.sonatype.org/
-                    
-                
-
                 
                     com.facebook.presto
                     presto-maven-plugin
@@ -3028,7 +3016,7 @@
                         true
                         
                             ossrh
-                            true
+                            ${release.autoPublish}
                             validated
                         
                     
diff --git a/presto-benchmark-driver/pom.xml b/presto-benchmark-driver/pom.xml
index 1ea86335581be..76725407eb1d5 100644
--- a/presto-benchmark-driver/pom.xml
+++ b/presto-benchmark-driver/pom.xml
@@ -93,48 +93,58 @@
         
     
 
-    
-        
-            
-                org.apache.maven.plugins
-                maven-shade-plugin
-                
-                    
-                        package
-                        
-                            shade
-                        
+    
+        
+            executable-jar
+            
+                
+                    !skipExecutableJar
+                
+            
+            
+                
+                    
+                        org.apache.maven.plugins
+                        maven-shade-plugin
+                        
+                            
+                                package
+                                
+                                    shade
+                                
+                                
+                                    true
+                                    executable
+                                    
+                                        
+                                            
+                                                ${main-class}
+                                            
+                                        
+                                    
+                                
+                            
+                        
+                    
+
+                    
+                        org.skife.maven
+                        really-executable-jar-maven-plugin
                         
-                            true
-                            executable
-                            
-                                
-                                    
-                                        ${main-class}
-                                    
-                                
-                            
+                            -Xmx1G
+                            executable
                         
-                    
-                
-            
-
-            
-                org.skife.maven
-                really-executable-jar-maven-plugin
-                
-                    -Xmx1G
-                    executable
-                
-                
-                    
-                        package
-                        
-                            really-executable-jar
-                        
-                    
-                
-            
-        
-    
+                        
+                            
+                                package
+                                
+                                    really-executable-jar
+                                
+                            
+                        
+                    
+                
+            
+        
+    
 
diff --git a/presto-benchto-benchmarks/src/main/java/com/facebook/presto/benchto/benchmarks/Dummy.java b/presto-benchto-benchmarks/src/main/java/com/facebook/presto/benchto/benchmarks/Dummy.java
new file mode 100644
index 0000000000000..e5bc3b155b6b8
--- /dev/null
+++ b/presto-benchto-benchmarks/src/main/java/com/facebook/presto/benchto/benchmarks/Dummy.java
@@ -0,0 +1,21 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.benchto.benchmarks;
+
+/**
+ * This class exists to force the creation of a jar for the presto-benchto-benchmarks module. This is needed to deploy the presto-benchto-benchmarks module to nexus.
+ */
+public class Dummy
+{
+}
diff --git a/presto-cli/pom.xml b/presto-cli/pom.xml
index 8d2bcb99858b6..c2298b07486c2 100644
--- a/presto-cli/pom.xml
+++ b/presto-cli/pom.xml
@@ -130,48 +130,6 @@
 
     
         
-            
-                org.apache.maven.plugins
-                maven-shade-plugin
-                
-                    
-                        package
-                        
-                            shade
-                        
-                        
-                            true
-                            executable
-                            false
-                            
-                                
-                                    
-                                        ${main-class}
-                                    
-                                
-                            
-                        
-                    
-                
-            
-
-            
-                org.skife.maven
-                really-executable-jar-maven-plugin
-                
-                    -Xmx1G
-                    executable
-                
-                
-                    
-                        package
-                        
-                            really-executable-jar
-                        
-                    
-                
-            
-
             
                 org.basepom.maven
                 duplicate-finder-maven-plugin
@@ -195,4 +153,60 @@
             
         
     
+
+    
+        
+            executable-jar
+            
+                
+                    !skipExecutableJar
+                
+            
+            
+                
+                    
+                        org.apache.maven.plugins
+                        maven-shade-plugin
+                        
+                            
+                                package
+                                
+                                    shade
+                                
+                                
+                                    true
+                                    executable
+                                    false
+                                    
+                                        
+                                            
+                                                ${main-class}
+                                            
+                                        
+                                    
+                                
+                            
+                        
+                    
+
+                    
+                        org.skife.maven
+                        really-executable-jar-maven-plugin
+                        
+                            -Xmx1G
+                            executable
+                        
+                        
+                            
+                                package
+                                
+                                    really-executable-jar
+                                
+                            
+                        
+                    
+                
+            
+        
+    
 
diff --git a/presto-testing-server-launcher/pom.xml b/presto-testing-server-launcher/pom.xml
index 67acac907c506..dc51f0626e6a0 100644
--- a/presto-testing-server-launcher/pom.xml
+++ b/presto-testing-server-launcher/pom.xml
@@ -45,49 +45,58 @@
         
     
 
-    
-        
-            
-                org.apache.maven.plugins
-                maven-shade-plugin
-                
-                    
-                        package
-                        
-                            shade
-                        
+    
+        
+            executable-jar
+            
+                
+                    !skipExecutableJar
+                
+            
+            
+                
+                    
+                        org.apache.maven.plugins
+                        maven-shade-plugin
+                        
+                            
+                                package
+                                
+                                    shade
+                                
+                                
+                                    true
+                                    executable
+                                    
+                                        
+                                            
+                                                ${main-class}
+                                            
+                                        
+                                    
+                                
+                            
+                        
+                    
+
+                    
+                        org.skife.maven
+                        really-executable-jar-maven-plugin
                         
-                            true
-                            executable
-                            
-                                
-                                    
-                                        ${main-class}
-                                    
-                                
-                            
+                            -Xmx1G
+                            executable
                         
-                    
-                
-            
-
-            
-                org.skife.maven
-                really-executable-jar-maven-plugin
-                
-                    -Xmx1G
-                    executable
-                
-                
-                    
-                        package
-                        
-                            really-executable-jar
-                        
-                    
-                
-            
-        
-    
-
+                        
+                            
+                                package
+                                
+                                    really-executable-jar
+                                
+                            
+                        
+                    
+                
+            
+        
+    
 

From 638736888560d5bbaecc34fa774e517238c39881 Mon Sep 17 00:00:00 2001
From: Rebecca Whitworth 
Date: Tue, 12 Aug 2025 06:46:13 +0100
Subject: [PATCH 027/113] [docs] Add doc for Type Mapping to
 connector/deltalake.rst #25357

Fix #25357
Added type mapping table for Delta Lake to PrestoDB

Co-Authored-By: Steve Burnett 
Co-Authored-By: Jalpreet Singh Nanda 
---
 .../src/main/sphinx/connector/deltalake.rst   | 44 +++++++++++++++++++
 1 file changed, 44 insertions(+)

diff --git a/presto-docs/src/main/sphinx/connector/deltalake.rst b/presto-docs/src/main/sphinx/connector/deltalake.rst
index 96c7899e56972..3b44a60f1deee 100644
--- a/presto-docs/src/main/sphinx/connector/deltalake.rst
+++ b/presto-docs/src/main/sphinx/connector/deltalake.rst
@@ -133,3 +133,47 @@ in the table ``sales.apac.sales_data``.
 
 Above query drops the external table ``sales.apac.sales_data_new``. This only drops the
 metadata for the table. The referenced data directory is not deleted.
+
+Delta Lake to PrestoDB type mapping
+-----------------------------------
+
+Map of Delta Lake types to the relevant PrestoDB types:
+
+.. list-table:: Delta Lake to PrestoDB type mapping
+  :widths: 50, 50
+  :header-rows: 1
+
+  * - Delta Lake type
+    - PrestoDB type
+  * - ``BOOLEAN``
+    - ``BOOLEAN``
+  * - ``SMALLINT``
+    - ``SMALLINT`` 
+  * - ``TINYINT``
+    - ``TINYINT``
+  * - ``INT``
+    - ``INTEGER``
+  * - ``LONG``
+    - ``BIGINT``
+  * - ``FLOAT``
+    - ``REAL``
+  * - ``DOUBLE``
+    - ``DOUBLE``
+  * - ``DECIMAL``
+    - ``DECIMAL``
+  * - ``STRING``
+    - ``VARCHAR``
+  * - ``BINARY``
+    - ``VARBINARY``
+  * - ``DATE``
+    - ``DATE``
+  * - ``TIMESTAMP_NTZ``
+    - ``TIMESTAMP``
+  * - ``TIMESTAMP``
+    - ``TIMESTAMP WITH TIME ZONE``
+  * - ``ARRAY``
+    - ``ARRAY``
+  * - ``MAP``
+    - ``MAP``
+  * - ``STRUCT``
+    - ``ROW``

From 915f96a492339f11c1f07d2d4352f8e6335dd988 Mon Sep 17 00:00:00 2001
From: Jialiang Tan 
Date: Thu, 28 Aug 2025 20:56:01 -0700
Subject: [PATCH 028/113] Add output stats for pos sink operators (#25915)

Summary:

Adds output row stats for sapphire-velox related sink operators
Properly close write file on broadcast write

Reviewed By: singcha

Differential Revision: D81271224
---
 .../presto_cpp/main/operators/BroadcastFactory.cpp         | 7 ++++++-
 .../presto_cpp/main/operators/BroadcastWrite.cpp           | 3 +++
 .../presto_cpp/main/operators/ShuffleWrite.cpp             | 2 ++
 3 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/presto-native-execution/presto_cpp/main/operators/BroadcastFactory.cpp b/presto-native-execution/presto_cpp/main/operators/BroadcastFactory.cpp
index 7dcfee2cbfad3..b82ea7192e6b1 100644
--- a/presto-native-execution/presto_cpp/main/operators/BroadcastFactory.cpp
+++ b/presto-native-execution/presto_cpp/main/operators/BroadcastFactory.cpp
@@ -82,7 +82,12 @@ void BroadcastFileWriter::collect(const RowVectorPtr& input) {
   write(input);
 }
 
-void BroadcastFileWriter::noMoreData() {}
+void BroadcastFileWriter::noMoreData() {
+  if (writeFile_ != nullptr) {
+    writeFile_->flush();
+    writeFile_->close();
+  }
+}
 
 RowVectorPtr BroadcastFileWriter::fileStats() {
   // No rows written.
diff --git a/presto-native-execution/presto_cpp/main/operators/BroadcastWrite.cpp b/presto-native-execution/presto_cpp/main/operators/BroadcastWrite.cpp
index cffcbd4a64320..fc835bbe5c73b 100644
--- a/presto-native-execution/presto_cpp/main/operators/BroadcastWrite.cpp
+++ b/presto-native-execution/presto_cpp/main/operators/BroadcastWrite.cpp
@@ -70,6 +70,9 @@ class BroadcastWriteOperator : public Operator {
     }
 
     fileBroadcastWriter_->collect(reorderedInput);
+    auto lockedStats = stats_.wlock();
+    lockedStats->addOutputVector(
+        reorderedInput->estimateFlatSize(), reorderedInput->size());
   }
 
   void noMoreInput() override {
diff --git a/presto-native-execution/presto_cpp/main/operators/ShuffleWrite.cpp b/presto-native-execution/presto_cpp/main/operators/ShuffleWrite.cpp
index 434353a1e5d2e..dbb99c0e62711 100644
--- a/presto-native-execution/presto_cpp/main/operators/ShuffleWrite.cpp
+++ b/presto-native-execution/presto_cpp/main/operators/ShuffleWrite.cpp
@@ -104,6 +104,8 @@ class ShuffleWriteOperator : public Operator {
             "collect");
       }
     }
+    auto lockedStats = stats_.wlock();
+    lockedStats->addOutputVector(input->estimateFlatSize(), input->size());
   }
 
   void noMoreInput() override {

From 7357711dc6433de2f627656171a6c3dce262abbf Mon Sep 17 00:00:00 2001
From: Xiaoxuan Meng 
Date: Fri, 29 Aug 2025 21:47:02 -0700
Subject: [PATCH 029/113] [native] fix: Fail the query if the aggregation stats
 node fails the sanity check

---
 .../main/types/PrestoToVeloxQueryPlan.cpp     | 21 +++++++++----------
 1 file changed, 10 insertions(+), 11 deletions(-)

diff --git a/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.cpp b/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.cpp
index caecfd18f7deb..94428644cea46 100644
--- a/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.cpp
+++ b/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.cpp
@@ -855,20 +855,19 @@ VeloxQueryPlanConverterBase::toColumnStatsSpec(
       sourceVeloxPlan);
 
   // Sanity checks on aggregation node.
-  if (aggregationNode->ignoreNullKeys() ||
-      aggregationNode->groupId().has_value() ||
-      aggregationNode->isPreGrouped() ||
-      !aggregationNode->globalGroupingSets().empty() ||
-      aggregationNode->aggregateNames().empty() ||
-      (aggregationNode->aggregateNames().size() !=
-       aggregationNode->aggregates().size())) {
-    return std::nullopt;
-  }
-  return std::make_optional(core::ColumnStatsSpec{
+  VELOX_CHECK(!aggregationNode->ignoreNullKeys());
+  VELOX_CHECK(!aggregationNode->groupId().has_value());
+  VELOX_CHECK(!aggregationNode->isPreGrouped());
+  VELOX_CHECK(aggregationNode->globalGroupingSets().empty());
+  VELOX_CHECK(!aggregationNode->aggregateNames().empty());
+  VELOX_CHECK_EQ(
+      aggregationNode->aggregateNames().size(),
+      aggregationNode->aggregates().size());
+  return core::ColumnStatsSpec{
       aggregationNode->groupingKeys(),
       aggregationNode->step(),
       aggregationNode->aggregateNames(),
-      aggregationNode->aggregates()});
+      aggregationNode->aggregates()};
 }
 
 std::vector

From ecf13445359ffba74b0650a8ddfe0fb5d4e40e8c Mon Sep 17 00:00:00 2001
From: Pramod Satya 
Date: Sat, 30 Aug 2025 15:16:19 -0700
Subject: [PATCH 030/113] [native] Advance velox

---
 .../runtime-metrics/PrometheusStatsReporter.h | 27 +++++++++++++++++++
 presto-native-execution/velox                 |  2 +-
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/presto-native-execution/presto_cpp/main/runtime-metrics/PrometheusStatsReporter.h b/presto-native-execution/presto_cpp/main/runtime-metrics/PrometheusStatsReporter.h
index e5e8031c1e36e..c5764d12581f6 100644
--- a/presto-native-execution/presto_cpp/main/runtime-metrics/PrometheusStatsReporter.h
+++ b/presto-native-execution/presto_cpp/main/runtime-metrics/PrometheusStatsReporter.h
@@ -76,6 +76,18 @@ class PrometheusStatsReporter : public facebook::velox::BaseStatsReporter {
       const std::vector& /* statTypes */,
       const std::vector& /* pcts */,
       const std::vector& /* slidingWindowsSeconds */) const override {};
+  
+  void registerDynamicQuantileMetricExportType(
+      const char* /* keyPattern */,
+      const std::vector& /* statTypes */,
+      const std::vector& /* pcts */,
+      const std::vector& /* slidingWindowsSeconds */) const override {};
+  
+  void registerDynamicQuantileMetricExportType(
+      folly::StringPiece /* keyPattern */,
+      const std::vector& /* statTypes */,
+      const std::vector& /* pcts */,
+      const std::vector& /* slidingWindowsSeconds */) const override {};
 
   void addMetricValue(const std::string& key, size_t value = 1) const override;
 
@@ -99,6 +111,21 @@ class PrometheusStatsReporter : public facebook::velox::BaseStatsReporter {
 
   void addQuantileMetricValue(folly::StringPiece /* key */, size_t /* value */)
       const override{};
+  
+  void addDynamicQuantileMetricValue(
+      const std::string& /* key */,
+      folly::Range /* subkeys */,
+      size_t /* value */) const override {};
+  
+  virtual void addDynamicQuantileMetricValue(
+      const char* /* key */,
+      folly::Range /* subkeys */,
+      size_t /* value */) const override {};
+
+  virtual void addDynamicQuantileMetricValue(
+      folly::StringPiece /* key */,
+      folly::Range /* subkeys */,
+      size_t /* value */) const override {};
 
   std::string fetchMetrics() override;
 
diff --git a/presto-native-execution/velox b/presto-native-execution/velox
index 9aea41b8a8c5f..31200dd7fccce 160000
--- a/presto-native-execution/velox
+++ b/presto-native-execution/velox
@@ -1 +1 @@
-Subproject commit 9aea41b8a8c5fc28afbbedc25f4488165fae6d21
+Subproject commit 31200dd7fccce657da5ca0a2409bbc6978e35a68

From 3b32df2e726d1f7cc1470d8394e9db1cf93e8c89 Mon Sep 17 00:00:00 2001
From: Natasha Sehgal 
Date: Thu, 28 Aug 2025 16:35:01 -0700
Subject: [PATCH 031/113] fix: Deserialized TDigest NaN checks (#25907)

Summary: Pull Request resolved: https://github.com/prestodb/presto/pull/25907

Reviewed By: Yuhta

Differential Revision:
D81173686

Privacy Context Container: L1202870
---
 .../java/com/facebook/presto/tdigest/TDigest.java   | 13 +++++++++++++
 .../com/facebook/presto/tdigest/TestTDigest.java    | 10 ++++++++++
 2 files changed, 23 insertions(+)

diff --git a/presto-main-base/src/main/java/com/facebook/presto/tdigest/TDigest.java b/presto-main-base/src/main/java/com/facebook/presto/tdigest/TDigest.java
index 80a750a887fee..0a7977cd6fa67 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/tdigest/TDigest.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/tdigest/TDigest.java
@@ -203,6 +203,13 @@ public static TDigest createTDigest(Slice slice)
             r.mean = new double[r.activeCentroids];
             sliceInput.readBytes(wrappedDoubleArray(r.weight), r.activeCentroids * SIZE_OF_DOUBLE);
             sliceInput.readBytes(wrappedDoubleArray(r.mean), r.activeCentroids * SIZE_OF_DOUBLE);
+
+            // Validate deserialized TDigest data
+            for (int i = 0; i < r.activeCentroids; i++) {
+                checkArgument(!isNaN(r.mean[i]), "Deserialized t-digest contains NaN mean value");
+                checkArgument(r.weight[i] > 0, "weight must be > 0");
+            }
+
             sliceInput.close();
             return r;
         }
@@ -712,6 +719,12 @@ public long estimatedInMemorySizeInBytes()
 
     public Slice serialize()
     {
+        // Validate data before serialization
+        for (int i = 0; i < activeCentroids; i++) {
+            checkArgument(!isNaN(mean[i]), "Cannot serialize t-digest with NaN mean value");
+            checkArgument(weight[i] > 0, "Cannot serialize t-digest with non-positive weight");
+        }
+
         SliceOutput sliceOutput = new DynamicSliceOutput(toIntExact(estimatedSerializedSizeInBytes()));
 
         sliceOutput.writeByte(1); // version 1 of T-Digest serialization
diff --git a/presto-main-base/src/test/java/com/facebook/presto/tdigest/TestTDigest.java b/presto-main-base/src/test/java/com/facebook/presto/tdigest/TestTDigest.java
index 773be051a5df5..04320aaa52fa8 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/tdigest/TestTDigest.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/tdigest/TestTDigest.java
@@ -397,6 +397,16 @@ public void testGeometricDistribution()
         }
     }
 
+    @Test(expectedExceptions = IllegalArgumentException.class, expectedExceptionsMessageRegExp = ".*Cannot serialize t-digest with NaN mean value.*")
+    public void testSerializationWithNaNMean()
+    {
+        double[] means = {1.0, Double.NaN, 3.0};
+        double[] weights = {1.0, 1.0, 1.0};
+        TDigest tDigest = createTDigest(means, weights, STANDARD_COMPRESSION_FACTOR, 1.0, 3.0, 6.0, 3);
+
+        tDigest.serialize();
+    }
+
     @Test(enabled = false)
     public void testPoissonDistribution()
     {

From b4bb7ce53ee48e8e1bbaeb831fc38121884164ba Mon Sep 17 00:00:00 2001
From: Amit Dutta 
Date: Mon, 1 Sep 2025 20:29:30 -0700
Subject: [PATCH 032/113] [native] Advance velox

---
 presto-native-execution/velox | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/presto-native-execution/velox b/presto-native-execution/velox
index 31200dd7fccce..eac4b4bc61445 160000
--- a/presto-native-execution/velox
+++ b/presto-native-execution/velox
@@ -1 +1 @@
-Subproject commit 31200dd7fccce657da5ca0a2409bbc6978e35a68
+Subproject commit eac4b4bc61445e58436dfaa6fa907353308db1e5

From 0b66c1eb55870dc029357646b7bfd73cfab97304 Mon Sep 17 00:00:00 2001
From: pratyakshsharma 
Date: Sat, 16 Aug 2025 16:48:08 +0530
Subject: [PATCH 033/113] Add `connector_name` to `metadata.catalogs` table

---
 .../src/main/sphinx/connector/system.rst      | 35 ++++++++++++-
 .../presto/connector/ConnectorManager.java    | 11 ++--
 .../connector/system/CatalogSystemTable.java  |  9 ++--
 .../com/facebook/presto/metadata/Catalog.java | 52 ++++++++++++++++++-
 .../facebook/presto/metadata/Metadata.java    |  7 +++
 .../presto/metadata/MetadataListing.java      | 15 ++++++
 .../presto/metadata/MetadataManager.java      |  7 +++
 .../InMemoryTransactionManager.java           | 21 ++++++++
 .../transaction/TransactionManager.java       |  7 +++
 .../presto/tests/SystemConnectorTests.java    |  4 +-
 .../selectInformationSchemaColumns.result     |  1 +
 11 files changed, 156 insertions(+), 13 deletions(-)

diff --git a/presto-docs/src/main/sphinx/connector/system.rst b/presto-docs/src/main/sphinx/connector/system.rst
index 9f0b8038e4479..0e388c38627fc 100644
--- a/presto-docs/src/main/sphinx/connector/system.rst
+++ b/presto-docs/src/main/sphinx/connector/system.rst
@@ -36,7 +36,40 @@ System Connector Tables
 ``metadata.catalogs``
 ^^^^^^^^^^^^^^^^^^^^^
 
-The catalogs table contains the list of available catalogs.
+The catalogs table contains the list of available catalogs. The columns in ``metadata.catalogs`` are:
+
+======================================= ======================================================================
+Column Name                             Description
+======================================= ======================================================================
+``catalog_name``                        The value of this column is derived from the names of
+                                        catalog.properties files present under ``etc/catalog`` path under
+                                        presto installation directory. Everything except the suffix
+                                        ``.properties`` is treated as the catalog name. For example, if there
+                                        is a file named ``my_catalog.properties``, then ``my_catalog`` will be
+                                        listed as the value for this column.
+
+``connector_id``                        The values in this column are a duplicate of the values in the
+                                        ``catalog_name`` column.
+
+``connector_name``                      This column represents the actual name of the underlying connector
+                                        that a particular catalog is using. This column contains the value of
+                                        ``connector.name`` property from the catalog.properties file.
+======================================= ======================================================================
+
+Example:
+
+Suppose a user configures a single catalog by creating a file named ``my_catalog.properties`` with the
+below contents::
+
+    connector.name=hive-hadoop2
+    hive.metastore.uri=thrift://localhost:9083
+
+``metadata.catalogs`` table will show below output::
+
+    presto> select * from system.metadata.catalogs;
+     catalog_name | connector_id | connector_name
+    --------------+--------------+----------------
+     my_catalog   | my_catalog   | hive-hadoop2
 
 ``metadata.schema_properties``
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
diff --git a/presto-main-base/src/main/java/com/facebook/presto/connector/ConnectorManager.java b/presto-main-base/src/main/java/com/facebook/presto/connector/ConnectorManager.java
index bbcbe517b50dc..dc85eea82f0e5 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/connector/ConnectorManager.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/connector/ConnectorManager.java
@@ -215,10 +215,10 @@ public synchronized ConnectorId createConnection(String catalogName, String conn
         requireNonNull(connectorName, "connectorName is null");
         ConnectorFactory connectorFactory = connectorFactories.get(connectorName);
         checkArgument(connectorFactory != null, "No factory for connector %s", connectorName);
-        return createConnection(catalogName, connectorFactory, properties);
+        return createConnection(catalogName, connectorFactory, properties, connectorName);
     }
 
-    private synchronized ConnectorId createConnection(String catalogName, ConnectorFactory connectorFactory, Map properties)
+    private synchronized ConnectorId createConnection(String catalogName, ConnectorFactory connectorFactory, Map properties, String connectorName)
     {
         checkState(!stopped.get(), "ConnectorManager is stopped");
         requireNonNull(catalogName, "catalogName is null");
@@ -229,12 +229,12 @@ private synchronized ConnectorId createConnection(String catalogName, ConnectorF
         ConnectorId connectorId = new ConnectorId(catalogName);
         checkState(!connectors.containsKey(connectorId), "A connector %s already exists", connectorId);
 
-        addCatalogConnector(catalogName, connectorId, connectorFactory, properties);
+        addCatalogConnector(catalogName, connectorId, connectorFactory, properties, connectorName);
 
         return connectorId;
     }
 
-    private synchronized void addCatalogConnector(String catalogName, ConnectorId connectorId, ConnectorFactory factory, Map properties)
+    private synchronized void addCatalogConnector(String catalogName, ConnectorId connectorId, ConnectorFactory factory, Map properties, String connectorName)
     {
         // create all connectors before adding, so a broken connector does not leave the system half updated
         MaterializedConnector connector = new MaterializedConnector(connectorId, createConnector(connectorId, factory, properties));
@@ -269,7 +269,8 @@ private synchronized void addCatalogConnector(String catalogName, ConnectorId co
                 informationSchemaConnector.getConnectorId(),
                 informationSchemaConnector.getConnector(),
                 systemConnector.getConnectorId(),
-                systemConnector.getConnector());
+                systemConnector.getConnector(),
+                connectorName);
 
         try {
             addConnectorInternal(connector);
diff --git a/presto-main-base/src/main/java/com/facebook/presto/connector/system/CatalogSystemTable.java b/presto-main-base/src/main/java/com/facebook/presto/connector/system/CatalogSystemTable.java
index 19cbea31463fc..e8fbcdf6a2303 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/connector/system/CatalogSystemTable.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/connector/system/CatalogSystemTable.java
@@ -15,8 +15,8 @@
 
 import com.facebook.presto.Session;
 import com.facebook.presto.common.predicate.TupleDomain;
+import com.facebook.presto.metadata.Catalog.CatalogContext;
 import com.facebook.presto.metadata.Metadata;
-import com.facebook.presto.spi.ConnectorId;
 import com.facebook.presto.spi.ConnectorSession;
 import com.facebook.presto.spi.ConnectorTableMetadata;
 import com.facebook.presto.spi.InMemoryRecordSet;
@@ -32,7 +32,7 @@
 
 import static com.facebook.presto.common.type.VarcharType.createUnboundedVarcharType;
 import static com.facebook.presto.connector.system.SystemConnectorSessionUtil.toSession;
-import static com.facebook.presto.metadata.MetadataListing.listCatalogs;
+import static com.facebook.presto.metadata.MetadataListing.listCatalogsWithConnectorContext;
 import static com.facebook.presto.metadata.MetadataUtil.TableMetadataBuilder.tableMetadataBuilder;
 import static com.facebook.presto.spi.SystemTable.Distribution.SINGLE_COORDINATOR;
 import static java.util.Objects.requireNonNull;
@@ -45,6 +45,7 @@ public class CatalogSystemTable
     public static final ConnectorTableMetadata CATALOG_TABLE = tableMetadataBuilder(CATALOG_TABLE_NAME)
             .column("catalog_name", createUnboundedVarcharType())
             .column("connector_id", createUnboundedVarcharType())
+            .column("connector_name", createUnboundedVarcharType())
             .build();
     private final Metadata metadata;
     private final AccessControl accessControl;
@@ -73,8 +74,8 @@ public RecordCursor cursor(ConnectorTransactionHandle transactionHandle, Connect
     {
         Session session = toSession(transactionHandle, connectorSession);
         Builder table = InMemoryRecordSet.builder(CATALOG_TABLE);
-        for (Map.Entry entry : listCatalogs(session, metadata, accessControl).entrySet()) {
-            table.addRow(entry.getKey(), entry.getValue().toString());
+        for (Map.Entry entry : listCatalogsWithConnectorContext(session, metadata, accessControl).entrySet()) {
+            table.addRow(entry.getKey(), entry.getValue().getCatalogName(), entry.getValue().getConnectorName());
         }
         return table.build().cursor();
     }
diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/Catalog.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/Catalog.java
index 4ee7ed05f963b..5cbdb9bfb1308 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/metadata/Catalog.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/Catalog.java
@@ -32,6 +32,8 @@ public class Catalog
     private final ConnectorId systemTablesId;
     private final Connector systemTables;
 
+    private final CatalogContext catalogContext;
+
     public Catalog(
             String catalogName,
             ConnectorId connectorId,
@@ -40,6 +42,26 @@ public Catalog(
             Connector informationSchema,
             ConnectorId systemTablesId,
             Connector systemTables)
+    {
+        this(catalogName,
+                connectorId,
+                connector,
+                informationSchemaId,
+                informationSchema,
+                systemTablesId,
+                systemTables,
+                catalogName);
+    }
+
+    public Catalog(
+            String catalogName,
+            ConnectorId connectorId,
+            Connector connector,
+            ConnectorId informationSchemaId,
+            Connector informationSchema,
+            ConnectorId systemTablesId,
+            Connector systemTables,
+            String connectorName)
     {
         this.catalogName = checkCatalogName(catalogName);
         this.connectorId = requireNonNull(connectorId, "connectorId is null");
@@ -48,8 +70,9 @@ public Catalog(
         this.informationSchema = requireNonNull(informationSchema, "informationSchema is null");
         this.systemTablesId = requireNonNull(systemTablesId, "systemTablesId is null");
         this.systemTables = requireNonNull(systemTables, "systemTables is null");
+        requireNonNull(connectorName, "connectorName is null");
+        this.catalogContext = new CatalogContext(catalogName, connectorName);
     }
-
     public String getCatalogName()
     {
         return catalogName;
@@ -60,6 +83,11 @@ public ConnectorId getConnectorId()
         return connectorId;
     }
 
+    public CatalogContext getCatalogContext()
+    {
+        return catalogContext;
+    }
+
     public ConnectorId getInformationSchemaId()
     {
         return informationSchemaId;
@@ -92,4 +120,26 @@ public String toString()
                 .add("connectorId", connectorId)
                 .toString();
     }
+
+    public class CatalogContext
+    {
+        private final String catalogName;
+        private final String connectorName;
+
+        public CatalogContext(String catalogName, String connectorName)
+        {
+            this.catalogName = catalogName;
+            this.connectorName = connectorName;
+        }
+
+        public String getCatalogName()
+        {
+            return catalogName;
+        }
+
+        public String getConnectorName()
+        {
+            return connectorName;
+        }
+    }
 }
diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/Metadata.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/Metadata.java
index f7283b436c65f..2a0c62d7d4b0e 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/metadata/Metadata.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/Metadata.java
@@ -20,6 +20,7 @@
 import com.facebook.presto.common.predicate.TupleDomain;
 import com.facebook.presto.common.type.Type;
 import com.facebook.presto.common.type.TypeSignature;
+import com.facebook.presto.metadata.Catalog.CatalogContext;
 import com.facebook.presto.spi.ColumnHandle;
 import com.facebook.presto.spi.ColumnMetadata;
 import com.facebook.presto.spi.ConnectorId;
@@ -52,6 +53,7 @@
 import com.facebook.presto.spi.statistics.ComputedStatistics;
 import com.facebook.presto.spi.statistics.TableStatistics;
 import com.facebook.presto.spi.statistics.TableStatisticsMetadata;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.util.concurrent.ListenableFuture;
 import io.airlift.slice.Slice;
 
@@ -360,6 +362,11 @@ public interface Metadata
      */
     Map getCatalogNames(Session session);
 
+    default Map getCatalogNamesWithConnectorContext(Session session)
+    {
+        return ImmutableMap.of();
+    }
+
     /**
      * Get the names that match the specified table prefix (never null).
      */
diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/MetadataListing.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/MetadataListing.java
index ac0e661bdc04c..c5a2532466a53 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/metadata/MetadataListing.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/MetadataListing.java
@@ -15,6 +15,7 @@
 
 import com.facebook.presto.Session;
 import com.facebook.presto.common.QualifiedObjectName;
+import com.facebook.presto.metadata.Catalog.CatalogContext;
 import com.facebook.presto.spi.ColumnMetadata;
 import com.facebook.presto.spi.ConnectorId;
 import com.facebook.presto.spi.SchemaTableName;
@@ -54,6 +55,20 @@ public static SortedMap listCatalogs(Session session, Metad
         return result.build();
     }
 
+    public static SortedMap listCatalogsWithConnectorContext(Session session, Metadata metadata, AccessControl accessControl)
+    {
+        Map catalogNamesWithConnectorContext = metadata.getCatalogNamesWithConnectorContext(session);
+        Set allowedCatalogs = accessControl.filterCatalogs(session.getIdentity(), session.getAccessControlContext(), catalogNamesWithConnectorContext.keySet());
+
+        ImmutableSortedMap.Builder result = ImmutableSortedMap.naturalOrder();
+        for (Map.Entry entry : catalogNamesWithConnectorContext.entrySet()) {
+            if (allowedCatalogs.contains(entry.getKey())) {
+                result.put(entry);
+            }
+        }
+        return result.build();
+    }
+
     public static SortedSet listSchemas(Session session, Metadata metadata, AccessControl accessControl, String catalogName)
     {
         Set schemaNames = ImmutableSet.copyOf(metadata.listSchemaNames(session, catalogName));
diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/MetadataManager.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/MetadataManager.java
index a8be73143072a..bfcafe2e98339 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/metadata/MetadataManager.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/MetadataManager.java
@@ -26,6 +26,7 @@
 import com.facebook.presto.common.predicate.TupleDomain;
 import com.facebook.presto.common.type.Type;
 import com.facebook.presto.common.type.TypeSignature;
+import com.facebook.presto.metadata.Catalog.CatalogContext;
 import com.facebook.presto.spi.ColumnHandle;
 import com.facebook.presto.spi.ColumnMetadata;
 import com.facebook.presto.spi.ConnectorDeleteTableHandle;
@@ -987,6 +988,12 @@ public Map getCatalogNames(Session session)
         return transactionManager.getCatalogNames(session.getRequiredTransactionId());
     }
 
+    @Override
+    public Map getCatalogNamesWithConnectorContext(Session session)
+    {
+        return transactionManager.getCatalogNamesWithConnectorContext(session.getRequiredTransactionId());
+    }
+
     @Override
     public List listViews(Session session, QualifiedTablePrefix prefix)
     {
diff --git a/presto-main-base/src/main/java/com/facebook/presto/transaction/InMemoryTransactionManager.java b/presto-main-base/src/main/java/com/facebook/presto/transaction/InMemoryTransactionManager.java
index 8afc8c07370e4..bd3a4dbd4f1fd 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/transaction/InMemoryTransactionManager.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/transaction/InMemoryTransactionManager.java
@@ -19,6 +19,7 @@
 import com.facebook.airlift.units.Duration;
 import com.facebook.presto.common.transaction.TransactionId;
 import com.facebook.presto.metadata.Catalog;
+import com.facebook.presto.metadata.Catalog.CatalogContext;
 import com.facebook.presto.metadata.CatalogManager;
 import com.facebook.presto.metadata.CatalogMetadata;
 import com.facebook.presto.spi.ConnectorId;
@@ -214,6 +215,12 @@ public Map getCatalogNames(TransactionId transactionId)
         return getTransactionMetadata(transactionId).getCatalogNames();
     }
 
+    @Override
+    public Map getCatalogNamesWithConnectorContext(TransactionId transactionId)
+    {
+        return getTransactionMetadata(transactionId).getCatalogNamesWithConnectorContext();
+    }
+
     @Override
     public Optional getOptionalCatalogMetadata(TransactionId transactionId, String catalogName)
     {
@@ -462,6 +469,20 @@ private synchronized Map getCatalogNames()
             return ImmutableMap.copyOf(catalogNames);
         }
 
+        private synchronized Map getCatalogNamesWithConnectorContext()
+        {
+            Map catalogNamesWithConnectorContext = new HashMap<>();
+            catalogByName.values().stream()
+                    .filter(Optional::isPresent)
+                    .map(Optional::get)
+                    .forEach(catalog -> catalogNamesWithConnectorContext.put(catalog.getCatalogName(), catalog.getCatalogContext()));
+
+            catalogManager.getCatalogs().stream()
+                    .forEach(catalog -> catalogNamesWithConnectorContext.putIfAbsent(catalog.getCatalogName(), catalog.getCatalogContext()));
+
+            return ImmutableMap.copyOf(catalogNamesWithConnectorContext);
+        }
+
         private synchronized Optional getConnectorId(String catalogName)
         {
             Optional catalog = catalogByName.get(catalogName);
diff --git a/presto-main-base/src/main/java/com/facebook/presto/transaction/TransactionManager.java b/presto-main-base/src/main/java/com/facebook/presto/transaction/TransactionManager.java
index b7d3e0ccdca4a..ecacaf515e0ca 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/transaction/TransactionManager.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/transaction/TransactionManager.java
@@ -15,6 +15,7 @@
 
 import com.facebook.presto.Session;
 import com.facebook.presto.common.transaction.TransactionId;
+import com.facebook.presto.metadata.Catalog.CatalogContext;
 import com.facebook.presto.metadata.CatalogMetadata;
 import com.facebook.presto.spi.ConnectorId;
 import com.facebook.presto.spi.connector.ConnectorTransactionHandle;
@@ -22,6 +23,7 @@
 import com.facebook.presto.spi.function.FunctionNamespaceTransactionHandle;
 import com.facebook.presto.spi.security.AccessControl;
 import com.facebook.presto.spi.transaction.IsolationLevel;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.util.concurrent.ListenableFuture;
 
 import java.util.List;
@@ -47,6 +49,11 @@ public interface TransactionManager
 
     Map getCatalogNames(TransactionId transactionId);
 
+    default Map getCatalogNamesWithConnectorContext(TransactionId transactionId)
+    {
+        return ImmutableMap.of();
+    }
+
     Optional getOptionalCatalogMetadata(TransactionId transactionId, String catalogName);
 
     void enableRollback(TransactionId transactionId);
diff --git a/presto-product-tests/src/main/java/com/facebook/presto/tests/SystemConnectorTests.java b/presto-product-tests/src/main/java/com/facebook/presto/tests/SystemConnectorTests.java
index 3317d07e90f8c..4f14b99e2d931 100644
--- a/presto-product-tests/src/main/java/com/facebook/presto/tests/SystemConnectorTests.java
+++ b/presto-product-tests/src/main/java/com/facebook/presto/tests/SystemConnectorTests.java
@@ -108,9 +108,9 @@ public void selectRuntimeTasks()
     @Test(groups = {SYSTEM_CONNECTOR, JDBC})
     public void selectMetadataCatalogs()
     {
-        String sql = "select catalog_name, connector_id from system.metadata.catalogs";
+        String sql = "select catalog_name, connector_id, connector_name from system.metadata.catalogs";
         assertThat(query(sql))
-                .hasColumns(VARCHAR, VARCHAR)
+                .hasColumns(VARCHAR, VARCHAR, VARCHAR)
                 .hasAnyRows();
     }
 
diff --git a/presto-product-tests/src/main/resources/sql-tests/testcases/system/selectInformationSchemaColumns.result b/presto-product-tests/src/main/resources/sql-tests/testcases/system/selectInformationSchemaColumns.result
index c442a420ebead..4103c03301025 100644
--- a/presto-product-tests/src/main/resources/sql-tests/testcases/system/selectInformationSchemaColumns.result
+++ b/presto-product-tests/src/main/resources/sql-tests/testcases/system/selectInformationSchemaColumns.result
@@ -46,6 +46,7 @@ system| metadata| analyze_properties| type| varchar| YES| null| null|
 system| metadata| analyze_properties| description| varchar| YES| null| null|
 system| metadata| catalogs| catalog_name| varchar| YES| null| null|
 system| metadata| catalogs| connector_id| varchar| YES| null| null|
+system| metadata| catalogs| connector_name| varchar| YES| null| null|
 system| metadata| column_properties| catalog_name| varchar| YES| null| null|
 system| metadata| column_properties| property_name| varchar| YES| null| null|
 system| metadata| column_properties| default_value| varchar| YES| null| null|

From 591e49f44732266c53b1a08a96201208df45f7be Mon Sep 17 00:00:00 2001
From: Reetika Agrawal 
Date: Fri, 22 Aug 2025 13:21:26 +0530
Subject: [PATCH 034/113] Enable and fix all Mongodb connector tests in CI

---
 .github/workflows/tests.yml                               | 2 +-
 presto-mongodb/pom.xml                                    | 6 ++----
 .../java/com/facebook/presto/mongodb/MongoSession.java    | 6 +-----
 .../presto/mongodb/TestMongoIntegrationSmokeTest.java     | 8 +++++---
 4 files changed, 9 insertions(+), 13 deletions(-)

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 0afc4c0fbfa82..3432d50e03e9b 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -59,7 +59,7 @@ jobs:
           - ":presto-hive -P test-hive-parquet"
           - ":presto-main-base"
           - ":presto-main"
-          - ":presto-mongodb -P test-mongo-distributed-queries"
+          - ":presto-mongodb -P ci-full-tests"
           - ":presto-redis -P test-redis-integration-smoke-test"
           - ":presto-elasticsearch"
           - ":presto-orc"
diff --git a/presto-mongodb/pom.xml b/presto-mongodb/pom.xml
index e8a6429ea6b4e..8644f91fb7809 100644
--- a/presto-mongodb/pom.xml
+++ b/presto-mongodb/pom.xml
@@ -222,17 +222,15 @@
                 
             
         
+        
         
-            test-mongo-distributed-queries
+            ci-full-tests
             
                 
                     
                         org.apache.maven.plugins
                         maven-surefire-plugin
                         
-                            
-                                **/TestMongoDistributedQueries.java
-                            
                         
                     
                 
diff --git a/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoSession.java b/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoSession.java
index c3fb4f1913575..be5ec84d46871 100644
--- a/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoSession.java
+++ b/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoSession.java
@@ -457,14 +457,10 @@ private boolean deleteTableMetadata(SchemaTableName schemaTableName)
         String tableName = schemaTableName.getTableName();
 
         MongoDatabase db = client.getDatabase(schemaName);
-        if (!collectionExists(db, tableName)) {
-            return false;
-        }
-
         DeleteResult result = db.getCollection(schemaCollection)
                 .deleteOne(new Document(TABLE_NAME_KEY, tableName));
 
-        return result.getDeletedCount() == 1;
+        return result.getDeletedCount() == 1 || !collectionExists(db, tableName);
     }
 
     private List guessTableFields(SchemaTableName schemaTableName)
diff --git a/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoIntegrationSmokeTest.java b/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoIntegrationSmokeTest.java
index a5e820d4a1dfb..fd35e5ea284ee 100644
--- a/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoIntegrationSmokeTest.java
+++ b/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoIntegrationSmokeTest.java
@@ -21,6 +21,7 @@
 import com.google.common.collect.ImmutableMap;
 import com.mongodb.client.MongoCollection;
 import org.bson.Document;
+import org.bson.types.ObjectId;
 import org.testng.annotations.AfterClass;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
@@ -146,7 +147,8 @@ public void testInsertWithEveryType()
         assertEquals(row.getField(5), LocalDate.of(1980, 5, 7));
         assertEquals(row.getField(6), LocalDateTime.of(1980, 5, 7, 11, 22, 33, 456_000_000));
         assertEquals(row.getField(7), LocalTime.of(11, 22, 33, 456_000_000));
-        assertEquals(row.getField(8), "{\"name\":\"alice\"}");
+        assertEquals(new ObjectId((byte[]) row.getField(8)), new ObjectId("ffffffffffffffffffffffff"));
+        assertEquals(row.getField(9).toString(), "{\"name\":\"alice\"}");
         assertUpdate("DROP TABLE test_insert_types_table");
         assertFalse(getQueryRunner().tableExists(getSession(), "test_insert_types_table"));
     }
@@ -218,13 +220,13 @@ public void testMaps()
         assertUpdate("CREATE TABLE test.tmp_map9 (col VARCHAR)");
         mongoQueryRunner.getMongoClient().getDatabase("test").getCollection("tmp_map9").insertOne(new Document(
                 ImmutableMap.of("col", new Document(ImmutableMap.of("key1", "value1", "key2", "value2")))));
-        assertQuery("SELECT col FROM test.tmp_map9", "SELECT '{ \"key1\" : \"value1\", \"key2\" : \"value2\" }'");
+        assertQuery("SELECT col FROM test.tmp_map9", "SELECT '{\"key1\": \"value1\", \"key2\": \"value2\"}'");
 
         assertUpdate("CREATE TABLE test.tmp_map10 (col VARCHAR)");
         mongoQueryRunner.getMongoClient().getDatabase("test").getCollection("tmp_map10").insertOne(new Document(
                 ImmutableMap.of("col", ImmutableList.of(new Document(ImmutableMap.of("key1", "value1", "key2", "value2")),
                         new Document(ImmutableMap.of("key3", "value3", "key4", "value4"))))));
-        assertQuery("SELECT col FROM test.tmp_map10", "SELECT '[{ \"key1\" : \"value1\", \"key2\" : \"value2\" }, { \"key3\" : \"value3\", \"key4\" : \"value4\" }]'");
+        assertQuery("SELECT col FROM test.tmp_map10", "SELECT '[{\"key1\": \"value1\", \"key2\": \"value2\"}, {\"key3\": \"value3\", \"key4\": \"value4\"}]'");
 
         assertUpdate("CREATE TABLE test.tmp_map11 (col VARCHAR)");
         mongoQueryRunner.getMongoClient().getDatabase("test").getCollection("tmp_map11").insertOne(new Document(

From 9da619b91d095ef590d70a59aec12e4a3182203f Mon Sep 17 00:00:00 2001
From: Reetika Agrawal 
Date: Fri, 22 Aug 2025 13:22:10 +0530
Subject: [PATCH 035/113] Enable case-senstive identifer support for Mongodb
 connector

---
 .../src/main/sphinx/connector/mongodb.rst     |   4 +
 .../presto/mongodb/MongoClientConfig.java     |  13 +
 .../presto/mongodb/MongoConnector.java        |   7 +-
 .../presto/mongodb/MongoMetadata.java         |  14 +-
 .../facebook/presto/mongodb/MongoSession.java |   2 +-
 .../presto/mongodb/MongoQueryRunner.java      |  13 +-
 .../presto/mongodb/TestMongoClientConfig.java |   7 +-
 .../TestMongoDbIntegrationMixedCaseTest.java  | 223 ++++++++++++++++++
 .../mongodb/TestMongoDistributedQueries.java  |   3 +-
 9 files changed, 271 insertions(+), 15 deletions(-)
 create mode 100644 presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoDbIntegrationMixedCaseTest.java

diff --git a/presto-docs/src/main/sphinx/connector/mongodb.rst b/presto-docs/src/main/sphinx/connector/mongodb.rst
index f40e7b0c50554..63a9bdedeb3af 100644
--- a/presto-docs/src/main/sphinx/connector/mongodb.rst
+++ b/presto-docs/src/main/sphinx/connector/mongodb.rst
@@ -51,6 +51,10 @@ Property Name                         Description
 ``mongodb.write-concern``             The write concern
 ``mongodb.required-replica-set``      The required replica set name
 ``mongodb.cursor-batch-size``         The number of elements to return in a batch
+``case-sensitive-name-matching``      Enable case-sensitive identifier support for schema,
+                                      table, and column names for the connector. When disabled,
+                                      names are matched case-insensitively using lowercase
+                                      normalization. Default is ``false``
 ===================================== ==============================================================
 
 ``mongodb.seeds``
diff --git a/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoClientConfig.java b/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoClientConfig.java
index 07987c290c365..b23570fdad83d 100644
--- a/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoClientConfig.java
+++ b/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoClientConfig.java
@@ -59,6 +59,7 @@ public class MongoClientConfig
     private WriteConcernType writeConcern = WriteConcernType.ACKNOWLEDGED;
     private String requiredReplicaSetName;
     private String implicitRowFieldPrefix = "_pos";
+    private boolean caseSensitiveNameMatchingEnabled;
 
     @NotNull
     public String getSchemaCollection()
@@ -328,4 +329,16 @@ public MongoClientConfig setSslEnabled(boolean sslEnabled)
         this.sslEnabled = sslEnabled;
         return this;
     }
+
+    public boolean isCaseSensitiveNameMatchingEnabled()
+    {
+        return caseSensitiveNameMatchingEnabled;
+    }
+
+    @Config("case-sensitive-name-matching")
+    public MongoClientConfig setCaseSensitiveNameMatchingEnabled(boolean caseSensitiveNameMatchingEnabled)
+    {
+        this.caseSensitiveNameMatchingEnabled = caseSensitiveNameMatchingEnabled;
+        return this;
+    }
 }
diff --git a/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoConnector.java b/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoConnector.java
index b21f6fdf4511c..f7feee5e7eaf3 100644
--- a/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoConnector.java
+++ b/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoConnector.java
@@ -41,18 +41,21 @@ public class MongoConnector
     private final MongoPageSinkProvider pageSinkProvider;
 
     private final ConcurrentMap transactions = new ConcurrentHashMap<>();
+    private final MongoClientConfig mongoClientConfig;
 
     @Inject
     public MongoConnector(
             MongoSession mongoSession,
             MongoSplitManager splitManager,
             MongoPageSourceProvider pageSourceProvider,
-            MongoPageSinkProvider pageSinkProvider)
+            MongoPageSinkProvider pageSinkProvider,
+            MongoClientConfig mongoClientConfig)
     {
         this.mongoSession = mongoSession;
         this.splitManager = requireNonNull(splitManager, "splitManager is null");
         this.pageSourceProvider = requireNonNull(pageSourceProvider, "pageSourceProvider is null");
         this.pageSinkProvider = requireNonNull(pageSinkProvider, "pageSinkProvider is null");
+        this.mongoClientConfig = requireNonNull(mongoClientConfig, "mongoClientConfig is null");
     }
 
     @Override
@@ -60,7 +63,7 @@ public ConnectorTransactionHandle beginTransaction(IsolationLevel isolationLevel
     {
         checkConnectorSupports(READ_UNCOMMITTED, isolationLevel);
         MongoTransactionHandle transaction = new MongoTransactionHandle();
-        transactions.put(transaction, new MongoMetadata(mongoSession));
+        transactions.put(transaction, new MongoMetadata(mongoSession, mongoClientConfig));
         return transaction;
     }
 
diff --git a/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoMetadata.java b/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoMetadata.java
index 1a75bcfe2feea..6133db8f0ac53 100644
--- a/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoMetadata.java
+++ b/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoMetadata.java
@@ -49,7 +49,7 @@
 import java.util.concurrent.atomic.AtomicReference;
 
 import static com.google.common.base.Preconditions.checkState;
-import static java.util.Locale.ENGLISH;
+import static java.util.Locale.ROOT;
 import static java.util.Objects.requireNonNull;
 import static java.util.stream.Collectors.toList;
 
@@ -61,10 +61,12 @@ public class MongoMetadata
     private final MongoSession mongoSession;
 
     private final AtomicReference rollbackAction = new AtomicReference<>();
+    private final MongoClientConfig mongoClientConfig;
 
-    public MongoMetadata(MongoSession mongoSession)
+    public MongoMetadata(MongoSession mongoSession, MongoClientConfig mongoClientConfig)
     {
         this.mongoSession = requireNonNull(mongoSession, "mongoSession is null");
+        this.mongoClientConfig = mongoClientConfig;
     }
 
     @Override
@@ -101,7 +103,7 @@ public List listTables(ConnectorSession session, String schemaN
 
         for (String schemaName : listSchemas(session, schemaNameOrNull)) {
             for (String tableName : mongoSession.getAllTables(schemaName)) {
-                tableNames.add(new SchemaTableName(schemaName, tableName.toLowerCase(ENGLISH)));
+                tableNames.add(new SchemaTableName(schemaName, normalizeIdentifier(session, tableName)));
             }
         }
         return tableNames.build();
@@ -319,4 +321,10 @@ public void dropColumn(ConnectorSession session, ConnectorTableHandle tableHandl
     {
         mongoSession.dropColumn(((MongoTableHandle) tableHandle), ((MongoColumnHandle) column).getName());
     }
+
+    @Override
+    public String normalizeIdentifier(ConnectorSession session, String identifier)
+    {
+        return mongoClientConfig.isCaseSensitiveNameMatchingEnabled() ? identifier : identifier.toLowerCase(ROOT);
+    }
 }
diff --git a/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoSession.java b/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoSession.java
index be5ec84d46871..c73f949ac551d 100644
--- a/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoSession.java
+++ b/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoSession.java
@@ -404,7 +404,7 @@ private Document getTableMetadata(SchemaTableName schemaTableName)
     public boolean collectionExists(MongoDatabase db, String collectionName)
     {
         for (String name : db.listCollectionNames()) {
-            if (name.equalsIgnoreCase(collectionName)) {
+            if (name.equals(collectionName)) {
                 return true;
             }
         }
diff --git a/presto-mongodb/src/test/java/com/facebook/presto/mongodb/MongoQueryRunner.java b/presto-mongodb/src/test/java/com/facebook/presto/mongodb/MongoQueryRunner.java
index 9c07bed90d15a..1b079c6c93645 100644
--- a/presto-mongodb/src/test/java/com/facebook/presto/mongodb/MongoQueryRunner.java
+++ b/presto-mongodb/src/test/java/com/facebook/presto/mongodb/MongoQueryRunner.java
@@ -24,6 +24,7 @@
 import io.airlift.tpch.TpchTable;
 
 import java.net.InetSocketAddress;
+import java.util.HashMap;
 import java.util.Map;
 
 import static com.facebook.airlift.testing.Closeables.closeAllSuppress;
@@ -53,10 +54,10 @@ private MongoQueryRunner(Session session, int workers)
     public static MongoQueryRunner createMongoQueryRunner(TpchTable... tables)
             throws Exception
     {
-        return createMongoQueryRunner(ImmutableList.copyOf(tables));
+        return createMongoQueryRunner(ImmutableList.copyOf(tables), ImmutableMap.of());
     }
 
-    public static MongoQueryRunner createMongoQueryRunner(Iterable> tables)
+    public static MongoQueryRunner createMongoQueryRunner(Iterable> tables, Map connectorProperties)
             throws Exception
     {
         MongoQueryRunner queryRunner = null;
@@ -66,12 +67,12 @@ public static MongoQueryRunner createMongoQueryRunner(Iterable> tab
             queryRunner.installPlugin(new TpchPlugin());
             queryRunner.createCatalog("tpch", "tpch");
 
-            Map properties = ImmutableMap.of(
-                    "mongodb.seeds", queryRunner.getAddress().getHostString() + ":" + queryRunner.getAddress().getPort(),
-                    "mongodb.socket-keep-alive", "true");
+            connectorProperties = new HashMap<>(connectorProperties);
+            connectorProperties.putIfAbsent("mongodb.seeds", queryRunner.getAddress().getHostString() + ":" + queryRunner.getAddress().getPort());
+            connectorProperties.putIfAbsent("mongodb.socket-keep-alive", "true");
 
             queryRunner.installPlugin(new MongoPlugin());
-            queryRunner.createCatalog("mongodb", "mongodb", properties);
+            queryRunner.createCatalog("mongodb", "mongodb", connectorProperties);
 
             copyTpchTables(queryRunner, "tpch", TINY_SCHEMA_NAME, createSession(), tables);
 
diff --git a/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoClientConfig.java b/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoClientConfig.java
index 42e4c35c2edfe..8c341c1d407d9 100644
--- a/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoClientConfig.java
+++ b/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoClientConfig.java
@@ -43,7 +43,8 @@ public void testDefaults()
                 .setReadPreferenceTags("")
                 .setWriteConcern(WriteConcernType.ACKNOWLEDGED)
                 .setRequiredReplicaSetName(null)
-                .setImplicitRowFieldPrefix("_pos"));
+                .setImplicitRowFieldPrefix("_pos")
+                .setCaseSensitiveNameMatchingEnabled(false));
     }
 
     @Test
@@ -66,6 +67,7 @@ public void testExplicitPropertyMappings()
                 .put("mongodb.write-concern", "UNACKNOWLEDGED")
                 .put("mongodb.required-replica-set", "replica_set")
                 .put("mongodb.implicit-row-field-prefix", "_prefix")
+                .put("case-sensitive-name-matching", "true")
                 .build();
 
         MongoClientConfig expected = new MongoClientConfig()
@@ -84,7 +86,8 @@ public void testExplicitPropertyMappings()
                 .setReadPreferenceTags("tag_name:tag_value")
                 .setWriteConcern(WriteConcernType.UNACKNOWLEDGED)
                 .setRequiredReplicaSetName("replica_set")
-                .setImplicitRowFieldPrefix("_prefix");
+                .setImplicitRowFieldPrefix("_prefix")
+                .setCaseSensitiveNameMatchingEnabled(true);
 
         ConfigAssertions.assertFullMapping(properties, expected);
     }
diff --git a/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoDbIntegrationMixedCaseTest.java b/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoDbIntegrationMixedCaseTest.java
new file mode 100644
index 0000000000000..24ba8fb81e1bd
--- /dev/null
+++ b/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoDbIntegrationMixedCaseTest.java
@@ -0,0 +1,223 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.mongodb;
+
+import com.facebook.presto.Session;
+import com.facebook.presto.testing.MaterializedResult;
+import com.facebook.presto.testing.QueryRunner;
+import com.facebook.presto.tests.AbstractTestQueryFramework;
+import com.google.common.collect.ImmutableMap;
+import com.mongodb.MongoClient;
+import io.airlift.tpch.TpchTable;
+import org.bson.Document;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import java.util.List;
+import java.util.stream.Collectors;
+
+import static com.facebook.presto.common.type.BigintType.BIGINT;
+import static com.facebook.presto.common.type.VarcharType.VARCHAR;
+import static com.facebook.presto.mongodb.MongoQueryRunner.createMongoQueryRunner;
+import static com.facebook.presto.testing.TestingSession.testSessionBuilder;
+import static com.facebook.presto.testing.assertions.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertTrue;
+
+@Test
+public class TestMongoDbIntegrationMixedCaseTest
+        extends AbstractTestQueryFramework
+{
+    private MongoQueryRunner mongoQueryRunner;
+
+    @Override
+    protected QueryRunner createQueryRunner()
+            throws Exception
+    {
+        return createMongoQueryRunner(TpchTable.getTables(), ImmutableMap.of("case-sensitive-name-matching", "true"));
+    }
+
+    @BeforeClass
+    public void setUp()
+    {
+        mongoQueryRunner = (MongoQueryRunner) getQueryRunner();
+    }
+
+    @AfterClass(alwaysRun = true)
+    public final void destroy()
+    {
+        if (mongoQueryRunner != null) {
+            mongoQueryRunner.shutdown();
+        }
+    }
+
+    public void testDescribeTableWithDifferentCaseInSameSchema()
+    {
+        try {
+            getQueryRunner().execute("CREATE TABLE ORDERS AS SELECT * FROM orders");
+
+            assertTrue(getQueryRunner().tableExists(getSession(), "orders"));
+            assertTrue(getQueryRunner().tableExists(getSession(), "ORDERS"));
+
+            MaterializedResult actualColumns = computeActual("DESC ORDERS").toTestTypes();
+
+            MaterializedResult expectedColumns = MaterializedResult.resultBuilder(getQueryRunner().getDefaultSession(), VARCHAR, VARCHAR, VARCHAR, VARCHAR, BIGINT, BIGINT, BIGINT)
+                    .row("orderkey", "bigint", "", "", 19L, null, null)
+                    .row("custkey", "bigint", "", "", 19L, null, null)
+                    .row("orderstatus", "varchar(1)", "", "", null, null, 1L)
+                    .row("totalprice", "double", "", "", 53L, null, null)
+                    .row("orderdate", "date", "", "", null, null, null)
+                    .row("orderpriority", "varchar(15)", "", "", null, null, 15L)
+                    .row("clerk", "varchar(15)", "", "", null, null, 15L)
+                    .row("shippriority", "integer", "", "", 10L, null, null)
+                    .row("comment", "varchar(79)", "", "", null, null, 79L)
+                    .build();
+            assertEquals(actualColumns, expectedColumns);
+        }
+        finally {
+            assertUpdate("DROP TABLE tpch.ORDERS");
+        }
+    }
+
+    @Test
+    public void testCreateAndDropTable()
+    {
+        Session session = testSessionBuilder()
+                .setCatalog("mongodb")
+                .setSchema("Mixed_Test_Database")
+                .build();
+
+        try {
+            getQueryRunner().execute(session, "CREATE TABLE TEST_CREATE(name VARCHAR(50), id int)");
+            assertTrue(getQueryRunner().tableExists(session, "TEST_CREATE"));
+            assertFalse(getQueryRunner().tableExists(session, "test_create"));
+
+            getQueryRunner().execute(session, "CREATE TABLE test_create(name VARCHAR(50), id int)");
+            assertTrue(getQueryRunner().tableExists(session, "test_create"));
+            assertFalse(getQueryRunner().tableExists(session, "Test_Create"));
+        }
+
+        finally {
+            assertUpdate(session, "DROP TABLE TEST_CREATE");
+            assertFalse(getQueryRunner().tableExists(session, "TEST_CREATE"));
+
+            assertUpdate(session, "DROP TABLE test_create");
+            assertFalse(getQueryRunner().tableExists(session, "test_create"));
+        }
+    }
+
+    @Test
+    public void testCreateTableAs()
+    {
+        Session session = testSessionBuilder()
+                .setCatalog("mongodb")
+                .setSchema("Mixed_Test_Database")
+                .build();
+
+        try {
+            getQueryRunner().execute(session, "CREATE TABLE TEST_CTAS AS SELECT * FROM tpch.region");
+            assertTrue(getQueryRunner().tableExists(session, "TEST_CTAS"));
+
+            getQueryRunner().execute(session, "CREATE TABLE IF NOT EXISTS test_ctas AS SELECT * FROM tpch.region");
+            assertTrue(getQueryRunner().tableExists(session, "test_ctas"));
+
+            getQueryRunner().execute(session, "CREATE TABLE TEST_CTAS_Join AS SELECT c.custkey, o.orderkey FROM " +
+                    "tpch.customer c INNER JOIN tpch.orders o ON c.custkey = o.custkey WHERE c.mktsegment = 'BUILDING'");
+            assertTrue(getQueryRunner().tableExists(session, "TEST_CTAS_Join"));
+
+            assertQueryFails("CREATE TABLE Mixed_Test_Database.TEST_CTAS_FAIL_Join AS SELECT c.custkey, o.orderkey FROM " +
+                    "tpch.customer c INNER JOIN tpch.ORDERS1 o ON c.custkey = o.custkey WHERE c.mktsegment = 'BUILDING'", "Table mongodb.tpch.ORDERS1 does not exist"); //failure scenario since tpch.ORDERS1 doesn't exist
+            assertFalse(getQueryRunner().tableExists(session, "TEST_CTAS_FAIL_Join"));
+
+            getQueryRunner().execute(session, "CREATE TABLE Test_CTAS_Mixed_Join AS SELECT Cus.custkey, Ord.orderkey FROM " +
+                    "tpch.customer Cus INNER JOIN tpch.orders Ord ON Cus.custkey = Ord.custkey WHERE Cus.mktsegment = 'BUILDING'");
+            assertTrue(getQueryRunner().tableExists(session, "Test_CTAS_Mixed_Join"));
+        }
+        finally {
+            getQueryRunner().execute(session, "DROP TABLE IF EXISTS TEST_CTAS");
+            getQueryRunner().execute(session, "DROP TABLE IF EXISTS test_ctas");
+            getQueryRunner().execute(session, "DROP TABLE IF EXISTS TEST_CTAS_Join");
+            getQueryRunner().execute(session, "DROP TABLE IF EXISTS Test_CTAS_Mixed_Join");
+        }
+    }
+
+    @Test
+    public void testInsert()
+    {
+        Session session = testSessionBuilder()
+                .setCatalog("mongodb")
+                .setSchema("Mixed_Test_Database")
+                .build();
+
+        try {
+            getQueryRunner().execute(session, "CREATE TABLE Test_Insert (x bigint, y varchar(100))");
+            getQueryRunner().execute(session, "INSERT INTO Test_Insert VALUES (123, 'test')");
+            assertTrue(getQueryRunner().tableExists(session, "Test_Insert"));
+            assertQuery("SELECT * FROM Mixed_Test_Database.Test_Insert", "SELECT 123 x, 'test' y");
+
+            getQueryRunner().execute(session, "CREATE TABLE TEST_INSERT (x bigint, Y varchar(100))");
+            getQueryRunner().execute(session, "INSERT INTO TEST_INSERT VALUES (1234, 'test1')");
+            assertTrue(getQueryRunner().tableExists(session, "TEST_INSERT"));
+            assertQuery("SELECT * FROM Mixed_Test_Database.TEST_INSERT", "SELECT 1234 x, 'test1' Y");
+        }
+        finally {
+            getQueryRunner().execute(session, "DROP TABLE IF EXISTS Test_Insert");
+            getQueryRunner().execute(session, "DROP TABLE IF EXISTS TEST_INSERT");
+        }
+    }
+
+    @Test
+    public void testMixedCaseColumns()
+    {
+        try {
+            assertUpdate("CREATE TABLE test (a integer, B integer)");
+            assertUpdate("CREATE TABLE TEST (a integer, aA integer)");
+            assertUpdate("INSERT INTO TEST VALUES (123, 12)", 1);
+            assertTableColumnNames("TEST", "a", "aA");
+            assertUpdate("ALTER TABLE TEST ADD COLUMN EMAIL varchar");
+            assertUpdate("ALTER TABLE TEST RENAME COLUMN a TO a_New");
+            assertTableColumnNames("TEST", "a_New", "aA", "EMAIL");
+            assertUpdate("ALTER TABLE TEST DROP COLUMN aA");
+            assertTableColumnNames("TEST", "a_New", "EMAIL");
+        }
+        finally {
+            assertUpdate("DROP TABLE test");
+            assertUpdate("DROP TABLE TEST");
+        }
+    }
+
+    @Test
+    public void testShowSchemas()
+    {
+        // Create two MongoDB databases directly, since Presto doesn't support create schema for mongodb
+        MongoClient mongoClient = mongoQueryRunner.getMongoClient();
+        try {
+            mongoClient.getDatabase("TESTDB1").getCollection("dummy").insertOne(new Document("x", 1));
+            mongoClient.getDatabase("testdb1").getCollection("dummy").insertOne(new Document("x", 1));
+
+            MaterializedResult result = computeActual("SHOW SCHEMAS");
+            List schemaNames = result.getMaterializedRows().stream()
+                    .map(row -> row.getField(0).toString())
+                    .collect(Collectors.toList());
+
+            assertTrue(schemaNames.contains("TESTDB1"));
+            assertTrue(schemaNames.contains("testdb1"));
+        }
+        finally {
+            mongoClient.getDatabase("TESTDB1").drop();
+            mongoClient.getDatabase("testdb1").drop();
+        }
+    }
+}
diff --git a/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoDistributedQueries.java b/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoDistributedQueries.java
index 08e6e9f04feac..23c58a846d886 100644
--- a/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoDistributedQueries.java
+++ b/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoDistributedQueries.java
@@ -15,6 +15,7 @@
 
 import com.facebook.presto.testing.QueryRunner;
 import com.facebook.presto.tests.AbstractTestQueries;
+import com.google.common.collect.ImmutableMap;
 import io.airlift.tpch.TpchTable;
 import org.testng.annotations.AfterClass;
 import org.testng.annotations.BeforeClass;
@@ -32,7 +33,7 @@ public class TestMongoDistributedQueries
     protected QueryRunner createQueryRunner()
             throws Exception
     {
-        return createMongoQueryRunner(TpchTable.getTables());
+        return createMongoQueryRunner(TpchTable.getTables(), ImmutableMap.of());
     }
 
     @BeforeClass

From b7609053da6a38336a1107d4339fdfe138a77672 Mon Sep 17 00:00:00 2001
From: dnskr 
Date: Sun, 31 Aug 2025 21:23:01 +0200
Subject: [PATCH 036/113] Fix minor documentation issues

---
 presto-docs/src/main/sphinx/connector/hive.rst               | 1 +
 presto-docs/src/main/sphinx/connector/postgresql.rst         | 1 +
 presto-docs/src/main/sphinx/plugin/native-sidecar-plugin.rst | 3 ++-
 presto-docs/src/main/sphinx/presto_cpp/properties.rst        | 1 +
 4 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/presto-docs/src/main/sphinx/connector/hive.rst b/presto-docs/src/main/sphinx/connector/hive.rst
index 7208faaff0b39..cfdb309eeaf9f 100644
--- a/presto-docs/src/main/sphinx/connector/hive.rst
+++ b/presto-docs/src/main/sphinx/connector/hive.rst
@@ -229,6 +229,7 @@ Property Name                                            Description
                                                          metadata, instead of their ordinal position. Also toggleable 
                                                          through the ``hive.orc_use_column_names`` session property.
 ======================================================== ============================================================ ============
+
 .. _constructor: https://github.com/apache/hadoop/blob/02a9190af5f8264e25966a80c8f9ea9bb6677899/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/conf/Configuration.java#L844-L875
 
 Avro Configuration Properties
diff --git a/presto-docs/src/main/sphinx/connector/postgresql.rst b/presto-docs/src/main/sphinx/connector/postgresql.rst
index 256f7f56e72bb..065b54ec11baa 100644
--- a/presto-docs/src/main/sphinx/connector/postgresql.rst
+++ b/presto-docs/src/main/sphinx/connector/postgresql.rst
@@ -149,6 +149,7 @@ The connector maps PostgreSQL types to the corresponding PrestoDB types:
     - ``VARCHAR``
   * - ``GEOGRAPHY``
     - ``VARCHAR``
+
 No other types are supported.
 
 PrestoDB to PostgreSQL type mapping
diff --git a/presto-docs/src/main/sphinx/plugin/native-sidecar-plugin.rst b/presto-docs/src/main/sphinx/plugin/native-sidecar-plugin.rst
index e9cca5dc14438..49440a816288b 100644
--- a/presto-docs/src/main/sphinx/plugin/native-sidecar-plugin.rst
+++ b/presto-docs/src/main/sphinx/plugin/native-sidecar-plugin.rst
@@ -29,6 +29,7 @@ Property Name                                Description
 ============================================ ===================================================================== ==============================
 
 .. _sidecar-worker-properties:
+
 Sidecar worker properties
 ^^^^^^^^^^^^^^^^^^^^^^^^^
 Enable sidecar functionality with:
@@ -68,7 +69,7 @@ Property Name                                Description
 ============================================ ===================================================================== ==============================
 
 Session properties
------------------
+------------------
 
 These properties must be configured in ``etc/session-property-providers/native-worker.properties`` to use the session property provider of the ``NativeSidecarPlugin``.
 
diff --git a/presto-docs/src/main/sphinx/presto_cpp/properties.rst b/presto-docs/src/main/sphinx/presto_cpp/properties.rst
index cdb078a35a6da..2f1819029c21d 100644
--- a/presto-docs/src/main/sphinx/presto_cpp/properties.rst
+++ b/presto-docs/src/main/sphinx/presto_cpp/properties.rst
@@ -75,6 +75,7 @@ alphabetical order.
     - Custom type names are peeled in the coordinator. Only the actual base type is preserved.
     - ``CAST(col AS EnumType)`` is rewritten as ``CAST(col AS )``.
     - ``ENUM_KEY(EnumType)`` is rewritten as ``ELEMENT_AT(MAP(, VARCHAR))``.
+
   This property can only be enabled with native execution.
 
 ``optimizer.optimize-hash-generation``

From 977cabec66c7fb80ca750e7adb8430d4d75ea4d6 Mon Sep 17 00:00:00 2001
From: dnskr 
Date: Sun, 31 Aug 2025 19:01:47 +0200
Subject: [PATCH 037/113] Enable copy button feature explicitly

---
 presto-docs/src/main/sphinx/conf.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/presto-docs/src/main/sphinx/conf.py b/presto-docs/src/main/sphinx/conf.py
index 4f1e7342c5a16..ea71636098d47 100644
--- a/presto-docs/src/main/sphinx/conf.py
+++ b/presto-docs/src/main/sphinx/conf.py
@@ -135,6 +135,7 @@ def get_version():
     'features': [
         'toc.follow',
         'toc.sticky',
+        'content.code.copy',
     ],
     'palette': [
         {

From a5256abcdb323644c61c3550dc2d0d33a5b59a5d Mon Sep 17 00:00:00 2001
From: Pratik Joseph Dabre 
Date: Fri, 29 Aug 2025 14:06:05 -0700
Subject: [PATCH 038/113] [native] Validate sidecar function signatures against
 plugin loaded function signatures at startup

---
 ...BuiltInPluginFunctionNamespaceManager.java |  5 ++
 ...uiltInSpecialFunctionNamespaceManager.java | 10 ++--
 .../metadata/FunctionAndTypeManager.java      |  7 ++-
 .../facebook/presto/server/PrestoServer.java  |  3 ++
 .../server/testing/TestingPrestoServer.java   |  6 +++
 .../NativeFunctionDefinitionProvider.java     | 27 ++++------
 .../NativeFunctionNamespaceManagerConfig.java | 52 +++++++++++++++++++
 .../NativeFunctionNamespaceManagerModule.java |  2 +-
 ...tNativeFunctionNamespaceManagerConfig.java | 51 ++++++++++++++++++
 ...ginLoadedDuplicateSqlInvokedFunctions.java | 24 ++-------
 10 files changed, 144 insertions(+), 43 deletions(-)
 create mode 100644 presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/functionNamespace/NativeFunctionNamespaceManagerConfig.java
 create mode 100644 presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/functionNamespace/TestNativeFunctionNamespaceManagerConfig.java

diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInPluginFunctionNamespaceManager.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInPluginFunctionNamespaceManager.java
index 75e0d3415d935..e0cf7dc25472f 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInPluginFunctionNamespaceManager.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInPluginFunctionNamespaceManager.java
@@ -31,6 +31,11 @@ public BuiltInPluginFunctionNamespaceManager(FunctionAndTypeManager functionAndT
         super(functionAndTypeManager);
     }
 
+    public void triggerConflictCheckWithBuiltInFunctions()
+    {
+        checkForNamingConflicts(this.getFunctionsFromDefaultNamespace());
+    }
+
     @Override
     public synchronized void registerBuiltInSpecialFunctions(List functions)
     {
diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInSpecialFunctionNamespaceManager.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInSpecialFunctionNamespaceManager.java
index 451e209b69da6..964b0e9a0891c 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInSpecialFunctionNamespaceManager.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInSpecialFunctionNamespaceManager.java
@@ -58,7 +58,7 @@ public abstract class BuiltInSpecialFunctionNamespaceManager
 {
     protected volatile FunctionMap functions = new FunctionMap();
     private final FunctionAndTypeManager functionAndTypeManager;
-    private final Supplier cachedFunctions =
+    protected final Supplier cachedFunctions =
             Suppliers.memoize(this::createFunctionMap);
     private final LoadingCache specializedFunctionKeyCache;
     private final LoadingCache specializedScalarCache;
@@ -212,12 +212,16 @@ public ScalarFunctionImplementation getScalarFunctionImplementation(FunctionHand
         return getScalarFunctionImplementation(((BuiltInFunctionHandle) functionHandle).getSignature());
     }
 
-    private synchronized FunctionMap createFunctionMap()
+    protected Collection getFunctionsFromDefaultNamespace()
     {
         Optional> functionNamespaceManager =
                 functionAndTypeManager.getServingFunctionNamespaceManager(functionAndTypeManager.getDefaultNamespace());
         checkArgument(functionNamespaceManager.isPresent(), "Cannot find function namespace for catalog '%s'", functionAndTypeManager.getDefaultNamespace().getCatalogName());
-        checkForNamingConflicts(functionNamespaceManager.get().listFunctions(Optional.empty(), Optional.empty()));
+        return functionNamespaceManager.get().listFunctions(Optional.empty(), Optional.empty());
+    }
+
+    private synchronized FunctionMap createFunctionMap()
+    {
         return functions;
     }
 
diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionAndTypeManager.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionAndTypeManager.java
index 85dffe5487b7c..bb0c6cf311eb6 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionAndTypeManager.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionAndTypeManager.java
@@ -979,7 +979,12 @@ public SpecializedFunctionKey getSpecializedFunctionKey(Signature signature, Col
         return builtInTypeAndFunctionNamespaceManager.doGetSpecializedFunctionKeyForMagicLiteralFunctions(signature, this);
     }
 
-    public CatalogSchemaName configureDefaultNamespace(String defaultNamespacePrefixString)
+    public BuiltInPluginFunctionNamespaceManager getBuiltInPluginFunctionNamespaceManager()
+    {
+        return builtInPluginFunctionNamespaceManager;
+    }
+
+    private CatalogSchemaName configureDefaultNamespace(String defaultNamespacePrefixString)
     {
         if (!defaultNamespacePrefixString.matches(DEFAULT_NAMESPACE_PREFIX_PATTERN.pattern())) {
             throw new PrestoException(GENERIC_USER_ERROR, format("Default namespace prefix string should be in the form of 'catalog.schema', found: %s", defaultNamespacePrefixString));
diff --git a/presto-main/src/main/java/com/facebook/presto/server/PrestoServer.java b/presto-main/src/main/java/com/facebook/presto/server/PrestoServer.java
index 6462e9b8fb57b..ab6f45c8174ae 100644
--- a/presto-main/src/main/java/com/facebook/presto/server/PrestoServer.java
+++ b/presto-main/src/main/java/com/facebook/presto/server/PrestoServer.java
@@ -209,6 +209,9 @@ public void run()
             injector.getInstance(ClientRequestFilterManager.class).loadClientRequestFilters();
             injector.getInstance(ExpressionOptimizerManager.class).loadExpressionOptimizerFactories();
 
+            injector.getInstance(FunctionAndTypeManager.class)
+                    .getBuiltInPluginFunctionNamespaceManager().triggerConflictCheckWithBuiltInFunctions();
+
             startAssociatedProcesses(injector);
 
             injector.getInstance(Announcer.class).start();
diff --git a/presto-main/src/main/java/com/facebook/presto/server/testing/TestingPrestoServer.java b/presto-main/src/main/java/com/facebook/presto/server/testing/TestingPrestoServer.java
index 43c574058f2bb..e58509da40263 100644
--- a/presto-main/src/main/java/com/facebook/presto/server/testing/TestingPrestoServer.java
+++ b/presto-main/src/main/java/com/facebook/presto/server/testing/TestingPrestoServer.java
@@ -548,6 +548,12 @@ public void installCoordinatorPlugin(CoordinatorPlugin plugin)
         pluginManager.installCoordinatorPlugin(plugin);
     }
 
+    public void triggerConflictCheckWithBuiltInFunctions()
+    {
+        metadata.getFunctionAndTypeManager()
+                .getBuiltInPluginFunctionNamespaceManager().triggerConflictCheckWithBuiltInFunctions();
+    }
+
     public DispatchManager getDispatchManager()
     {
         return dispatchManager;
diff --git a/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/functionNamespace/NativeFunctionDefinitionProvider.java b/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/functionNamespace/NativeFunctionDefinitionProvider.java
index a0cd355cfe175..677752dd4d2c7 100644
--- a/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/functionNamespace/NativeFunctionDefinitionProvider.java
+++ b/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/functionNamespace/NativeFunctionDefinitionProvider.java
@@ -14,26 +14,24 @@
 package com.facebook.presto.sidecar.functionNamespace;
 
 import com.facebook.airlift.http.client.HttpClient;
-import com.facebook.airlift.http.client.HttpUriBuilder;
 import com.facebook.airlift.http.client.Request;
 import com.facebook.airlift.json.JsonCodec;
 import com.facebook.airlift.log.Logger;
 import com.facebook.presto.functionNamespace.JsonBasedUdfFunctionMetadata;
 import com.facebook.presto.functionNamespace.UdfFunctionSignatureMap;
 import com.facebook.presto.sidecar.ForSidecarInfo;
-import com.facebook.presto.spi.Node;
 import com.facebook.presto.spi.NodeManager;
 import com.facebook.presto.spi.PrestoException;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableMap;
 import com.google.inject.Inject;
 
-import java.net.URI;
 import java.util.List;
 import java.util.Map;
 
 import static com.facebook.airlift.http.client.JsonResponseHandler.createJsonResponseHandler;
 import static com.facebook.airlift.http.client.Request.Builder.prepareGet;
+import static com.facebook.presto.builtin.tools.NativeSidecarFunctionRegistryTool.getSidecarLocationOnStartup;
 import static com.facebook.presto.spi.StandardErrorCode.INVALID_ARGUMENTS;
 import static java.util.Objects.requireNonNull;
 
@@ -42,27 +40,29 @@ public class NativeFunctionDefinitionProvider
 {
     private static final Logger log = Logger.get(NativeFunctionDefinitionProvider.class);
     private final JsonCodec>> nativeFunctionSignatureMapJsonCodec;
-    private final NodeManager nodeManager;
     private final HttpClient httpClient;
-    private static final String FUNCTION_SIGNATURES_ENDPOINT = "/v1/functions";
+    private final NativeFunctionNamespaceManagerConfig config;
 
     @Inject
     public NativeFunctionDefinitionProvider(
             @ForSidecarInfo HttpClient httpClient,
             JsonCodec>> nativeFunctionSignatureMapJsonCodec,
-            NodeManager nodeManager)
+            NativeFunctionNamespaceManagerConfig config)
     {
         this.nativeFunctionSignatureMapJsonCodec =
                 requireNonNull(nativeFunctionSignatureMapJsonCodec, "nativeFunctionSignatureMapJsonCodec is null");
-        this.nodeManager = requireNonNull(nodeManager, "nodeManager is null");
-        this.httpClient = requireNonNull(httpClient, "typeManager is null");
+        this.httpClient = requireNonNull(httpClient, "httpClient is null");
+        this.config = requireNonNull(config, "config is null");
     }
 
     @Override
     public UdfFunctionSignatureMap getUdfDefinition(NodeManager nodeManager)
     {
         try {
-            Request request = prepareGet().setUri(getSidecarLocation()).build();
+            Request request =
+                    prepareGet().setUri(
+                            getSidecarLocationOnStartup(
+                                    nodeManager, config.getSidecarNumRetries(), config.getSidecarRetryDelay().toMillis())).build();
             Map> nativeFunctionSignatureMap = httpClient.execute(request, createJsonResponseHandler(nativeFunctionSignatureMapJsonCodec));
             return new UdfFunctionSignatureMap(ImmutableMap.copyOf(nativeFunctionSignatureMap));
         }
@@ -71,15 +71,6 @@ public UdfFunctionSignatureMap getUdfDefinition(NodeManager nodeManager)
         }
     }
 
-    private URI getSidecarLocation()
-    {
-        Node sidecarNode = nodeManager.getSidecarNode();
-        return HttpUriBuilder
-                .uriBuilderFrom(sidecarNode.getHttpUri())
-                .appendPath(FUNCTION_SIGNATURES_ENDPOINT)
-                .build();
-    }
-
     @VisibleForTesting
     public HttpClient getHttpClient()
     {
diff --git a/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/functionNamespace/NativeFunctionNamespaceManagerConfig.java b/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/functionNamespace/NativeFunctionNamespaceManagerConfig.java
new file mode 100644
index 0000000000000..38aeb75d1e5aa
--- /dev/null
+++ b/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/functionNamespace/NativeFunctionNamespaceManagerConfig.java
@@ -0,0 +1,52 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.sidecar.functionNamespace;
+
+import com.facebook.airlift.configuration.Config;
+import com.facebook.airlift.configuration.ConfigDescription;
+import com.facebook.airlift.units.Duration;
+
+import static java.util.concurrent.TimeUnit.MINUTES;
+
+public class NativeFunctionNamespaceManagerConfig
+{
+    private int sidecarNumRetries = 8;
+    private Duration sidecarRetryDelay = new Duration(1, MINUTES);
+
+    public int getSidecarNumRetries()
+    {
+        return sidecarNumRetries;
+    }
+
+    @Config("sidecar.num-retries")
+    @ConfigDescription("Max times to retry fetching sidecar node")
+    public NativeFunctionNamespaceManagerConfig setSidecarNumRetries(int sidecarNumRetries)
+    {
+        this.sidecarNumRetries = sidecarNumRetries;
+        return this;
+    }
+
+    public Duration getSidecarRetryDelay()
+    {
+        return sidecarRetryDelay;
+    }
+
+    @Config("sidecar.retry-delay")
+    @ConfigDescription("Cooldown period to retry when fetching sidecar node fails")
+    public NativeFunctionNamespaceManagerConfig setSidecarRetryDelay(Duration sidecarRetryDelay)
+    {
+        this.sidecarRetryDelay = sidecarRetryDelay;
+        return this;
+    }
+}
diff --git a/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/functionNamespace/NativeFunctionNamespaceManagerModule.java b/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/functionNamespace/NativeFunctionNamespaceManagerModule.java
index 50f36269456d3..bff0a8d91d820 100644
--- a/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/functionNamespace/NativeFunctionNamespaceManagerModule.java
+++ b/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/functionNamespace/NativeFunctionNamespaceManagerModule.java
@@ -37,7 +37,6 @@ public class NativeFunctionNamespaceManagerModule
         implements Module
 {
     private final String catalogName;
-
     private final NodeManager nodeManager;
     private final FunctionMetadataManager functionMetadataManager;
 
@@ -55,6 +54,7 @@ public void configure(Binder binder)
 
         configBinder(binder).bindConfig(SqlInvokedFunctionNamespaceManagerConfig.class);
         configBinder(binder).bindConfig(SqlFunctionLanguageConfig.class);
+        configBinder(binder).bindConfig(NativeFunctionNamespaceManagerConfig.class);
         binder.bind(FunctionDefinitionProvider.class).to(NativeFunctionDefinitionProvider.class).in(SINGLETON);
         binder.bind(new TypeLiteral>>>() {})
                 .toInstance(new JsonCodecFactory().mapJsonCodec(String.class, listJsonCodec(JsonBasedUdfFunctionMetadata.class)));
diff --git a/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/functionNamespace/TestNativeFunctionNamespaceManagerConfig.java b/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/functionNamespace/TestNativeFunctionNamespaceManagerConfig.java
new file mode 100644
index 0000000000000..feee5f6ce8c9c
--- /dev/null
+++ b/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/functionNamespace/TestNativeFunctionNamespaceManagerConfig.java
@@ -0,0 +1,51 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.sidecar.functionNamespace;
+
+import com.facebook.airlift.units.Duration;
+import com.google.common.collect.ImmutableMap;
+import org.testng.annotations.Test;
+
+import java.util.Map;
+
+import static com.facebook.airlift.configuration.testing.ConfigAssertions.assertFullMapping;
+import static com.facebook.airlift.configuration.testing.ConfigAssertions.assertRecordedDefaults;
+import static com.facebook.airlift.configuration.testing.ConfigAssertions.recordDefaults;
+import static java.util.concurrent.TimeUnit.MINUTES;
+
+public class TestNativeFunctionNamespaceManagerConfig
+{
+    @Test
+    public void testDefaults()
+    {
+        assertRecordedDefaults(recordDefaults(NativeFunctionNamespaceManagerConfig.class)
+                .setSidecarNumRetries(8)
+                .setSidecarRetryDelay(new Duration(1, MINUTES)));
+    }
+
+    @Test
+    public void testExplicitPropertyMappings()
+    {
+        Map properties = new ImmutableMap.Builder()
+                .put("sidecar.num-retries", "15")
+                .put("sidecar.retry-delay", "5m")
+                .build();
+
+        NativeFunctionNamespaceManagerConfig expected = new NativeFunctionNamespaceManagerConfig()
+                .setSidecarNumRetries(15)
+                .setSidecarRetryDelay(new Duration(5, MINUTES));
+
+        assertFullMapping(properties, expected);
+    }
+}
diff --git a/presto-tests/src/test/java/com/facebook/presto/functions/TestPluginLoadedDuplicateSqlInvokedFunctions.java b/presto-tests/src/test/java/com/facebook/presto/functions/TestPluginLoadedDuplicateSqlInvokedFunctions.java
index 4ab7e80a82906..eba96c2001ce3 100644
--- a/presto-tests/src/test/java/com/facebook/presto/functions/TestPluginLoadedDuplicateSqlInvokedFunctions.java
+++ b/presto-tests/src/test/java/com/facebook/presto/functions/TestPluginLoadedDuplicateSqlInvokedFunctions.java
@@ -22,12 +22,9 @@
 import org.testng.annotations.Test;
 
 import java.util.Set;
-import java.util.regex.Pattern;
 
-import static com.facebook.presto.metadata.BuiltInTypeAndFunctionNamespaceManager.JAVA_BUILTIN_NAMESPACE;
 import static com.facebook.presto.testing.TestingSession.testSessionBuilder;
-import static java.lang.String.format;
-import static org.testng.Assert.fail;
+import static org.testng.Assert.assertThrows;
 
 public class TestPluginLoadedDuplicateSqlInvokedFunctions
 {
@@ -45,20 +42,6 @@ public void setup()
                 .build());
     }
 
-    public void assertInvalidFunction(String expr, String exceptionPattern)
-    {
-        try {
-            client.execute("SELECT " + expr);
-            fail("Function expected to fail but not");
-        }
-        catch (Exception e) {
-            if (!(e.getMessage().matches(exceptionPattern))) {
-                fail(format("Expected exception message '%s' to match '%s' but not",
-                        e.getMessage(), exceptionPattern));
-            }
-        }
-    }
-
     private static class TestDuplicateFunctionsPlugin
             implements Plugin
     {
@@ -71,10 +54,11 @@ public Set> getSqlInvokedFunctions()
         }
     }
 
+    // As soon as we trigger the conflict check with the built-in functions, an error will be thrown if duplicate signatures are found.
+    // In `PrestoServer.java` this conflict check is triggered implicitly.
     @Test
     public void testDuplicateFunctionsLoaded()
     {
-        assertInvalidFunction(JAVA_BUILTIN_NAMESPACE + ".modulo(10,3)",
-                Pattern.quote(format("java.lang.IllegalArgumentException: Function already registered: %s.array_intersect(array(array(T))):array(T)", JAVA_BUILTIN_NAMESPACE)));
+        assertThrows(IllegalArgumentException.class, () -> server.triggerConflictCheckWithBuiltInFunctions());
     }
 }

From 4cf74fc3469bae5b8e0ad3a9e6a42006074434b7 Mon Sep 17 00:00:00 2001
From: dnskr 
Date: Sun, 31 Aug 2025 17:39:29 +0200
Subject: [PATCH 039/113] Generalize the doc page explaining how to deploy
 Presto with Homebrew

---
 .../main/sphinx/installation/deploy-brew.rst  | 177 ++++--------------
 1 file changed, 35 insertions(+), 142 deletions(-)

diff --git a/presto-docs/src/main/sphinx/installation/deploy-brew.rst b/presto-docs/src/main/sphinx/installation/deploy-brew.rst
index 51cb3008da70e..bed202f4bd4d0 100644
--- a/presto-docs/src/main/sphinx/installation/deploy-brew.rst
+++ b/presto-docs/src/main/sphinx/installation/deploy-brew.rst
@@ -1,37 +1,31 @@
-=====================================
-Deploy Presto on a Mac using Homebrew
-=====================================
+============================
+Deploy Presto using Homebrew
+============================
 
-- If you are deploying Presto on an Intel Mac, see `Deploy Presto on an Intel Mac using Homebrew`_.
+This guide explains how to install and get started with Presto on macOS, Linux or WSL2 using the Homebrew package manager.
 
-- If you are deploying Presto on an Apple Silicon Mac that has an M1 or M2 chip, see `Deploy Presto on an Apple Silicon Mac using Homebrew`_. 
+Prerequisites
+-------------
 
-Deploy Presto on an Intel Mac using Homebrew
---------------------------------------------
-*Note*: These steps were developed and tested on Mac OS X on Intel. These steps will not work with Apple Silicon (M1 or M2) chips.
-
-Following these steps, you will:
-
-- install the Presto service and CLI on an Intel Mac using `Homebrew `_
-- start and stop the Presto service
-- start the Presto CLI
+`Homebrew `_ installed.
 
 Install Presto
-^^^^^^^^^^^^^^
+--------------
 
-Follow these steps to install Presto on an Intel Mac using `Homebrew `_. 
+Run the following command to install the latest version of Presto using the `Homebrew Formulae `_:
 
-1. If you do not have brew installed, run the following command:
+.. code-block:: none
 
-   ``/bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/master/install.sh)"``
+   brew install prestodb
 
-2. To install Presto, run the following command:
+Homebrew installs packages in the ``Cellar`` directory, which can be found with this command:
 
-   ``brew install prestodb``
+.. code-block:: none
 
-   Presto is installed in the directory */usr/local/Cellar/prestodb/.* 
+   brew --cellar
 
-The following files are created in the *libexec/etc* directory in the Presto install directory:
+The directory ``$(brew --cellar)/prestodb//libexec`` contains the Presto files used to run and configure the service.
+For example, the ``etc`` directory within the Presto installation contains the following default configuration files:
 
 - node.properties
 - jvm.config
@@ -39,35 +33,28 @@ The following files are created in the *libexec/etc* directory in the Presto ins
 - log.properties
 - catalog/jmx.properties
 
-For example, the full path to the node.properties file is */usr/local/Cellar/prestodb//libexec/etc/node.properties*. 
-
-The Presto CLI is installed in the *bin* directory of the Presto install directory: */usr/local/Cellar/prestodb//bin*.
-
-The executables are added to */usr/local/bin* path and should be available as part of $PATH.
-
 Start and Stop Presto
-^^^^^^^^^^^^^^^^^^^^^
+---------------------
 
-To start Presto, use the ``presto-server`` helper script. 
+Presto is installed with the ``presto-server`` helper script, which simplifies managing the cluster.
+For example, run the following command to start the Presto service in the foreground:
 
-To start the Presto service in the background, run the following command: 
-
-``presto-server start``
+.. code-block:: none
 
-To start the Presto service in the foreground, run the following command:
+   presto-server run
 
-``presto-server run``
+To stop Presto from running in the foreground, press ``Ctrl + C`` until the terminal prompt appears, or close the terminal.
 
-To stop the Presto service in the background, run the following command:
+For more available commands and options, use help:
 
-``presto-server stop``
+.. code-block:: none
 
-To stop the Presto service in the foreground, close the terminal or select Ctrl + C until the terminal prompt is shown. 
+   presto-server --help
 
 Open the Presto Console
-^^^^^^^^^^^^^^^^^^^^^^^
+-----------------------
 
-After starting Presto, you can access the web UI at the default port ``8080`` using the following link in a browser:
+After starting the service, Presto Console can be accessible at the default port ``8080`` using the following link in a browser:
 
 .. code-block:: none
 
@@ -79,117 +66,23 @@ After starting Presto, you can access the web UI at the default port ``8080`` us
 For more information about the Presto Console, see :doc:`/clients/presto-console`.
 
 Start the Presto CLI
-^^^^^^^^^^^^^^^^^^^^
+--------------------
 
 The Presto CLI is a terminal-based interactive shell for running queries, and is a
 `self-executing `_
 JAR file that acts like a normal UNIX executable.
 
-The Presto CLI is installed in the *bin* directory of the Presto install directory: */usr/local/Cellar/prestodb//bin*.
-
-To run the Presto CLI, run the following command:
-
-``presto``
-
-The Presto CLI starts and displays the prompt ``presto>``. 
-
-For more information, see :doc:`/clients/presto-cli`.
-
-Deploy Presto on an Apple Silicon Mac using Homebrew 
-----------------------------------------------------
-*Note*: These steps were developed and tested on Mac OS X on Apple Silicon. These steps will not work with Intel chips.
-
-Following these steps, you will:
-
-- install the Presto service and CLI on an Apple Silicon Mac using `Homebrew `_
-- start and stop the Presto service
-- start the Presto CLI
-
-Install Presto
-^^^^^^^^^^^^^^
-
-Follow these steps to install Presto on an Apple Silicon Mac using `Homebrew `_. 
-
-1. If you do not have brew installed, run the following command:
-
-   ``arch -x86_64 /bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/master/install.sh)"``
-
-   This installs Homebrew into ``/usr/local/bin``. 
-   
-   *Note*: The default installation of Homebrew on Apple Silicon is to ``/opt/homebrew``.
-
-2. To allow the shell to look for Homebrew in ``/usr/local/bin`` before it looks for Homebrew in ``/opt/homebrew``, run the following command:
-
-   ``export PATH=/usr/local/bin:$PATH``
-
-3. To install Presto, run the following command:
-
-   ``arch -x86_64 brew install prestodb``
-
-   Presto is installed in the directory */usr/local/Cellar/prestodb/.* The executables ``presto`` 
-   and ``presto-server`` are installed in ``/usr/local/bin/``.
-
-The following files are created in the *libexec/etc* directory in the Presto install directory:
-
-- node.properties
-- jvm.config
-- config.properties
-- log.properties
-- catalog/jmx.properties
-
-For example, the full path to the node.properties file is */usr/local/Cellar/prestodb//libexec/etc/node.properties*. 
-
-The Presto CLI is installed in the *bin* directory of the Presto install directory: */usr/local/Cellar/prestodb//bin*.
-
-The executables are added to */usr/local/bin* path and should be available as part of $PATH.
-
-Start and Stop Presto
-^^^^^^^^^^^^^^^^^^^^^
-
-To start Presto, use the ``presto-server`` helper script. 
-
-To start the Presto service in the background, run the following command: 
-
-``arch -x86_64 presto-server start``
-
-To start the Presto service in the foreground, run the following command:
-
-``arch -x86_64 presto-server run``
-
-To stop the Presto service in the background, run the following command:
-
-``presto-server stop``
-
-To stop the Presto service in the foreground, close the terminal or select Ctrl + C until the terminal prompt is shown. 
-
-Open the Presto Console
-^^^^^^^^^^^^^^^^^^^^^^^
-
-After starting Presto, you can access the web UI at the default port ``8080`` using the following link in a browser:
+The Presto CLI is installed in the directory ``$(brew --cellar)/prestodb//bin``.
+To run the Presto CLI, use the following command:
 
 .. code-block:: none
 
-    http://localhost:8080
-
-.. figure:: ../images/presto_console.png
-   :align: center
-
-For more information about the Presto Console, see :doc:`/clients/presto-console`.
-
-Start the Presto CLI
-^^^^^^^^^^^^^^^^^^^^
-
-The Presto CLI is a terminal-based interactive shell for running queries, and is a
-`self-executing `_
-JAR file that acts like a normal UNIX executable.
-
-The Presto CLI is installed in the *bin* directory of the Presto install directory: */usr/local/Cellar/prestodb//bin*.
-The executable ``presto`` is installed in ``/usr/local/bin/``.
+   presto
 
-To run the Presto CLI, run the following command:
+The Presto CLI starts and displays its prompt:
 
-``presto``
+.. code-block:: none
 
-The Presto CLI starts and displays the prompt ``presto>``. 
+   presto>
 
-For more information, see :doc:`/clients/presto-cli`.
\ No newline at end of file
+For more information, see :doc:`/clients/presto-cli`.

From 27280fcab433068bb6d028c1e98bd27822120cb2 Mon Sep 17 00:00:00 2001
From: Pratik Joseph Dabre 
Date: Tue, 2 Sep 2025 12:29:02 -0700
Subject: [PATCH 040/113] Extract all inlined sql invoked functions into a new
 plugin presto-sql-invoked-functions-plugin

---
 pom.xml                                       |  1 +
 presto-hive/pom.xml                           |  6 ++++
 .../TestDistributedQueriesSingleNode.java     |  5 ++-
 .../hive/TestHiveDistributedNanQueries.java   |  6 +++-
 .../hive/TestHiveDistributedQueries.java      |  5 ++-
 ...tedQueriesWithExchangeMaterialization.java |  5 ++-
 ...tedQueriesWithOptimizedRepartitioning.java |  5 ++-
 ...stHiveDistributedQueriesWithThriftRpc.java |  5 ++-
 .../TestHivePushdownDistributedQueries.java   |  5 ++-
 .../hive/TestLambdaSubfieldPruning.java       |  2 ++
 .../hive/TestParquetDistributedQueries.java   |  5 ++-
 .../i18n/functions/TestMyanmarFunctions.java  |  8 +++++
 presto-main-base/pom.xml                      |  9 +++++
 ...uiltInSpecialFunctionNamespaceManager.java | 13 ++-----
 ...uiltInTypeAndFunctionNamespaceManager.java | 11 ------
 .../presto/metadata/FunctionListBuilder.java  | 14 --------
 .../scalar/ArrayIntersectFunction.java        | 12 -------
 .../facebook/presto/testing/QueryRunner.java  |  5 +++
 .../scalar/AbstractTestFunctions.java         |  8 ++++-
 .../operator/scalar/FunctionAssertions.java   | 15 ++++++--
 ...ossJoinWithArrayNotContainsToAntiJoin.java | 11 ++++++
 ...nWithArrayContainsToEquiJoinCondition.java | 11 ++++++
 presto-main/etc/config.properties             |  3 +-
 .../presto/ml/AbstractTestMLFunctions.java    |  8 +++++
 presto-native-execution/pom.xml               | 13 +++++++
 .../TestPrestoNativeArrayFunctionQueries.java | 15 +++++---
 .../TestPrestoNativeGeneralQueriesJSON.java   |  9 +++--
 .../TestPrestoNativeGeneralQueriesThrift.java |  9 +++--
 .../TestDistributedEngineOnlyQueries.java     | 12 +++++++
 .../presto/pinot/udf/TestPinotFunctions.java  |  8 +++++
 presto-server/src/main/provisio/presto.xml    |  6 ++++
 presto-sql-invoked-functions-plugin/pom.xml   | 29 +++++++++++++++
 .../scalar/sql/ArrayIntersectFunction.java    | 35 ++++++++++++++++++
 .../presto}/scalar/sql/ArraySqlFunctions.java |  2 +-
 .../scalar/sql/MapNormalizeFunction.java      |  2 +-
 .../presto}/scalar/sql/MapSqlFunctions.java   |  2 +-
 .../scalar/sql/SimpleSamplingPercent.java     |  2 +-
 .../scalar/sql/SqlInvokedFunctionsPlugin.java | 36 +++++++++++++++++++
 .../scalar/sql/StringSqlFunctions.java        |  2 +-
 .../functions/TestTeradataDateFunctions.java  |  4 ++-
 .../functions/TestTeradataFunctions.java      |  8 +++++
 presto-tests/pom.xml                          |  7 ++++
 .../presto/tests/DistributedQueryRunner.java  |  8 +++++
 .../TestBuiltInConnectorFunctions.java        |  4 +--
 ...ginLoadedDuplicateSqlInvokedFunctions.java | 11 ++++++
 .../TestDistributedEngineOnlyQueries.java     |  5 ++-
 .../tests/TestDistributedSpilledQueries.java  |  2 ++
 ...tributedSpilledQueriesWithTempStorage.java |  3 ++
 .../presto/tests/TestLocalQueries.java        |  3 ++
 .../tests/TestQueryPlanDeterminism.java       |  2 ++
 .../presto/tests/TestSqlFunctions.java        |  3 +-
 .../tests/TestVerboseOptimizerInfo.java       |  3 +-
 52 files changed, 343 insertions(+), 80 deletions(-)
 create mode 100644 presto-sql-invoked-functions-plugin/pom.xml
 create mode 100644 presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/ArrayIntersectFunction.java
 rename {presto-main-base/src/main/java/com/facebook/presto/operator => presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto}/scalar/sql/ArraySqlFunctions.java (99%)
 rename {presto-main-base/src/main/java/com/facebook/presto/operator => presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto}/scalar/sql/MapNormalizeFunction.java (96%)
 rename {presto-main-base/src/main/java/com/facebook/presto/operator => presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto}/scalar/sql/MapSqlFunctions.java (99%)
 rename {presto-main-base/src/main/java/com/facebook/presto/operator => presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto}/scalar/sql/SimpleSamplingPercent.java (96%)
 create mode 100644 presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/SqlInvokedFunctionsPlugin.java
 rename {presto-main-base/src/main/java/com/facebook/presto/operator => presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto}/scalar/sql/StringSqlFunctions.java (97%)

diff --git a/pom.xml b/pom.xml
index f979056f76891..c792fc783d343 100644
--- a/pom.xml
+++ b/pom.xml
@@ -221,6 +221,7 @@
         presto-function-server
         presto-router-example-plugin-scheduler
         presto-plan-checker-router-plugin
+        presto-sql-invoked-functions-plugin
         presto-spark-classloader-spark${dep.pos.classloader.module-name.suffix}
     
 
diff --git a/presto-hive/pom.xml b/presto-hive/pom.xml
index 0e6cdf23cb5b5..4b855fb5df084 100644
--- a/presto-hive/pom.xml
+++ b/presto-hive/pom.xml
@@ -494,6 +494,12 @@
                 
             
         
+        
+            com.facebook.presto
+            presto-sql-invoked-functions-plugin
+            ${project.version}
+            test
+        
     
 
     
diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestDistributedQueriesSingleNode.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestDistributedQueriesSingleNode.java
index bfaec5f22852d..e3328a940c71b 100644
--- a/presto-hive/src/test/java/com/facebook/presto/hive/TestDistributedQueriesSingleNode.java
+++ b/presto-hive/src/test/java/com/facebook/presto/hive/TestDistributedQueriesSingleNode.java
@@ -13,6 +13,7 @@
  */
 package com.facebook.presto.hive;
 
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.testing.QueryRunner;
 import com.facebook.presto.tests.AbstractTestDistributedQueries;
 import com.google.common.collect.ImmutableMap;
@@ -32,11 +33,13 @@ protected QueryRunner createQueryRunner()
     {
         ImmutableMap.Builder coordinatorProperties = ImmutableMap.builder();
         coordinatorProperties.put("single-node-execution-enabled", "true");
-        return HiveQueryRunner.createQueryRunner(
+        QueryRunner queryRunner = HiveQueryRunner.createQueryRunner(
                 getTables(),
                 ImmutableMap.of(),
                 coordinatorProperties.build(),
                 Optional.empty());
+        queryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
+        return queryRunner;
     }
 
     @Override
diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedNanQueries.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedNanQueries.java
index fa2738c63fd72..fc74dbb2f4433 100644
--- a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedNanQueries.java
+++ b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedNanQueries.java
@@ -14,6 +14,7 @@
 
 package com.facebook.presto.hive;
 
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.testing.QueryRunner;
 import com.facebook.presto.tests.AbstractTestNanQueries;
 import com.google.common.collect.ImmutableList;
@@ -28,6 +29,9 @@ public class TestHiveDistributedNanQueries
     protected QueryRunner createQueryRunner()
             throws Exception
     {
-        return HiveQueryRunner.createQueryRunner(ImmutableList.of(), ImmutableMap.of("use-new-nan-definition", "true"), ImmutableMap.of(), Optional.empty());
+        QueryRunner queryRunner =
+                HiveQueryRunner.createQueryRunner(ImmutableList.of(), ImmutableMap.of("use-new-nan-definition", "true"), ImmutableMap.of(), Optional.empty());
+        queryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
+        return queryRunner;
     }
 }
diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueries.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueries.java
index 7552535ae9a4f..134cf4b101f1f 100644
--- a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueries.java
+++ b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueries.java
@@ -15,6 +15,7 @@
 
 import com.facebook.presto.Session;
 import com.facebook.presto.hive.TestHiveEventListenerPlugin.TestingHiveEventListener;
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.spi.QueryId;
 import com.facebook.presto.spi.eventlistener.EventListener;
 import com.facebook.presto.testing.MaterializedResult;
@@ -59,7 +60,9 @@ public class TestHiveDistributedQueries
     protected QueryRunner createQueryRunner()
             throws Exception
     {
-        return HiveQueryRunner.createQueryRunner(getTables());
+        QueryRunner queryRunner = HiveQueryRunner.createQueryRunner(getTables());
+        queryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
+        return queryRunner;
     }
 
     @Override
diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueriesWithExchangeMaterialization.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueriesWithExchangeMaterialization.java
index 226bbf3888b53..a878243e32b91 100644
--- a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueriesWithExchangeMaterialization.java
+++ b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueriesWithExchangeMaterialization.java
@@ -14,6 +14,7 @@
 package com.facebook.presto.hive;
 
 import com.facebook.presto.Session;
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.testing.MaterializedResult;
 import com.facebook.presto.testing.QueryRunner;
 import com.facebook.presto.tests.AbstractTestDistributedQueries;
@@ -40,7 +41,9 @@ public class TestHiveDistributedQueriesWithExchangeMaterialization
     protected QueryRunner createQueryRunner()
             throws Exception
     {
-        return createMaterializingQueryRunner(getTables());
+        QueryRunner queryRunner = createMaterializingQueryRunner(getTables());
+        queryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
+        return queryRunner;
     }
 
     @Test
diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueriesWithOptimizedRepartitioning.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueriesWithOptimizedRepartitioning.java
index d0bf8f6499428..2477f9c31cf46 100644
--- a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueriesWithOptimizedRepartitioning.java
+++ b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueriesWithOptimizedRepartitioning.java
@@ -13,6 +13,7 @@
  */
 package com.facebook.presto.hive;
 
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.testing.QueryRunner;
 import com.facebook.presto.tests.AbstractTestDistributedQueries;
 import com.google.common.collect.ImmutableMap;
@@ -30,13 +31,15 @@ public class TestHiveDistributedQueriesWithOptimizedRepartitioning
     protected QueryRunner createQueryRunner()
             throws Exception
     {
-        return HiveQueryRunner.createQueryRunner(
+        QueryRunner queryRunner = HiveQueryRunner.createQueryRunner(
                 getTables(),
                 ImmutableMap.of(
                         "experimental.optimized-repartitioning", "true",
                         // Use small SerializedPages to force flushing
                         "driver.max-page-partitioning-buffer-size", "10000B"),
                 Optional.empty());
+        queryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
+        return queryRunner;
     }
 
     @Override
diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueriesWithThriftRpc.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueriesWithThriftRpc.java
index 6f877f35a4ce3..1bbc9590dcd50 100644
--- a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueriesWithThriftRpc.java
+++ b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveDistributedQueriesWithThriftRpc.java
@@ -13,6 +13,7 @@
  */
 package com.facebook.presto.hive;
 
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.testing.QueryRunner;
 import com.facebook.presto.tests.AbstractTestDistributedQueries;
 import com.google.common.collect.ImmutableMap;
@@ -30,13 +31,15 @@ public class TestHiveDistributedQueriesWithThriftRpc
     protected QueryRunner createQueryRunner()
             throws Exception
     {
-        return HiveQueryRunner.createQueryRunner(
+        QueryRunner queryRunner = HiveQueryRunner.createQueryRunner(
                 getTables(),
                 ImmutableMap.of(
                         "internal-communication.task-communication-protocol", "THRIFT",
                         "internal-communication.server-info-communication-protocol", "THRIFT"),
                 ImmutableMap.of(),
                 Optional.empty());
+        queryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
+        return queryRunner;
     }
 
     @Override
diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestHivePushdownDistributedQueries.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestHivePushdownDistributedQueries.java
index 28ec3c9ef3491..dc816175ffb12 100644
--- a/presto-hive/src/test/java/com/facebook/presto/hive/TestHivePushdownDistributedQueries.java
+++ b/presto-hive/src/test/java/com/facebook/presto/hive/TestHivePushdownDistributedQueries.java
@@ -13,6 +13,7 @@
  */
 package com.facebook.presto.hive;
 
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.testing.MaterializedResult;
 import com.facebook.presto.testing.QueryRunner;
 import com.facebook.presto.tests.AbstractTestDistributedQueries;
@@ -34,7 +35,7 @@ public class TestHivePushdownDistributedQueries
     protected QueryRunner createQueryRunner()
             throws Exception
     {
-        return HiveQueryRunner.createQueryRunner(
+        QueryRunner queryRunner = HiveQueryRunner.createQueryRunner(
                 getTables(),
                 ImmutableMap.of("experimental.pushdown-subfields-enabled", "true",
                         "experimental.pushdown-dereference-enabled", "true"),
@@ -44,6 +45,8 @@ protected QueryRunner createQueryRunner()
                         "hive.partial_aggregation_pushdown_enabled", "true",
                         "hive.partial_aggregation_pushdown_for_variable_length_datatypes_enabled", "true"),
                 Optional.empty());
+        queryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
+        return queryRunner;
     }
 
     @Override
diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestLambdaSubfieldPruning.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestLambdaSubfieldPruning.java
index a0559b03b0302..b99eff80743ab 100644
--- a/presto-hive/src/test/java/com/facebook/presto/hive/TestLambdaSubfieldPruning.java
+++ b/presto-hive/src/test/java/com/facebook/presto/hive/TestLambdaSubfieldPruning.java
@@ -14,6 +14,7 @@
 package com.facebook.presto.hive;
 
 import com.facebook.presto.Session;
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.testing.QueryRunner;
 import com.facebook.presto.tests.AbstractTestQueryFramework;
 import com.facebook.presto.tests.DistributedQueryRunner;
@@ -126,6 +127,7 @@ private static DistributedQueryRunner createLineItemExTable(DistributedQueryRunn
 
                             "FROM lineitem  \n");
         }
+        queryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
         return queryRunner;
     }
 
diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestParquetDistributedQueries.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestParquetDistributedQueries.java
index d8bb271cebcd0..0aef10a5b4bf0 100644
--- a/presto-hive/src/test/java/com/facebook/presto/hive/TestParquetDistributedQueries.java
+++ b/presto-hive/src/test/java/com/facebook/presto/hive/TestParquetDistributedQueries.java
@@ -14,6 +14,7 @@
 package com.facebook.presto.hive;
 
 import com.facebook.presto.Session;
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.testing.MaterializedResult;
 import com.facebook.presto.testing.QueryRunner;
 import com.facebook.presto.tests.AbstractTestDistributedQueries;
@@ -46,7 +47,7 @@ protected QueryRunner createQueryRunner()
                 .put("hive.partial_aggregation_pushdown_enabled", "true")
                 .put("hive.partial_aggregation_pushdown_for_variable_length_datatypes_enabled", "true")
                 .build();
-        return HiveQueryRunner.createQueryRunner(
+        QueryRunner queryRunner = HiveQueryRunner.createQueryRunner(
                 getTables(),
                 ImmutableMap.of(
                         "experimental.pushdown-subfields-enabled", "true",
@@ -54,6 +55,8 @@ protected QueryRunner createQueryRunner()
                 "sql-standard",
                 parquetProperties,
                 Optional.empty());
+        queryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
+        return queryRunner;
     }
 
     @Test
diff --git a/presto-i18n-functions/src/test/java/com/facebook/presto/i18n/functions/TestMyanmarFunctions.java b/presto-i18n-functions/src/test/java/com/facebook/presto/i18n/functions/TestMyanmarFunctions.java
index 1f5bc00908408..f0259c46f8c5c 100644
--- a/presto-i18n-functions/src/test/java/com/facebook/presto/i18n/functions/TestMyanmarFunctions.java
+++ b/presto-i18n-functions/src/test/java/com/facebook/presto/i18n/functions/TestMyanmarFunctions.java
@@ -14,15 +14,23 @@
 package com.facebook.presto.i18n.functions;
 
 import com.facebook.presto.operator.scalar.AbstractTestFunctions;
+import com.facebook.presto.sql.analyzer.FeaturesConfig;
+import com.facebook.presto.sql.analyzer.FunctionsConfig;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
 
+import static com.facebook.presto.SessionTestUtils.TEST_SESSION;
 import static com.facebook.presto.common.type.VarcharType.VARCHAR;
 import static com.facebook.presto.metadata.FunctionExtractor.extractFunctions;
 
 public class TestMyanmarFunctions
         extends AbstractTestFunctions
 {
+    public TestMyanmarFunctions()
+    {
+        super(TEST_SESSION, new FeaturesConfig(), new FunctionsConfig(), false);
+    }
+
     @BeforeClass
     public void setUp()
     {
diff --git a/presto-main-base/pom.xml b/presto-main-base/pom.xml
index c51c38e73f2f7..b9ff6288fe744 100644
--- a/presto-main-base/pom.xml
+++ b/presto-main-base/pom.xml
@@ -461,9 +461,18 @@
             io.netty
             netty-transport
         
+
         
             com.facebook.presto
             presto-built-in-worker-function-tools
+            ${project.version}
+            test
+        
+
+        
+            com.facebook.presto
+            presto-sql-invoked-functions-plugin
+            ${project.version}
             test
         
     
diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInSpecialFunctionNamespaceManager.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInSpecialFunctionNamespaceManager.java
index 964b0e9a0891c..dcd00723dc890 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInSpecialFunctionNamespaceManager.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInSpecialFunctionNamespaceManager.java
@@ -33,8 +33,6 @@
 import com.facebook.presto.spi.function.SqlFunction;
 import com.facebook.presto.spi.function.SqlInvokedFunction;
 import com.facebook.presto.spi.function.SqlInvokedScalarFunctionImplementation;
-import com.google.common.base.Supplier;
-import com.google.common.base.Suppliers;
 import com.google.common.cache.CacheBuilder;
 import com.google.common.cache.CacheLoader;
 import com.google.common.cache.LoadingCache;
@@ -49,7 +47,6 @@
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Throwables.throwIfInstanceOf;
 import static com.google.common.collect.ImmutableList.toImmutableList;
-import static java.util.Collections.emptyList;
 import static java.util.Objects.requireNonNull;
 import static java.util.concurrent.TimeUnit.HOURS;
 
@@ -58,8 +55,6 @@ public abstract class BuiltInSpecialFunctionNamespaceManager
 {
     protected volatile FunctionMap functions = new FunctionMap();
     private final FunctionAndTypeManager functionAndTypeManager;
-    protected final Supplier cachedFunctions =
-            Suppliers.memoize(this::createFunctionMap);
     private final LoadingCache specializedFunctionKeyCache;
     private final LoadingCache specializedScalarCache;
 
@@ -85,11 +80,7 @@ public BuiltInSpecialFunctionNamespaceManager(FunctionAndTypeManager functionAnd
     @Override
     public Collection getFunctions(Optional transactionHandle, QualifiedObjectName functionName)
     {
-        if (functions.list().isEmpty() ||
-                (!functionName.getCatalogSchemaName().equals(functionAndTypeManager.getDefaultNamespace()))) {
-            return emptyList();
-        }
-        return cachedFunctions.get().get(functionName);
+        return functions.get(functionName);
     }
 
     /**
@@ -98,7 +89,7 @@ public Collection getFunctions(Optional listFunctions(Optional likePattern, Optional escape)
     {
-        return cachedFunctions.get().list();
+        return functions.list();
     }
 
     @Override
diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInTypeAndFunctionNamespaceManager.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInTypeAndFunctionNamespaceManager.java
index 1c3afcf3fa115..40c493fc887c6 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInTypeAndFunctionNamespaceManager.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInTypeAndFunctionNamespaceManager.java
@@ -203,11 +203,6 @@
 import com.facebook.presto.operator.scalar.WilsonInterval;
 import com.facebook.presto.operator.scalar.WordStemFunction;
 import com.facebook.presto.operator.scalar.queryplan.JsonPrestoQueryPlanFunctions;
-import com.facebook.presto.operator.scalar.sql.ArraySqlFunctions;
-import com.facebook.presto.operator.scalar.sql.MapNormalizeFunction;
-import com.facebook.presto.operator.scalar.sql.MapSqlFunctions;
-import com.facebook.presto.operator.scalar.sql.SimpleSamplingPercent;
-import com.facebook.presto.operator.scalar.sql.StringSqlFunctions;
 import com.facebook.presto.operator.window.CumulativeDistributionFunction;
 import com.facebook.presto.operator.window.DenseRankFunction;
 import com.facebook.presto.operator.window.FirstValueFunction;
@@ -994,12 +989,6 @@ private List getBuiltInFunctions(FunctionsConfig function
                 .aggregate(ThetaSketchAggregationFunction.class)
                 .scalars(ThetaSketchFunctions.class)
                 .function(MergeTDigestFunction.MERGE)
-                .sqlInvokedScalar(MapNormalizeFunction.class)
-                .sqlInvokedScalars(ArraySqlFunctions.class)
-                .sqlInvokedScalars(ArrayIntersectFunction.class)
-                .sqlInvokedScalars(MapSqlFunctions.class)
-                .sqlInvokedScalars(SimpleSamplingPercent.class)
-                .sqlInvokedScalars(StringSqlFunctions.class)
                 .scalar(DynamicFilterPlaceholderFunction.class)
                 .scalars(EnumCasts.class)
                 .scalars(LongEnumOperators.class)
diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionListBuilder.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionListBuilder.java
index f3dab976902a9..799dd5f8074fe 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionListBuilder.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionListBuilder.java
@@ -15,7 +15,6 @@
 
 import com.facebook.presto.operator.scalar.annotations.CodegenScalarFromAnnotationsParser;
 import com.facebook.presto.operator.scalar.annotations.ScalarFromAnnotationsParser;
-import com.facebook.presto.operator.scalar.annotations.SqlInvokedScalarFromAnnotationsParser;
 import com.facebook.presto.operator.window.WindowAnnotationsParser;
 import com.facebook.presto.spi.function.SqlFunction;
 import com.facebook.presto.spi.function.WindowFunction;
@@ -24,7 +23,6 @@
 import java.util.ArrayList;
 import java.util.List;
 
-import static com.facebook.presto.metadata.BuiltInTypeAndFunctionNamespaceManager.JAVA_BUILTIN_NAMESPACE;
 import static java.util.Objects.requireNonNull;
 
 public class FunctionListBuilder
@@ -61,18 +59,6 @@ public FunctionListBuilder scalars(Class clazz)
         return this;
     }
 
-    public FunctionListBuilder sqlInvokedScalar(Class clazz)
-    {
-        functions.addAll(SqlInvokedScalarFromAnnotationsParser.parseFunctionDefinition(clazz, JAVA_BUILTIN_NAMESPACE));
-        return this;
-    }
-
-    public FunctionListBuilder sqlInvokedScalars(Class clazz)
-    {
-        functions.addAll(SqlInvokedScalarFromAnnotationsParser.parseFunctionDefinitions(clazz, JAVA_BUILTIN_NAMESPACE));
-        return this;
-    }
-
     public FunctionListBuilder codegenScalars(Class clazz)
     {
         functions.addAll(CodegenScalarFromAnnotationsParser.parseFunctionDefinitions(clazz));
diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/ArrayIntersectFunction.java b/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/ArrayIntersectFunction.java
index d977420d28339..4ee510a86c467 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/ArrayIntersectFunction.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/ArrayIntersectFunction.java
@@ -19,8 +19,6 @@
 import com.facebook.presto.spi.function.Description;
 import com.facebook.presto.spi.function.OperatorDependency;
 import com.facebook.presto.spi.function.ScalarFunction;
-import com.facebook.presto.spi.function.SqlInvokedScalarFunction;
-import com.facebook.presto.spi.function.SqlParameter;
 import com.facebook.presto.spi.function.SqlType;
 import com.facebook.presto.spi.function.TypeParameter;
 
@@ -60,14 +58,4 @@ public static Block intersect(
 
         return typedSet.getBlock();
     }
-
-    @SqlInvokedScalarFunction(value = "array_intersect", deterministic = true, calledOnNullInput = false)
-    @Description("Intersects elements of all arrays in the given array")
-    @TypeParameter("T")
-    @SqlParameter(name = "input", type = "array>")
-    @SqlType("array")
-    public static String arrayIntersectArray()
-    {
-        return "RETURN reduce(input, IF((cardinality(input) = 0), ARRAY[], input[1]), (s, x) -> array_intersect(s, x), (s) -> s)";
-    }
 }
diff --git a/presto-main-base/src/main/java/com/facebook/presto/testing/QueryRunner.java b/presto-main-base/src/main/java/com/facebook/presto/testing/QueryRunner.java
index 5ab9fcbb94a71..f6c1a0b9b1405 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/testing/QueryRunner.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/testing/QueryRunner.java
@@ -126,6 +126,11 @@ default void loadPlanCheckerProviderManager(String planCheckerProviderName, Map<
         throw new UnsupportedOperationException();
     }
 
+    default void triggerConflictCheckWithBuiltInFunctions()
+    {
+        throw new UnsupportedOperationException();
+    }
+
     class MaterializedResultWithPlan
     {
         private final MaterializedResult materializedResult;
diff --git a/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/AbstractTestFunctions.java b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/AbstractTestFunctions.java
index d31ae0b09eb02..0644a6ebbdb68 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/AbstractTestFunctions.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/AbstractTestFunctions.java
@@ -59,6 +59,7 @@ public abstract class AbstractTestFunctions
     private final FeaturesConfig featuresConfig;
     private final FunctionsConfig functionsConfig;
     protected FunctionAssertions functionAssertions;
+    private final boolean loadInlinedSqlInvokedFunctionsPlugin;
 
     protected AbstractTestFunctions()
     {
@@ -81,18 +82,23 @@ protected AbstractTestFunctions(FunctionsConfig functionsConfig)
     }
 
     protected AbstractTestFunctions(Session session, FeaturesConfig featuresConfig, FunctionsConfig functionsConfig)
+    {
+        this(session, featuresConfig, functionsConfig, true);
+    }
+    protected AbstractTestFunctions(Session session, FeaturesConfig featuresConfig, FunctionsConfig functionsConfig, boolean loadInlinedSqlInvokedFunctionsPlugin)
     {
         this.session = requireNonNull(session, "session is null");
         this.featuresConfig = requireNonNull(featuresConfig, "featuresConfig is null");
         this.functionsConfig = requireNonNull(functionsConfig, "config is null")
                 .setLegacyLogFunction(true)
                 .setUseNewNanDefinition(true);
+        this.loadInlinedSqlInvokedFunctionsPlugin = loadInlinedSqlInvokedFunctionsPlugin;
     }
 
     @BeforeClass
     public final void initTestFunctions()
     {
-        functionAssertions = new FunctionAssertions(session, featuresConfig, functionsConfig, false);
+        functionAssertions = new FunctionAssertions(session, featuresConfig, functionsConfig, false, loadInlinedSqlInvokedFunctionsPlugin);
     }
 
     @AfterClass(alwaysRun = true)
diff --git a/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/FunctionAssertions.java b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/FunctionAssertions.java
index 17ba3c70823db..b2965a54329c5 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/FunctionAssertions.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/FunctionAssertions.java
@@ -41,6 +41,7 @@
 import com.facebook.presto.operator.project.CursorProcessor;
 import com.facebook.presto.operator.project.PageProcessor;
 import com.facebook.presto.operator.project.PageProjectionWithOutputs;
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.spi.ColumnHandle;
 import com.facebook.presto.spi.ConnectorId;
 import com.facebook.presto.spi.ConnectorPageSource;
@@ -221,20 +222,25 @@ public FunctionAssertions()
 
     public FunctionAssertions(Session session)
     {
-        this(session, new FeaturesConfig(), new FunctionsConfig(), false);
+        this(session, new FeaturesConfig(), new FunctionsConfig(), false, true);
     }
 
     public FunctionAssertions(Session session, FeaturesConfig featuresConfig)
     {
-        this(session, featuresConfig, new FunctionsConfig(), false);
+        this(session, featuresConfig, new FunctionsConfig(), false, true);
     }
 
     public FunctionAssertions(Session session, FunctionsConfig functionsConfig)
     {
-        this(session, new FeaturesConfig(), functionsConfig, false);
+        this(session, new FeaturesConfig(), functionsConfig, false, true);
     }
 
     public FunctionAssertions(Session session, FeaturesConfig featuresConfig, FunctionsConfig functionsConfig, boolean refreshSession)
+    {
+        this(session, featuresConfig, functionsConfig, refreshSession, true);
+    }
+
+    public FunctionAssertions(Session session, FeaturesConfig featuresConfig, FunctionsConfig functionsConfig, boolean refreshSession, boolean loadInlinedSqlInvokedFunctionsPlugin)
     {
         requireNonNull(session, "session is null");
         runner = new LocalQueryRunner(session, featuresConfig, functionsConfig);
@@ -244,6 +250,9 @@ public FunctionAssertions(Session session, FeaturesConfig featuresConfig, Functi
         else {
             this.session = session;
         }
+        if (loadInlinedSqlInvokedFunctionsPlugin) {
+            runner.installPlugin(new SqlInvokedFunctionsPlugin());
+        }
         metadata = runner.getMetadata();
         compiler = runner.getExpressionCompiler();
     }
diff --git a/presto-main-base/src/test/java/com/facebook/presto/sql/planner/iterative/rule/TestCrossJoinWithArrayNotContainsToAntiJoin.java b/presto-main-base/src/test/java/com/facebook/presto/sql/planner/iterative/rule/TestCrossJoinWithArrayNotContainsToAntiJoin.java
index 4bd3f167ac11a..194cf489d3998 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/sql/planner/iterative/rule/TestCrossJoinWithArrayNotContainsToAntiJoin.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/sql/planner/iterative/rule/TestCrossJoinWithArrayNotContainsToAntiJoin.java
@@ -14,6 +14,7 @@
 package com.facebook.presto.sql.planner.iterative.rule;
 
 import com.facebook.presto.common.type.ArrayType;
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.spi.plan.Assignments;
 import com.facebook.presto.spi.plan.FilterNode;
 import com.facebook.presto.spi.plan.JoinNode;
@@ -22,9 +23,11 @@
 import com.facebook.presto.spi.plan.UnnestNode;
 import com.facebook.presto.spi.plan.ValuesNode;
 import com.facebook.presto.sql.planner.iterative.rule.test.BaseRuleTest;
+import com.facebook.presto.sql.planner.iterative.rule.test.RuleTester;
 import com.facebook.presto.sql.planner.plan.EnforceSingleRowNode;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
+import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
 
 import static com.facebook.presto.SystemSessionProperties.REWRITE_CROSS_JOIN_ARRAY_NOT_CONTAINS_TO_ANTI_JOIN;
@@ -33,10 +36,18 @@
 import static com.facebook.presto.common.type.VarcharType.VARCHAR;
 import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.node;
 import static com.facebook.presto.sql.planner.iterative.rule.test.PlanBuilder.constantExpressions;
+import static java.util.Collections.singletonList;
 
 public class TestCrossJoinWithArrayNotContainsToAntiJoin
         extends BaseRuleTest
 {
+    @BeforeClass
+    @Override
+    public void setUp()
+    {
+        tester = new RuleTester(singletonList(new SqlInvokedFunctionsPlugin()));
+    }
+
     @Test
     public void testTriggerForBigInt()
     {
diff --git a/presto-main-base/src/test/java/com/facebook/presto/sql/planner/iterative/rule/TestLeftJoinWithArrayContainsToEquiJoinCondition.java b/presto-main-base/src/test/java/com/facebook/presto/sql/planner/iterative/rule/TestLeftJoinWithArrayContainsToEquiJoinCondition.java
index 657093aec5be2..cdf060f86d390 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/sql/planner/iterative/rule/TestLeftJoinWithArrayContainsToEquiJoinCondition.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/sql/planner/iterative/rule/TestLeftJoinWithArrayContainsToEquiJoinCondition.java
@@ -14,11 +14,14 @@
 package com.facebook.presto.sql.planner.iterative.rule;
 
 import com.facebook.presto.common.type.ArrayType;
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.spi.plan.JoinType;
 import com.facebook.presto.sql.planner.iterative.rule.test.BaseRuleTest;
+import com.facebook.presto.sql.planner.iterative.rule.test.RuleTester;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
 
 import java.util.Optional;
@@ -32,10 +35,18 @@
 import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.project;
 import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.unnest;
 import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.values;
+import static java.util.Collections.singletonList;
 
 public class TestLeftJoinWithArrayContainsToEquiJoinCondition
         extends BaseRuleTest
 {
+    @BeforeClass
+    @Override
+    public void setUp()
+    {
+        tester = new RuleTester(singletonList(new SqlInvokedFunctionsPlugin()));
+    }
+
     @Test
     public void testTriggerForBigIntArrayRightSide()
     {
diff --git a/presto-main/etc/config.properties b/presto-main/etc/config.properties
index a120422c4b6f1..164015db096d4 100644
--- a/presto-main/etc/config.properties
+++ b/presto-main/etc/config.properties
@@ -51,7 +51,8 @@ plugin.bundles=\
   ../presto-node-ttl-fetchers/pom.xml,\
   ../presto-hive-function-namespace/pom.xml,\
   ../presto-delta/pom.xml,\
-  ../presto-hudi/pom.xml
+  ../presto-hudi/pom.xml, \
+  ../presto-sql-invoked-functions-plugin/pom.xml
 
 presto.version=testversion
 node-scheduler.include-coordinator=true
diff --git a/presto-ml/src/test/java/com/facebook/presto/ml/AbstractTestMLFunctions.java b/presto-ml/src/test/java/com/facebook/presto/ml/AbstractTestMLFunctions.java
index dffbefd9189aa..4ee697957e540 100644
--- a/presto-ml/src/test/java/com/facebook/presto/ml/AbstractTestMLFunctions.java
+++ b/presto-ml/src/test/java/com/facebook/presto/ml/AbstractTestMLFunctions.java
@@ -15,13 +15,21 @@
 package com.facebook.presto.ml;
 
 import com.facebook.presto.operator.scalar.AbstractTestFunctions;
+import com.facebook.presto.sql.analyzer.FeaturesConfig;
+import com.facebook.presto.sql.analyzer.FunctionsConfig;
 import org.testng.annotations.BeforeClass;
 
+import static com.facebook.presto.SessionTestUtils.TEST_SESSION;
 import static com.facebook.presto.metadata.FunctionExtractor.extractFunctions;
 
 abstract class AbstractTestMLFunctions
         extends AbstractTestFunctions
 {
+    public AbstractTestMLFunctions()
+    {
+        super(TEST_SESSION, new FeaturesConfig(), new FunctionsConfig(), false);
+    }
+
     @BeforeClass
     protected void registerFunctions()
     {
diff --git a/presto-native-execution/pom.xml b/presto-native-execution/pom.xml
index 8f849f866335a..4f48109d76937 100644
--- a/presto-native-execution/pom.xml
+++ b/presto-native-execution/pom.xml
@@ -328,6 +328,19 @@
             presto-jdbc
             test
         
+
+        
+            org.apache.commons
+            commons-lang3
+            test
+        
+
+        
+            com.facebook.presto
+            presto-sql-invoked-functions-plugin
+            ${project.version}
+            test
+        
     
 
     
diff --git a/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/TestPrestoNativeArrayFunctionQueries.java b/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/TestPrestoNativeArrayFunctionQueries.java
index bfa21641206b2..f45f20c545d00 100644
--- a/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/TestPrestoNativeArrayFunctionQueries.java
+++ b/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/TestPrestoNativeArrayFunctionQueries.java
@@ -13,6 +13,7 @@
  */
 package com.facebook.presto.nativeworker;
 
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.testing.ExpectedQueryRunner;
 import com.facebook.presto.testing.QueryRunner;
 
@@ -20,18 +21,24 @@ public class TestPrestoNativeArrayFunctionQueries
         extends AbstractTestNativeArrayFunctionQueries
 {
     @Override
-    protected QueryRunner createQueryRunner() throws Exception
+    protected QueryRunner createQueryRunner()
+            throws Exception
     {
-        return PrestoNativeQueryRunnerUtils.nativeHiveQueryRunnerBuilder()
+        QueryRunner queryRunner = PrestoNativeQueryRunnerUtils.nativeHiveQueryRunnerBuilder()
                 .setAddStorageFormatToPath(true)
                 .build();
+        queryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
+        return queryRunner;
     }
 
     @Override
-    protected ExpectedQueryRunner createExpectedQueryRunner() throws Exception
+    protected ExpectedQueryRunner createExpectedQueryRunner()
+            throws Exception
     {
-        return PrestoNativeQueryRunnerUtils.javaHiveQueryRunnerBuilder()
+        QueryRunner queryRunner = PrestoNativeQueryRunnerUtils.javaHiveQueryRunnerBuilder()
                 .setAddStorageFormatToPath(true)
                 .build();
+        queryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
+        return queryRunner;
     }
 }
diff --git a/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/TestPrestoNativeGeneralQueriesJSON.java b/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/TestPrestoNativeGeneralQueriesJSON.java
index 4b1d27f30426e..bcba83d00da65 100644
--- a/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/TestPrestoNativeGeneralQueriesJSON.java
+++ b/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/TestPrestoNativeGeneralQueriesJSON.java
@@ -13,6 +13,7 @@
  */
 package com.facebook.presto.nativeworker;
 
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.testing.ExpectedQueryRunner;
 import com.facebook.presto.testing.QueryRunner;
 
@@ -23,17 +24,21 @@ public class TestPrestoNativeGeneralQueriesJSON
     protected QueryRunner createQueryRunner()
             throws Exception
     {
-        return PrestoNativeQueryRunnerUtils.nativeHiveQueryRunnerBuilder()
+        QueryRunner queryRunner = PrestoNativeQueryRunnerUtils.nativeHiveQueryRunnerBuilder()
                 .setAddStorageFormatToPath(true)
                 .build();
+        queryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
+        return queryRunner;
     }
 
     @Override
     protected ExpectedQueryRunner createExpectedQueryRunner()
             throws Exception
     {
-        return PrestoNativeQueryRunnerUtils.javaHiveQueryRunnerBuilder()
+        QueryRunner queryRunner = PrestoNativeQueryRunnerUtils.javaHiveQueryRunnerBuilder()
                 .setAddStorageFormatToPath(true)
                 .build();
+        queryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
+        return queryRunner;
     }
 }
diff --git a/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/TestPrestoNativeGeneralQueriesThrift.java b/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/TestPrestoNativeGeneralQueriesThrift.java
index a1c33cd09f51c..fd82dcde9209e 100644
--- a/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/TestPrestoNativeGeneralQueriesThrift.java
+++ b/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/TestPrestoNativeGeneralQueriesThrift.java
@@ -13,6 +13,7 @@
  */
 package com.facebook.presto.nativeworker;
 
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.testing.ExpectedQueryRunner;
 import com.facebook.presto.testing.QueryRunner;
 
@@ -23,18 +24,22 @@ public class TestPrestoNativeGeneralQueriesThrift
     protected QueryRunner createQueryRunner()
             throws Exception
     {
-        return PrestoNativeQueryRunnerUtils.nativeHiveQueryRunnerBuilder()
+        QueryRunner queryRunner = PrestoNativeQueryRunnerUtils.nativeHiveQueryRunnerBuilder()
                 .setAddStorageFormatToPath(true)
                 .setUseThrift(true)
                 .build();
+        queryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
+        return queryRunner;
     }
 
     @Override
     protected ExpectedQueryRunner createExpectedQueryRunner()
             throws Exception
     {
-        return PrestoNativeQueryRunnerUtils.javaHiveQueryRunnerBuilder()
+        QueryRunner queryRunner = PrestoNativeQueryRunnerUtils.javaHiveQueryRunnerBuilder()
                 .setAddStorageFormatToPath(true)
                 .build();
+        queryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
+        return queryRunner;
     }
 }
diff --git a/presto-native-tests/src/test/java/com/facebook/presto/nativetests/TestDistributedEngineOnlyQueries.java b/presto-native-tests/src/test/java/com/facebook/presto/nativetests/TestDistributedEngineOnlyQueries.java
index 71c884e828554..b54a9711940b9 100644
--- a/presto-native-tests/src/test/java/com/facebook/presto/nativetests/TestDistributedEngineOnlyQueries.java
+++ b/presto-native-tests/src/test/java/com/facebook/presto/nativetests/TestDistributedEngineOnlyQueries.java
@@ -123,4 +123,16 @@ public void testKeyBasedSampling()
     public void testDefaultSamplingPercent()
     {
     }
+
+    @Override
+    @Test(enabled = false)
+    public void testLeftJoinWithArrayContainsCondition()
+    {
+    }
+
+    @Override
+    @Test(enabled = false)
+    public void testTry()
+    {
+    }
 }
diff --git a/presto-pinot/src/test/java/com/facebook/presto/pinot/udf/TestPinotFunctions.java b/presto-pinot/src/test/java/com/facebook/presto/pinot/udf/TestPinotFunctions.java
index acfd681d0e48b..3dec94b4fe4ce 100644
--- a/presto-pinot/src/test/java/com/facebook/presto/pinot/udf/TestPinotFunctions.java
+++ b/presto-pinot/src/test/java/com/facebook/presto/pinot/udf/TestPinotFunctions.java
@@ -15,15 +15,23 @@
 
 import com.facebook.presto.operator.scalar.AbstractTestFunctions;
 import com.facebook.presto.pinot.PinotPlugin;
+import com.facebook.presto.sql.analyzer.FeaturesConfig;
+import com.facebook.presto.sql.analyzer.FunctionsConfig;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
 
+import static com.facebook.presto.SessionTestUtils.TEST_SESSION;
 import static com.facebook.presto.common.type.DoubleType.DOUBLE;
 import static com.facebook.presto.metadata.FunctionExtractor.extractFunctions;
 
 public class TestPinotFunctions
         extends AbstractTestFunctions
 {
+    public TestPinotFunctions()
+    {
+        super(TEST_SESSION, new FeaturesConfig(), new FunctionsConfig(), false);
+    }
+
     @BeforeClass
     public void setUp()
     {
diff --git a/presto-server/src/main/provisio/presto.xml b/presto-server/src/main/provisio/presto.xml
index 14bcfd5b3c31a..10359099bce2a 100644
--- a/presto-server/src/main/provisio/presto.xml
+++ b/presto-server/src/main/provisio/presto.xml
@@ -280,4 +280,10 @@
             
         
     
+
+    
+        
+            
+        
+    
 
diff --git a/presto-sql-invoked-functions-plugin/pom.xml b/presto-sql-invoked-functions-plugin/pom.xml
new file mode 100644
index 0000000000000..5ab338547b4ef
--- /dev/null
+++ b/presto-sql-invoked-functions-plugin/pom.xml
@@ -0,0 +1,29 @@
+
+    4.0.0
+    
+        com.facebook.presto
+        presto-root
+        0.295-SNAPSHOT
+    
+
+    presto-sql-invoked-functions-plugin
+    Presto - Sql invoked functions plugin
+    presto-plugin
+
+    
+        ${project.parent.basedir}
+    
+
+    
+        
+            com.facebook.presto
+            presto-spi
+            provided
+        
+        
+            com.google.guava
+            guava
+        
+    
+
diff --git a/presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/ArrayIntersectFunction.java b/presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/ArrayIntersectFunction.java
new file mode 100644
index 0000000000000..860faede7942d
--- /dev/null
+++ b/presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/ArrayIntersectFunction.java
@@ -0,0 +1,35 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.scalar.sql;
+
+import com.facebook.presto.spi.function.Description;
+import com.facebook.presto.spi.function.SqlInvokedScalarFunction;
+import com.facebook.presto.spi.function.SqlParameter;
+import com.facebook.presto.spi.function.SqlType;
+import com.facebook.presto.spi.function.TypeParameter;
+
+public class ArrayIntersectFunction
+{
+    private ArrayIntersectFunction() {}
+
+    @SqlInvokedScalarFunction(value = "array_intersect", deterministic = true, calledOnNullInput = false)
+    @Description("Intersects elements of all arrays in the given array")
+    @TypeParameter("T")
+    @SqlParameter(name = "input", type = "array>")
+    @SqlType("array")
+    public static String arrayIntersectArray()
+    {
+        return "RETURN reduce(input, IF((cardinality(input) = 0), ARRAY[], input[1]), (s, x) -> array_intersect(s, x), (s) -> s)";
+    }
+}
diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/sql/ArraySqlFunctions.java b/presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/ArraySqlFunctions.java
similarity index 99%
rename from presto-main-base/src/main/java/com/facebook/presto/operator/scalar/sql/ArraySqlFunctions.java
rename to presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/ArraySqlFunctions.java
index 2fc1218ea2408..b5ab56eb6f312 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/sql/ArraySqlFunctions.java
+++ b/presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/ArraySqlFunctions.java
@@ -11,7 +11,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.facebook.presto.operator.scalar.sql;
+package com.facebook.presto.scalar.sql;
 
 import com.facebook.presto.spi.function.Description;
 import com.facebook.presto.spi.function.SqlInvokedScalarFunction;
diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/sql/MapNormalizeFunction.java b/presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/MapNormalizeFunction.java
similarity index 96%
rename from presto-main-base/src/main/java/com/facebook/presto/operator/scalar/sql/MapNormalizeFunction.java
rename to presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/MapNormalizeFunction.java
index 1618dac2c4356..99104f3027b8d 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/sql/MapNormalizeFunction.java
+++ b/presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/MapNormalizeFunction.java
@@ -11,7 +11,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.facebook.presto.operator.scalar.sql;
+package com.facebook.presto.scalar.sql;
 
 import com.facebook.presto.spi.function.Description;
 import com.facebook.presto.spi.function.SqlInvokedScalarFunction;
diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/sql/MapSqlFunctions.java b/presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/MapSqlFunctions.java
similarity index 99%
rename from presto-main-base/src/main/java/com/facebook/presto/operator/scalar/sql/MapSqlFunctions.java
rename to presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/MapSqlFunctions.java
index 0cf9558a22989..498c068dedf39 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/sql/MapSqlFunctions.java
+++ b/presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/MapSqlFunctions.java
@@ -11,7 +11,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.facebook.presto.operator.scalar.sql;
+package com.facebook.presto.scalar.sql;
 
 import com.facebook.presto.spi.function.Description;
 import com.facebook.presto.spi.function.SqlInvokedScalarFunction;
diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/sql/SimpleSamplingPercent.java b/presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/SimpleSamplingPercent.java
similarity index 96%
rename from presto-main-base/src/main/java/com/facebook/presto/operator/scalar/sql/SimpleSamplingPercent.java
rename to presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/SimpleSamplingPercent.java
index 473bbedafd637..5aa6b6b5d966e 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/sql/SimpleSamplingPercent.java
+++ b/presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/SimpleSamplingPercent.java
@@ -11,7 +11,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.facebook.presto.operator.scalar.sql;
+package com.facebook.presto.scalar.sql;
 
 import com.facebook.presto.spi.function.Description;
 import com.facebook.presto.spi.function.SqlInvokedScalarFunction;
diff --git a/presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/SqlInvokedFunctionsPlugin.java b/presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/SqlInvokedFunctionsPlugin.java
new file mode 100644
index 0000000000000..9a1e3650f72fe
--- /dev/null
+++ b/presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/SqlInvokedFunctionsPlugin.java
@@ -0,0 +1,36 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.scalar.sql;
+
+import com.facebook.presto.spi.Plugin;
+import com.google.common.collect.ImmutableSet;
+
+import java.util.Set;
+
+public class SqlInvokedFunctionsPlugin
+        implements Plugin
+{
+    @Override
+    public Set> getSqlInvokedFunctions()
+    {
+        return ImmutableSet.>builder()
+                .add(ArraySqlFunctions.class)
+                .add(MapNormalizeFunction.class)
+                .add(MapSqlFunctions.class)
+                .add(SimpleSamplingPercent.class)
+                .add(StringSqlFunctions.class)
+                .add(ArrayIntersectFunction.class)
+                .build();
+    }
+}
diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/sql/StringSqlFunctions.java b/presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/StringSqlFunctions.java
similarity index 97%
rename from presto-main-base/src/main/java/com/facebook/presto/operator/scalar/sql/StringSqlFunctions.java
rename to presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/StringSqlFunctions.java
index 1b0e6e15661ed..9201e4bb1cf94 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/sql/StringSqlFunctions.java
+++ b/presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/StringSqlFunctions.java
@@ -12,7 +12,7 @@
  * limitations under the License.
  */
 
-package com.facebook.presto.operator.scalar.sql;
+package com.facebook.presto.scalar.sql;
 
 import com.facebook.presto.spi.function.Description;
 import com.facebook.presto.spi.function.SqlInvokedScalarFunction;
diff --git a/presto-teradata-functions/src/test/java/com/facebook/presto/teradata/functions/TestTeradataDateFunctions.java b/presto-teradata-functions/src/test/java/com/facebook/presto/teradata/functions/TestTeradataDateFunctions.java
index 9b4a72aacc792..16bd47af0d5c1 100644
--- a/presto-teradata-functions/src/test/java/com/facebook/presto/teradata/functions/TestTeradataDateFunctions.java
+++ b/presto-teradata-functions/src/test/java/com/facebook/presto/teradata/functions/TestTeradataDateFunctions.java
@@ -18,6 +18,8 @@
 import com.facebook.presto.common.type.SqlDate;
 import com.facebook.presto.common.type.TimestampType;
 import com.facebook.presto.operator.scalar.AbstractTestFunctions;
+import com.facebook.presto.sql.analyzer.FeaturesConfig;
+import com.facebook.presto.sql.analyzer.FunctionsConfig;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
 
@@ -40,7 +42,7 @@ public class TestTeradataDateFunctions
 
     protected TestTeradataDateFunctions()
     {
-        super(SESSION);
+        super(SESSION, new FeaturesConfig(), new FunctionsConfig(), false);
     }
 
     @BeforeClass
diff --git a/presto-teradata-functions/src/test/java/com/facebook/presto/teradata/functions/TestTeradataFunctions.java b/presto-teradata-functions/src/test/java/com/facebook/presto/teradata/functions/TestTeradataFunctions.java
index be6efbffec1ef..e724ca94566e7 100644
--- a/presto-teradata-functions/src/test/java/com/facebook/presto/teradata/functions/TestTeradataFunctions.java
+++ b/presto-teradata-functions/src/test/java/com/facebook/presto/teradata/functions/TestTeradataFunctions.java
@@ -14,9 +14,12 @@
 package com.facebook.presto.teradata.functions;
 
 import com.facebook.presto.operator.scalar.AbstractTestFunctions;
+import com.facebook.presto.sql.analyzer.FeaturesConfig;
+import com.facebook.presto.sql.analyzer.FunctionsConfig;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
 
+import static com.facebook.presto.SessionTestUtils.TEST_SESSION;
 import static com.facebook.presto.common.type.BigintType.BIGINT;
 import static com.facebook.presto.common.type.VarcharType.VARCHAR;
 import static com.facebook.presto.common.type.VarcharType.createVarcharType;
@@ -25,6 +28,11 @@
 public class TestTeradataFunctions
         extends AbstractTestFunctions
 {
+    public TestTeradataFunctions()
+    {
+        super(TEST_SESSION, new FeaturesConfig(), new FunctionsConfig(), false);
+    }
+
     @BeforeClass
     public void setUp()
     {
diff --git a/presto-tests/pom.xml b/presto-tests/pom.xml
index 3623dd0991913..963a82a22adad 100644
--- a/presto-tests/pom.xml
+++ b/presto-tests/pom.xml
@@ -378,6 +378,13 @@
             ratis-metrics-default
             test
         
+
+        
+            com.facebook.presto
+            presto-sql-invoked-functions-plugin
+            ${project.version}
+            test
+        
     
 
     
diff --git a/presto-tests/src/main/java/com/facebook/presto/tests/DistributedQueryRunner.java b/presto-tests/src/main/java/com/facebook/presto/tests/DistributedQueryRunner.java
index 67bf65f67eb55..6949fc4e9c8b3 100644
--- a/presto-tests/src/main/java/com/facebook/presto/tests/DistributedQueryRunner.java
+++ b/presto-tests/src/main/java/com/facebook/presto/tests/DistributedQueryRunner.java
@@ -1040,6 +1040,14 @@ public void loadPlanCheckerProviderManager(String planCheckerProviderName, Map> getSqlInvokedFunctions()
                     .add(TestDuplicateSqlInvokedFunctions.class)
                     .build();
         }
+
+        @Override
+        public Set> getFunctions()
+        {
+            return ImmutableSet.>builder()
+                    .add(TestFunctions.class)
+                    // Adding a SQL Invoked function in the built-in functions to mimic a conflict.
+                    .add(ArrayIntersectFunction.class)
+                    .build();
+        }
     }
 
     // As soon as we trigger the conflict check with the built-in functions, an error will be thrown if duplicate signatures are found.
diff --git a/presto-tests/src/test/java/com/facebook/presto/tests/TestDistributedEngineOnlyQueries.java b/presto-tests/src/test/java/com/facebook/presto/tests/TestDistributedEngineOnlyQueries.java
index b18173ff0bf20..8d6910ccaa59d 100644
--- a/presto-tests/src/test/java/com/facebook/presto/tests/TestDistributedEngineOnlyQueries.java
+++ b/presto-tests/src/test/java/com/facebook/presto/tests/TestDistributedEngineOnlyQueries.java
@@ -13,6 +13,7 @@
  */
 package com.facebook.presto.tests;
 
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.testing.QueryRunner;
 import com.facebook.presto.tests.tpch.TpchQueryRunnerBuilder;
 
@@ -23,6 +24,8 @@ public class TestDistributedEngineOnlyQueries
     protected QueryRunner createQueryRunner()
             throws Exception
     {
-        return TpchQueryRunnerBuilder.builder().build();
+        QueryRunner queryRunner = TpchQueryRunnerBuilder.builder().build();
+        queryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
+        return queryRunner;
     }
 }
diff --git a/presto-tests/src/test/java/com/facebook/presto/tests/TestDistributedSpilledQueries.java b/presto-tests/src/test/java/com/facebook/presto/tests/TestDistributedSpilledQueries.java
index e945493c2969e..1ab8522eb50a3 100644
--- a/presto-tests/src/test/java/com/facebook/presto/tests/TestDistributedSpilledQueries.java
+++ b/presto-tests/src/test/java/com/facebook/presto/tests/TestDistributedSpilledQueries.java
@@ -15,6 +15,7 @@
 
 import com.facebook.presto.Session;
 import com.facebook.presto.SystemSessionProperties;
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.testing.QueryRunner;
 import com.facebook.presto.tpch.TpchPlugin;
 import com.google.common.collect.ImmutableMap;
@@ -63,6 +64,7 @@ public static QueryRunner localCreateQueryRunner()
         try {
             queryRunner.installPlugin(new TpchPlugin());
             queryRunner.createCatalog("tpch", "tpch");
+            queryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
             return queryRunner;
         }
         catch (Exception e) {
diff --git a/presto-tests/src/test/java/com/facebook/presto/tests/TestDistributedSpilledQueriesWithTempStorage.java b/presto-tests/src/test/java/com/facebook/presto/tests/TestDistributedSpilledQueriesWithTempStorage.java
index c4ca56a3f78c0..739c5dc826631 100644
--- a/presto-tests/src/test/java/com/facebook/presto/tests/TestDistributedSpilledQueriesWithTempStorage.java
+++ b/presto-tests/src/test/java/com/facebook/presto/tests/TestDistributedSpilledQueriesWithTempStorage.java
@@ -15,6 +15,7 @@
 
 import com.facebook.presto.Session;
 import com.facebook.presto.SystemSessionProperties;
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.testing.QueryRunner;
 import com.facebook.presto.tpch.TpchPlugin;
 import com.google.common.collect.ImmutableMap;
@@ -57,6 +58,8 @@ public static DistributedQueryRunner localCreateQueryRunner()
         try {
             queryRunner.installPlugin(new TpchPlugin());
             queryRunner.createCatalog("tpch", "tpch");
+            queryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
+
             return queryRunner;
         }
         catch (Exception e) {
diff --git a/presto-tests/src/test/java/com/facebook/presto/tests/TestLocalQueries.java b/presto-tests/src/test/java/com/facebook/presto/tests/TestLocalQueries.java
index b381b01f5ada3..e7f2cf5a9a781 100644
--- a/presto-tests/src/test/java/com/facebook/presto/tests/TestLocalQueries.java
+++ b/presto-tests/src/test/java/com/facebook/presto/tests/TestLocalQueries.java
@@ -15,6 +15,7 @@
 
 import com.facebook.presto.Session;
 import com.facebook.presto.metadata.SessionPropertyManager;
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.spi.CatalogSchemaTableName;
 import com.facebook.presto.spi.ConnectorId;
 import com.facebook.presto.sql.planner.planPrinter.IOPlanPrinter.ColumnConstraint;
@@ -78,6 +79,8 @@ public static LocalQueryRunner createLocalQueryRunner()
         sessionPropertyManager.addSystemSessionProperties(TEST_SYSTEM_PROPERTIES);
         sessionPropertyManager.addConnectorSessionProperties(new ConnectorId(TESTING_CATALOG), TEST_CATALOG_PROPERTIES);
 
+        localQueryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
+
         return localQueryRunner;
     }
 
diff --git a/presto-tests/src/test/java/com/facebook/presto/tests/TestQueryPlanDeterminism.java b/presto-tests/src/test/java/com/facebook/presto/tests/TestQueryPlanDeterminism.java
index 98fdee5fe7ebf..8fbcc731fec0a 100644
--- a/presto-tests/src/test/java/com/facebook/presto/tests/TestQueryPlanDeterminism.java
+++ b/presto-tests/src/test/java/com/facebook/presto/tests/TestQueryPlanDeterminism.java
@@ -16,6 +16,7 @@
 import com.facebook.presto.Session;
 import com.facebook.presto.common.type.Type;
 import com.facebook.presto.metadata.SessionPropertyManager;
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.spi.ConnectorId;
 import com.facebook.presto.testing.LocalQueryRunner;
 import com.facebook.presto.testing.MaterializedResult;
@@ -74,6 +75,7 @@ protected QueryRunner createQueryRunner()
         sessionPropertyManager.addSystemSessionProperties(TEST_SYSTEM_PROPERTIES);
         sessionPropertyManager.addConnectorSessionProperties(new ConnectorId(TESTING_CATALOG), TEST_CATALOG_PROPERTIES);
 
+        localQueryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
         return localQueryRunner;
     }
 
diff --git a/presto-tests/src/test/java/com/facebook/presto/tests/TestSqlFunctions.java b/presto-tests/src/test/java/com/facebook/presto/tests/TestSqlFunctions.java
index cafb8e59e7fc1..e4ec0ad5e39a8 100644
--- a/presto-tests/src/test/java/com/facebook/presto/tests/TestSqlFunctions.java
+++ b/presto-tests/src/test/java/com/facebook/presto/tests/TestSqlFunctions.java
@@ -18,6 +18,7 @@
 import com.facebook.presto.common.type.TypeSignature;
 import com.facebook.presto.common.type.TypeSignatureParameter;
 import com.facebook.presto.common.type.UserDefinedType;
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.spi.function.Parameter;
 import com.facebook.presto.spi.function.RoutineCharacteristics;
 import com.facebook.presto.spi.function.SqlFunctionId;
@@ -99,7 +100,7 @@ protected QueryRunner createQueryRunner()
             queryRunner.getMetadata().getFunctionAndTypeManager().addUserDefinedType(COUNTRY_ENUM);
 
             queryRunner.execute("CREATE TYPE testing.type.person AS (first_name varchar, last_name varchar, age tinyint, country testing.enum.country)");
-
+            queryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
             return queryRunner;
         }
         catch (Exception e) {
diff --git a/presto-tests/src/test/java/com/facebook/presto/tests/TestVerboseOptimizerInfo.java b/presto-tests/src/test/java/com/facebook/presto/tests/TestVerboseOptimizerInfo.java
index 3cb9f26f1b8fd..add9d0be8cd07 100644
--- a/presto-tests/src/test/java/com/facebook/presto/tests/TestVerboseOptimizerInfo.java
+++ b/presto-tests/src/test/java/com/facebook/presto/tests/TestVerboseOptimizerInfo.java
@@ -16,6 +16,7 @@
 
 import com.facebook.presto.Session;
 import com.facebook.presto.metadata.SessionPropertyManager;
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.spi.ConnectorId;
 import com.facebook.presto.testing.LocalQueryRunner;
 import com.facebook.presto.testing.MaterializedResult;
@@ -71,7 +72,7 @@ public static LocalQueryRunner createLocalQueryRunner()
         SessionPropertyManager sessionPropertyManager = localQueryRunner.getMetadata().getSessionPropertyManager();
         sessionPropertyManager.addSystemSessionProperties(TEST_SYSTEM_PROPERTIES);
         sessionPropertyManager.addConnectorSessionProperties(new ConnectorId(TESTING_CATALOG), TEST_CATALOG_PROPERTIES);
-
+        localQueryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
         return localQueryRunner;
     }
 

From 6d69b415813b9160a7abd01d4d5603fd1deeac9c Mon Sep 17 00:00:00 2001
From: HeidiHan0000 <44453261+HeidiHan0000@users.noreply.github.com>
Date: Tue, 2 Sep 2025 17:15:38 -0700
Subject: [PATCH 041/113] refactor: TypeParser to use
 velox/.../prestosql/...TypeParser (#25877)

Summary: https://github.com/facebookincubator/velox/pull/14208 moves
velox/type/parser -> velox/functions/prestosql/types/parser. This is
because this type parser was intended to be used for presto specific
purposes, but was placed in a folder which seems too general. With this
merged, now moving references to Velox's type parser to the new folder.

Differential Revision: D80919225


```
== NO RELEASE NOTE ==
```
---
 presto-native-execution/presto_cpp/main/CMakeLists.txt        | 2 +-
 presto-native-execution/presto_cpp/main/types/CMakeLists.txt  | 2 +-
 presto-native-execution/presto_cpp/main/types/TypeParser.cpp  | 4 ++--
 .../presto_cpp/main/types/tests/CMakeLists.txt                | 2 +-
 4 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/presto-native-execution/presto_cpp/main/CMakeLists.txt b/presto-native-execution/presto_cpp/main/CMakeLists.txt
index 38770def39269..29836ba760385 100644
--- a/presto-native-execution/presto_cpp/main/CMakeLists.txt
+++ b/presto-native-execution/presto_cpp/main/CMakeLists.txt
@@ -78,10 +78,10 @@ target_link_libraries(
   velox_hive_iceberg_splitreader
   velox_hive_partition_function
   velox_presto_serializer
+  velox_presto_type_parser
   velox_s3fs
   velox_serialization
   velox_time
-  velox_type_parser
   velox_type
   velox_type_fbhive
   velox_type_tz
diff --git a/presto-native-execution/presto_cpp/main/types/CMakeLists.txt b/presto-native-execution/presto_cpp/main/types/CMakeLists.txt
index d264bc95f80c2..8b168254334aa 100644
--- a/presto-native-execution/presto_cpp/main/types/CMakeLists.txt
+++ b/presto-native-execution/presto_cpp/main/types/CMakeLists.txt
@@ -11,7 +11,7 @@
 # limitations under the License.
 
 add_library(presto_type_converter OBJECT TypeParser.cpp)
-target_link_libraries(presto_type_converter velox_type_parser)
+target_link_libraries(presto_type_converter velox_presto_type_parser)
 
 add_library(presto_velox_expr_conversion OBJECT PrestoToVeloxExpr.cpp)
 target_link_libraries(presto_velox_expr_conversion velox_presto_types
diff --git a/presto-native-execution/presto_cpp/main/types/TypeParser.cpp b/presto-native-execution/presto_cpp/main/types/TypeParser.cpp
index f8c757adf6e10..4a19066c38fb0 100644
--- a/presto-native-execution/presto_cpp/main/types/TypeParser.cpp
+++ b/presto-native-execution/presto_cpp/main/types/TypeParser.cpp
@@ -15,7 +15,7 @@
 #include 
 
 #include "presto_cpp/main/types/TypeParser.h"
-#include "velox/type/parser/TypeParser.h"
+#include "velox/functions/prestosql/types/parser/TypeParser.h"
 
 #include "presto_cpp/main/common/Configs.h"
 
@@ -33,7 +33,7 @@ velox::TypePtr TypeParser::parse(const std::string& text) const {
     return it->second;
   }
 
-  auto result = velox::parseType(text);
+  auto result = velox::functions::prestosql::parseType(text);
   cache_.insert({text, result});
   return result;
 }
diff --git a/presto-native-execution/presto_cpp/main/types/tests/CMakeLists.txt b/presto-native-execution/presto_cpp/main/types/tests/CMakeLists.txt
index 9560036e501d4..48b22e9a63328 100644
--- a/presto-native-execution/presto_cpp/main/types/tests/CMakeLists.txt
+++ b/presto-native-execution/presto_cpp/main/types/tests/CMakeLists.txt
@@ -67,9 +67,9 @@ target_link_libraries(
   velox_tpch_connector
   velox_hive_partition_function
   velox_presto_serializer
+  velox_presto_type_parser
   velox_serialization
   velox_type
-  velox_type_parser
   Boost::filesystem
   ${RE2}
   ${FOLLY_WITH_DEPENDENCIES}

From b89a041322c972c98efd3b09d720699f20915362 Mon Sep 17 00:00:00 2001
From: Mahadevuni Naveen Kumar 
Date: Tue, 2 Sep 2025 12:18:06 +0530
Subject: [PATCH 042/113] Move annotation LiteralParameter to SPI for use in
 other modules

---
 .../presto/operator/annotations/ImplementationDependency.java   | 2 +-
 .../facebook/presto/operator/scalar/CharacterStringCasts.java   | 2 +-
 .../com/facebook/presto/operator/scalar/JoniRegexpCasts.java    | 2 +-
 .../java/com/facebook/presto/operator/scalar/JsonFunctions.java | 2 +-
 .../java/com/facebook/presto/operator/scalar/MathFunctions.java | 2 +-
 .../com/facebook/presto/operator/scalar/StringFunctions.java    | 2 +-
 .../src/main/java/com/facebook/presto/type/BigintOperators.java | 1 +
 .../src/main/java/com/facebook/presto/type/LikeFunctions.java   | 1 +
 .../presto/operator/TestAnnotationEngineForAggregates.java      | 2 +-
 .../presto/operator/TestAnnotationEngineForScalars.java         | 2 +-
 .../facebook/presto/operator/scalar/TestStringFunctions.java    | 2 +-
 .../com/facebook/presto/spi/function}/LiteralParameter.java     | 2 +-
 12 files changed, 12 insertions(+), 10 deletions(-)
 rename {presto-main-base/src/main/java/com/facebook/presto/type => presto-spi/src/main/java/com/facebook/presto/spi/function}/LiteralParameter.java (95%)

diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/annotations/ImplementationDependency.java b/presto-main-base/src/main/java/com/facebook/presto/operator/annotations/ImplementationDependency.java
index 4f30469473589..06e897574a1d4 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/operator/annotations/ImplementationDependency.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/operator/annotations/ImplementationDependency.java
@@ -20,9 +20,9 @@
 import com.facebook.presto.spi.function.Convention;
 import com.facebook.presto.spi.function.FunctionDependency;
 import com.facebook.presto.spi.function.InvocationConvention;
+import com.facebook.presto.spi.function.LiteralParameter;
 import com.facebook.presto.spi.function.OperatorDependency;
 import com.facebook.presto.spi.function.TypeParameter;
-import com.facebook.presto.type.LiteralParameter;
 
 import java.lang.annotation.Annotation;
 import java.lang.reflect.AnnotatedElement;
diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/CharacterStringCasts.java b/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/CharacterStringCasts.java
index 0ecf2917fd5cb..b5198a36699ae 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/CharacterStringCasts.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/CharacterStringCasts.java
@@ -14,10 +14,10 @@
 package com.facebook.presto.operator.scalar;
 
 import com.facebook.presto.common.function.OperatorType;
+import com.facebook.presto.spi.function.LiteralParameter;
 import com.facebook.presto.spi.function.LiteralParameters;
 import com.facebook.presto.spi.function.ScalarOperator;
 import com.facebook.presto.spi.function.SqlType;
-import com.facebook.presto.type.LiteralParameter;
 import com.google.common.collect.ImmutableList;
 import io.airlift.slice.Slice;
 import io.airlift.slice.SliceUtf8;
diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/JoniRegexpCasts.java b/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/JoniRegexpCasts.java
index c0f551099672c..eb6264e1a0a1d 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/JoniRegexpCasts.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/JoniRegexpCasts.java
@@ -15,11 +15,11 @@
 
 import com.facebook.presto.common.function.OperatorType;
 import com.facebook.presto.spi.PrestoException;
+import com.facebook.presto.spi.function.LiteralParameter;
 import com.facebook.presto.spi.function.LiteralParameters;
 import com.facebook.presto.spi.function.ScalarOperator;
 import com.facebook.presto.spi.function.SqlType;
 import com.facebook.presto.type.JoniRegexpType;
-import com.facebook.presto.type.LiteralParameter;
 import io.airlift.jcodings.specific.NonStrictUTF8Encoding;
 import io.airlift.joni.Option;
 import io.airlift.joni.Regex;
diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/JsonFunctions.java b/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/JsonFunctions.java
index 339ff5c10dba9..41242a9fcbf50 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/JsonFunctions.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/JsonFunctions.java
@@ -21,13 +21,13 @@
 import com.facebook.presto.common.type.StandardTypes;
 import com.facebook.presto.common.type.Type;
 import com.facebook.presto.spi.PrestoException;
+import com.facebook.presto.spi.function.LiteralParameter;
 import com.facebook.presto.spi.function.LiteralParameters;
 import com.facebook.presto.spi.function.ScalarFunction;
 import com.facebook.presto.spi.function.ScalarOperator;
 import com.facebook.presto.spi.function.SqlNullable;
 import com.facebook.presto.spi.function.SqlType;
 import com.facebook.presto.type.JsonPathType;
-import com.facebook.presto.type.LiteralParameter;
 import com.fasterxml.jackson.core.JsonFactory;
 import com.fasterxml.jackson.core.JsonParser;
 import com.fasterxml.jackson.core.JsonToken;
diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/MathFunctions.java b/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/MathFunctions.java
index e778108d74303..1bdd52292b2f7 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/MathFunctions.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/MathFunctions.java
@@ -21,13 +21,13 @@
 import com.facebook.presto.operator.aggregation.TypedSet;
 import com.facebook.presto.spi.PrestoException;
 import com.facebook.presto.spi.function.Description;
+import com.facebook.presto.spi.function.LiteralParameter;
 import com.facebook.presto.spi.function.LiteralParameters;
 import com.facebook.presto.spi.function.ScalarFunction;
 import com.facebook.presto.spi.function.Signature;
 import com.facebook.presto.spi.function.SqlNullable;
 import com.facebook.presto.spi.function.SqlType;
 import com.facebook.presto.type.Constraint;
-import com.facebook.presto.type.LiteralParameter;
 import com.facebook.presto.util.SecureRandomGeneration;
 import com.google.common.primitives.Doubles;
 import io.airlift.slice.Slice;
diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/StringFunctions.java b/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/StringFunctions.java
index 2a55fe9485f29..ba7722722f609 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/StringFunctions.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/StringFunctions.java
@@ -19,6 +19,7 @@
 import com.facebook.presto.common.type.StandardTypes;
 import com.facebook.presto.spi.PrestoException;
 import com.facebook.presto.spi.function.Description;
+import com.facebook.presto.spi.function.LiteralParameter;
 import com.facebook.presto.spi.function.LiteralParameters;
 import com.facebook.presto.spi.function.ScalarFunction;
 import com.facebook.presto.spi.function.ScalarOperator;
@@ -26,7 +27,6 @@
 import com.facebook.presto.spi.function.SqlType;
 import com.facebook.presto.type.CodePointsType;
 import com.facebook.presto.type.Constraint;
-import com.facebook.presto.type.LiteralParameter;
 import com.google.common.primitives.Ints;
 import io.airlift.slice.InvalidCodePointException;
 import io.airlift.slice.InvalidUtf8Exception;
diff --git a/presto-main-base/src/main/java/com/facebook/presto/type/BigintOperators.java b/presto-main-base/src/main/java/com/facebook/presto/type/BigintOperators.java
index 92dc55dd997db..4689354736795 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/type/BigintOperators.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/type/BigintOperators.java
@@ -20,6 +20,7 @@
 import com.facebook.presto.spi.function.BlockIndex;
 import com.facebook.presto.spi.function.BlockPosition;
 import com.facebook.presto.spi.function.IsNull;
+import com.facebook.presto.spi.function.LiteralParameter;
 import com.facebook.presto.spi.function.LiteralParameters;
 import com.facebook.presto.spi.function.ScalarOperator;
 import com.facebook.presto.spi.function.SqlNullable;
diff --git a/presto-main-base/src/main/java/com/facebook/presto/type/LikeFunctions.java b/presto-main-base/src/main/java/com/facebook/presto/type/LikeFunctions.java
index 7b89abc3bf48d..672d8b3542ff8 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/type/LikeFunctions.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/type/LikeFunctions.java
@@ -16,6 +16,7 @@
 import com.facebook.presto.common.function.OperatorType;
 import com.facebook.presto.common.type.StandardTypes;
 import com.facebook.presto.spi.PrestoException;
+import com.facebook.presto.spi.function.LiteralParameter;
 import com.facebook.presto.spi.function.LiteralParameters;
 import com.facebook.presto.spi.function.ScalarFunction;
 import com.facebook.presto.spi.function.ScalarOperator;
diff --git a/presto-main-base/src/test/java/com/facebook/presto/operator/TestAnnotationEngineForAggregates.java b/presto-main-base/src/test/java/com/facebook/presto/operator/TestAnnotationEngineForAggregates.java
index 5927f1fc65ee6..220b91ba548d8 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/operator/TestAnnotationEngineForAggregates.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/operator/TestAnnotationEngineForAggregates.java
@@ -43,6 +43,7 @@
 import com.facebook.presto.spi.function.Description;
 import com.facebook.presto.spi.function.FunctionKind;
 import com.facebook.presto.spi.function.InputFunction;
+import com.facebook.presto.spi.function.LiteralParameter;
 import com.facebook.presto.spi.function.LiteralParameters;
 import com.facebook.presto.spi.function.LongVariableConstraint;
 import com.facebook.presto.spi.function.OperatorDependency;
@@ -53,7 +54,6 @@
 import com.facebook.presto.spi.function.TypeParameterSpecialization;
 import com.facebook.presto.spi.function.aggregation.AggregationMetadata;
 import com.facebook.presto.type.Constraint;
-import com.facebook.presto.type.LiteralParameter;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import io.airlift.slice.Slice;
diff --git a/presto-main-base/src/test/java/com/facebook/presto/operator/TestAnnotationEngineForScalars.java b/presto-main-base/src/test/java/com/facebook/presto/operator/TestAnnotationEngineForScalars.java
index 80ab4502de2e6..787a2a5f8da27 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/operator/TestAnnotationEngineForScalars.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/operator/TestAnnotationEngineForScalars.java
@@ -30,13 +30,13 @@
 import com.facebook.presto.spi.function.Description;
 import com.facebook.presto.spi.function.FunctionKind;
 import com.facebook.presto.spi.function.IsNull;
+import com.facebook.presto.spi.function.LiteralParameter;
 import com.facebook.presto.spi.function.LiteralParameters;
 import com.facebook.presto.spi.function.ScalarFunction;
 import com.facebook.presto.spi.function.Signature;
 import com.facebook.presto.spi.function.SqlNullable;
 import com.facebook.presto.spi.function.SqlType;
 import com.facebook.presto.spi.function.TypeParameter;
-import com.facebook.presto.type.LiteralParameter;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import io.airlift.slice.Slice;
diff --git a/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestStringFunctions.java b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestStringFunctions.java
index 9197fae5aa811..0023402dd9e91 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestStringFunctions.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestStringFunctions.java
@@ -18,10 +18,10 @@
 import com.facebook.presto.common.type.SqlVarbinary;
 import com.facebook.presto.common.type.StandardTypes;
 import com.facebook.presto.spi.function.Description;
+import com.facebook.presto.spi.function.LiteralParameter;
 import com.facebook.presto.spi.function.LiteralParameters;
 import com.facebook.presto.spi.function.ScalarFunction;
 import com.facebook.presto.spi.function.SqlType;
-import com.facebook.presto.type.LiteralParameter;
 import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
diff --git a/presto-main-base/src/main/java/com/facebook/presto/type/LiteralParameter.java b/presto-spi/src/main/java/com/facebook/presto/spi/function/LiteralParameter.java
similarity index 95%
rename from presto-main-base/src/main/java/com/facebook/presto/type/LiteralParameter.java
rename to presto-spi/src/main/java/com/facebook/presto/spi/function/LiteralParameter.java
index 548419647d58b..f57ba796fdc53 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/type/LiteralParameter.java
+++ b/presto-spi/src/main/java/com/facebook/presto/spi/function/LiteralParameter.java
@@ -11,7 +11,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.facebook.presto.type;
+package com.facebook.presto.spi.function;
 
 import java.lang.annotation.ElementType;
 import java.lang.annotation.Retention;

From 76cac253f04f7f03492c7c22d65dae920e3df753 Mon Sep 17 00:00:00 2001
From: wangd 
Date: Tue, 2 Sep 2025 14:37:58 +0800
Subject: [PATCH 043/113] [Native]Fix typo in PrestoServer.h and
 IcebergOutputTableHandle.hpp.inc

---
 presto-native-execution/presto_cpp/main/PrestoServer.h        | 4 ++--
 .../iceberg/special/IcebergOutputTableHandle.hpp.inc          | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/presto-native-execution/presto_cpp/main/PrestoServer.h b/presto-native-execution/presto_cpp/main/PrestoServer.h
index afeb45fc12092..8e658c89a157e 100644
--- a/presto-native-execution/presto_cpp/main/PrestoServer.h
+++ b/presto-native-execution/presto_cpp/main/PrestoServer.h
@@ -172,8 +172,8 @@ class PrestoServer {
 
   /// Invoked to get the ip address of the process. In certain deployment
   /// setup, each process has different ip address. Deployment environment
-  /// may provide there own library to get process specific ip address.
-  /// In such cases, getLocalIp can be overriden to pass process specific
+  /// may provide their own library to get process specific ip address.
+  /// In such cases, getLocalIp can be overridden to pass process specific
   /// ip address.
   virtual std::string getLocalIp() const;
 
diff --git a/presto-native-execution/presto_cpp/presto_protocol/connector/iceberg/special/IcebergOutputTableHandle.hpp.inc b/presto-native-execution/presto_cpp/presto_protocol/connector/iceberg/special/IcebergOutputTableHandle.hpp.inc
index 406238e860df3..e168a870cea33 100644
--- a/presto-native-execution/presto_cpp/presto_protocol/connector/iceberg/special/IcebergOutputTableHandle.hpp.inc
+++ b/presto-native-execution/presto_cpp/presto_protocol/connector/iceberg/special/IcebergOutputTableHandle.hpp.inc
@@ -12,7 +12,7 @@
  * limitations under the License.
  */
 
-// IcebergInsertTableHandle is special since it needs an usage of
+// IcebergOutputTableHandle is special since it needs an usage of
 // hive::.
 
 namespace facebook::presto::protocol::iceberg {

From 8c8e940c1ebdf3f066b70387aff6ffed27181b6c Mon Sep 17 00:00:00 2001
From: wangd 
Date: Mon, 25 Aug 2025 14:50:14 +0800
Subject: [PATCH 044/113] [Native]Remove redundant parentheses in
 PrestoToVeloxQueryPlan.cpp

---
 .../main/types/PrestoToVeloxQueryPlan.cpp      | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.cpp b/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.cpp
index 94428644cea46..ccd4b1d0f6294 100644
--- a/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.cpp
+++ b/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.cpp
@@ -2018,7 +2018,7 @@ core::PlanFragment VeloxQueryPlanConverterBase::toVeloxQueryPlan(
         planFragment.planNode = core::PartitionedOutputNode::single(
             partitionedOutputNodeId,
             outputType,
-            toVeloxSerdeKind((partitioningScheme.encoding)),
+            toVeloxSerdeKind(partitioningScheme.encoding),
             sourceNode);
         return planFragment;
       case protocol::SystemPartitioning::FIXED: {
@@ -2032,7 +2032,7 @@ core::PlanFragment VeloxQueryPlanConverterBase::toVeloxQueryPlan(
               planFragment.planNode = core::PartitionedOutputNode::single(
                   partitionedOutputNodeId,
                   outputType,
-                  toVeloxSerdeKind((partitioningScheme.encoding)),
+                  toVeloxSerdeKind(partitioningScheme.encoding),
                   sourceNode);
               return planFragment;
             }
@@ -2045,7 +2045,7 @@ core::PlanFragment VeloxQueryPlanConverterBase::toVeloxQueryPlan(
                     partitioningScheme.replicateNullsAndAny,
                     std::make_shared(),
                     outputType,
-                    toVeloxSerdeKind((partitioningScheme.encoding)),
+                    toVeloxSerdeKind(partitioningScheme.encoding),
                     sourceNode);
             return planFragment;
           }
@@ -2058,7 +2058,7 @@ core::PlanFragment VeloxQueryPlanConverterBase::toVeloxQueryPlan(
               planFragment.planNode = core::PartitionedOutputNode::single(
                   partitionedOutputNodeId,
                   outputType,
-                  toVeloxSerdeKind((partitioningScheme.encoding)),
+                  toVeloxSerdeKind(partitioningScheme.encoding),
                   sourceNode);
               return planFragment;
             }
@@ -2072,7 +2072,7 @@ core::PlanFragment VeloxQueryPlanConverterBase::toVeloxQueryPlan(
                     std::make_shared(
                         inputType, keyChannels, constValues),
                     outputType,
-                    toVeloxSerdeKind((partitioningScheme.encoding)),
+                    toVeloxSerdeKind(partitioningScheme.encoding),
                     sourceNode);
             return planFragment;
           }
@@ -2081,7 +2081,7 @@ core::PlanFragment VeloxQueryPlanConverterBase::toVeloxQueryPlan(
                 partitionedOutputNodeId,
                 1,
                 outputType,
-                toVeloxSerdeKind((partitioningScheme.encoding)),
+                toVeloxSerdeKind(partitioningScheme.encoding),
                 sourceNode);
             return planFragment;
           }
@@ -2100,7 +2100,7 @@ core::PlanFragment VeloxQueryPlanConverterBase::toVeloxQueryPlan(
         planFragment.planNode = core::PartitionedOutputNode::arbitrary(
             partitionedOutputNodeId,
             std::move(outputType),
-            toVeloxSerdeKind((partitioningScheme.encoding)),
+            toVeloxSerdeKind(partitioningScheme.encoding),
             std::move(sourceNode));
         return planFragment;
       }
@@ -2119,7 +2119,7 @@ core::PlanFragment VeloxQueryPlanConverterBase::toVeloxQueryPlan(
     planFragment.planNode = core::PartitionedOutputNode::single(
         partitionedOutputNodeId,
         outputType,
-        toVeloxSerdeKind((partitioningScheme.encoding)),
+        toVeloxSerdeKind(partitioningScheme.encoding),
         sourceNode);
     return planFragment;
   }
@@ -2135,7 +2135,7 @@ core::PlanFragment VeloxQueryPlanConverterBase::toVeloxQueryPlan(
       partitioningScheme.replicateNullsAndAny,
       std::shared_ptr(std::move(spec)),
       toRowType(partitioningScheme.outputLayout, typeParser_),
-      toVeloxSerdeKind((partitioningScheme.encoding)),
+      toVeloxSerdeKind(partitioningScheme.encoding),
       sourceNode);
   return planFragment;
 }

From 8cf71c16b61523ed04281fbab37717990a389c94 Mon Sep 17 00:00:00 2001
From: Hazmi 
Date: Thu, 17 Jul 2025 18:24:43 +0300
Subject: [PATCH 045/113] Add a DB based session property manager

Co-authored-by: Ashish Tadose 
Co-authored-by: Ariel Weisberg 
Co-authored-by: Jalpreet Singh Nanda (:imjalpreet) 
---
 README.md                                     |   2 +
 pom.xml                                       |   6 +
 .../admin/session-property-managers.rst       |  84 ++++++++++-
 presto-function-namespace-managers/pom.xml    |   1 -
 presto-main/etc/session-property-config.json  |   0
 .../etc/session-property-config.properties    |   0
 presto-session-property-managers/pom.xml      |  63 +++++++++
 .../AbstractSessionPropertyManager.java       |  65 +++++++++
 .../presto/session/SessionMatchSpec.java      | 103 +++++++++++++-
 ...onPropertyConfigurationManagerPlugin.java} |   8 +-
 .../session/db/DbSessionPropertyManager.java  |  47 +++++++
 .../db/DbSessionPropertyManagerConfig.java    |  54 ++++++++
 .../db/DbSessionPropertyManagerFactory.java   |  53 +++++++
 .../db/DbSessionPropertyManagerModule.java    |  33 +++++
 .../presto/session/db/DbSpecsProvider.java    |  29 ++++
 .../session/db/RefreshingDbSpecsProvider.java |  94 +++++++++++++
 .../session/db/SessionPropertiesDao.java      | 130 ++++++++++++++++++
 .../db/SessionPropertiesDaoProvider.java      |  46 +++++++
 .../FileSessionPropertyManager.java           |  45 ++----
 .../FileSessionPropertyManagerConfig.java     |   2 +-
 .../FileSessionPropertyManagerFactory.java    |   2 +-
 .../FileSessionPropertyManagerModule.java     |   2 +-
 ...> AbstractTestSessionPropertyManager.java} |  77 ++++++-----
 .../db/TestDbSessionPropertyManager.java      | 112 +++++++++++++++
 .../TestDbSessionPropertyManagerConfig.java   |  52 +++++++
 .../TestDbSessionPropertyManagerMariadb.java  |  29 ++++
 .../db/TestDbSessionPropertyManagerMysql.java |  23 ++++
 .../file/TestFileSessionPropertyManager.java  |  62 +++++++++
 .../TestFileSessionPropertyManagerConfig.java |   2 +-
 ...estoSparkLauncherIntegrationSmokeTest.java |   5 +-
 30 files changed, 1144 insertions(+), 87 deletions(-)
 create mode 100644 presto-main/etc/session-property-config.json
 create mode 100644 presto-main/etc/session-property-config.properties
 create mode 100644 presto-session-property-managers/src/main/java/com/facebook/presto/session/AbstractSessionPropertyManager.java
 rename presto-session-property-managers/src/main/java/com/facebook/presto/session/{FileSessionPropertyManagerPlugin.java => SessionPropertyConfigurationManagerPlugin.java} (73%)
 create mode 100644 presto-session-property-managers/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManager.java
 create mode 100644 presto-session-property-managers/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManagerConfig.java
 create mode 100644 presto-session-property-managers/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManagerFactory.java
 create mode 100644 presto-session-property-managers/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManagerModule.java
 create mode 100644 presto-session-property-managers/src/main/java/com/facebook/presto/session/db/DbSpecsProvider.java
 create mode 100644 presto-session-property-managers/src/main/java/com/facebook/presto/session/db/RefreshingDbSpecsProvider.java
 create mode 100644 presto-session-property-managers/src/main/java/com/facebook/presto/session/db/SessionPropertiesDao.java
 create mode 100644 presto-session-property-managers/src/main/java/com/facebook/presto/session/db/SessionPropertiesDaoProvider.java
 rename presto-session-property-managers/src/main/java/com/facebook/presto/session/{ => file}/FileSessionPropertyManager.java (60%)
 rename presto-session-property-managers/src/main/java/com/facebook/presto/session/{ => file}/FileSessionPropertyManagerConfig.java (96%)
 rename presto-session-property-managers/src/main/java/com/facebook/presto/session/{ => file}/FileSessionPropertyManagerFactory.java (97%)
 rename presto-session-property-managers/src/main/java/com/facebook/presto/session/{ => file}/FileSessionPropertyManagerModule.java (96%)
 rename presto-session-property-managers/src/test/java/com/facebook/presto/session/{TestFileSessionPropertyManager.java => AbstractTestSessionPropertyManager.java} (72%)
 create mode 100644 presto-session-property-managers/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManager.java
 create mode 100644 presto-session-property-managers/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManagerConfig.java
 create mode 100644 presto-session-property-managers/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManagerMariadb.java
 create mode 100644 presto-session-property-managers/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManagerMysql.java
 create mode 100644 presto-session-property-managers/src/test/java/com/facebook/presto/session/file/TestFileSessionPropertyManager.java
 rename presto-session-property-managers/src/test/java/com/facebook/presto/session/{ => file}/TestFileSessionPropertyManagerConfig.java (97%)

diff --git a/README.md b/README.md
index 2da36a8a31aea..2b4aa6d998c1d 100644
--- a/README.md
+++ b/README.md
@@ -59,6 +59,8 @@ Additionally, the Hive plugin must be configured with location of your Hive meta
 
     -Dhive.metastore.uri=thrift://localhost:9083
 
+To modify the loaded plugins in IntelliJ, modify the `config.properties` located in `presto-main/etc`. You can modify `plugin.bundles` with the location of the plugin pom.xml
+
 ### Additional configuration for Java 17
 
 When running with Java 17, additional `--add-opens` flags are required to allow reflective access used by certain catalogs based on which catalogs are configured.  
diff --git a/pom.xml b/pom.xml
index c792fc783d343..5bdb0ad74d96e 100644
--- a/pom.xml
+++ b/pom.xml
@@ -89,6 +89,7 @@
         3.18.0
         6.0.0
         17.0.0
+        3.5.4
 
         2
 
@@ -2679,6 +2680,11 @@
                 ${dep.arrow.version}
             
 
+            
+                org.mariadb.jdbc
+                mariadb-java-client
+                ${dep.mariadb.version}
+            
         
     
 
diff --git a/presto-docs/src/main/sphinx/admin/session-property-managers.rst b/presto-docs/src/main/sphinx/admin/session-property-managers.rst
index 5ebdfb1a88bc3..8147d9741f051 100644
--- a/presto-docs/src/main/sphinx/admin/session-property-managers.rst
+++ b/presto-docs/src/main/sphinx/admin/session-property-managers.rst
@@ -5,10 +5,11 @@ Session Property Managers
 Administrators can add session properties to control the behavior for subsets of their workload.
 These properties are defaults and can be overridden by users (if authorized to do so). Session
 properties can be used to control resource usage, enable or disable features, and change query
-characteristics. Session property managers are pluggable.
+characteristics. Session property managers are pluggable. A session property manager can either
+be database-based or file-based.
 
-Add an ``etc/session-property-config.properties`` file with the following contents to enable
-the built-in manager that reads a JSON config file:
+To enable a built-in manager that reads a JSON configuration file, add an
+``etc/session-property-config.properties`` file with the following contents:
 
 .. code-block:: none
 
@@ -24,6 +25,66 @@ by default. All matching rules contribute to constructing a list of session prop
 are applied in the order they are specified. Rules specified later in the file override values
 for properties that have been previously encountered.
 
+
+For the database-based built-in manager, add an
+``etc/session-property-config.properties`` file with the following contents:
+
+.. code-block:: text
+
+    session-property-config.configuration-manager=db
+    session-property-manager.db.url=jdbc:mysql://localhost:3306/session_properties?user=user&password=pass&createDatabaseIfNotExist=true
+    session-property-manager.db.refresh-period=50s
+
+Change the value of ``session-property-manager.db.url`` to the JDBC URL of a database.
+
+``session-property-manager.db.refresh-period`` should be set to how often Presto refreshes
+to fetch the latest session properties from the database.
+
+This database consists of three tables: ``session_specs``, ``session_client_tags`` and ``session_property_values``.
+Presto will create the database on startup if you set ``createDatabaseIfNotExist`` to ``true`` in your JDBC URL.
+If the tables do not exist, Presto will create them on startup.
+
+.. code-block:: text
+
+    mysql> DESCRIBE session_specs;
+    +-----------------------------+--------------+------+-----+---------+----------------+
+    | Field                       | Type         | Null | Key | Default | Extra          |
+    +-----------------------------+--------------+------+-----+---------+----------------+
+    | spec_id                     | bigint       | NO   | PRI | NULL    | auto_increment |
+    | user_regex                  | varchar(512) | YES  |     | NULL    |                |
+    | source_regex                | varchar(512) | YES  |     | NULL    |                |
+    | query_type                  | varchar(512) | YES  |     | NULL    |                |
+    | group_regex                 | varchar(512) | YES  |     | NULL    |                |
+    | client_info_regex           | varchar(512) | YES  |     | NULL    |                |
+    | override_session_properties | tinyint(1)   | YES  |     | NULL    |                |
+    | priority                    | int          | NO   |     | NULL    |                |
+    +-----------------------------+--------------+------+-----+---------+----------------+
+    8 rows in set (0.016 sec)
+
+.. code-block:: text
+
+    mysql> DESCRIBE session_client_tags;
+    +-------------+--------------+------+-----+---------+-------+
+    | Field       | Type         | Null | Key | Default | Extra |
+    +-------------+--------------+------+-----+---------+-------+
+    | tag_spec_id | bigint       | NO   | PRI | NULL    |       |
+    | client_tag  | varchar(512) | NO   | PRI | NULL    |       |
+    +-------------+--------------+------+-----+---------+-------+
+    2 rows in set (0.062 sec)
+
+.. code-block:: text
+
+    mysql> DESCRIBE session_property_values;
+    +--------------------------+--------------+------+-----+---------+-------+
+    | Field                    | Type         | Null | Key | Default | Extra |
+    +--------------------------+--------------+------+-----+---------+-------+
+    | property_spec_id         | bigint       | NO   | PRI | NULL    |       |
+    | session_property_name    | varchar(512) | NO   | PRI | NULL    |       |
+    | session_property_value   | varchar(512) | YES  |     | NULL    |       |
+    | session_property_catalog | varchar(512) | YES  |     | NULL    |       |
+    +--------------------------+--------------+------+-----+---------+-------+
+    3 rows in set (0.009 sec)
+
 Match Rules
 -----------
 
@@ -52,9 +113,15 @@ Match Rules
   Note that once a session property has been overridden by ANY rule it remains overridden even if later higher precedence rules change the
   value, but don't specify override.
 
-* ``sessionProperties``: map with string keys and values. Each entry is a system or catalog property name and
+* ``sessionProperties``: map with string keys and values. Each entry is a system property name and
   corresponding value. Values must be specified as strings, no matter the actual data type.
 
+* ``catalogSessionProperties``: map with string keys corresponding to the catalog name, and a map with string keys
+  and values as the value. Each entry is a catalog name and corresponding map of session property values.
+
+* For the database session property manager, catalog & system session properties are located in the same table.
+  ``session_property_catalog`` should be null for system session properties.
+
 Example
 -------
 
@@ -71,6 +138,8 @@ Consider the following set of requirements:
 * All high memory ETL queries (tagged with 'high_mem_etl') are routed to subgroups under the ``global.pipeline`` group,
   and must be configured to enable :doc:`/admin/exchange-materialization`.
 
+* All iceberg catalog queries should override the ``delete-as-join-rewrite-enabled`` property
+
 These requirements can be expressed with the following rules:
 
 .. code-block:: json
@@ -104,5 +173,12 @@ These requirements can be expressed with the following rules:
           "partitioning_provider_catalog": "hive",
           "hash_partition_count": 4096
         }
+      },
+      {
+        "catalogSessionProperties": {
+          "iceberg": {
+            "delete_as_join_rewrite_enabled": "true"
+          }
+        }
       }
     ]
diff --git a/presto-function-namespace-managers/pom.xml b/presto-function-namespace-managers/pom.xml
index 7e3a3e95425df..a882b50e35df6 100644
--- a/presto-function-namespace-managers/pom.xml
+++ b/presto-function-namespace-managers/pom.xml
@@ -148,7 +148,6 @@
         
             org.mariadb.jdbc
             mariadb-java-client
-            3.5.4
             runtime
         
 
diff --git a/presto-main/etc/session-property-config.json b/presto-main/etc/session-property-config.json
new file mode 100644
index 0000000000000..e69de29bb2d1d
diff --git a/presto-main/etc/session-property-config.properties b/presto-main/etc/session-property-config.properties
new file mode 100644
index 0000000000000..e69de29bb2d1d
diff --git a/presto-session-property-managers/pom.xml b/presto-session-property-managers/pom.xml
index 6386c244b92e5..0ba95f4c6ee88 100644
--- a/presto-session-property-managers/pom.xml
+++ b/presto-session-property-managers/pom.xml
@@ -24,6 +24,11 @@
             bootstrap
         
 
+        
+            com.facebook.airlift
+            log
+        
+
         
             com.facebook.airlift
             json
@@ -44,6 +49,11 @@
             guice
         
 
+        
+            javax.annotation
+            javax.annotation-api
+        
+
         
             javax.inject
             javax.inject
@@ -54,6 +64,11 @@
             jakarta.inject-api
         
 
+        
+            com.facebook.airlift
+            concurrent
+        
+
         
             jakarta.validation
             jakarta.validation-api
@@ -69,6 +84,30 @@
             jackson-core
         
 
+        
+            org.jdbi
+            jdbi3-core
+        
+
+        
+            org.jdbi
+            jdbi3-sqlobject
+        
+
+        
+            com.mysql
+            mysql-connector-j
+            true
+            runtime
+        
+
+        
+            org.mariadb.jdbc
+            mariadb-java-client
+            true
+            runtime
+        
+
         
         
             com.facebook.presto
@@ -88,12 +127,24 @@
             provided
         
 
+        
+            io.airlift
+            slice
+            provided
+        
+
         
             com.fasterxml.jackson.core
             jackson-annotations
             provided
         
 
+        
+            org.openjdk.jol
+            jol-core
+            provided
+        
+
         
         
             com.facebook.presto
@@ -112,5 +163,17 @@
             testing
             test
         
+
+        
+            com.facebook.presto
+            testing-mysql-server-8
+            test
+        
+
+        
+            com.facebook.presto
+            testing-mysql-server-base
+            test
+        
     
 
diff --git a/presto-session-property-managers/src/main/java/com/facebook/presto/session/AbstractSessionPropertyManager.java b/presto-session-property-managers/src/main/java/com/facebook/presto/session/AbstractSessionPropertyManager.java
new file mode 100644
index 0000000000000..7780652db038a
--- /dev/null
+++ b/presto-session-property-managers/src/main/java/com/facebook/presto/session/AbstractSessionPropertyManager.java
@@ -0,0 +1,65 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.session;
+
+import com.facebook.presto.spi.session.SessionConfigurationContext;
+import com.facebook.presto.spi.session.SessionPropertyConfigurationManager;
+
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+public abstract class AbstractSessionPropertyManager
+        implements SessionPropertyConfigurationManager
+{
+    private static final String DEFAULT_PROPERTIES = "defaultProperties";
+    private static final String OVERRIDE_PROPERTIES = "overrideProperties";
+
+    @Override
+    public final SystemSessionPropertyConfiguration getSystemSessionProperties(SessionConfigurationContext context)
+    {
+        Map defaultProperties = new HashMap<>();
+        Set overridePropertyNames = new HashSet<>();
+        for (SessionMatchSpec sessionMatchSpec : getSessionMatchSpecs()) {
+            Map newProperties = sessionMatchSpec.match(sessionMatchSpec.getSessionProperties(), context);
+            defaultProperties.putAll(newProperties);
+            if (sessionMatchSpec.getOverrideSessionProperties().orElse(false)) {
+                overridePropertyNames.addAll(newProperties.keySet());
+            }
+        }
+
+        // Once a property has been overridden it stays that way and the value is updated by any rule
+        Map overrideProperties = new HashMap<>();
+        for (String propertyName : overridePropertyNames) {
+            overrideProperties.put(propertyName, defaultProperties.get(propertyName));
+        }
+
+        return new SystemSessionPropertyConfiguration(defaultProperties, overrideProperties);
+    }
+
+    @Override
+    public final Map> getCatalogSessionProperties(SessionConfigurationContext context)
+    {
+        Map> catalogProperties = new HashMap<>();
+        for (SessionMatchSpec sessionMatchSpec : getSessionMatchSpecs()) {
+            Map> newProperties = sessionMatchSpec.match(sessionMatchSpec.getCatalogSessionProperties(), context);
+            catalogProperties.putAll(newProperties);
+        }
+        return catalogProperties;
+    }
+
+    protected abstract List getSessionMatchSpecs();
+}
diff --git a/presto-session-property-managers/src/main/java/com/facebook/presto/session/SessionMatchSpec.java b/presto-session-property-managers/src/main/java/com/facebook/presto/session/SessionMatchSpec.java
index 18f34a80adf9d..f5e655c705eef 100644
--- a/presto-session-property-managers/src/main/java/com/facebook/presto/session/SessionMatchSpec.java
+++ b/presto-session-property-managers/src/main/java/com/facebook/presto/session/SessionMatchSpec.java
@@ -17,16 +17,24 @@
 import com.facebook.presto.spi.session.SessionConfigurationContext;
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
+import com.google.common.base.Splitter;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+import org.jdbi.v3.core.mapper.RowMapper;
+import org.jdbi.v3.core.statement.StatementContext;
 
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.regex.Pattern;
 
+import static com.facebook.presto.session.db.SessionPropertiesDao.EMPTY_CATALOG;
+import static com.google.common.base.Preconditions.checkArgument;
 import static java.util.Objects.requireNonNull;
 
 public class SessionMatchSpec
@@ -39,6 +47,7 @@ public class SessionMatchSpec
     private final Optional resourceGroupRegex;
     private final Optional overrideSessionProperties;
     private final Map sessionProperties;
+    private final Map> catalogSessionProperties;
 
     @JsonCreator
     public SessionMatchSpec(
@@ -49,7 +58,8 @@ public SessionMatchSpec(
             @JsonProperty("group") Optional resourceGroupRegex,
             @JsonProperty("clientInfo") Optional clientInfoRegex,
             @JsonProperty("overrideSessionProperties") Optional overrideSessionProperties,
-            @JsonProperty("sessionProperties") Map sessionProperties)
+            @JsonProperty("sessionProperties") Map sessionProperties,
+            @JsonProperty("catalogSessionProperties") Map> catalogSessionProperties)
     {
         this.userRegex = requireNonNull(userRegex, "userRegex is null");
         this.sourceRegex = requireNonNull(sourceRegex, "sourceRegex is null");
@@ -61,9 +71,10 @@ public SessionMatchSpec(
         this.overrideSessionProperties = requireNonNull(overrideSessionProperties, "overrideSessionProperties is null");
         requireNonNull(sessionProperties, "sessionProperties is null");
         this.sessionProperties = ImmutableMap.copyOf(sessionProperties);
+        this.catalogSessionProperties = ImmutableMap.copyOf(catalogSessionProperties);
     }
 
-    public Map match(SessionConfigurationContext context)
+    public  Map match(Map object, SessionConfigurationContext context)
     {
         if (userRegex.isPresent() && !userRegex.get().matcher(context.getUser()).matches()) {
             return ImmutableMap.of();
@@ -99,7 +110,7 @@ public Map match(SessionConfigurationContext context)
             }
         }
 
-        return sessionProperties;
+        return object;
     }
 
     @JsonProperty("user")
@@ -149,4 +160,90 @@ public Map getSessionProperties()
     {
         return sessionProperties;
     }
+
+    @JsonProperty
+    public Map> getCatalogSessionProperties()
+    {
+        return catalogSessionProperties;
+    }
+
+    public static class Mapper
+            implements RowMapper
+    {
+        @Override
+        public SessionMatchSpec map(ResultSet resultSet, StatementContext context)
+                throws SQLException
+        {
+            SessionInfo sessionInfo = getProperties(
+                    Optional.ofNullable(resultSet.getString("session_property_names")),
+                    Optional.ofNullable(resultSet.getString("session_property_values")),
+                    Optional.ofNullable(resultSet.getString("session_property_catalogs")));
+
+            return new SessionMatchSpec(
+                    Optional.ofNullable(resultSet.getString("user_regex")).map(Pattern::compile),
+                    Optional.ofNullable(resultSet.getString("source_regex")).map(Pattern::compile),
+                    Optional.ofNullable(resultSet.getString("client_tags")).map(tag -> Splitter.on(",").splitToList(tag)),
+                    Optional.ofNullable(resultSet.getString("query_type")),
+                    Optional.ofNullable(resultSet.getString("group_regex")).map(Pattern::compile),
+                    Optional.ofNullable(resultSet.getString("client_info_regex")).map(Pattern::compile),
+                    Optional.of(resultSet.getBoolean("override_session_properties")),
+                    sessionInfo.getSessionProperties(),
+                    sessionInfo.getCatalogProperties());
+        }
+
+        private SessionInfo getProperties(Optional names, Optional values, Optional catalogs)
+        {
+            if (!names.isPresent()) {
+                return new SessionInfo(ImmutableMap.of(), ImmutableMap.of());
+            }
+
+            checkArgument(catalogs.isPresent(), "names are present, but catalogs are not");
+            checkArgument(values.isPresent(), "names are present, but values are not");
+            List sessionPropertyNames = Splitter.on(",").splitToList(names.get());
+            List sessionPropertyValues = Splitter.on(",").splitToList(values.get());
+            List sessionPropertyCatalogs = Splitter.on(",").splitToList(catalogs.get());
+            checkArgument(sessionPropertyNames.size() == sessionPropertyValues.size(),
+                    "The number of property names and values should be the same");
+
+            checkArgument(sessionPropertyNames.size() == sessionPropertyCatalogs.size(),
+                    "The number of property names and catalog values should be the same");
+
+            Map> catalogSessionProperties = new HashMap<>();
+            Map systemSessionProperties = new HashMap<>();
+            for (int i = 0; i < sessionPropertyNames.size(); i++) {
+                if (sessionPropertyCatalogs.get(i).equals(EMPTY_CATALOG)) {
+                    systemSessionProperties.put(sessionPropertyNames.get(i), sessionPropertyValues.get(i));
+                }
+                else {
+                    catalogSessionProperties
+                            .computeIfAbsent(sessionPropertyCatalogs.get(i), k -> new HashMap<>())
+                            .put(sessionPropertyNames.get(i), sessionPropertyValues.get(i));
+                }
+            }
+
+            return new SessionInfo(systemSessionProperties, catalogSessionProperties);
+        }
+        static class SessionInfo
+        {
+            private final Map sessionProperties;
+            private final Map> catalogProperties;
+
+            public SessionInfo(Map sessionProperties,
+                               Map> catalogProperties)
+            {
+                this.sessionProperties = sessionProperties;
+                this.catalogProperties = catalogProperties;
+            }
+
+            public Map getSessionProperties()
+            {
+                return sessionProperties;
+            }
+
+            public Map> getCatalogProperties()
+            {
+                return catalogProperties;
+            }
+        }
+    }
 }
diff --git a/presto-session-property-managers/src/main/java/com/facebook/presto/session/FileSessionPropertyManagerPlugin.java b/presto-session-property-managers/src/main/java/com/facebook/presto/session/SessionPropertyConfigurationManagerPlugin.java
similarity index 73%
rename from presto-session-property-managers/src/main/java/com/facebook/presto/session/FileSessionPropertyManagerPlugin.java
rename to presto-session-property-managers/src/main/java/com/facebook/presto/session/SessionPropertyConfigurationManagerPlugin.java
index bde86e9400f31..b89048a8733f5 100644
--- a/presto-session-property-managers/src/main/java/com/facebook/presto/session/FileSessionPropertyManagerPlugin.java
+++ b/presto-session-property-managers/src/main/java/com/facebook/presto/session/SessionPropertyConfigurationManagerPlugin.java
@@ -13,16 +13,20 @@
  */
 package com.facebook.presto.session;
 
+import com.facebook.presto.session.db.DbSessionPropertyManagerFactory;
+import com.facebook.presto.session.file.FileSessionPropertyManagerFactory;
 import com.facebook.presto.spi.Plugin;
 import com.facebook.presto.spi.session.SessionPropertyConfigurationManagerFactory;
 import com.google.common.collect.ImmutableList;
 
-public class FileSessionPropertyManagerPlugin
+public class SessionPropertyConfigurationManagerPlugin
         implements Plugin
 {
     @Override
     public Iterable getSessionPropertyConfigurationManagerFactories()
     {
-        return ImmutableList.of(new FileSessionPropertyManagerFactory());
+        return ImmutableList.of(
+                new FileSessionPropertyManagerFactory(),
+                new DbSessionPropertyManagerFactory());
     }
 }
diff --git a/presto-session-property-managers/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManager.java b/presto-session-property-managers/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManager.java
new file mode 100644
index 0000000000000..528d51c395c6e
--- /dev/null
+++ b/presto-session-property-managers/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManager.java
@@ -0,0 +1,47 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.session.db;
+
+import com.facebook.presto.session.AbstractSessionPropertyManager;
+import com.facebook.presto.session.SessionMatchSpec;
+import com.facebook.presto.spi.session.SessionConfigurationContext;
+import com.facebook.presto.spi.session.SessionPropertyConfigurationManager;
+
+import javax.inject.Inject;
+
+import java.util.List;
+
+import static java.util.Objects.requireNonNull;
+
+/**
+ * A {@link SessionPropertyConfigurationManager} implementation that connects to a database for fetching information
+ * about session property overrides given {@link SessionConfigurationContext}.
+ */
+public class DbSessionPropertyManager
+        extends AbstractSessionPropertyManager
+{
+    private final DbSpecsProvider specsProvider;
+
+    @Inject
+    public DbSessionPropertyManager(DbSpecsProvider specsProvider)
+    {
+        this.specsProvider = requireNonNull(specsProvider, "specsProvider is null");
+    }
+
+    @Override
+    protected List getSessionMatchSpecs()
+    {
+        return this.specsProvider.get();
+    }
+}
diff --git a/presto-session-property-managers/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManagerConfig.java b/presto-session-property-managers/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManagerConfig.java
new file mode 100644
index 0000000000000..4a94ac0d9e61a
--- /dev/null
+++ b/presto-session-property-managers/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManagerConfig.java
@@ -0,0 +1,54 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.session.db;
+
+import com.facebook.airlift.configuration.Config;
+import com.facebook.airlift.units.Duration;
+import com.facebook.airlift.units.MinDuration;
+import jakarta.validation.constraints.NotNull;
+
+import static java.util.concurrent.TimeUnit.SECONDS;
+
+public class DbSessionPropertyManagerConfig
+{
+    private String configDbUrl;
+    private Duration specsRefreshPeriod = new Duration(10, SECONDS);
+
+    @NotNull
+    public String getConfigDbUrl()
+    {
+        return configDbUrl;
+    }
+
+    @Config("session-property-manager.db.url")
+    public DbSessionPropertyManagerConfig setConfigDbUrl(String configDbUrl)
+    {
+        this.configDbUrl = configDbUrl;
+        return this;
+    }
+
+    @NotNull
+    @MinDuration("1ms")
+    public Duration getSpecsRefreshPeriod()
+    {
+        return specsRefreshPeriod;
+    }
+
+    @Config("session-property-manager.db.refresh-period")
+    public DbSessionPropertyManagerConfig setSpecsRefreshPeriod(Duration specsRefreshPeriod)
+    {
+        this.specsRefreshPeriod = specsRefreshPeriod;
+        return this;
+    }
+}
diff --git a/presto-session-property-managers/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManagerFactory.java b/presto-session-property-managers/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManagerFactory.java
new file mode 100644
index 0000000000000..7a6eba7821ef0
--- /dev/null
+++ b/presto-session-property-managers/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManagerFactory.java
@@ -0,0 +1,53 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.session.db;
+
+import com.facebook.airlift.bootstrap.Bootstrap;
+import com.facebook.airlift.json.JsonModule;
+import com.facebook.presto.spi.resourceGroups.SessionPropertyConfigurationManagerContext;
+import com.facebook.presto.spi.session.SessionPropertyConfigurationManager;
+import com.facebook.presto.spi.session.SessionPropertyConfigurationManagerFactory;
+import com.google.inject.Injector;
+
+import java.util.Map;
+
+import static com.google.common.base.Throwables.throwIfUnchecked;
+
+public class DbSessionPropertyManagerFactory
+        implements SessionPropertyConfigurationManagerFactory
+{
+    @Override
+    public String getName()
+    {
+        return "db";
+    }
+
+    @Override
+    public SessionPropertyConfigurationManager create(Map config, SessionPropertyConfigurationManagerContext context)
+    {
+        try {
+            Bootstrap app = new Bootstrap(new JsonModule(), new DbSessionPropertyManagerModule());
+            Injector injector = app
+                    .doNotInitializeLogging()
+                    .setRequiredConfigurationProperties(config)
+                    .initialize();
+
+            return injector.getInstance(DbSessionPropertyManager.class);
+        }
+        catch (Exception e) {
+            throwIfUnchecked(e);
+            throw new RuntimeException(e);
+        }
+    }
+}
diff --git a/presto-session-property-managers/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManagerModule.java b/presto-session-property-managers/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManagerModule.java
new file mode 100644
index 0000000000000..48f480b7603aa
--- /dev/null
+++ b/presto-session-property-managers/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManagerModule.java
@@ -0,0 +1,33 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.session.db;
+
+import com.google.inject.Binder;
+import com.google.inject.Module;
+import com.google.inject.Scopes;
+
+import static com.facebook.airlift.configuration.ConfigBinder.configBinder;
+
+public class DbSessionPropertyManagerModule
+        implements Module
+{
+    @Override
+    public void configure(Binder binder)
+    {
+        configBinder(binder).bindConfig(DbSessionPropertyManagerConfig.class);
+        binder.bind(DbSessionPropertyManager.class).in(Scopes.SINGLETON);
+        binder.bind(SessionPropertiesDao.class).toProvider(SessionPropertiesDaoProvider.class).in(Scopes.SINGLETON);
+        binder.bind(DbSpecsProvider.class).to(RefreshingDbSpecsProvider.class).in(Scopes.SINGLETON);
+    }
+}
diff --git a/presto-session-property-managers/src/main/java/com/facebook/presto/session/db/DbSpecsProvider.java b/presto-session-property-managers/src/main/java/com/facebook/presto/session/db/DbSpecsProvider.java
new file mode 100644
index 0000000000000..278b52daadb67
--- /dev/null
+++ b/presto-session-property-managers/src/main/java/com/facebook/presto/session/db/DbSpecsProvider.java
@@ -0,0 +1,29 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.session.db;
+
+import com.facebook.presto.session.SessionMatchSpec;
+
+import java.util.List;
+import java.util.function.Supplier;
+
+/**
+ * This interface was created to separate the scheduling logic for {@link SessionMatchSpec} loading. This also helps
+ * us test the core logic of {@link DbSessionPropertyManager} in a modular fashion by letting us use a test
+ * implementation of this interface.
+ */
+public interface DbSpecsProvider
+        extends Supplier>
+{
+}
diff --git a/presto-session-property-managers/src/main/java/com/facebook/presto/session/db/RefreshingDbSpecsProvider.java b/presto-session-property-managers/src/main/java/com/facebook/presto/session/db/RefreshingDbSpecsProvider.java
new file mode 100644
index 0000000000000..2694fd5e54e9e
--- /dev/null
+++ b/presto-session-property-managers/src/main/java/com/facebook/presto/session/db/RefreshingDbSpecsProvider.java
@@ -0,0 +1,94 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.session.db;
+
+import com.facebook.airlift.log.Logger;
+import com.facebook.presto.session.SessionMatchSpec;
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.ImmutableList;
+
+import javax.annotation.PostConstruct;
+import javax.annotation.PreDestroy;
+import javax.inject.Inject;
+
+import java.util.List;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static com.facebook.airlift.concurrent.Threads.daemonThreadsNamed;
+import static java.util.Objects.requireNonNull;
+import static java.util.concurrent.Executors.newSingleThreadScheduledExecutor;
+
+/**
+ * Periodically schedules the loading of specs from the database during initialization. Returns the most recent successfully
+ * loaded specs on every get() invocation.
+ */
+public class RefreshingDbSpecsProvider
+        implements DbSpecsProvider
+{
+    private static final Logger log = Logger.get(RefreshingDbSpecsProvider.class);
+
+    private final AtomicReference> sessionMatchSpecs = new AtomicReference<>(ImmutableList.of());
+    private final SessionPropertiesDao dao;
+
+    private final ScheduledExecutorService executor = newSingleThreadScheduledExecutor(daemonThreadsNamed("RefreshingDbSpecsProvider"));
+    private final AtomicBoolean started = new AtomicBoolean();
+    private final long refreshPeriodMillis;
+
+    @Inject
+    public RefreshingDbSpecsProvider(DbSessionPropertyManagerConfig config, SessionPropertiesDao dao)
+    {
+        requireNonNull(config, "config is null");
+        this.dao = requireNonNull(dao, "dao is null");
+        this.refreshPeriodMillis = config.getSpecsRefreshPeriod().toMillis();
+        dao.createSessionSpecsTable();
+        dao.createSessionClientTagsTable();
+        dao.createSessionPropertiesTable();
+    }
+
+    @PostConstruct
+    public void initialize()
+    {
+        if (!started.getAndSet(true)) {
+            executor.scheduleWithFixedDelay(this::refresh, 0, refreshPeriodMillis, TimeUnit.MILLISECONDS);
+        }
+    }
+
+    @VisibleForTesting
+    void refresh()
+    {
+        try {
+            sessionMatchSpecs.set(ImmutableList.copyOf(dao.getSessionMatchSpecs()));
+        }
+        catch (Throwable e) {
+            // Catch all exceptions here since throwing an exception from executor#scheduleWithFixedDelay method
+            // suppresses all future scheduled invocations
+            log.error(e, "Error loading configuration from database");
+        }
+    }
+
+    @PreDestroy
+    public void destroy()
+    {
+        executor.shutdownNow();
+    }
+
+    @Override
+    public List get()
+    {
+        return sessionMatchSpecs.get();
+    }
+}
diff --git a/presto-session-property-managers/src/main/java/com/facebook/presto/session/db/SessionPropertiesDao.java b/presto-session-property-managers/src/main/java/com/facebook/presto/session/db/SessionPropertiesDao.java
new file mode 100644
index 0000000000000..71a93f0838c4a
--- /dev/null
+++ b/presto-session-property-managers/src/main/java/com/facebook/presto/session/db/SessionPropertiesDao.java
@@ -0,0 +1,130 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.session.db;
+
+import com.facebook.presto.session.SessionMatchSpec;
+import com.google.common.annotations.VisibleForTesting;
+import org.jdbi.v3.sqlobject.customizer.Bind;
+import org.jdbi.v3.sqlobject.statement.SqlQuery;
+import org.jdbi.v3.sqlobject.statement.SqlUpdate;
+import org.jdbi.v3.sqlobject.statement.UseRowMapper;
+
+import java.util.List;
+
+/**
+ * Dao should guarantee that the list of SessionMatchSpecs is returned in increasing order of priority. i.e. if two
+ * rows in the ResultSet specify different values for the same property, the row coming in later will override the
+ * value set by the row coming in earlier.
+ */
+public interface SessionPropertiesDao
+{
+    String SESSION_SPECS_TABLE = "session_specs";
+    String CLIENT_TAGS_TABLE = "session_client_tags";
+    String PROPERTIES_TABLE = "session_property_values";
+    public static String EMPTY_CATALOG = "__NULL__";
+
+    @SqlUpdate("CREATE TABLE IF NOT EXISTS " + SESSION_SPECS_TABLE + "(\n" +
+            "spec_id BIGINT NOT NULL AUTO_INCREMENT,\n" +
+            "user_regex VARCHAR(512),\n" +
+            "source_regex VARCHAR(512),\n" +
+            "query_type VARCHAR(512),\n" +
+            "group_regex VARCHAR(512),\n" +
+            "client_info_regex VARCHAR(512),\n" +
+            "override_session_properties TINYINT(1),\n" +
+            "priority INT NOT NULL,\n" +
+            "PRIMARY KEY (spec_id)\n" +
+            ")")
+    void createSessionSpecsTable();
+
+    @SqlUpdate("CREATE TABLE IF NOT EXISTS " + CLIENT_TAGS_TABLE + "(\n" +
+            "tag_spec_id BIGINT NOT NULL,\n" +
+            "client_tag VARCHAR(512) NOT NULL,\n" +
+            "PRIMARY KEY (tag_spec_id, client_tag),\n" +
+            "FOREIGN KEY (tag_spec_id) REFERENCES session_specs (spec_id)\n" +
+            ")")
+    void createSessionClientTagsTable();
+
+    @SqlUpdate("CREATE TABLE IF NOT EXISTS " + PROPERTIES_TABLE + "(\n" +
+            "property_spec_id BIGINT NOT NULL,\n" +
+            "session_property_name VARCHAR(512),\n" +
+            "session_property_value VARCHAR(512),\n" +
+            "catalog VARCHAR(512),\n" +
+            "PRIMARY KEY (property_spec_id, session_property_name),\n" +
+            "FOREIGN KEY (property_spec_id) REFERENCES session_specs (spec_id)\n" +
+            ")")
+    void createSessionPropertiesTable();
+
+    @SqlUpdate("DROP TABLE IF EXISTS " + SESSION_SPECS_TABLE)
+    void dropSessionSpecsTable();
+
+    @SqlUpdate("DROP TABLE IF EXISTS " + CLIENT_TAGS_TABLE)
+    void dropSessionClientTagsTable();
+
+    @SqlUpdate("DROP TABLE IF EXISTS " + PROPERTIES_TABLE)
+    void dropSessionPropertiesTable();
+
+    @SqlQuery("SELECT " +
+            "S.spec_id,\n" +
+            "S.user_regex,\n" +
+            "S.source_regex,\n" +
+            "S.query_type,\n" +
+            "S.group_regex,\n" +
+            "S.client_info_regex,\n" +
+            "S.override_session_properties,\n" +
+            "S.client_tags,\n" +
+            "GROUP_CONCAT(P.session_property_name ORDER BY P.session_property_name) session_property_names,\n" +
+            "GROUP_CONCAT(P.session_property_value ORDER BY P.session_property_name) session_property_values,\n" +
+            "GROUP_CONCAT(COALESCE(P.catalog, '" + EMPTY_CATALOG + "') ORDER BY P.session_property_name) session_property_catalogs\n" +
+            "FROM\n" +
+            "(SELECT\n" +
+            "A.spec_id, A.user_regex, A.source_regex, A.query_type, A.group_regex, A.client_info_regex, A.override_session_properties, A.priority,\n" +
+            "GROUP_CONCAT(DISTINCT B.client_tag) client_tags\n" +
+            "FROM " + SESSION_SPECS_TABLE + " A\n" +
+            "LEFT JOIN " + CLIENT_TAGS_TABLE + " B\n" +
+            "ON A.spec_id = B.tag_spec_id\n" +
+            "GROUP BY A.spec_id, A.user_regex, A.source_regex, A.query_type, A.group_regex, A.client_info_regex, A.override_session_properties, A.priority)\n" +
+            " S JOIN\n" +
+            PROPERTIES_TABLE + " P\n" +
+            "ON S.spec_id = P.property_spec_id\n" +
+            "GROUP BY S.spec_id, S.user_regex, S.source_regex, S.query_type, S.group_regex, S.client_info_regex, S.override_session_properties, S.priority, S.client_tags\n" +
+            "ORDER BY S.priority asc")
+    @UseRowMapper(SessionMatchSpec.Mapper.class)
+    List getSessionMatchSpecs();
+
+    @VisibleForTesting
+    @SqlUpdate("INSERT INTO " + SESSION_SPECS_TABLE + " (spec_id, user_regex, source_regex, query_type, group_regex, client_info_regex, override_session_properties, priority)\n" +
+            "VALUES (:spec_id, :user_regex, :source_regex, :query_type, :group_regex, :client_info_regex, :override_session_properties, :priority)")
+    void insertSpecRow(
+            @Bind("spec_id") long specId,
+            @Bind("user_regex") String userRegex,
+            @Bind("source_regex") String sourceRegex,
+            @Bind("query_type") String queryType,
+            @Bind("group_regex") String groupRegex,
+            @Bind("client_info_regex") String clientInfoRegex,
+            @Bind("override_session_properties") Integer overrideSessionProperties,
+            @Bind("priority") int priority);
+
+    @VisibleForTesting
+    @SqlUpdate("INSERT INTO " + CLIENT_TAGS_TABLE + " (tag_spec_id, client_tag) VALUES (:spec_id, :client_tag)")
+    void insertClientTag(@Bind("spec_id") long specId, @Bind("client_tag") String clientTag);
+
+    @VisibleForTesting
+    @SqlUpdate("INSERT INTO " + PROPERTIES_TABLE + " (property_spec_id, session_property_name, session_property_value, catalog)\n" +
+            "VALUES (:property_spec_id, :session_property_name, :session_property_value, :catalog)")
+    void insertSessionProperty(
+            @Bind("property_spec_id") long propertySpecId,
+            @Bind("session_property_name") String sessionPropertyName,
+            @Bind("session_property_value") String sessionPropertyValue,
+            @Bind("catalog") String catalog);
+}
diff --git a/presto-session-property-managers/src/main/java/com/facebook/presto/session/db/SessionPropertiesDaoProvider.java b/presto-session-property-managers/src/main/java/com/facebook/presto/session/db/SessionPropertiesDaoProvider.java
new file mode 100644
index 0000000000000..06617a1885396
--- /dev/null
+++ b/presto-session-property-managers/src/main/java/com/facebook/presto/session/db/SessionPropertiesDaoProvider.java
@@ -0,0 +1,46 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.session.db;
+
+import org.jdbi.v3.core.Jdbi;
+import org.jdbi.v3.sqlobject.SqlObjectPlugin;
+
+import javax.inject.Inject;
+import javax.inject.Provider;
+
+import java.sql.DriverManager;
+
+import static java.util.Objects.requireNonNull;
+
+public class SessionPropertiesDaoProvider
+        implements Provider
+{
+    private final SessionPropertiesDao dao;
+
+    @Inject
+    public SessionPropertiesDaoProvider(DbSessionPropertyManagerConfig config)
+    {
+        requireNonNull(config, "config is null");
+        requireNonNull(config.getConfigDbUrl(), "db url is null");
+        this.dao = Jdbi.create(() -> DriverManager.getConnection(config.getConfigDbUrl()))
+                .installPlugin(new SqlObjectPlugin())
+                .onDemand(SessionPropertiesDao.class);
+    }
+
+    @Override
+    public SessionPropertiesDao get()
+    {
+        return dao;
+    }
+}
diff --git a/presto-session-property-managers/src/main/java/com/facebook/presto/session/FileSessionPropertyManager.java b/presto-session-property-managers/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManager.java
similarity index 60%
rename from presto-session-property-managers/src/main/java/com/facebook/presto/session/FileSessionPropertyManager.java
rename to presto-session-property-managers/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManager.java
index 9904baad66869..6f559a00f2b91 100644
--- a/presto-session-property-managers/src/main/java/com/facebook/presto/session/FileSessionPropertyManager.java
+++ b/presto-session-property-managers/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManager.java
@@ -11,40 +11,36 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.facebook.presto.session;
+package com.facebook.presto.session.file;
 
 import com.facebook.airlift.json.JsonCodec;
 import com.facebook.airlift.json.JsonCodecFactory;
 import com.facebook.airlift.json.JsonObjectMapperProvider;
-import com.facebook.presto.spi.session.SessionConfigurationContext;
-import com.facebook.presto.spi.session.SessionPropertyConfigurationManager;
+import com.facebook.presto.session.AbstractSessionPropertyManager;
+import com.facebook.presto.session.SessionMatchSpec;
 import com.fasterxml.jackson.databind.JsonMappingException;
 import com.fasterxml.jackson.databind.exc.UnrecognizedPropertyException;
-import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableList;
 import jakarta.inject.Inject;
 
 import java.io.IOException;
 import java.io.UncheckedIOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
-import java.util.HashMap;
-import java.util.HashSet;
 import java.util.List;
-import java.util.Map;
-import java.util.Set;
 
 import static com.fasterxml.jackson.databind.DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES;
 import static java.lang.String.format;
 import static java.util.Objects.requireNonNull;
 
 public class FileSessionPropertyManager
-        implements SessionPropertyConfigurationManager
+        extends AbstractSessionPropertyManager
 {
     public static final JsonCodec> CODEC = new JsonCodecFactory(
             () -> new JsonObjectMapperProvider().get().enable(FAIL_ON_UNKNOWN_PROPERTIES))
             .listJsonCodec(SessionMatchSpec.class);
 
-    private final List sessionMatchSpecs;
+    private final ImmutableList sessionMatchSpecs;
 
     @Inject
     public FileSessionPropertyManager(FileSessionPropertyManagerConfig config)
@@ -53,7 +49,7 @@ public FileSessionPropertyManager(FileSessionPropertyManagerConfig config)
 
         Path configurationFile = config.getConfigFile().toPath();
         try {
-            sessionMatchSpecs = CODEC.fromJson(Files.readAllBytes(configurationFile));
+            sessionMatchSpecs = ImmutableList.copyOf(CODEC.fromJson(Files.readAllBytes(configurationFile)));
         }
         catch (IOException e) {
             throw new UncheckedIOException(e);
@@ -80,31 +76,8 @@ public FileSessionPropertyManager(FileSessionPropertyManagerConfig config)
     }
 
     @Override
-    public SystemSessionPropertyConfiguration getSystemSessionProperties(SessionConfigurationContext context)
+    protected List getSessionMatchSpecs()
     {
-        // later properties override earlier properties
-        Map defaultProperties = new HashMap<>();
-        Set overridePropertyNames = new HashSet<>();
-        for (SessionMatchSpec sessionMatchSpec : sessionMatchSpecs) {
-            Map newProperties = sessionMatchSpec.match(context);
-            defaultProperties.putAll(newProperties);
-            if (sessionMatchSpec.getOverrideSessionProperties().orElse(false)) {
-                overridePropertyNames.addAll(newProperties.keySet());
-            }
-        }
-
-        // Once a property has been overridden it stays that way and the value is updated by any rule
-        Map overrideProperties = new HashMap<>();
-        for (String propertyName : overridePropertyNames) {
-            overrideProperties.put(propertyName, defaultProperties.get(propertyName));
-        }
-
-        return new SystemSessionPropertyConfiguration(ImmutableMap.copyOf(defaultProperties), ImmutableMap.copyOf(overrideProperties));
-    }
-
-    @Override
-    public Map> getCatalogSessionProperties(SessionConfigurationContext context)
-    {
-        return ImmutableMap.of();
+        return sessionMatchSpecs;
     }
 }
diff --git a/presto-session-property-managers/src/main/java/com/facebook/presto/session/FileSessionPropertyManagerConfig.java b/presto-session-property-managers/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManagerConfig.java
similarity index 96%
rename from presto-session-property-managers/src/main/java/com/facebook/presto/session/FileSessionPropertyManagerConfig.java
rename to presto-session-property-managers/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManagerConfig.java
index b1e72bed0aedd..54164b968b105 100644
--- a/presto-session-property-managers/src/main/java/com/facebook/presto/session/FileSessionPropertyManagerConfig.java
+++ b/presto-session-property-managers/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManagerConfig.java
@@ -11,7 +11,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.facebook.presto.session;
+package com.facebook.presto.session.file;
 
 import com.facebook.airlift.configuration.Config;
 import jakarta.validation.constraints.NotNull;
diff --git a/presto-session-property-managers/src/main/java/com/facebook/presto/session/FileSessionPropertyManagerFactory.java b/presto-session-property-managers/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManagerFactory.java
similarity index 97%
rename from presto-session-property-managers/src/main/java/com/facebook/presto/session/FileSessionPropertyManagerFactory.java
rename to presto-session-property-managers/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManagerFactory.java
index aec39239cfbb5..08a49b74f54f5 100644
--- a/presto-session-property-managers/src/main/java/com/facebook/presto/session/FileSessionPropertyManagerFactory.java
+++ b/presto-session-property-managers/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManagerFactory.java
@@ -11,7 +11,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.facebook.presto.session;
+package com.facebook.presto.session.file;
 
 import com.facebook.airlift.bootstrap.Bootstrap;
 import com.facebook.airlift.json.JsonModule;
diff --git a/presto-session-property-managers/src/main/java/com/facebook/presto/session/FileSessionPropertyManagerModule.java b/presto-session-property-managers/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManagerModule.java
similarity index 96%
rename from presto-session-property-managers/src/main/java/com/facebook/presto/session/FileSessionPropertyManagerModule.java
rename to presto-session-property-managers/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManagerModule.java
index 03fee8d5207dc..9d50e6144df4c 100644
--- a/presto-session-property-managers/src/main/java/com/facebook/presto/session/FileSessionPropertyManagerModule.java
+++ b/presto-session-property-managers/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManagerModule.java
@@ -11,7 +11,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.facebook.presto.session;
+package com.facebook.presto.session.file;
 
 import com.google.inject.Binder;
 import com.google.inject.Module;
diff --git a/presto-session-property-managers/src/test/java/com/facebook/presto/session/TestFileSessionPropertyManager.java b/presto-session-property-managers/src/test/java/com/facebook/presto/session/AbstractTestSessionPropertyManager.java
similarity index 72%
rename from presto-session-property-managers/src/test/java/com/facebook/presto/session/TestFileSessionPropertyManager.java
rename to presto-session-property-managers/src/test/java/com/facebook/presto/session/AbstractTestSessionPropertyManager.java
index efd2c0f6b6411..3f09af61b1562 100644
--- a/presto-session-property-managers/src/test/java/com/facebook/presto/session/TestFileSessionPropertyManager.java
+++ b/presto-session-property-managers/src/test/java/com/facebook/presto/session/AbstractTestSessionPropertyManager.java
@@ -11,34 +11,25 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-
 package com.facebook.presto.session;
 
-import com.facebook.airlift.testing.TempFile;
 import com.facebook.presto.common.resourceGroups.QueryType;
 import com.facebook.presto.spi.resourceGroups.ResourceGroupId;
 import com.facebook.presto.spi.session.SessionConfigurationContext;
-import com.facebook.presto.spi.session.SessionPropertyConfigurationManager;
-import com.facebook.presto.spi.session.SessionPropertyConfigurationManager.SystemSessionPropertyConfiguration;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import org.testng.annotations.Test;
 
 import java.io.IOException;
-import java.nio.file.Files;
-import java.nio.file.Path;
-import java.util.Arrays;
 import java.util.Map;
 import java.util.Optional;
 import java.util.regex.Pattern;
 
-import static com.facebook.presto.session.FileSessionPropertyManager.CODEC;
-import static org.testng.Assert.assertEquals;
-
-public class TestFileSessionPropertyManager
+@Test(singleThreaded = true)
+public abstract class AbstractTestSessionPropertyManager
 {
-    private static final SessionConfigurationContext CONTEXT = new SessionConfigurationContext(
+    protected static final SessionConfigurationContext CONTEXT = new SessionConfigurationContext(
             "user",
             Optional.empty(),
             Optional.of("source"),
@@ -48,6 +39,13 @@ public class TestFileSessionPropertyManager
             Optional.of("bar"),
             "testversion");
 
+    protected abstract void assertProperties(Map properties, SessionMatchSpec... spec)
+            throws IOException;
+    protected abstract void assertProperties(Map defaultProperties, Map overrideProperties, SessionMatchSpec... spec)
+            throws IOException;
+    protected abstract void assertProperties(Map defaultProperties, Map overrideProperties, Map> catalogProperties, SessionMatchSpec... spec)
+            throws IOException;
+
     @Test
     public void testResourceGroupMatch()
             throws IOException
@@ -61,7 +59,8 @@ public void testResourceGroupMatch()
                 Optional.of(Pattern.compile("global.pipeline.user_.*")),
                 Optional.empty(),
                 Optional.empty(),
-                properties);
+                properties,
+                ImmutableMap.of());
 
         assertProperties(properties, spec);
     }
@@ -79,7 +78,8 @@ public void testClientTagMatch()
                 Optional.empty(),
                 Optional.empty(),
                 Optional.empty(),
-                properties);
+                properties,
+                ImmutableMap.of());
 
         assertProperties(properties, spec);
     }
@@ -96,7 +96,8 @@ public void testMultipleMatch()
                 Optional.empty(),
                 Optional.empty(),
                 Optional.empty(),
-                ImmutableMap.of("PROPERTY1", "VALUE1"));
+                ImmutableMap.of("PROPERTY1", "VALUE1", "PROPERTY3", "VALUE3"),
+                ImmutableMap.of());
         SessionMatchSpec spec2 = new SessionMatchSpec(
                 Optional.empty(),
                 Optional.empty(),
@@ -105,9 +106,10 @@ public void testMultipleMatch()
                 Optional.empty(),
                 Optional.empty(),
                 Optional.empty(),
-                ImmutableMap.of("PROPERTY1", "VALUE1", "PROPERTY2", "VALUE2"));
+                ImmutableMap.of("PROPERTY1", "VALUE1", "PROPERTY2", "VALUE2"),
+                ImmutableMap.of());
 
-        assertProperties(ImmutableMap.of("PROPERTY1", "VALUE1", "PROPERTY2", "VALUE2"), spec1, spec2);
+        assertProperties(ImmutableMap.of("PROPERTY1", "VALUE1", "PROPERTY2", "VALUE2", "PROPERTY3", "VALUE3"), spec1, spec2);
     }
 
     @Test
@@ -122,7 +124,8 @@ public void testNoMatch()
                 Optional.of(Pattern.compile("global.interactive.user_.*")),
                 Optional.empty(),
                 Optional.empty(),
-                ImmutableMap.of("PROPERTY", "VALUE"));
+                ImmutableMap.of("PROPERTY", "VALUE"),
+                ImmutableMap.of());
 
         assertProperties(ImmutableMap.of(), spec);
     }
@@ -140,7 +143,8 @@ public void testClientInfoMatch()
                 Optional.empty(),
                 Optional.of(Pattern.compile("bar")),
                 Optional.empty(),
-                properties);
+                properties,
+                ImmutableMap.of());
 
         assertProperties(properties, spec);
     }
@@ -160,7 +164,8 @@ public void testOverride()
                 Optional.empty(),
                 Optional.of(Pattern.compile("bar")),
                 Optional.of(true),
-                overrideProperties);
+                overrideProperties,
+                ImmutableMap.of());
 
         SessionMatchSpec specDefault = new SessionMatchSpec(
                 Optional.empty(),
@@ -170,7 +175,8 @@ public void testOverride()
                 Optional.empty(),
                 Optional.of(Pattern.compile("bar")),
                 Optional.empty(),
-                defaultProperties);
+                defaultProperties,
+                ImmutableMap.of());
 
         // PROPERTY1 should be an override property with the value from the default (non-override, higher precedence)
         // spec.
@@ -178,22 +184,23 @@ public void testOverride()
         assertProperties(defaultProperties, ImmutableMap.of("PROPERTY1", "VALUE2"), specOverride, specDefault);
     }
 
-    private static void assertProperties(Map defaultProperties, SessionMatchSpec... spec)
+    @Test
+    public void testCatalogProperty()
             throws IOException
     {
-        assertProperties(defaultProperties, ImmutableMap.of(), spec);
-    }
+        ImmutableMap defaultProperties = ImmutableMap.of("PROPERTY1", "VALUE1");
+        ImmutableMap> catalogProperties = ImmutableMap.of("CATALOG", ImmutableMap.of("PROPERTY", "VALUE"));
+        SessionMatchSpec spec = new SessionMatchSpec(
+                Optional.empty(),
+                Optional.empty(),
+                Optional.empty(),
+                Optional.empty(),
+                Optional.empty(),
+                Optional.empty(),
+                Optional.empty(),
+                defaultProperties,
+                catalogProperties);
 
-    private static void assertProperties(Map defaultProperties, Map overrideProperties, SessionMatchSpec... spec)
-            throws IOException
-    {
-        try (TempFile tempFile = new TempFile()) {
-            Path configurationFile = tempFile.path();
-            Files.write(configurationFile, CODEC.toJsonBytes(Arrays.asList(spec)));
-            SessionPropertyConfigurationManager manager = new FileSessionPropertyManager(new FileSessionPropertyManagerConfig().setConfigFile(configurationFile.toFile()));
-            SystemSessionPropertyConfiguration propertyConfiguration = manager.getSystemSessionProperties(CONTEXT);
-            assertEquals(propertyConfiguration.systemPropertyDefaults, defaultProperties);
-            assertEquals(propertyConfiguration.systemPropertyOverrides, overrideProperties);
-        }
+        assertProperties(defaultProperties, ImmutableMap.of(), catalogProperties, spec);
     }
 }
diff --git a/presto-session-property-managers/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManager.java b/presto-session-property-managers/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManager.java
new file mode 100644
index 0000000000000..865e8f840edf4
--- /dev/null
+++ b/presto-session-property-managers/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManager.java
@@ -0,0 +1,112 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.session.db;
+
+import com.facebook.presto.session.AbstractTestSessionPropertyManager;
+import com.facebook.presto.session.SessionMatchSpec;
+import com.facebook.presto.spi.session.SessionPropertyConfigurationManager;
+import com.facebook.presto.testing.mysql.MySqlOptions;
+import com.facebook.presto.testing.mysql.TestingMySqlServer;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import org.testng.annotations.AfterClass;
+
+import java.util.Map;
+import java.util.regex.Pattern;
+
+import static org.testng.Assert.assertEquals;
+
+public abstract class TestDbSessionPropertyManager
+        extends AbstractTestSessionPropertyManager
+{
+    private static final MySqlOptions MY_SQL_OPTIONS = MySqlOptions.builder()
+            .build();
+
+    private final String driver;
+
+    private final TestingMySqlServer mysqlServer;
+
+    public TestDbSessionPropertyManager(String driver)
+            throws Exception
+    {
+        this.driver = driver;
+        this.mysqlServer = new TestingMySqlServer("testuser", "testpass", ImmutableList.of(), MY_SQL_OPTIONS);
+    }
+
+    @AfterClass(alwaysRun = true)
+    public void tearDown()
+            throws Exception
+    {
+        mysqlServer.close();
+    }
+
+    @Override
+    protected void assertProperties(Map defaultProperties, SessionMatchSpec... spec)
+    {
+        assertProperties(defaultProperties, ImmutableMap.of(), ImmutableMap.of(), spec);
+    }
+
+    @Override
+    protected void assertProperties(Map defaultProperties, Map overrideProperties, SessionMatchSpec... specs)
+    {
+        assertProperties(defaultProperties, overrideProperties, ImmutableMap.of(), specs);
+    }
+
+    @Override
+    protected void assertProperties(Map defaultProperties, Map overrideProperties, Map> catalogProperties, SessionMatchSpec... specs)
+    {
+        DbSessionPropertyManagerConfig config = new DbSessionPropertyManagerConfig()
+                .setConfigDbUrl(overrideJdbcUrl(mysqlServer.getJdbcUrl("session") + "&createDatabaseIfNotExist=true"));
+
+        SessionPropertiesDaoProvider sessionPropertiesDaoProvider = new SessionPropertiesDaoProvider(config);
+        SessionPropertiesDao dao = sessionPropertiesDaoProvider.get();
+        RefreshingDbSpecsProvider dbSpecsProvider = new RefreshingDbSpecsProvider(config, sessionPropertiesDaoProvider.get());
+        SessionPropertyConfigurationManager manager = new DbSessionPropertyManager(dbSpecsProvider);
+        int id = 1;
+        try {
+            for (SessionMatchSpec spec : specs) {
+                int finalId = id;
+                dao.insertSpecRow(
+                        finalId,
+                        spec.getUserRegex().map(Pattern::pattern).orElse(null),
+                        spec.getSourceRegex().map(Pattern::pattern).orElse(null),
+                        spec.getQueryType().orElse(null),
+                        spec.getResourceGroupRegex().map(Pattern::pattern).orElse(null),
+                        spec.getClientInfoRegex().map(Pattern::pattern).orElse(null),
+                        spec.getOverrideSessionProperties().map(val -> val ? 1 : 0).orElse(null),
+                        finalId);
+                spec.getClientTags().forEach(tag -> dao.insertClientTag(finalId, tag));
+                spec.getSessionProperties().forEach((key, value) -> dao.insertSessionProperty(finalId, key, value, null));
+                spec.getCatalogSessionProperties().forEach((catalog, property) -> property.forEach((key, value) -> dao.insertSessionProperty(finalId, key, value, catalog)));
+                id++;
+            }
+            dbSpecsProvider.refresh();
+            SessionPropertyConfigurationManager.SystemSessionPropertyConfiguration propertyConfiguration = manager.getSystemSessionProperties(CONTEXT);
+            assertEquals(propertyConfiguration.systemPropertyDefaults, defaultProperties);
+            assertEquals(propertyConfiguration.systemPropertyOverrides, overrideProperties);
+            assertEquals(manager.getCatalogSessionProperties(CONTEXT), catalogProperties);
+        }
+        finally {
+            dao.dropSessionPropertiesTable();
+            dao.dropSessionClientTagsTable();
+            dao.dropSessionSpecsTable();
+            dbSpecsProvider.destroy();
+        }
+    }
+
+    protected String overrideJdbcUrl(String url)
+    {
+        return url;
+    }
+}
diff --git a/presto-session-property-managers/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManagerConfig.java b/presto-session-property-managers/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManagerConfig.java
new file mode 100644
index 0000000000000..d976eb722390c
--- /dev/null
+++ b/presto-session-property-managers/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManagerConfig.java
@@ -0,0 +1,52 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.session.db;
+
+import com.facebook.airlift.units.Duration;
+import com.google.common.collect.ImmutableMap;
+import org.testng.annotations.Test;
+
+import java.util.Map;
+import java.util.concurrent.TimeUnit;
+
+import static com.facebook.airlift.configuration.testing.ConfigAssertions.assertFullMapping;
+import static com.facebook.airlift.configuration.testing.ConfigAssertions.assertRecordedDefaults;
+import static com.facebook.airlift.configuration.testing.ConfigAssertions.recordDefaults;
+import static java.util.concurrent.TimeUnit.SECONDS;
+
+public class TestDbSessionPropertyManagerConfig
+{
+    @Test
+    public void testDefaults()
+    {
+        assertRecordedDefaults(recordDefaults(DbSessionPropertyManagerConfig.class)
+                .setConfigDbUrl(null)
+                .setSpecsRefreshPeriod(new Duration(10, SECONDS)));
+    }
+
+    @Test
+    public void testExplicitPropertyMappings()
+    {
+        Map properties = new ImmutableMap.Builder()
+                .put("session-property-manager.db.url", "foo")
+                .put("session-property-manager.db.refresh-period", "50s")
+                .build();
+
+        DbSessionPropertyManagerConfig expected = new DbSessionPropertyManagerConfig()
+                .setConfigDbUrl("foo")
+                .setSpecsRefreshPeriod(new Duration(50, TimeUnit.SECONDS));
+
+        assertFullMapping(properties, expected);
+    }
+}
diff --git a/presto-session-property-managers/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManagerMariadb.java b/presto-session-property-managers/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManagerMariadb.java
new file mode 100644
index 0000000000000..3e440b95d3070
--- /dev/null
+++ b/presto-session-property-managers/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManagerMariadb.java
@@ -0,0 +1,29 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.session.db;
+
+public class TestDbSessionPropertyManagerMariadb
+        extends TestDbSessionPropertyManager
+{
+    public TestDbSessionPropertyManagerMariadb() throws Exception
+    {
+        super("mariadb");
+    }
+
+    @Override
+    public String overrideJdbcUrl(String url)
+    {
+        return url.replaceFirst("jdbc:mysql:", "jdbc:mariadb:");
+    }
+}
diff --git a/presto-session-property-managers/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManagerMysql.java b/presto-session-property-managers/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManagerMysql.java
new file mode 100644
index 0000000000000..806f5c8da9cac
--- /dev/null
+++ b/presto-session-property-managers/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManagerMysql.java
@@ -0,0 +1,23 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.session.db;
+
+public class TestDbSessionPropertyManagerMysql
+        extends TestDbSessionPropertyManager
+{
+    public TestDbSessionPropertyManagerMysql() throws Exception
+    {
+        super("mysql");
+    }
+}
diff --git a/presto-session-property-managers/src/test/java/com/facebook/presto/session/file/TestFileSessionPropertyManager.java b/presto-session-property-managers/src/test/java/com/facebook/presto/session/file/TestFileSessionPropertyManager.java
new file mode 100644
index 0000000000000..77f5f79ffdb1d
--- /dev/null
+++ b/presto-session-property-managers/src/test/java/com/facebook/presto/session/file/TestFileSessionPropertyManager.java
@@ -0,0 +1,62 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.facebook.presto.session.file;
+
+import com.facebook.airlift.testing.TempFile;
+import com.facebook.presto.session.AbstractTestSessionPropertyManager;
+import com.facebook.presto.session.SessionMatchSpec;
+import com.facebook.presto.spi.session.SessionPropertyConfigurationManager;
+import com.google.common.collect.ImmutableMap;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.Arrays;
+import java.util.Map;
+
+import static com.facebook.presto.session.file.FileSessionPropertyManager.CODEC;
+import static org.testng.Assert.assertEquals;
+
+public class TestFileSessionPropertyManager
+        extends AbstractTestSessionPropertyManager
+{
+    @Override
+    protected void assertProperties(Map defaultProperties, SessionMatchSpec... specs)
+            throws IOException
+    {
+        assertProperties(defaultProperties, ImmutableMap.of(), ImmutableMap.of(), specs);
+    }
+
+    @Override
+    protected void assertProperties(Map defaultProperties, Map overrideProperties, SessionMatchSpec... specs)
+            throws IOException
+    {
+        assertProperties(defaultProperties, overrideProperties, ImmutableMap.of(), specs);
+    }
+
+    protected void assertProperties(Map defaultProperties, Map overrideProperties, Map> catalogProperties, SessionMatchSpec... specs)
+            throws IOException
+    {
+        try (TempFile tempFile = new TempFile()) {
+            Path configurationFile = tempFile.path();
+            Files.write(configurationFile, CODEC.toJsonBytes(Arrays.asList(specs)));
+            SessionPropertyConfigurationManager manager = new FileSessionPropertyManager(new FileSessionPropertyManagerConfig().setConfigFile(configurationFile.toFile()));
+            SessionPropertyConfigurationManager.SystemSessionPropertyConfiguration propertyConfiguration = manager.getSystemSessionProperties(CONTEXT);
+            assertEquals(propertyConfiguration.systemPropertyDefaults, defaultProperties);
+            assertEquals(propertyConfiguration.systemPropertyOverrides, overrideProperties);
+            assertEquals(manager.getCatalogSessionProperties(CONTEXT), catalogProperties);
+        }
+    }
+}
diff --git a/presto-session-property-managers/src/test/java/com/facebook/presto/session/TestFileSessionPropertyManagerConfig.java b/presto-session-property-managers/src/test/java/com/facebook/presto/session/file/TestFileSessionPropertyManagerConfig.java
similarity index 97%
rename from presto-session-property-managers/src/test/java/com/facebook/presto/session/TestFileSessionPropertyManagerConfig.java
rename to presto-session-property-managers/src/test/java/com/facebook/presto/session/file/TestFileSessionPropertyManagerConfig.java
index b555593fa09a0..3c6c509a4490e 100644
--- a/presto-session-property-managers/src/test/java/com/facebook/presto/session/TestFileSessionPropertyManagerConfig.java
+++ b/presto-session-property-managers/src/test/java/com/facebook/presto/session/file/TestFileSessionPropertyManagerConfig.java
@@ -11,7 +11,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.facebook.presto.session;
+package com.facebook.presto.session.file;
 
 import com.google.common.collect.ImmutableMap;
 import org.testng.annotations.Test;
diff --git a/presto-spark-testing/src/test/java/com/facebook/presto/spark/testing/TestPrestoSparkLauncherIntegrationSmokeTest.java b/presto-spark-testing/src/test/java/com/facebook/presto/spark/testing/TestPrestoSparkLauncherIntegrationSmokeTest.java
index 4d73d96734caf..e826d055e57e8 100644
--- a/presto-spark-testing/src/test/java/com/facebook/presto/spark/testing/TestPrestoSparkLauncherIntegrationSmokeTest.java
+++ b/presto-spark-testing/src/test/java/com/facebook/presto/spark/testing/TestPrestoSparkLauncherIntegrationSmokeTest.java
@@ -51,7 +51,7 @@
 import java.util.concurrent.TimeoutException;
 import java.util.regex.Pattern;
 
-import static com.facebook.presto.session.FileSessionPropertyManager.CODEC;
+import static com.facebook.presto.session.file.FileSessionPropertyManager.CODEC;
 import static com.facebook.presto.spark.testing.Processes.destroyProcess;
 import static com.facebook.presto.testing.TestingSession.testSessionBuilder;
 import static com.facebook.presto.tests.QueryAssertions.assertEqualsIgnoreOrder;
@@ -181,7 +181,8 @@ public void setUp()
                 Optional.of(Pattern.compile("global.*")),
                 Optional.empty(),
                 Optional.empty(),
-                properties);
+                properties,
+                ImmutableMap.of());
         sessionPropertyConfigJsonFile = new File(tempDir, "session-property-config.json");
         Files.write(sessionPropertyConfigJsonFile.toPath(), CODEC.toJsonBytes(Collections.singletonList(spec)));
         sessionPropertyConfig = new File(tempDir, "session-property-configuration.properties");

From 2535ff424d0dd01590eb118578cf0f70d676acb7 Mon Sep 17 00:00:00 2001
From: Hazmi 
Date: Thu, 17 Jul 2025 18:24:43 +0300
Subject: [PATCH 046/113] Split presto-session-property-managers module

Currently, the presto-session-property-managers module has gotten large due to the different implementations for file based
and db based session property managers. If implemented, this commit will split the module into 3,
``presto-session-property-managers-common`` will contain the API implementation for a session property manager
along with some tests. ``presto-file-session-property-manager`` will contain the file based implementation, and
``presto-db-session-property-manager`` will contain the db based implementation.
---
 CODEOWNERS                                    |   4 +-
 pom.xml                                       |  25 +++-
 .../pom.xml                                   |  34 ++---
 ...ionPropertyConfigurationManagerPlugin.java |   7 +-
 .../session/db/DbSessionPropertyManager.java  |   0
 .../db/DbSessionPropertyManagerConfig.java    |   0
 .../db/DbSessionPropertyManagerFactory.java   |   0
 .../db/DbSessionPropertyManagerModule.java    |   0
 .../presto/session/db/DbSpecsProvider.java    |   0
 .../session/db/RefreshingDbSpecsProvider.java |   0
 .../session/db/SessionPropertiesDao.java      |   0
 .../db/SessionPropertiesDaoProvider.java      |   0
 .../db/TestDbSessionPropertyManager.java      |   0
 .../TestDbSessionPropertyManagerConfig.java   |   0
 .../TestDbSessionPropertyManagerMariadb.java  |   0
 .../db/TestDbSessionPropertyManagerMysql.java |   0
 .../admin/session-property-managers.rst       |   3 +-
 presto-file-session-property-manager/pom.xml  | 139 ++++++++++++++++++
 ...ionPropertyConfigurationManagerPlugin.java |  29 ++++
 .../file/FileSessionPropertyManager.java      |   0
 .../FileSessionPropertyManagerConfig.java     |   0
 .../FileSessionPropertyManagerFactory.java    |   0
 .../FileSessionPropertyManagerModule.java     |   0
 .../file/TestFileSessionPropertyManager.java  |   0
 .../TestFileSessionPropertyManagerConfig.java |   4 +-
 presto-main/etc/session-property-config.json  |   0
 .../etc/session-property-config.properties    |   0
 presto-server/src/main/provisio/presto.xml    |  10 +-
 .../pom.xml                                   |  81 ++++++++++
 .../AbstractSessionPropertyManager.java       |   0
 .../presto/session/SessionMatchSpec.java      |   4 +-
 .../AbstractTestSessionPropertyManager.java   |   0
 presto-spark-package/pom.xml                  |  10 +-
 .../src/main/assembly/presto.xml              |   8 +-
 presto-spark-testing/pom.xml                  |   8 +-
 presto-test-coverage/pom.xml                  |  16 +-
 36 files changed, 343 insertions(+), 39 deletions(-)
 rename {presto-session-property-managers => presto-db-session-property-manager}/pom.xml (88%)
 rename presto-session-property-managers/src/main/java/com/facebook/presto/session/SessionPropertyConfigurationManagerPlugin.java => presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/DbSessionPropertyConfigurationManagerPlugin.java (77%)
 rename {presto-session-property-managers => presto-db-session-property-manager}/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManager.java (100%)
 rename {presto-session-property-managers => presto-db-session-property-manager}/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManagerConfig.java (100%)
 rename {presto-session-property-managers => presto-db-session-property-manager}/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManagerFactory.java (100%)
 rename {presto-session-property-managers => presto-db-session-property-manager}/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManagerModule.java (100%)
 rename {presto-session-property-managers => presto-db-session-property-manager}/src/main/java/com/facebook/presto/session/db/DbSpecsProvider.java (100%)
 rename {presto-session-property-managers => presto-db-session-property-manager}/src/main/java/com/facebook/presto/session/db/RefreshingDbSpecsProvider.java (100%)
 rename {presto-session-property-managers => presto-db-session-property-manager}/src/main/java/com/facebook/presto/session/db/SessionPropertiesDao.java (100%)
 rename {presto-session-property-managers => presto-db-session-property-manager}/src/main/java/com/facebook/presto/session/db/SessionPropertiesDaoProvider.java (100%)
 rename {presto-session-property-managers => presto-db-session-property-manager}/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManager.java (100%)
 rename {presto-session-property-managers => presto-db-session-property-manager}/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManagerConfig.java (100%)
 rename {presto-session-property-managers => presto-db-session-property-manager}/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManagerMariadb.java (100%)
 rename {presto-session-property-managers => presto-db-session-property-manager}/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManagerMysql.java (100%)
 create mode 100644 presto-file-session-property-manager/pom.xml
 create mode 100644 presto-file-session-property-manager/src/main/java/com/facebook/presto/session/file/FileSessionPropertyConfigurationManagerPlugin.java
 rename {presto-session-property-managers => presto-file-session-property-manager}/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManager.java (100%)
 rename {presto-session-property-managers => presto-file-session-property-manager}/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManagerConfig.java (100%)
 rename {presto-session-property-managers => presto-file-session-property-manager}/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManagerFactory.java (100%)
 rename {presto-session-property-managers => presto-file-session-property-manager}/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManagerModule.java (100%)
 rename {presto-session-property-managers => presto-file-session-property-manager}/src/test/java/com/facebook/presto/session/file/TestFileSessionPropertyManager.java (100%)
 rename {presto-session-property-managers => presto-file-session-property-manager}/src/test/java/com/facebook/presto/session/file/TestFileSessionPropertyManagerConfig.java (94%)
 delete mode 100644 presto-main/etc/session-property-config.json
 delete mode 100644 presto-main/etc/session-property-config.properties
 create mode 100644 presto-session-property-managers-common/pom.xml
 rename {presto-session-property-managers => presto-session-property-managers-common}/src/main/java/com/facebook/presto/session/AbstractSessionPropertyManager.java (100%)
 rename {presto-session-property-managers => presto-session-property-managers-common}/src/main/java/com/facebook/presto/session/SessionMatchSpec.java (98%)
 rename {presto-session-property-managers => presto-session-property-managers-common}/src/test/java/com/facebook/presto/session/AbstractTestSessionPropertyManager.java (100%)

diff --git a/CODEOWNERS b/CODEOWNERS
index faf25f724f484..0b2831ba7ee31 100644
--- a/CODEOWNERS
+++ b/CODEOWNERS
@@ -76,7 +76,9 @@
 /presto-resource-group-managers @prestodb/committers
 /presto-router @prestodb/committers
 /presto-server @prestodb/committers
-/presto-session-property-managers @prestodb/committers
+/presto-session-property-managers-common @prestodb/committers
+/presto-file-session-property-manager @prestodb/committers
+/presto-db-session-property-manager @prestodb/committers
 /presto-singlestore @prestodb/committers
 /presto-spi @prestodb/committers
 /presto-sqlserver @prestodb/committers
diff --git a/pom.xml b/pom.xml
index 5bdb0ad74d96e..15a6f0ec95a20 100644
--- a/pom.xml
+++ b/pom.xml
@@ -172,7 +172,9 @@
         presto-plugin-toolkit
         presto-resource-group-managers
         presto-password-authenticators
-        presto-session-property-managers
+        presto-session-property-managers-common
+        presto-db-session-property-manager
+        presto-file-session-property-manager
         presto-benchto-benchmarks
         presto-thrift-api
         presto-thrift-testing-server
@@ -343,7 +345,26 @@
 
             
                 com.facebook.presto
-                presto-session-property-managers
+                presto-session-property-managers-common
+                ${project.version}
+            
+
+            
+                com.facebook.presto
+                presto-session-property-managers-common
+                ${project.version}
+                test-jar
+            
+
+            
+                com.facebook.presto
+                presto-db-session-property-manager
+                ${project.version}
+            
+
+            
+                com.facebook.presto
+                presto-file-session-property-manager
                 ${project.version}
             
 
diff --git a/presto-session-property-managers/pom.xml b/presto-db-session-property-manager/pom.xml
similarity index 88%
rename from presto-session-property-managers/pom.xml
rename to presto-db-session-property-manager/pom.xml
index 0ba95f4c6ee88..d0ef4ca921014 100644
--- a/presto-session-property-managers/pom.xml
+++ b/presto-db-session-property-manager/pom.xml
@@ -8,17 +8,21 @@
         0.295-SNAPSHOT
     
 
-    presto-session-property-managers
-    presto-session-property-managers
-    Presto - Session Property Managers
+    presto-db-session-property-manager
+    presto-db-session-property-manager
+    Presto - DB Session Property Manager
     presto-plugin
 
     
         ${project.parent.basedir}
-        true
     
 
     
+        
+            com.facebook.presto
+            presto-session-property-managers-common
+        
+
         
             com.facebook.airlift
             bootstrap
@@ -59,11 +63,6 @@
             javax.inject
         
 
-        
-            jakarta.inject
-            jakarta.inject-api
-        
-
         
             com.facebook.airlift
             concurrent
@@ -74,16 +73,6 @@
             jakarta.validation-api
         
 
-        
-            com.fasterxml.jackson.core
-            jackson-databind
-        
-
-        
-            com.fasterxml.jackson.core
-            jackson-core
-        
-
         
             org.jdbi
             jdbi3-core
@@ -175,5 +164,12 @@
             testing-mysql-server-base
             test
         
+
+        
+            com.facebook.presto
+            presto-session-property-managers-common
+            test-jar
+            test
+        
     
 
diff --git a/presto-session-property-managers/src/main/java/com/facebook/presto/session/SessionPropertyConfigurationManagerPlugin.java b/presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/DbSessionPropertyConfigurationManagerPlugin.java
similarity index 77%
rename from presto-session-property-managers/src/main/java/com/facebook/presto/session/SessionPropertyConfigurationManagerPlugin.java
rename to presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/DbSessionPropertyConfigurationManagerPlugin.java
index b89048a8733f5..bbbc653abd2c9 100644
--- a/presto-session-property-managers/src/main/java/com/facebook/presto/session/SessionPropertyConfigurationManagerPlugin.java
+++ b/presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/DbSessionPropertyConfigurationManagerPlugin.java
@@ -11,22 +11,19 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.facebook.presto.session;
+package com.facebook.presto.session.db;
 
-import com.facebook.presto.session.db.DbSessionPropertyManagerFactory;
-import com.facebook.presto.session.file.FileSessionPropertyManagerFactory;
 import com.facebook.presto.spi.Plugin;
 import com.facebook.presto.spi.session.SessionPropertyConfigurationManagerFactory;
 import com.google.common.collect.ImmutableList;
 
-public class SessionPropertyConfigurationManagerPlugin
+public class DbSessionPropertyConfigurationManagerPlugin
         implements Plugin
 {
     @Override
     public Iterable getSessionPropertyConfigurationManagerFactories()
     {
         return ImmutableList.of(
-                new FileSessionPropertyManagerFactory(),
                 new DbSessionPropertyManagerFactory());
     }
 }
diff --git a/presto-session-property-managers/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManager.java b/presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManager.java
similarity index 100%
rename from presto-session-property-managers/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManager.java
rename to presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManager.java
diff --git a/presto-session-property-managers/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManagerConfig.java b/presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManagerConfig.java
similarity index 100%
rename from presto-session-property-managers/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManagerConfig.java
rename to presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManagerConfig.java
diff --git a/presto-session-property-managers/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManagerFactory.java b/presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManagerFactory.java
similarity index 100%
rename from presto-session-property-managers/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManagerFactory.java
rename to presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManagerFactory.java
diff --git a/presto-session-property-managers/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManagerModule.java b/presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManagerModule.java
similarity index 100%
rename from presto-session-property-managers/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManagerModule.java
rename to presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/DbSessionPropertyManagerModule.java
diff --git a/presto-session-property-managers/src/main/java/com/facebook/presto/session/db/DbSpecsProvider.java b/presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/DbSpecsProvider.java
similarity index 100%
rename from presto-session-property-managers/src/main/java/com/facebook/presto/session/db/DbSpecsProvider.java
rename to presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/DbSpecsProvider.java
diff --git a/presto-session-property-managers/src/main/java/com/facebook/presto/session/db/RefreshingDbSpecsProvider.java b/presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/RefreshingDbSpecsProvider.java
similarity index 100%
rename from presto-session-property-managers/src/main/java/com/facebook/presto/session/db/RefreshingDbSpecsProvider.java
rename to presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/RefreshingDbSpecsProvider.java
diff --git a/presto-session-property-managers/src/main/java/com/facebook/presto/session/db/SessionPropertiesDao.java b/presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/SessionPropertiesDao.java
similarity index 100%
rename from presto-session-property-managers/src/main/java/com/facebook/presto/session/db/SessionPropertiesDao.java
rename to presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/SessionPropertiesDao.java
diff --git a/presto-session-property-managers/src/main/java/com/facebook/presto/session/db/SessionPropertiesDaoProvider.java b/presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/SessionPropertiesDaoProvider.java
similarity index 100%
rename from presto-session-property-managers/src/main/java/com/facebook/presto/session/db/SessionPropertiesDaoProvider.java
rename to presto-db-session-property-manager/src/main/java/com/facebook/presto/session/db/SessionPropertiesDaoProvider.java
diff --git a/presto-session-property-managers/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManager.java b/presto-db-session-property-manager/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManager.java
similarity index 100%
rename from presto-session-property-managers/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManager.java
rename to presto-db-session-property-manager/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManager.java
diff --git a/presto-session-property-managers/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManagerConfig.java b/presto-db-session-property-manager/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManagerConfig.java
similarity index 100%
rename from presto-session-property-managers/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManagerConfig.java
rename to presto-db-session-property-manager/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManagerConfig.java
diff --git a/presto-session-property-managers/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManagerMariadb.java b/presto-db-session-property-manager/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManagerMariadb.java
similarity index 100%
rename from presto-session-property-managers/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManagerMariadb.java
rename to presto-db-session-property-manager/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManagerMariadb.java
diff --git a/presto-session-property-managers/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManagerMysql.java b/presto-db-session-property-manager/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManagerMysql.java
similarity index 100%
rename from presto-session-property-managers/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManagerMysql.java
rename to presto-db-session-property-manager/src/test/java/com/facebook/presto/session/db/TestDbSessionPropertyManagerMysql.java
diff --git a/presto-docs/src/main/sphinx/admin/session-property-managers.rst b/presto-docs/src/main/sphinx/admin/session-property-managers.rst
index 8147d9741f051..4035e8dbb474f 100644
--- a/presto-docs/src/main/sphinx/admin/session-property-managers.rst
+++ b/presto-docs/src/main/sphinx/admin/session-property-managers.rst
@@ -6,7 +6,8 @@ Administrators can add session properties to control the behavior for subsets of
 These properties are defaults and can be overridden by users (if authorized to do so). Session
 properties can be used to control resource usage, enable or disable features, and change query
 characteristics. Session property managers are pluggable. A session property manager can either
-be database-based or file-based.
+be database-based or file-based. For production environments, the database-based manager is
+recommended as the properties can be updated without requiring a cluster restart.
 
 To enable a built-in manager that reads a JSON configuration file, add an
 ``etc/session-property-config.properties`` file with the following contents:
diff --git a/presto-file-session-property-manager/pom.xml b/presto-file-session-property-manager/pom.xml
new file mode 100644
index 0000000000000..5c3a4ceb2dddd
--- /dev/null
+++ b/presto-file-session-property-manager/pom.xml
@@ -0,0 +1,139 @@
+
+
+    4.0.0
+
+    
+        com.facebook.presto
+        presto-root
+        0.295-SNAPSHOT
+    
+
+    presto-file-session-property-manager
+    presto-file-session-property-manager
+    Presto - File Session Property Manager
+    presto-plugin
+
+    
+        ${project.parent.basedir}
+    
+
+    
+        
+            com.facebook.presto
+            presto-session-property-managers-common
+        
+
+        
+            com.facebook.airlift
+            bootstrap
+        
+
+        
+            com.facebook.airlift
+            json
+        
+
+        
+            com.facebook.airlift
+            configuration
+        
+
+        
+            com.google.guava
+            guava
+        
+
+        
+            com.google.inject
+            guice
+        
+
+        
+            javax.inject
+            javax.inject
+        
+
+        
+            jakarta.inject
+            jakarta.inject-api
+        
+
+        
+            jakarta.validation
+            jakarta.validation-api
+        
+
+        
+            com.fasterxml.jackson.core
+            jackson-databind
+        
+
+        
+            com.fasterxml.jackson.core
+            jackson-core
+        
+
+        
+        
+            com.facebook.presto
+            presto-spi
+            provided
+        
+
+        
+            com.facebook.presto
+            presto-common
+            provided
+        
+
+        
+            com.facebook.airlift
+            units
+            provided
+        
+
+        
+            io.airlift
+            slice
+            provided
+        
+
+        
+            com.fasterxml.jackson.core
+            jackson-annotations
+            provided
+        
+
+        
+            org.openjdk.jol
+            jol-core
+            provided
+        
+
+        
+        
+            com.facebook.presto
+            presto-testng-services
+            test
+        
+
+        
+            org.testng
+            testng
+            test
+        
+
+        
+            com.facebook.airlift
+            testing
+            test
+        
+
+        
+            com.facebook.presto
+            presto-session-property-managers-common
+            test-jar
+            test
+        
+    
+
diff --git a/presto-file-session-property-manager/src/main/java/com/facebook/presto/session/file/FileSessionPropertyConfigurationManagerPlugin.java b/presto-file-session-property-manager/src/main/java/com/facebook/presto/session/file/FileSessionPropertyConfigurationManagerPlugin.java
new file mode 100644
index 0000000000000..573c62295572c
--- /dev/null
+++ b/presto-file-session-property-manager/src/main/java/com/facebook/presto/session/file/FileSessionPropertyConfigurationManagerPlugin.java
@@ -0,0 +1,29 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.session.file;
+
+import com.facebook.presto.spi.Plugin;
+import com.facebook.presto.spi.session.SessionPropertyConfigurationManagerFactory;
+import com.google.common.collect.ImmutableList;
+
+public class FileSessionPropertyConfigurationManagerPlugin
+        implements Plugin
+{
+    @Override
+    public Iterable getSessionPropertyConfigurationManagerFactories()
+    {
+        return ImmutableList.of(
+                new FileSessionPropertyManagerFactory());
+    }
+}
diff --git a/presto-session-property-managers/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManager.java b/presto-file-session-property-manager/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManager.java
similarity index 100%
rename from presto-session-property-managers/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManager.java
rename to presto-file-session-property-manager/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManager.java
diff --git a/presto-session-property-managers/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManagerConfig.java b/presto-file-session-property-manager/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManagerConfig.java
similarity index 100%
rename from presto-session-property-managers/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManagerConfig.java
rename to presto-file-session-property-manager/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManagerConfig.java
diff --git a/presto-session-property-managers/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManagerFactory.java b/presto-file-session-property-manager/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManagerFactory.java
similarity index 100%
rename from presto-session-property-managers/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManagerFactory.java
rename to presto-file-session-property-manager/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManagerFactory.java
diff --git a/presto-session-property-managers/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManagerModule.java b/presto-file-session-property-manager/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManagerModule.java
similarity index 100%
rename from presto-session-property-managers/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManagerModule.java
rename to presto-file-session-property-manager/src/main/java/com/facebook/presto/session/file/FileSessionPropertyManagerModule.java
diff --git a/presto-session-property-managers/src/test/java/com/facebook/presto/session/file/TestFileSessionPropertyManager.java b/presto-file-session-property-manager/src/test/java/com/facebook/presto/session/file/TestFileSessionPropertyManager.java
similarity index 100%
rename from presto-session-property-managers/src/test/java/com/facebook/presto/session/file/TestFileSessionPropertyManager.java
rename to presto-file-session-property-manager/src/test/java/com/facebook/presto/session/file/TestFileSessionPropertyManager.java
diff --git a/presto-session-property-managers/src/test/java/com/facebook/presto/session/file/TestFileSessionPropertyManagerConfig.java b/presto-file-session-property-manager/src/test/java/com/facebook/presto/session/file/TestFileSessionPropertyManagerConfig.java
similarity index 94%
rename from presto-session-property-managers/src/test/java/com/facebook/presto/session/file/TestFileSessionPropertyManagerConfig.java
rename to presto-file-session-property-manager/src/test/java/com/facebook/presto/session/file/TestFileSessionPropertyManagerConfig.java
index 3c6c509a4490e..43403ef267bce 100644
--- a/presto-session-property-managers/src/test/java/com/facebook/presto/session/file/TestFileSessionPropertyManagerConfig.java
+++ b/presto-file-session-property-manager/src/test/java/com/facebook/presto/session/file/TestFileSessionPropertyManagerConfig.java
@@ -16,7 +16,7 @@
 import com.google.common.collect.ImmutableMap;
 import org.testng.annotations.Test;
 
-import java.io.File;
+import java.nio.file.Paths;
 import java.util.Map;
 
 import static com.facebook.airlift.configuration.testing.ConfigAssertions.assertFullMapping;
@@ -40,7 +40,7 @@ public void testExplicitPropertyMappings()
                 .build();
 
         FileSessionPropertyManagerConfig expected = new FileSessionPropertyManagerConfig()
-                .setConfigFile(new File("/test.json"));
+                .setConfigFile(Paths.get("/test.json").toFile());
 
         assertFullMapping(properties, expected);
     }
diff --git a/presto-main/etc/session-property-config.json b/presto-main/etc/session-property-config.json
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/presto-main/etc/session-property-config.properties b/presto-main/etc/session-property-config.properties
deleted file mode 100644
index e69de29bb2d1d..0000000000000
diff --git a/presto-server/src/main/provisio/presto.xml b/presto-server/src/main/provisio/presto.xml
index 10359099bce2a..b14b36a768e69 100644
--- a/presto-server/src/main/provisio/presto.xml
+++ b/presto-server/src/main/provisio/presto.xml
@@ -41,8 +41,14 @@
         
     
 
-    
-        
+    
+        
+            
+        
+    
+
+    
+        
             
         
     
diff --git a/presto-session-property-managers-common/pom.xml b/presto-session-property-managers-common/pom.xml
new file mode 100644
index 0000000000000..098b68e7ea58c
--- /dev/null
+++ b/presto-session-property-managers-common/pom.xml
@@ -0,0 +1,81 @@
+
+
+    4.0.0
+
+    
+        com.facebook.presto
+        presto-root
+        0.295-SNAPSHOT
+    
+
+    presto-session-property-managers-common
+    presto-session-property-managers-common
+    Presto - Session Property Managers Common
+    jar
+
+    
+        ${project.parent.basedir}
+    
+
+    
+        
+            com.google.guava
+            guava
+        
+
+        
+            org.jdbi
+            jdbi3-core
+        
+
+        
+        
+            com.facebook.presto
+            presto-spi
+            provided
+        
+
+        
+            com.facebook.presto
+            presto-common
+            provided
+        
+
+        
+            com.facebook.airlift
+            units
+            provided
+        
+
+        
+            io.airlift
+            slice
+            provided
+        
+
+        
+            com.fasterxml.jackson.core
+            jackson-annotations
+            provided
+        
+
+        
+            org.openjdk.jol
+            jol-core
+            provided
+        
+
+        
+        
+            com.facebook.presto
+            presto-testng-services
+            test
+        
+
+        
+            org.testng
+            testng
+            test
+        
+    
+
diff --git a/presto-session-property-managers/src/main/java/com/facebook/presto/session/AbstractSessionPropertyManager.java b/presto-session-property-managers-common/src/main/java/com/facebook/presto/session/AbstractSessionPropertyManager.java
similarity index 100%
rename from presto-session-property-managers/src/main/java/com/facebook/presto/session/AbstractSessionPropertyManager.java
rename to presto-session-property-managers-common/src/main/java/com/facebook/presto/session/AbstractSessionPropertyManager.java
diff --git a/presto-session-property-managers/src/main/java/com/facebook/presto/session/SessionMatchSpec.java b/presto-session-property-managers-common/src/main/java/com/facebook/presto/session/SessionMatchSpec.java
similarity index 98%
rename from presto-session-property-managers/src/main/java/com/facebook/presto/session/SessionMatchSpec.java
rename to presto-session-property-managers-common/src/main/java/com/facebook/presto/session/SessionMatchSpec.java
index f5e655c705eef..8d7a88826d90d 100644
--- a/presto-session-property-managers/src/main/java/com/facebook/presto/session/SessionMatchSpec.java
+++ b/presto-session-property-managers-common/src/main/java/com/facebook/presto/session/SessionMatchSpec.java
@@ -33,12 +33,12 @@
 import java.util.Set;
 import java.util.regex.Pattern;
 
-import static com.facebook.presto.session.db.SessionPropertiesDao.EMPTY_CATALOG;
 import static com.google.common.base.Preconditions.checkArgument;
 import static java.util.Objects.requireNonNull;
 
 public class SessionMatchSpec
 {
+    public static String emptyCatalog = "__NULL__";
     private final Optional userRegex;
     private final Optional sourceRegex;
     private final Set clientTags;
@@ -211,7 +211,7 @@ private SessionInfo getProperties(Optional names, Optional value
             Map> catalogSessionProperties = new HashMap<>();
             Map systemSessionProperties = new HashMap<>();
             for (int i = 0; i < sessionPropertyNames.size(); i++) {
-                if (sessionPropertyCatalogs.get(i).equals(EMPTY_CATALOG)) {
+                if (sessionPropertyCatalogs.get(i).equals(emptyCatalog)) {
                     systemSessionProperties.put(sessionPropertyNames.get(i), sessionPropertyValues.get(i));
                 }
                 else {
diff --git a/presto-session-property-managers/src/test/java/com/facebook/presto/session/AbstractTestSessionPropertyManager.java b/presto-session-property-managers-common/src/test/java/com/facebook/presto/session/AbstractTestSessionPropertyManager.java
similarity index 100%
rename from presto-session-property-managers/src/test/java/com/facebook/presto/session/AbstractTestSessionPropertyManager.java
rename to presto-session-property-managers-common/src/test/java/com/facebook/presto/session/AbstractTestSessionPropertyManager.java
diff --git a/presto-spark-package/pom.xml b/presto-spark-package/pom.xml
index 84ee9009f0bd7..44222dece8b37 100644
--- a/presto-spark-package/pom.xml
+++ b/presto-spark-package/pom.xml
@@ -152,7 +152,15 @@
 
         
             com.facebook.presto
-            presto-session-property-managers
+            presto-file-session-property-manager
+            ${project.version}
+            zip
+            provided
+        
+
+        
+            com.facebook.presto
+            presto-db-session-property-manager
             ${project.version}
             zip
             provided
diff --git a/presto-spark-package/src/main/assembly/presto.xml b/presto-spark-package/src/main/assembly/presto.xml
index 165bac55826e0..79446c79d59ba 100644
--- a/presto-spark-package/src/main/assembly/presto.xml
+++ b/presto-spark-package/src/main/assembly/presto.xml
@@ -85,8 +85,12 @@
             plugin/sqlserver
         
         
-            ${project.build.directory}/dependency/presto-session-property-managers-${project.version}
-            plugin/session-property-managers
+            ${project.build.directory}/dependency/presto-file-session-property-manager-${project.version}
+            plugin/file-session-property-manager
+        
+        
+            ${project.build.directory}/dependency/presto-db-session-property-manager-${project.version}
+            plugin/db-session-property-manager
         
     
 
diff --git a/presto-spark-testing/pom.xml b/presto-spark-testing/pom.xml
index 6ecb44edc0a25..14d2dcf21d548 100644
--- a/presto-spark-testing/pom.xml
+++ b/presto-spark-testing/pom.xml
@@ -88,7 +88,13 @@
 
         
             com.facebook.presto
-            presto-session-property-managers
+            presto-session-property-managers-common
+            test
+        
+
+        
+            com.facebook.presto
+            presto-file-session-property-manager
             test
         
 
diff --git a/presto-test-coverage/pom.xml b/presto-test-coverage/pom.xml
index dc812c9675e32..849974018f171 100644
--- a/presto-test-coverage/pom.xml
+++ b/presto-test-coverage/pom.xml
@@ -480,7 +480,21 @@
 
         
             com.facebook.presto
-            presto-session-property-managers
+            presto-session-property-managers-common
+            ${project.version}
+            compile
+        
+
+        
+            com.facebook.presto
+            presto-db-session-property-manager
+            ${project.version}
+            compile
+        
+
+        
+            com.facebook.presto
+            presto-file-session-property-manager
             ${project.version}
             compile
         

From 00f3749103af7b70e8541847342b09de17e185fb Mon Sep 17 00:00:00 2001
From: Abhash Jain 
Date: Tue, 2 Sep 2025 23:25:48 -0700
Subject: [PATCH 047/113] [presto] Fixing the Writer count from the system
 config (#25941)

Summary:

In this diff I update the taskWriterCount and taskPartitionedWriterCount value to system config which was deleted in D80124169 diff.
kTaskWriterCount is important config for the Impulse connector created table to increase the parallism to ingestion the data by creating multiple writer drivers.

Differential Revision: D81522582
---
 .../presto_cpp/main/PrestoToVeloxQueryConfig.cpp            | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/presto-native-execution/presto_cpp/main/PrestoToVeloxQueryConfig.cpp b/presto-native-execution/presto_cpp/main/PrestoToVeloxQueryConfig.cpp
index 8f2fbb5d202ec..e4e216450287a 100644
--- a/presto-native-execution/presto_cpp/main/PrestoToVeloxQueryConfig.cpp
+++ b/presto-native-execution/presto_cpp/main/PrestoToVeloxQueryConfig.cpp
@@ -139,6 +139,12 @@ void updateFromSystemConfigs(
       {std::string(SystemConfig::kUseLegacyArrayAgg),
        velox::core::QueryConfig::kPrestoArrayAggIgnoreNulls},
 
+      {std::string{SystemConfig::kTaskWriterCount},
+        velox::core::QueryConfig::kTaskWriterCount},
+
+      {std::string{SystemConfig::kTaskPartitionedWriterCount},
+        velox::core::QueryConfig::kTaskPartitionedWriterCount},
+
       {std::string(SystemConfig::kSinkMaxBufferSize),
        velox::core::QueryConfig::kMaxOutputBufferSize,
        [](const auto& value) {

From 686d1c5253927a8adc79685164d87e8ace9009a3 Mon Sep 17 00:00:00 2001
From: Pradeep Vaka 
Date: Wed, 3 Sep 2025 15:53:14 -0700
Subject: [PATCH 048/113] Support multiple catalogs for Presto spark native
 execution (#25943)

## Description
This PR implements native catalog properties support for Presto Spark,
enabling proper configuration and management of catalog properties for
native execution processes.

**Key Changes:**

* **Added `NativeExecutionCatalogProperties` class**: A new class that
holds catalog properties for native execution processes, where each
catalog generates a separate `.properties` file.

* **Enhanced `WorkerProperty` and `PrestoSparkWorkerProperty`**:
Extended the worker property classes to support catalog properties
configuration and proper property file generation.

* **Updated Native Execution Module**: Modified `NativeExecutionModule`
and `NativeExecutionProcess` to integrate catalog properties into the
native execution workflow.

* **Improved Configuration Integration**: Updated `PrestoSparkModule`,
`PrestoSparkServiceFactory` to properly wire the new catalog properties
functionality.

* **Enhanced Test Coverage**: Added and updated tests in
`TestNativeExecutionSystemConfig`, `TestNativeExecutionProcess`, and
other test classes to ensure proper catalog properties handling.

## Motivation and Context
This change is required to support proper catalog configuration in
Presto Spark's native execution mode. Previously, catalog properties
were not properly managed for native execution processes, which limited
the ability to configure connectors effectively in native mode.

## Test Plan
*   **Unit Tests**: Added test `NativeExecutionCatalogProperties`
* **Integration Tests**: Updated `TestNativeExecutionProcess` and
`TestPrestoSparkHttpClient` to verify catalog properties integration

## Contributor checklist

- [x] Please make sure your submission complies with our [contributing
guide](https://github.com/prestodb/presto/blob/master/CONTRIBUTING.md),
in particular [code
style](https://github.com/prestodb/presto/blob/master/CONTRIBUTING.md#code-style)
and [commit
standards](https://github.com/prestodb/presto/blob/master/CONTRIBUTING.md#commit-standards).
- [x] PR description addresses the issue accurately and concisely. If
the change is non-trivial, a GitHub Issue is referenced.
- [x] Adequate tests were added if applicable.
- [x] CI passed.

## Release Notes
```
== NO RELEASE NOTE ==
```
---
 .../PrestoSparkNativeQueryRunnerUtils.java    |  4 +-
 .../presto/spark/PrestoSparkModule.java       |  2 -
 .../spark/PrestoSparkServiceFactory.java      |  3 +-
 .../nativeprocess/NativeExecutionModule.java  | 28 +++---
 .../nativeprocess/NativeExecutionProcess.java |  3 +-
 .../NativeExecutionCatalogProperties.java     | 37 ++++++++
 .../NativeExecutionConnectorConfig.java       | 86 -------------------
 .../property/PrestoSparkWorkerProperty.java   |  6 +-
 .../execution/property/WorkerProperty.java    | 32 +++++--
 .../execution/TestNativeExecutionProcess.java |  4 +-
 .../http/TestPrestoSparkHttpClient.java       |  4 +-
 .../TestNativeExecutionSystemConfig.java      | 40 ++++-----
 12 files changed, 104 insertions(+), 145 deletions(-)
 create mode 100644 presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/NativeExecutionCatalogProperties.java
 delete mode 100644 presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/NativeExecutionConnectorConfig.java

diff --git a/presto-native-execution/src/test/java/com/facebook/presto/spark/PrestoSparkNativeQueryRunnerUtils.java b/presto-native-execution/src/test/java/com/facebook/presto/spark/PrestoSparkNativeQueryRunnerUtils.java
index f51acf99a6155..14113b22d4fb3 100644
--- a/presto-native-execution/src/test/java/com/facebook/presto/spark/PrestoSparkNativeQueryRunnerUtils.java
+++ b/presto-native-execution/src/test/java/com/facebook/presto/spark/PrestoSparkNativeQueryRunnerUtils.java
@@ -18,7 +18,6 @@
 import com.facebook.presto.hive.metastore.ExtendedHiveMetastore;
 import com.facebook.presto.nativeworker.PrestoNativeQueryRunnerUtils;
 import com.facebook.presto.spark.execution.nativeprocess.NativeExecutionModule;
-import com.facebook.presto.spark.execution.property.NativeExecutionConnectorConfig;
 import com.facebook.presto.spi.security.PrincipalType;
 import com.facebook.presto.testing.QueryRunner;
 import com.google.common.collect.ImmutableList;
@@ -125,7 +124,8 @@ public static PrestoSparkQueryRunner createTpchRunner()
         return createRunner(
                 "tpchstandard",
                 new NativeExecutionModule(
-                        Optional.of(new NativeExecutionConnectorConfig().setConnectorName("tpch"))));
+                        Optional.of(
+                                ImmutableMap.of("hive", ImmutableMap.of("connector.name", "tpch")))));
     }
 
     public static PrestoSparkQueryRunner createRunner(String defaultCatalog, Optional baseDir, Map additionalConfigProperties, Map additionalSparkProperties, ImmutableList nativeModules)
diff --git a/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkModule.java b/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkModule.java
index 82716e8a840c2..520d1d8c5bf7d 100644
--- a/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkModule.java
+++ b/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkModule.java
@@ -120,7 +120,6 @@
 import com.facebook.presto.spark.execution.PrestoSparkBroadcastTableCacheManager;
 import com.facebook.presto.spark.execution.PrestoSparkExecutionExceptionFactory;
 import com.facebook.presto.spark.execution.http.BatchTaskUpdateRequest;
-import com.facebook.presto.spark.execution.property.NativeExecutionConnectorConfig;
 import com.facebook.presto.spark.execution.property.NativeExecutionNodeConfig;
 import com.facebook.presto.spark.execution.shuffle.PrestoSparkLocalShuffleReadInfo;
 import com.facebook.presto.spark.execution.shuffle.PrestoSparkLocalShuffleWriteInfo;
@@ -284,7 +283,6 @@ protected void setup(Binder binder)
         configBinder(binder).bindConfig(PrestoSparkConfig.class);
         configBinder(binder).bindConfig(TracingConfig.class);
         configBinder(binder).bindConfig(NativeExecutionNodeConfig.class);
-        configBinder(binder).bindConfig(NativeExecutionConnectorConfig.class);
         configBinder(binder).bindConfig(PlanCheckerProviderManagerConfig.class);
         configBinder(binder).bindConfig(SecurityConfig.class);
 
diff --git a/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkServiceFactory.java b/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkServiceFactory.java
index 741b018631c62..c86d3b30f94ef 100644
--- a/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkServiceFactory.java
+++ b/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkServiceFactory.java
@@ -27,6 +27,7 @@
 import com.google.inject.Module;
 
 import java.util.List;
+import java.util.Optional;
 
 import static com.facebook.presto.spark.classloader_interface.PrestoSparkConfiguration.METADATA_STORAGE_TYPE_LOCAL;
 import static com.google.common.base.Preconditions.checkArgument;
@@ -70,7 +71,7 @@ protected List getAdditionalModules(PrestoSparkConfiguration configurati
         return ImmutableList.of(
                 new PrestoSparkLocalMetadataStorageModule(),
                 // TODO: Need to let NativeExecutionModule addition be controlled by configuration as well.
-                new NativeExecutionModule());
+                new NativeExecutionModule(Optional.of(configuration.getCatalogProperties())));
     }
 
     protected SqlParserOptions getSqlParserOptions()
diff --git a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/nativeprocess/NativeExecutionModule.java b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/nativeprocess/NativeExecutionModule.java
index d3e9fb1bc3786..5636681a85da5 100644
--- a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/nativeprocess/NativeExecutionModule.java
+++ b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/nativeprocess/NativeExecutionModule.java
@@ -13,9 +13,7 @@
  */
 package com.facebook.presto.spark.execution.nativeprocess;
 
-import com.facebook.presto.spark.execution.property.NativeExecutionConnectorConfig;
-import com.facebook.presto.spark.execution.property.NativeExecutionNodeConfig;
-import com.facebook.presto.spark.execution.property.NativeExecutionSystemConfig;
+import com.facebook.presto.spark.execution.property.NativeExecutionCatalogProperties;
 import com.facebook.presto.spark.execution.property.PrestoSparkWorkerProperty;
 import com.facebook.presto.spark.execution.property.WorkerProperty;
 import com.facebook.presto.spark.execution.shuffle.PrestoSparkLocalShuffleInfoTranslator;
@@ -32,6 +30,7 @@
 import com.google.inject.TypeLiteral;
 import okhttp3.OkHttpClient;
 
+import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.TimeUnit;
 
@@ -40,20 +39,20 @@
 public class NativeExecutionModule
         implements Module
 {
-    private Optional connectorConfig;
+    private Optional>> catalogProperties;
 
     // For use by production system where the configurations can only be tuned via configurations.
     public NativeExecutionModule()
     {
-        this.connectorConfig = Optional.empty();
+        this.catalogProperties = Optional.empty();
     }
 
     // In the future, we would make more bindings injected into NativeExecutionModule
     // to be able to test various configuration parameters
     @VisibleForTesting
-    public NativeExecutionModule(Optional connectorConfig)
+    public NativeExecutionModule(Optional>> catalogProperties)
     {
-        this.connectorConfig = connectorConfig;
+        this.catalogProperties = catalogProperties;
     }
 
     @Override
@@ -74,17 +73,14 @@ protected void bindShuffle(Binder binder)
 
     protected void bindWorkerProperties(Binder binder)
     {
+        // Bind NativeExecutionCatalogProperties - this is not bound elsewhere
+        binder.bind(NativeExecutionCatalogProperties.class).toInstance(
+                new NativeExecutionCatalogProperties(catalogProperties.orElse(ImmutableMap.of())));
+
+        // Bind worker property classes
         newOptionalBinder(binder, new TypeLiteral>() {
         }).setDefault().to(PrestoSparkWorkerProperty.class).in(Scopes.SINGLETON);
-        if (connectorConfig.isPresent()) {
-            binder.bind(PrestoSparkWorkerProperty.class).toInstance(
-                    new PrestoSparkWorkerProperty(connectorConfig.get(),
-                            new NativeExecutionNodeConfig(), new NativeExecutionSystemConfig(
-                            ImmutableMap.of())));
-        }
-        else {
-            binder.bind(PrestoSparkWorkerProperty.class).in(Scopes.SINGLETON);
-        }
+        binder.bind(PrestoSparkWorkerProperty.class).in(Scopes.SINGLETON);
     }
 
     protected void bindHttpClient(Binder binder)
diff --git a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/nativeprocess/NativeExecutionProcess.java b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/nativeprocess/NativeExecutionProcess.java
index 794574b46c7f3..d0c7937c0fe52 100644
--- a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/nativeprocess/NativeExecutionProcess.java
+++ b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/nativeprocess/NativeExecutionProcess.java
@@ -343,8 +343,7 @@ private void populateConfigurationFiles(String configBasePath)
         workerProperty.populateAllProperties(
                 Paths.get(configBasePath, WORKER_CONFIG_FILE),
                 Paths.get(configBasePath, WORKER_NODE_CONFIG_FILE),
-                Paths.get(configBasePath, format("%s%s.properties", WORKER_CONNECTOR_CONFIG_FILE,
-                    getNativeExecutionCatalogName(session))));
+                Paths.get(configBasePath, WORKER_CONNECTOR_CONFIG_FILE));  // Directory path for catalogs
     }
 
     private void doGetServerInfo(SettableFuture future)
diff --git a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/NativeExecutionCatalogProperties.java b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/NativeExecutionCatalogProperties.java
new file mode 100644
index 0000000000000..a217c4e7f5ae8
--- /dev/null
+++ b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/NativeExecutionCatalogProperties.java
@@ -0,0 +1,37 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.spark.execution.property;
+
+import java.util.Map;
+
+import static java.util.Objects.requireNonNull;
+
+/**
+ * This class holds catalog properties for native execution process.
+ * Each catalog will generate a separate .properties file.
+ */
+public class NativeExecutionCatalogProperties
+{
+    private final Map> catalogProperties;
+
+    public NativeExecutionCatalogProperties(Map> catalogProperties)
+    {
+        this.catalogProperties = requireNonNull(catalogProperties, "catalogProperties is null");
+    }
+
+    public Map> getAllCatalogProperties()
+    {
+        return catalogProperties;
+    }
+}
diff --git a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/NativeExecutionConnectorConfig.java b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/NativeExecutionConnectorConfig.java
deleted file mode 100644
index 34fc929b6da79..0000000000000
--- a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/NativeExecutionConnectorConfig.java
+++ /dev/null
@@ -1,86 +0,0 @@
-/*
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package com.facebook.presto.spark.execution.property;
-
-import com.facebook.airlift.configuration.Config;
-import com.facebook.airlift.units.DataSize;
-import com.google.common.collect.ImmutableMap;
-
-import java.util.Map;
-
-/**
- * This config class corresponds to catalog/.properties for native execution process. Properties inside will be used in PrestoServer.cpp
- * Currently, presto-on-spark native only supports querying 1 catalog in a single spark session.
- */
-public class NativeExecutionConnectorConfig
-{
-    private static final String CACHE_ENABLED = "cache.enabled";
-    private static final String CACHE_MAX_CACHE_SIZE = "cache.max-cache-size";
-    private static final String CONNECTOR_NAME = "connector.name";
-
-    private boolean cacheEnabled;
-    // maxCacheSize in MB
-    private DataSize maxCacheSize = new DataSize(0, DataSize.Unit.MEGABYTE);
-    private String connectorName = "hive";
-
-    public Map getAllProperties()
-    {
-        ImmutableMap.Builder builder = ImmutableMap.builder();
-        return builder.put(CACHE_ENABLED, String.valueOf(isCacheEnabled()))
-                .put(CACHE_MAX_CACHE_SIZE, String.valueOf(getDataSizeInLong(getMaxCacheSize().convertTo(DataSize.Unit.MEGABYTE))))
-                .put(CONNECTOR_NAME, getConnectorName())
-                .build();
-    }
-
-    public boolean isCacheEnabled()
-    {
-        return cacheEnabled;
-    }
-
-    @Config(CACHE_ENABLED)
-    public NativeExecutionConnectorConfig setCacheEnabled(boolean cacheEnabled)
-    {
-        this.cacheEnabled = cacheEnabled;
-        return this;
-    }
-
-    public DataSize getMaxCacheSize()
-    {
-        return maxCacheSize;
-    }
-
-    @Config(CACHE_MAX_CACHE_SIZE)
-    public NativeExecutionConnectorConfig setMaxCacheSize(DataSize maxCacheSize)
-    {
-        this.maxCacheSize = maxCacheSize;
-        return this;
-    }
-
-    public String getConnectorName()
-    {
-        return connectorName;
-    }
-
-    @Config(CONNECTOR_NAME)
-    public NativeExecutionConnectorConfig setConnectorName(String connectorName)
-    {
-        this.connectorName = connectorName;
-        return this;
-    }
-
-    private Long getDataSizeInLong(DataSize size)
-    {
-        return Double.valueOf(size.getValue()).longValue();
-    }
-}
diff --git a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/PrestoSparkWorkerProperty.java b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/PrestoSparkWorkerProperty.java
index f83bcb1464e92..becc2c4b47d71 100644
--- a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/PrestoSparkWorkerProperty.java
+++ b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/PrestoSparkWorkerProperty.java
@@ -19,14 +19,14 @@
  * A utility class that helps with properties and its materialization.
  */
 public class PrestoSparkWorkerProperty
-        extends WorkerProperty
+        extends WorkerProperty
 {
     @Inject
     public PrestoSparkWorkerProperty(
-            NativeExecutionConnectorConfig connectorConfig,
+            NativeExecutionCatalogProperties catalogProperties,
             NativeExecutionNodeConfig nodeConfig,
             NativeExecutionSystemConfig systemConfig)
     {
-        super(connectorConfig, nodeConfig, systemConfig);
+        super(catalogProperties, nodeConfig, systemConfig);
     }
 }
diff --git a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/WorkerProperty.java b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/WorkerProperty.java
index 1a2faa3edf009..67b5aa22986fb 100644
--- a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/WorkerProperty.java
+++ b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/property/WorkerProperty.java
@@ -64,25 +64,25 @@
  * }
  * }
  */
-public class WorkerProperty
+public class WorkerProperty
 {
-    private final T1 connectorConfig;
+    private final T1 catalogProperties;
     private final T2 nodeConfig;
     private final T3 systemConfig;
 
     public WorkerProperty(
-            T1 connectorConfig,
+            T1 catalogProperties,
             T2 nodeConfig,
             T3 systemConfig)
     {
         this.systemConfig = requireNonNull(systemConfig, "systemConfig is null");
         this.nodeConfig = requireNonNull(nodeConfig, "nodeConfig is null");
-        this.connectorConfig = requireNonNull(connectorConfig, "connectorConfig is null");
+        this.catalogProperties = requireNonNull(catalogProperties, "catalogProperties is null");
     }
 
-    public T1 getConnectorConfig()
+    public T1 getCatalogProperties()
     {
-        return connectorConfig;
+        return catalogProperties;
     }
 
     public T2 getNodeConfig()
@@ -95,12 +95,12 @@ public T3 getSystemConfig()
         return systemConfig;
     }
 
-    public void populateAllProperties(Path systemConfigPath, Path nodeConfigPath, Path connectorConfigPath)
+    public void populateAllProperties(Path systemConfigPath, Path nodeConfigPath, Path catalogConfigsPath)
             throws IOException
     {
         populateProperty(systemConfig.getAllProperties(), systemConfigPath);
         populateProperty(nodeConfig.getAllProperties(), nodeConfigPath);
-        populateProperty(connectorConfig.getAllProperties(), connectorConfigPath);
+        populateCatalogProperties(catalogProperties.getAllCatalogProperties(), catalogConfigsPath);
     }
 
     private void populateProperty(Map properties, Path path)
@@ -122,4 +122,20 @@ private void populateProperty(Map properties, Path path)
             throw e;
         }
     }
+
+    private void populateCatalogProperties(Map> catalogProperties, Path path)
+            throws IOException
+    {
+        File catalogDir = path.toFile();
+        if (!catalogDir.exists()) {
+            catalogDir.mkdirs();
+        }
+
+        for (Map.Entry> catalogEntry : catalogProperties.entrySet()) {
+            String catalogName = catalogEntry.getKey();
+            Map properties = catalogEntry.getValue();
+            Path catalogConfigPath = path.resolve(catalogName + ".properties");
+            populateProperty(properties, catalogConfigPath);
+        }
+    }
 }
diff --git a/presto-spark-base/src/test/java/com/facebook/presto/spark/execution/TestNativeExecutionProcess.java b/presto-spark-base/src/test/java/com/facebook/presto/spark/execution/TestNativeExecutionProcess.java
index 1483270d325e2..4819dc6fd3e53 100644
--- a/presto-spark-base/src/test/java/com/facebook/presto/spark/execution/TestNativeExecutionProcess.java
+++ b/presto-spark-base/src/test/java/com/facebook/presto/spark/execution/TestNativeExecutionProcess.java
@@ -22,7 +22,7 @@
 import com.facebook.presto.spark.execution.http.TestPrestoSparkHttpClient;
 import com.facebook.presto.spark.execution.nativeprocess.NativeExecutionProcess;
 import com.facebook.presto.spark.execution.nativeprocess.NativeExecutionProcessFactory;
-import com.facebook.presto.spark.execution.property.NativeExecutionConnectorConfig;
+import com.facebook.presto.spark.execution.property.NativeExecutionCatalogProperties;
 import com.facebook.presto.spark.execution.property.NativeExecutionNodeConfig;
 import com.facebook.presto.spark.execution.property.NativeExecutionSystemConfig;
 import com.facebook.presto.spark.execution.property.PrestoSparkWorkerProperty;
@@ -88,7 +88,7 @@ private NativeExecutionProcessFactory createNativeExecutionProcessFactory()
         TaskId taskId = new TaskId("testid", 0, 0, 0, 0);
         ScheduledExecutorService errorScheduler = newScheduledThreadPool(4);
         PrestoSparkWorkerProperty workerProperty = new PrestoSparkWorkerProperty(
-                new NativeExecutionConnectorConfig(),
+                new NativeExecutionCatalogProperties(ImmutableMap.of()),
                 new NativeExecutionNodeConfig(),
                 new NativeExecutionSystemConfig(ImmutableMap.of()));
         NativeExecutionProcessFactory factory = new NativeExecutionProcessFactory(
diff --git a/presto-spark-base/src/test/java/com/facebook/presto/spark/execution/http/TestPrestoSparkHttpClient.java b/presto-spark-base/src/test/java/com/facebook/presto/spark/execution/http/TestPrestoSparkHttpClient.java
index 6927065bb0842..700d55ade4e80 100644
--- a/presto-spark-base/src/test/java/com/facebook/presto/spark/execution/http/TestPrestoSparkHttpClient.java
+++ b/presto-spark-base/src/test/java/com/facebook/presto/spark/execution/http/TestPrestoSparkHttpClient.java
@@ -33,7 +33,7 @@
 import com.facebook.presto.spark.execution.nativeprocess.HttpNativeExecutionTaskResultFetcher;
 import com.facebook.presto.spark.execution.nativeprocess.NativeExecutionProcess;
 import com.facebook.presto.spark.execution.nativeprocess.NativeExecutionProcessFactory;
-import com.facebook.presto.spark.execution.property.NativeExecutionConnectorConfig;
+import com.facebook.presto.spark.execution.property.NativeExecutionCatalogProperties;
 import com.facebook.presto.spark.execution.property.NativeExecutionNodeConfig;
 import com.facebook.presto.spark.execution.property.NativeExecutionSystemConfig;
 import com.facebook.presto.spark.execution.property.PrestoSparkWorkerProperty;
@@ -897,7 +897,7 @@ private NativeExecutionProcess createNativeExecutionProcess(
             TestingResponseManager responseManager)
     {
         PrestoSparkWorkerProperty workerProperty = new PrestoSparkWorkerProperty(
-                new NativeExecutionConnectorConfig(),
+                new NativeExecutionCatalogProperties(ImmutableMap.of()),
                 new NativeExecutionNodeConfig(),
                 new NativeExecutionSystemConfig(ImmutableMap.of()));
         NativeExecutionProcessFactory factory = new NativeExecutionProcessFactory(
diff --git a/presto-spark-base/src/test/java/com/facebook/presto/spark/execution/property/TestNativeExecutionSystemConfig.java b/presto-spark-base/src/test/java/com/facebook/presto/spark/execution/property/TestNativeExecutionSystemConfig.java
index 53200d798baba..0568193e7c815 100644
--- a/presto-spark-base/src/test/java/com/facebook/presto/spark/execution/property/TestNativeExecutionSystemConfig.java
+++ b/presto-spark-base/src/test/java/com/facebook/presto/spark/execution/property/TestNativeExecutionSystemConfig.java
@@ -14,7 +14,6 @@
 package com.facebook.presto.spark.execution.property;
 
 import com.facebook.airlift.configuration.testing.ConfigAssertions;
-import com.facebook.airlift.units.DataSize;
 import com.google.common.collect.ImmutableMap;
 import org.testng.annotations.Test;
 
@@ -202,31 +201,25 @@ public void testNativeExecutionNodeConfig()
     }
 
     @Test
-    public void testNativeExecutionConnectorConfig()
+    public void testNativeExecutionCatalogProperties()
     {
-        // Test defaults
-        assertRecordedDefaults(ConfigAssertions.recordDefaults(NativeExecutionConnectorConfig.class)
-                .setCacheEnabled(false)
-                .setMaxCacheSize(new DataSize(0, DataSize.Unit.MEGABYTE))
-                .setConnectorName("hive"));
+        // Test default constructor
+        NativeExecutionCatalogProperties config = new NativeExecutionCatalogProperties(ImmutableMap.of());
+        assertEquals(config.getAllCatalogProperties(), ImmutableMap.of());
 
-        // Test explicit property mapping. Also makes sure properties returned by getAllProperties() covers full property list.
-        NativeExecutionConnectorConfig expected = new NativeExecutionConnectorConfig()
-                .setConnectorName("custom")
-                .setMaxCacheSize(new DataSize(32, DataSize.Unit.MEGABYTE))
-                .setCacheEnabled(true);
-        Map properties = new java.util.HashMap<>(expected.getAllProperties());
-        // Since the cache.max-cache-size requires to be size without the unit which to be compatible with the C++,
-        // here we convert the size from Long type (in string format) back to DataSize for comparison
-        properties.put("cache.max-cache-size", String.valueOf(new DataSize(Double.parseDouble(properties.get("cache.max-cache-size")), DataSize.Unit.MEGABYTE)));
-        assertFullMapping(properties, expected);
+        // Test constructor with catalog properties
+        Map> catalogProperties = ImmutableMap.of(
+                "hive", ImmutableMap.of("hive.metastore.uri", "thrift://localhost:9083"),
+                "tpch", ImmutableMap.of("tpch.splits-per-node", "4"));
+        NativeExecutionCatalogProperties configWithProps = new NativeExecutionCatalogProperties(catalogProperties);
+        assertEquals(configWithProps.getAllCatalogProperties(), catalogProperties);
     }
 
     @Test
     public void testFilePropertiesPopulator()
     {
         PrestoSparkWorkerProperty workerProperty = new PrestoSparkWorkerProperty(
-                new NativeExecutionConnectorConfig(), new NativeExecutionNodeConfig(),
+                new NativeExecutionCatalogProperties(ImmutableMap.of()), new NativeExecutionNodeConfig(),
                 new NativeExecutionSystemConfig(ImmutableMap.of()));
         testPropertiesPopulate(workerProperty);
     }
@@ -238,12 +231,17 @@ private void testPropertiesPopulate(PrestoSparkWorkerProperty workerProperty)
             directory = Files.createTempDirectory("presto");
             Path configPropertiesPath = Paths.get(directory.toString(), "config.properties");
             Path nodePropertiesPath = Paths.get(directory.toString(), "node.properties");
-            Path connectorPropertiesPath = Paths.get(directory.toString(), "catalog/hive.properties");
-            workerProperty.populateAllProperties(configPropertiesPath, nodePropertiesPath, connectorPropertiesPath);
+            Path catalogDirectory = Paths.get(directory.toString(), "catalog");  // Directory path for catalogs
+            workerProperty.populateAllProperties(configPropertiesPath, nodePropertiesPath, catalogDirectory);
 
             verifyProperties(workerProperty.getSystemConfig().getAllProperties(), readPropertiesFromDisk(configPropertiesPath));
             verifyProperties(workerProperty.getNodeConfig().getAllProperties(), readPropertiesFromDisk(nodePropertiesPath));
-            verifyProperties(workerProperty.getConnectorConfig().getAllProperties(), readPropertiesFromDisk(connectorPropertiesPath));
+            // Verify each catalog file was created properly
+            workerProperty.getCatalogProperties().getAllCatalogProperties().forEach(
+                    (catalogName, catalogProperties) -> {
+                        Path catalogFilePath = catalogDirectory.resolve(catalogName + ".properties");
+                        verifyProperties(catalogProperties, readPropertiesFromDisk(catalogFilePath));
+                    });
         }
         catch (Exception exception) {
             exception.printStackTrace();

From dbc36d3758ba0d3d7ab6462b68ba0a03875f1f3e Mon Sep 17 00:00:00 2001
From: Zac Wen 
Date: Wed, 3 Sep 2025 10:49:43 -0700
Subject: [PATCH 049/113] Add native index join perf related session properties

---
 .../NativeWorkerSessionPropertyProvider.java  | 32 ++++++++++++++++---
 1 file changed, 28 insertions(+), 4 deletions(-)

diff --git a/presto-main-base/src/main/java/com/facebook/presto/sessionpropertyproviders/NativeWorkerSessionPropertyProvider.java b/presto-main-base/src/main/java/com/facebook/presto/sessionpropertyproviders/NativeWorkerSessionPropertyProvider.java
index 0e54809e87b80..300f2485b5e97 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/sessionpropertyproviders/NativeWorkerSessionPropertyProvider.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/sessionpropertyproviders/NativeWorkerSessionPropertyProvider.java
@@ -28,7 +28,6 @@
 import static com.facebook.presto.spi.session.PropertyMetadata.stringProperty;
 import static java.util.Objects.requireNonNull;
 
-@Deprecated
 public class NativeWorkerSessionPropertyProvider
         implements WorkerSessionPropertyProvider
 {
@@ -81,6 +80,10 @@ public class NativeWorkerSessionPropertyProvider
     public static final String NATIVE_REQUEST_DATA_SIZES_MAX_WAIT_SEC = "native_request_data_sizes_max_wait_sec";
     public static final String NATIVE_QUERY_MEMORY_RECLAIMER_PRIORITY = "native_query_memory_reclaimer_priority";
     public static final String NATIVE_MAX_NUM_SPLITS_LISTENED_TO = "native_max_num_splits_listened_to";
+    public static final String NATIVE_INDEX_LOOKUP_JOIN_MAX_PREFETCH_BATCHES = "native_index_lookup_join_max_prefetch_batches";
+    public static final String NATIVE_INDEX_LOOKUP_JOIN_SPLIT_OUTPUT = "native_index_lookup_join_split_output";
+    public static final String NATIVE_UNNEST_SPLIT_OUTPUT = "native_unnest_split_output";
+
     private final List> sessionProperties;
 
     @Inject
@@ -152,7 +155,7 @@ public NativeWorkerSessionPropertyProvider(FeaturesConfig featuresConfig)
                 longProperty(
                         NATIVE_WRITER_FLUSH_THRESHOLD_BYTES,
                         "Native Execution only. Minimum memory footprint size required to reclaim memory from a file " +
-                        "writer by flushing its buffered data to disk.",
+                                "writer by flushing its buffered data to disk.",
                         96L << 20,
                         false),
                 booleanProperty(
@@ -266,7 +269,7 @@ public NativeWorkerSessionPropertyProvider(FeaturesConfig featuresConfig)
                         "",
                         !nativeExecution),
                 stringProperty(NATIVE_QUERY_TRACE_FRAGMENT_ID,
-                            "The fragment id of the traced task.",
+                        "The fragment id of the traced task.",
                         "",
                         !nativeExecution),
                 stringProperty(NATIVE_QUERY_TRACE_SHARD_ID,
@@ -366,13 +369,34 @@ public NativeWorkerSessionPropertyProvider(FeaturesConfig featuresConfig)
                 integerProperty(
                         NATIVE_QUERY_MEMORY_RECLAIMER_PRIORITY,
                         "Native Execution only. Priority of memory recliamer when deciding on memory pool to abort." +
-                        "Lower value has higher priority and less likely to be choosen for memory pool abort",
+                                "Lower value has higher priority and less likely to be choosen for memory pool abort",
                         2147483647,
                         !nativeExecution),
                 integerProperty(
                         NATIVE_MAX_NUM_SPLITS_LISTENED_TO,
                         "Maximum number of splits to listen to per table scan node per worker.",
                         0,
+                        !nativeExecution),
+
+                integerProperty(
+                        NATIVE_INDEX_LOOKUP_JOIN_MAX_PREFETCH_BATCHES,
+                        "Specifies the max number of input batches to prefetch to do index lookup ahead. " +
+                                "If it is zero, then process one input batch at a time.",
+                        0,
+                        !nativeExecution),
+                booleanProperty(
+                        NATIVE_INDEX_LOOKUP_JOIN_SPLIT_OUTPUT,
+                        "If this is true, then the index join operator might split output for each input " +
+                                "batch based on the output batch size control. Otherwise, it tries to produce a " +
+                                "single output for each input batch.",
+                        true,
+                        !nativeExecution),
+                booleanProperty(
+                        NATIVE_UNNEST_SPLIT_OUTPUT,
+                        "If this is true, then the unnest operator might split output for each input " +
+                                "batch based on the output batch size control. Otherwise, it produces a single " +
+                                "output for each input batch.",
+                        true,
                         !nativeExecution));
     }
 

From 55df18542c710de6f9cb32a6c19cbc5f828489f8 Mon Sep 17 00:00:00 2001
From: Joe Abraham 
Date: Wed, 13 Aug 2025 11:36:40 +0530
Subject: [PATCH 050/113] Add documentation to run native worker with sidecar

---
 presto-native-execution/README.md             | 33 ++++++++++++++++---
 .../etc/catalog/tpchstandard.properties       |  2 +-
 presto-native-execution/etc/config.properties |  2 +-
 presto-native-execution/etc/velox.properties  |  1 -
 .../etc_sidecar/catalog/hive.properties       |  1 +
 .../etc_sidecar/catalog/iceberg.properties    |  1 +
 .../catalog/tpchstandard.properties           |  1 +
 .../etc_sidecar/config.properties             |  7 ++++
 .../etc_sidecar/node.properties               |  3 ++
 9 files changed, 43 insertions(+), 8 deletions(-)
 delete mode 100644 presto-native-execution/etc/velox.properties
 create mode 100644 presto-native-execution/etc_sidecar/catalog/hive.properties
 create mode 100644 presto-native-execution/etc_sidecar/catalog/iceberg.properties
 create mode 100644 presto-native-execution/etc_sidecar/catalog/tpchstandard.properties
 create mode 100644 presto-native-execution/etc_sidecar/config.properties
 create mode 100644 presto-native-execution/etc_sidecar/node.properties

diff --git a/presto-native-execution/README.md b/presto-native-execution/README.md
index f1e31632de56e..5339a1d995053 100644
--- a/presto-native-execution/README.md
+++ b/presto-native-execution/README.md
@@ -214,12 +214,23 @@ Run IcebergExternalWorkerQueryRunner,
       `$DATA_DIR/iceberg_data//`. Here `file_format` could be `PARQUET | ORC | AVRO` and `catalog_type` could be `HIVE | HADOOP | NESSIE | REST`.
   * Use classpath of module: choose `presto-native-execution` module.
 
+Run NativeSidecarPluginQueryRunner:
+* Edit/Create `NativeSidecarPluginQueryRunner` Application Run/Debug Configuration (alter paths accordingly).
+  * Main class: `com.facebook.presto.sidecar.NativeSidecarPluginQueryRunner`.
+  * VM options : `-ea -Xmx5G -XX:+ExitOnOutOfMemoryError -Duser.timezone=America/Bahia_Banderas -Dhive.security=legacy`.
+  * Working directory: `$MODULE_DIR$`
+  * Environment variables: `PRESTO_SERVER=/Users//git/presto/presto-native-execution/cmake-build-debug/presto_cpp/main/presto_server;DATA_DIR=/Users//Desktop/data;WORKER_COUNT=0`
+  * Use classpath of module: choose `presto-native-execution` module.
+
 Run CLion:
 * File->Close Project if any is open.
 * Open `presto/presto-native-execution` directory as CMake project and wait till CLion loads/generates cmake files, symbols, etc.
 * Edit configuration for `presto_server` module (alter paths accordingly).
   * Program arguments: `--logtostderr=1 --v=1 --etc_dir=/Users//git/presto/presto-native-execution/etc`
   * Working directory: `/Users//git/presto/presto-native-execution`
+* For sidecar, Edit configuration for `presto_server` module (alter paths accordingly).
+  * Program arguments: `--logtostderr=1 --v=1 --etc_dir=/Users//git/presto/presto-native-execution/etc_sidecar`
+  * Working directory: `/Users//git/presto/presto-native-execution`
 * Edit menu CLion->Preferences->Build, Execution, Deployment->CMake
   * CMake options: `-DVELOX_BUILD_TESTING=ON -DCMAKE_BUILD_TYPE=Debug`
   * Build options: `-- -j 12`
@@ -237,12 +248,24 @@ Run CLion:
   * For Hive, Run `HiveExternalWorkerQueryRunner` from IntelliJ and wait until it starts (`======== SERVER STARTED ========` is displayed in the log output).
   * For Iceberg, Run `IcebergExternalWorkerQueryRunner` from IntelliJ and wait until it starts (`======== SERVER STARTED ========` is displayed in the log output).
 * Scroll up the log output and find `Discovery URL http://127.0.0.1:50555`. The port is 'random' with every start.
-* Copy that port (or the whole URL) to the `discovery.uri` field in `presto/presto-native-execution/etc/config.properties` for the worker to discover the Coordinator.
+* Copy that port (or the whole URL) to the `discovery.uri` field in `presto/presto-native-execution/etc/config.properties` for the worker to announce itself to the Coordinator.
 * In CLion run "presto_server" module. Connection success will be indicated by `Announcement succeeded: 202` line in the log output.
-* Two ways to run Presto client to start executing queries on the running local setup:
-  1. In command line from presto root directory run the presto client:
-      * `java -jar presto-cli/target/presto-cli-*-executable.jar --catalog hive --schema tpch`
-  2. Run `Presto Client` Application (see above on how to create and setup the configuration) inside IntelliJ
+* See **Run Presto Client** to start executing queries on the running local setup.
+
+### Run Presto Coordinator + Sidecar
+* Note that everything below can be done without using IDEs by running command line commands (not in this readme).
+* Add a property `presto.default-namespace=native.default` to `presto-native-execution/etc/config.properties`.
+* Run `NativeSidecarPluginQueryRunner` from IntelliJ and wait until it starts (`======== SERVER STARTED ========` is displayed in the log output).
+* Scroll up the log output and find `Discovery URL http://127.0.0.1:50555`. The port is 'random' with every startup.
+* Copy that port (or the whole URL) to the `discovery.uri` field in`presto/presto-native-execution/etc_sidecar/config.properties` for the sidecar to announce itself to the Coordinator.
+* In CLion run "presto_server" module. Connection success will be indicated by `Announcement succeeded: 202` line in the log output.
+* See **Run Presto Client** to start executing queries on the running local setup.
+
+### Run Presto Client
+* Run the following command from the presto root directory to start the Presto client:
+    ```
+    java -jar presto-cli/target/presto-cli-*-executable.jar --catalog hive --schema tpch
+    ```
 * You can start from `show tables;` and `describe table;` queries and execute more queries as needed.
 
 ### Run Integration (End to End or E2E) Tests
diff --git a/presto-native-execution/etc/catalog/tpchstandard.properties b/presto-native-execution/etc/catalog/tpchstandard.properties
index 16e833ca8f436..75110c5acf145 100644
--- a/presto-native-execution/etc/catalog/tpchstandard.properties
+++ b/presto-native-execution/etc/catalog/tpchstandard.properties
@@ -1 +1 @@
-connector.name=tpch
\ No newline at end of file
+connector.name=tpch
diff --git a/presto-native-execution/etc/config.properties b/presto-native-execution/etc/config.properties
index 37db096a9402a..bec551886110b 100644
--- a/presto-native-execution/etc/config.properties
+++ b/presto-native-execution/etc/config.properties
@@ -1,4 +1,4 @@
-discovery.uri=http://127.0.0.1:58215
+discovery.uri=http://127.0.0.1:
 presto.version=testversion
 http-server.http.port=7777
 shutdown-onset-sec=1
diff --git a/presto-native-execution/etc/velox.properties b/presto-native-execution/etc/velox.properties
deleted file mode 100644
index 6c2506bd99a8e..0000000000000
--- a/presto-native-execution/etc/velox.properties
+++ /dev/null
@@ -1 +0,0 @@
-mutable-config=true
\ No newline at end of file
diff --git a/presto-native-execution/etc_sidecar/catalog/hive.properties b/presto-native-execution/etc_sidecar/catalog/hive.properties
new file mode 100644
index 0000000000000..466b7e664e44f
--- /dev/null
+++ b/presto-native-execution/etc_sidecar/catalog/hive.properties
@@ -0,0 +1 @@
+connector.name=hive
diff --git a/presto-native-execution/etc_sidecar/catalog/iceberg.properties b/presto-native-execution/etc_sidecar/catalog/iceberg.properties
new file mode 100644
index 0000000000000..f3a43dcb28126
--- /dev/null
+++ b/presto-native-execution/etc_sidecar/catalog/iceberg.properties
@@ -0,0 +1 @@
+connector.name=iceberg
diff --git a/presto-native-execution/etc_sidecar/catalog/tpchstandard.properties b/presto-native-execution/etc_sidecar/catalog/tpchstandard.properties
new file mode 100644
index 0000000000000..75110c5acf145
--- /dev/null
+++ b/presto-native-execution/etc_sidecar/catalog/tpchstandard.properties
@@ -0,0 +1 @@
+connector.name=tpch
diff --git a/presto-native-execution/etc_sidecar/config.properties b/presto-native-execution/etc_sidecar/config.properties
new file mode 100644
index 0000000000000..76ea4efa7c059
--- /dev/null
+++ b/presto-native-execution/etc_sidecar/config.properties
@@ -0,0 +1,7 @@
+discovery.uri=http://127.0.0.1:
+presto.version=testversion
+http-server.http.port=7778
+shutdown-onset-sec=1
+runtime-metrics-collection-enabled=true
+native-sidecar=true
+presto.default-namespace=native.default
diff --git a/presto-native-execution/etc_sidecar/node.properties b/presto-native-execution/etc_sidecar/node.properties
new file mode 100644
index 0000000000000..1d92b7ace8087
--- /dev/null
+++ b/presto-native-execution/etc_sidecar/node.properties
@@ -0,0 +1,3 @@
+node.environment=testing
+node.internal-address=127.0.0.1
+node.location=testing-location

From e7f7e26a327e3107d9cbf48a7f54c21b9af09af8 Mon Sep 17 00:00:00 2001
From: Steve Burnett 
Date: Wed, 3 Sep 2025 10:53:08 -0400
Subject: [PATCH 051/113] Add node selection strategy documentation

---
 presto-docs/src/main/sphinx/connector/iceberg.rst  |  3 ++-
 presto-docs/src/main/sphinx/develop/connectors.rst | 14 +++++++++++++-
 2 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/presto-docs/src/main/sphinx/connector/iceberg.rst b/presto-docs/src/main/sphinx/connector/iceberg.rst
index 9267f5a8d8df1..7b9adc3c0933a 100644
--- a/presto-docs/src/main/sphinx/connector/iceberg.rst
+++ b/presto-docs/src/main/sphinx/connector/iceberg.rst
@@ -541,7 +541,8 @@ Property Name                                         Description
                                                       assign a split to. Splits which read data from the same file within
                                                       the same chunk will hash to the same node. A smaller chunk size will
                                                       result in a higher probability splits being distributed evenly across
-                                                      the cluster, but reduce locality.
+                                                      the cluster, but reduce locality. 
+                                                      See :ref:`develop/connectors:Node Selection Strategy`.
 ``iceberg.parquet_dereference_pushdown_enabled``      Overrides the behavior of the connector property                        Yes                 No
                                                       ``iceberg.enable-parquet-dereference-pushdown`` in the current session.
 ===================================================== ======================================================================= =================== =============================================
diff --git a/presto-docs/src/main/sphinx/develop/connectors.rst b/presto-docs/src/main/sphinx/develop/connectors.rst
index c77877890e3d3..24713c68cc176 100644
--- a/presto-docs/src/main/sphinx/develop/connectors.rst
+++ b/presto-docs/src/main/sphinx/develop/connectors.rst
@@ -8,7 +8,8 @@ you adapt your data source to the API expected by Presto, you can write
 queries against this data.
 
 ConnectorSplit
-----------------
+--------------
+
 Instances of your connector splits.
 
 The ``getNodeSelectionStrategy`` method indicates the node affinity
@@ -81,3 +82,14 @@ Given a split and a list of columns, the record set provider is
 responsible for delivering data to the Presto execution engine.
 It creates a ``RecordSet``, which in turn creates a ``RecordCursor``
 that is used by Presto to read the column values for each row.
+
+Node Selection Strategy
+-----------------------
+
+The node selection strategy is specified by a connector on each split.  The possible values are:
+
+* HARD_AFFINITY - The Presto runtime must schedule this split on the nodes specified on ``ConnectorSplit#getPreferredNodes``.
+* SOFT_AFFINITY - The Presto runtime should prefer ``ConnectorSplit#getPreferredNodes`` nodes, but doesn't have to. Use this value primarily for caching.
+* NO_PREFERENCE - No preference. 
+
+Use the ``node_selection_strategy`` session property in Hive and Iceberg to override this. 
\ No newline at end of file

From 04a922fd9b3848012fa15d759a5fd1f3d0e0a015 Mon Sep 17 00:00:00 2001
From: Prashant Sharma 
Date: Thu, 4 Sep 2025 15:41:11 +0530
Subject: [PATCH 052/113] [native] README: fixed a typo in aws dep install
 command. (#25878)

## Description
Fixes a small typo on aws dependency install command.

## Motivation and Context



## Impact


## Test Plan


## Contributor checklist

- [ ] Please make sure your submission complies with our [contributing
guide](https://github.com/prestodb/presto/blob/master/CONTRIBUTING.md),
in particular [code
style](https://github.com/prestodb/presto/blob/master/CONTRIBUTING.md#code-style)
and [commit
standards](https://github.com/prestodb/presto/blob/master/CONTRIBUTING.md#commit-standards).
- [ ] PR description addresses the issue accurately and concisely. If
the change is non-trivial, a GitHub Issue is referenced.
- [ ] Documented new properties (with its default value), SQL syntax,
functions, or other functionality.
- [ ] If release notes are required, they follow the [release notes
guidelines](https://github.com/prestodb/presto/wiki/Release-Notes-Guidelines).
- [ ] Adequate tests were added if applicable.
- [ ] CI passed.

## Release Notes
Please follow [release notes
guidelines](https://github.com/prestodb/presto/wiki/Release-Notes-Guidelines)
and fill in the release notes below.

```
== NO RELEASE NOTE ==
```
---
 presto-native-execution/README.md                   |   6 ++++--
 .../images}/cl_clangformat_switcherenable.png       | Bin
 2 files changed, 4 insertions(+), 2 deletions(-)
 rename presto-native-execution/{ => docs/images}/cl_clangformat_switcherenable.png (100%)

diff --git a/presto-native-execution/README.md b/presto-native-execution/README.md
index 5339a1d995053..59f4fe00424ee 100644
--- a/presto-native-execution/README.md
+++ b/presto-native-execution/README.md
@@ -95,7 +95,9 @@ S3 support needs the [AWS SDK C++](https://github.com/aws/aws-sdk-cpp) library.
 This dependency can be installed by running the target platform build script
 from the `presto/presto-native-execution` directory.
 
-`./velox/scripts/setup-centos9.sh install_aws`
+`./velox/scripts/setup-centos9.sh install_aws_deps`
+    Or
+`./velox/scripts/setup-ubuntu.sh install_aws_deps`
 
 #### JWT Authentication
 To enable JWT authentication support, add `-DPRESTO_ENABLE_JWT=ON` to the
@@ -240,7 +242,7 @@ Run CLion:
 * To enable clang format you need
   * Open any h or cpp file in the editor and select `Enable ClangFormat` by clicking `4 spaces` rectangle in the status bar (bottom right) which is next to `UTF-8` bar.
 
-    ![ScreenShot](cl_clangformat_switcherenable.png)
+    ![ScreenShot](docs/images/cl_clangformat_switcherenable.png)
 
 ### Run Presto Coordinator + Worker
 * Note that everything below can be done without using IDEs by running command line commands (not in this readme).
diff --git a/presto-native-execution/cl_clangformat_switcherenable.png b/presto-native-execution/docs/images/cl_clangformat_switcherenable.png
similarity index 100%
rename from presto-native-execution/cl_clangformat_switcherenable.png
rename to presto-native-execution/docs/images/cl_clangformat_switcherenable.png

From 92a8ae51ae69e99904605456edd602d43a46ab35 Mon Sep 17 00:00:00 2001
From: Shrinidhi Joshi 
Date: Wed, 3 Sep 2025 15:18:28 -0700
Subject: [PATCH 053/113] CODEOWNERS: Expand codeownership of presto-spark
 owners to include presto-spark code in presto-native-execution module

---
 CODEOWNERS | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/CODEOWNERS b/CODEOWNERS
index 0b2831ba7ee31..e4a0e3bb56aed 100644
--- a/CODEOWNERS
+++ b/CODEOWNERS
@@ -123,6 +123,7 @@ CODEOWNERS @prestodb/team-tsc
 #####################################################################
 # Presto on Spark module
 /presto-spark* @shrinidhijoshi @prestodb/committers
+/presto-native-execution/*/com/facebook/presto/spark/* @shrinidhijoshi @prestodb/committers
 
 #####################################################################
 # Presto connectors and plugins
@@ -166,4 +167,3 @@ CODEOWNERS @prestodb/team-tsc
 # Presto CI and builds
 /.github @czentgr @unidevel @prestodb/committers
 /docker @czentgr @unidevel @prestodb/committers
-

From 9c5004f0a4ed6de3727696a431583a7294dc28ba Mon Sep 17 00:00:00 2001
From: Gary Helmling 
Date: Mon, 1 Sep 2025 23:09:28 -0700
Subject: [PATCH 054/113] Fix up conditional inclusion of Spark2/3 modules

---
 pom.xml                                    | 111 ++++++++++++++++-----
 presto-spark-base/pom.xml                  |  36 +++++--
 presto-spark-classloader-interface/pom.xml |  45 +++++----
 3 files changed, 133 insertions(+), 59 deletions(-)

diff --git a/pom.xml b/pom.xml
index 15a6f0ec95a20..d5ac80313ec19 100644
--- a/pom.xml
+++ b/pom.xml
@@ -90,8 +90,7 @@
         6.0.0
         17.0.0
         3.5.4
-
-        2
+        2.0.2-6
 
         
+                                            org.codehaus.plexus:plexus-utils
+                                            com.google.guava:guava
+                                            com.fasterxml.jackson.core:jackson-annotations
+                                            com.fasterxml.jackson.core:jackson-core
+                                            com.fasterxml.jackson.core:jackson-databind
+                                        
+                                    
+                                
+                            
+                        
+
+                        
+                            org.basepom.maven
+                            duplicate-finder-maven-plugin
+                            
+                                true
+                                
+                                    com.github.benmanes.caffeine.*
+                                    
+                                    META-INF.versions.9.module-info
+                                    
+                                    META-INF.versions.11.module-info
+                                    
+                                    META-INF.versions.9.org.apache.lucene.*
+                                
+
+                            
+                        
+
+                    
+                
+            
+
+        
         
             spark3
 
@@ -3173,19 +3237,12 @@
             
 
             
-                3
+                3.4.1-1
             
 
-            
-                
-                    
-                        com.facebook.presto.spark
-                        spark-core
-                        3.4.1-1
-                        provided
-                    
-                
-            
+            
+                presto-spark-classloader-spark3
+            
 
             
                 
diff --git a/presto-spark-base/pom.xml b/presto-spark-base/pom.xml
index 37dc1d4fb541b..7767219a7ec65 100644
--- a/presto-spark-base/pom.xml
+++ b/presto-spark-base/pom.xml
@@ -16,7 +16,6 @@
         9.4.55.v20240627
         4.12.0
         3.9.1
-        2
     
 
     
@@ -55,12 +54,6 @@
             provided
         
 
-        
-            com.facebook.presto
-            presto-spark-classloader-spark${dep.pos.classloader.module-name.suffix}
-            provided
-        
-
         
             com.facebook.presto
             presto-client
@@ -538,6 +531,25 @@
                 
             
         
+        
+            spark2
+
+            
+                true
+                
+                    !spark-version
+                
+            
+
+            
+                
+                    com.facebook.presto
+                    presto-spark-classloader-spark2
+                    ${project.version}
+                
+            
+        
+
         
             spark3
 
@@ -548,11 +560,13 @@
                 
             
 
-            
-                3
-            
-
             
+                
+                    com.facebook.presto
+                    presto-spark-classloader-spark3
+                    ${project.version}
+                
+
                 
                     com.facebook.presto.spark
                     spark-core
diff --git a/presto-spark-classloader-interface/pom.xml b/presto-spark-classloader-interface/pom.xml
index 8f93be12818a4..a95e45827c5f1 100644
--- a/presto-spark-classloader-interface/pom.xml
+++ b/presto-spark-classloader-interface/pom.xml
@@ -13,7 +13,6 @@
     
         ${project.parent.basedir}
         true
-        2
     
 
     
@@ -23,11 +22,6 @@
             provided
         
 
-        
-            com.facebook.presto
-            presto-spark-classloader-spark${dep.pos.classloader.module-name.suffix}
-        
-
         
             com.google.guava
             guava
@@ -40,6 +34,25 @@
     
 
     
+        
+            spark2
+
+            
+                true
+                
+                    !spark-version
+                
+            
+
+            
+                
+                    com.facebook.presto
+                    presto-spark-classloader-spark2
+                    ${project.version}
+                
+            
+
+        
         
             spark3
 
@@ -50,22 +63,12 @@
                 
             
 
-            
-                3
-            
-
-            
-                
-                    
-                        com.facebook.presto.spark
-                        spark-core
-                        3.4.1-1
-                        compile
-                    
-                
-            
-
             
+                
+                    com.facebook.presto
+                    presto-spark-classloader-spark3
+                    ${project.version}
+                
                 
                     org.scala-lang
                     scala-library

From 8c7f2a8487683d73f64b5c765ff073e85dee8fbf Mon Sep 17 00:00:00 2001
From: Artem Selishchev 
Date: Thu, 4 Sep 2025 18:28:57 -0700
Subject: [PATCH 055/113] [presto] Move out M2Y from RegressionState for
 regr_slope and regr_intercept functions (#25475) (#25748)

Summary:

## Context
Currently we don't enforce intermediate/return type are the same in
Coordinator and Prestissimo Worker.
Velox creates vectors for intermediate/return results based on a plan
that comes from Coordinator. Then Prestissimo tries to use those vector
and not crash.

In practise we had a crash some time ago due to such a mismatch
(D74199165).
And I added validation to Velox to catch such kind of mismatches early:
https://github.com/facebookincubator/velox/pull/13322
But we wasn't able to enable it in prod, because the validation failed
for "regr_slope" and "regr_intercept" functions.

## What's changed?
In this diff I'm fixing "regr_slope" and "regr_intercept" intermediate
types. Basically in Java `AggregationState` for all these functions is
the same:
```
    AggregationFunction("regr_slope")
    AggregationFunction("regr_intercept")
    AggregationFunction("regr_sxy")
    AggregationFunction("regr_sxx")
    AggregationFunction("regr_syy")
    AggregationFunction("regr_r2")
    AggregationFunction("regr_count")
    AggregationFunction("regr_avgy")
    AggregationFunction("regr_avgx")
```
But in Prestissimo the state storage is more optimal:
```
    AggregationFunction("regr_slope")
    AggregationFunction("regr_intercept")
```
These 2 aggregation functions don't have M2Y field. And this is more
efficient, because we don't waste memory and CPU on the field, that
aren't needed.

So I moved M2Y to extended class, the same as it works in Velox:
https://github.com/facebookincubator/velox/blob/main/velox/functions/prestosql/aggregates/CovarianceAggregates.cpp?fbclid=IwY2xjawLRTetleHRuA2FlbQIxMQBicmlkETFiT0N3UFR0M2VKOHl6MHRhAR6KRQ1VUQdCkZXzwj14sMQrVZ-R9QBH1utuGJb5U_lyGzDwt8PwV317QRVNJg_aem_-ePxZ-fHO5MNgfUmayVJFA#L326-L337

No major changes, mostly just reorganized the code.

## Test plan
I tested `REGR_SLOPE`, `REGR_INTERCEPT` and `REGR_R2` functions since
they are heavily used in prod and cover both cases: with and without M2Y
field.

What my test looked like. For all 3 `REGR_*` functions I found some prod
queries, then:
1. Ran them on prev Java build
2. Ran them on new (with this PR) Java build
3. Ran them on prev Prestissimo build
4. Ran them on new (with this PR) Prestissimo build

And compared the output results. They all were identical.
With this manual test we covered `Coordinator -> Java Worker` and
`Coordinator -> Prestissimo Worker` integrations.

## Next steps
In this diff I'm trying to apply the same optimization to Java. With
this fix, the signatures will become the same in Java and Prestissimo
and we will be able to enable the validation

Differential Revision: D77625566

== NO RELEASE NOTES ==
---
 ...uiltInTypeAndFunctionNamespaceManager.java |   4 +
 .../aggregation/AggregationUtils.java         |  25 ++-
 .../DoubleRegressionAggregation.java          | 103 ------------
 .../DoubleRegressionExtendedAggregation.java  | 149 ++++++++++++++++++
 .../RealRegressionAggregation.java            | 103 ------------
 .../RealRegressionExtendedAggregation.java    | 149 ++++++++++++++++++
 .../state/ExtendedRegressionState.java        |  22 +++
 .../aggregation/state/RegressionState.java    |   4 -
 8 files changed, 345 insertions(+), 214 deletions(-)
 create mode 100644 presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/DoubleRegressionExtendedAggregation.java
 create mode 100644 presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/RealRegressionExtendedAggregation.java
 create mode 100644 presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/state/ExtendedRegressionState.java

diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInTypeAndFunctionNamespaceManager.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInTypeAndFunctionNamespaceManager.java
index 40c493fc887c6..d37d658e69204 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInTypeAndFunctionNamespaceManager.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInTypeAndFunctionNamespaceManager.java
@@ -67,6 +67,7 @@
 import com.facebook.presto.operator.aggregation.DoubleCovarianceAggregation;
 import com.facebook.presto.operator.aggregation.DoubleHistogramAggregation;
 import com.facebook.presto.operator.aggregation.DoubleRegressionAggregation;
+import com.facebook.presto.operator.aggregation.DoubleRegressionExtendedAggregation;
 import com.facebook.presto.operator.aggregation.DoubleSumAggregation;
 import com.facebook.presto.operator.aggregation.EntropyAggregation;
 import com.facebook.presto.operator.aggregation.GeometricMeanAggregations;
@@ -84,6 +85,7 @@
 import com.facebook.presto.operator.aggregation.RealGeometricMeanAggregations;
 import com.facebook.presto.operator.aggregation.RealHistogramAggregation;
 import com.facebook.presto.operator.aggregation.RealRegressionAggregation;
+import com.facebook.presto.operator.aggregation.RealRegressionExtendedAggregation;
 import com.facebook.presto.operator.aggregation.RealSumAggregation;
 import com.facebook.presto.operator.aggregation.ReduceAggregationFunction;
 import com.facebook.presto.operator.aggregation.SumDataSizeForStats;
@@ -744,7 +746,9 @@ private List getBuiltInFunctions(FunctionsConfig function
                 .aggregates(DoubleCovarianceAggregation.class)
                 .aggregates(RealCovarianceAggregation.class)
                 .aggregates(DoubleRegressionAggregation.class)
+                .aggregates(DoubleRegressionExtendedAggregation.class)
                 .aggregates(RealRegressionAggregation.class)
+                .aggregates(RealRegressionExtendedAggregation.class)
                 .aggregates(DoubleCorrelationAggregation.class)
                 .aggregates(RealCorrelationAggregation.class)
                 .aggregates(BitwiseOrAggregation.class)
diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/AggregationUtils.java b/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/AggregationUtils.java
index 578e782bc8d91..c78186ebb2417 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/AggregationUtils.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/AggregationUtils.java
@@ -22,6 +22,7 @@
 import com.facebook.presto.operator.aggregation.state.CentralMomentsState;
 import com.facebook.presto.operator.aggregation.state.CorrelationState;
 import com.facebook.presto.operator.aggregation.state.CovarianceState;
+import com.facebook.presto.operator.aggregation.state.ExtendedRegressionState;
 import com.facebook.presto.operator.aggregation.state.RegressionState;
 import com.facebook.presto.operator.aggregation.state.VarianceState;
 import com.facebook.presto.spi.function.AggregationFunctionImplementation;
@@ -145,9 +146,14 @@ public static double getCorrelation(CorrelationState state)
     public static void updateRegressionState(RegressionState state, double x, double y)
     {
         double oldMeanX = state.getMeanX();
-        double oldMeanY = state.getMeanY();
         updateCovarianceState(state, x, y);
         state.setM2X(state.getM2X() + (x - oldMeanX) * (x - state.getMeanX()));
+    }
+
+    public static void updateExtendedRegressionState(ExtendedRegressionState state, double x, double y)
+    {
+        double oldMeanY = state.getMeanY();
+        updateRegressionState(state, x, y);
         state.setM2Y(state.getM2Y() + (y - oldMeanY) * (y - state.getMeanY()));
     }
 
@@ -189,12 +195,12 @@ public static double getRegressionSxy(RegressionState state)
         return state.getC2();
     }
 
-    public static double getRegressionSyy(RegressionState state)
+    public static double getRegressionSyy(ExtendedRegressionState state)
     {
         return state.getM2Y();
     }
 
-    public static double getRegressionR2(RegressionState state)
+    public static double getRegressionR2(ExtendedRegressionState state)
     {
         if (state.getM2X() != 0 && state.getM2Y() == 0) {
             return 1.0;
@@ -311,10 +317,21 @@ public static void mergeRegressionState(RegressionState state, RegressionState o
         long na = state.getCount();
         long nb = otherState.getCount();
         state.setM2X(state.getM2X() + otherState.getM2X() + na * nb * Math.pow(state.getMeanX() - otherState.getMeanX(), 2) / (double) (na + nb));
-        state.setM2Y(state.getM2Y() + otherState.getM2Y() + na * nb * Math.pow(state.getMeanY() - otherState.getMeanY(), 2) / (double) (na + nb));
         updateCovarianceState(state, otherState);
     }
 
+    public static void mergeExtendedRegressionState(ExtendedRegressionState state, ExtendedRegressionState otherState)
+    {
+        if (otherState.getCount() == 0) {
+            return;
+        }
+
+        long na = state.getCount();
+        long nb = otherState.getCount();
+        state.setM2Y(state.getM2Y() + otherState.getM2Y() + na * nb * Math.pow(state.getMeanY() - otherState.getMeanY(), 2) / (double) (na + nb));
+        mergeRegressionState(state, otherState);
+    }
+
     public static String generateAggregationName(String baseName, TypeSignature outputType, List inputTypes)
     {
         StringBuilder sb = new StringBuilder();
diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/DoubleRegressionAggregation.java b/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/DoubleRegressionAggregation.java
index 24d1c6e61fcf5..db3ad26ec5d6d 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/DoubleRegressionAggregation.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/DoubleRegressionAggregation.java
@@ -24,15 +24,8 @@
 import com.facebook.presto.spi.function.SqlType;
 
 import static com.facebook.presto.common.type.DoubleType.DOUBLE;
-import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionAvgx;
-import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionAvgy;
-import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionCount;
 import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionIntercept;
-import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionR2;
 import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionSlope;
-import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionSxx;
-import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionSxy;
-import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionSyy;
 import static com.facebook.presto.operator.aggregation.AggregationUtils.mergeRegressionState;
 import static com.facebook.presto.operator.aggregation.AggregationUtils.updateRegressionState;
 
@@ -78,100 +71,4 @@ public static void regrIntercept(@AggregationState RegressionState state, BlockB
             out.appendNull();
         }
     }
-
-    @AggregationFunction("regr_sxy")
-    @OutputFunction(StandardTypes.DOUBLE)
-    public static void regrSxy(@AggregationState RegressionState state, BlockBuilder out)
-    {
-        double result = getRegressionSxy(state);
-        double count = getRegressionCount(state);
-        if (Double.isFinite(result) && Double.isFinite(count) && count > 0) {
-            DOUBLE.writeDouble(out, result);
-        }
-        else {
-            out.appendNull();
-        }
-    }
-
-    @AggregationFunction("regr_sxx")
-    @OutputFunction(StandardTypes.DOUBLE)
-    public static void regrSxx(@AggregationState RegressionState state, BlockBuilder out)
-    {
-        double result = getRegressionSxx(state);
-        double count = getRegressionCount(state);
-        if (Double.isFinite(result) && Double.isFinite(count) && count > 0) {
-            DOUBLE.writeDouble(out, result);
-        }
-        else {
-            out.appendNull();
-        }
-    }
-
-    @AggregationFunction("regr_syy")
-    @OutputFunction(StandardTypes.DOUBLE)
-    public static void regrSyy(@AggregationState RegressionState state, BlockBuilder out)
-    {
-        double result = getRegressionSyy(state);
-        double count = getRegressionCount(state);
-        if (Double.isFinite(result) && Double.isFinite(count) && count > 0) {
-            DOUBLE.writeDouble(out, result);
-        }
-        else {
-            out.appendNull();
-        }
-    }
-
-    @AggregationFunction("regr_r2")
-    @OutputFunction(StandardTypes.DOUBLE)
-    public static void regrR2(@AggregationState RegressionState state, BlockBuilder out)
-    {
-        double result = getRegressionR2(state);
-        if (Double.isFinite(result)) {
-            DOUBLE.writeDouble(out, result);
-        }
-        else {
-            out.appendNull();
-        }
-    }
-
-    @AggregationFunction("regr_count")
-    @OutputFunction(StandardTypes.DOUBLE)
-    public static void regrCount(@AggregationState RegressionState state, BlockBuilder out)
-    {
-        double result = getRegressionCount(state);
-        if (Double.isFinite(result) && result > 0) {
-            DOUBLE.writeDouble(out, result);
-        }
-        else {
-            out.appendNull();
-        }
-    }
-
-    @AggregationFunction("regr_avgy")
-    @OutputFunction(StandardTypes.DOUBLE)
-    public static void regrAvgy(@AggregationState RegressionState state, BlockBuilder out)
-    {
-        double result = getRegressionAvgy(state);
-        double count = getRegressionCount(state);
-        if (Double.isFinite(result) && Double.isFinite(count) && count > 0) {
-            DOUBLE.writeDouble(out, result);
-        }
-        else {
-            out.appendNull();
-        }
-    }
-
-    @AggregationFunction("regr_avgx")
-    @OutputFunction(StandardTypes.DOUBLE)
-    public static void regrAvgx(@AggregationState RegressionState state, BlockBuilder out)
-    {
-        double result = getRegressionAvgx(state);
-        double count = getRegressionCount(state);
-        if (Double.isFinite(result) && Double.isFinite(count) && count > 0) {
-            DOUBLE.writeDouble(out, result);
-        }
-        else {
-            out.appendNull();
-        }
-    }
 }
diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/DoubleRegressionExtendedAggregation.java b/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/DoubleRegressionExtendedAggregation.java
new file mode 100644
index 0000000000000..3550cd0936949
--- /dev/null
+++ b/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/DoubleRegressionExtendedAggregation.java
@@ -0,0 +1,149 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.operator.aggregation;
+
+import com.facebook.presto.common.block.BlockBuilder;
+import com.facebook.presto.common.type.StandardTypes;
+import com.facebook.presto.operator.aggregation.state.ExtendedRegressionState;
+import com.facebook.presto.spi.function.AggregationFunction;
+import com.facebook.presto.spi.function.AggregationState;
+import com.facebook.presto.spi.function.CombineFunction;
+import com.facebook.presto.spi.function.InputFunction;
+import com.facebook.presto.spi.function.OutputFunction;
+import com.facebook.presto.spi.function.SqlType;
+
+import static com.facebook.presto.common.type.DoubleType.DOUBLE;
+import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionAvgx;
+import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionAvgy;
+import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionCount;
+import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionR2;
+import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionSxx;
+import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionSxy;
+import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionSyy;
+import static com.facebook.presto.operator.aggregation.AggregationUtils.mergeExtendedRegressionState;
+import static com.facebook.presto.operator.aggregation.AggregationUtils.updateExtendedRegressionState;
+
+@AggregationFunction
+public class DoubleRegressionExtendedAggregation
+{
+    private DoubleRegressionExtendedAggregation() {}
+
+    @InputFunction
+    public static void input(@AggregationState ExtendedRegressionState state, @SqlType(StandardTypes.DOUBLE) double dependentValue, @SqlType(StandardTypes.DOUBLE) double independentValue)
+    {
+        updateExtendedRegressionState(state, independentValue, dependentValue);
+    }
+
+    @CombineFunction
+    public static void combine(@AggregationState ExtendedRegressionState state, @AggregationState ExtendedRegressionState otherState)
+    {
+        mergeExtendedRegressionState(state, otherState);
+    }
+
+    @AggregationFunction("regr_sxy")
+    @OutputFunction(StandardTypes.DOUBLE)
+    public static void regrSxy(@AggregationState ExtendedRegressionState state, BlockBuilder out)
+    {
+        double result = getRegressionSxy(state);
+        double count = getRegressionCount(state);
+        if (Double.isFinite(result) && Double.isFinite(count) && count > 0) {
+            DOUBLE.writeDouble(out, result);
+        }
+        else {
+            out.appendNull();
+        }
+    }
+
+    @AggregationFunction("regr_sxx")
+    @OutputFunction(StandardTypes.DOUBLE)
+    public static void regrSxx(@AggregationState ExtendedRegressionState state, BlockBuilder out)
+    {
+        double result = getRegressionSxx(state);
+        double count = getRegressionCount(state);
+        if (Double.isFinite(result) && Double.isFinite(count) && count > 0) {
+            DOUBLE.writeDouble(out, result);
+        }
+        else {
+            out.appendNull();
+        }
+    }
+
+    @AggregationFunction("regr_syy")
+    @OutputFunction(StandardTypes.DOUBLE)
+    public static void regrSyy(@AggregationState ExtendedRegressionState state, BlockBuilder out)
+    {
+        double result = getRegressionSyy(state);
+        double count = getRegressionCount(state);
+        if (Double.isFinite(result) && Double.isFinite(count) && count > 0) {
+            DOUBLE.writeDouble(out, result);
+        }
+        else {
+            out.appendNull();
+        }
+    }
+
+    @AggregationFunction("regr_r2")
+    @OutputFunction(StandardTypes.DOUBLE)
+    public static void regrR2(@AggregationState ExtendedRegressionState state, BlockBuilder out)
+    {
+        double result = getRegressionR2(state);
+        if (Double.isFinite(result)) {
+            DOUBLE.writeDouble(out, result);
+        }
+        else {
+            out.appendNull();
+        }
+    }
+
+    @AggregationFunction("regr_count")
+    @OutputFunction(StandardTypes.DOUBLE)
+    public static void regrCount(@AggregationState ExtendedRegressionState state, BlockBuilder out)
+    {
+        double result = getRegressionCount(state);
+        if (Double.isFinite(result) && result > 0) {
+            DOUBLE.writeDouble(out, result);
+        }
+        else {
+            out.appendNull();
+        }
+    }
+
+    @AggregationFunction("regr_avgy")
+    @OutputFunction(StandardTypes.DOUBLE)
+    public static void regrAvgy(@AggregationState ExtendedRegressionState state, BlockBuilder out)
+    {
+        double result = getRegressionAvgy(state);
+        double count = getRegressionCount(state);
+        if (Double.isFinite(result) && Double.isFinite(count) && count > 0) {
+            DOUBLE.writeDouble(out, result);
+        }
+        else {
+            out.appendNull();
+        }
+    }
+
+    @AggregationFunction("regr_avgx")
+    @OutputFunction(StandardTypes.DOUBLE)
+    public static void regrAvgx(@AggregationState ExtendedRegressionState state, BlockBuilder out)
+    {
+        double result = getRegressionAvgx(state);
+        double count = getRegressionCount(state);
+        if (Double.isFinite(result) && Double.isFinite(count) && count > 0) {
+            DOUBLE.writeDouble(out, result);
+        }
+        else {
+            out.appendNull();
+        }
+    }
+}
diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/RealRegressionAggregation.java b/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/RealRegressionAggregation.java
index 1fe5d006da1a9..a75222bfa93c4 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/RealRegressionAggregation.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/RealRegressionAggregation.java
@@ -24,15 +24,8 @@
 import com.facebook.presto.spi.function.SqlType;
 
 import static com.facebook.presto.common.type.RealType.REAL;
-import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionAvgx;
-import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionAvgy;
-import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionCount;
 import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionIntercept;
-import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionR2;
 import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionSlope;
-import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionSxx;
-import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionSxy;
-import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionSyy;
 import static java.lang.Float.floatToRawIntBits;
 import static java.lang.Float.intBitsToFloat;
 
@@ -78,100 +71,4 @@ public static void regrIntercept(@AggregationState RegressionState state, BlockB
             out.appendNull();
         }
     }
-
-    @AggregationFunction("regr_sxy")
-    @OutputFunction(StandardTypes.REAL)
-    public static void regrSxy(@AggregationState RegressionState state, BlockBuilder out)
-    {
-        double result = getRegressionSxy(state);
-        double count = getRegressionCount(state);
-        if (Double.isFinite(result) && Double.isFinite(count) && count > 0) {
-            REAL.writeLong(out, floatToRawIntBits((float) result));
-        }
-        else {
-            out.appendNull();
-        }
-    }
-
-    @AggregationFunction("regr_sxx")
-    @OutputFunction(StandardTypes.REAL)
-    public static void regrSxx(@AggregationState RegressionState state, BlockBuilder out)
-    {
-        double result = getRegressionSxx(state);
-        double count = getRegressionCount(state);
-        if (Double.isFinite(result) && Double.isFinite(count) && count > 0) {
-            REAL.writeLong(out, floatToRawIntBits((float) result));
-        }
-        else {
-            out.appendNull();
-        }
-    }
-
-    @AggregationFunction("regr_syy")
-    @OutputFunction(StandardTypes.REAL)
-    public static void regrSyy(@AggregationState RegressionState state, BlockBuilder out)
-    {
-        double result = getRegressionSyy(state);
-        double count = getRegressionCount(state);
-        if (Double.isFinite(result) && Double.isFinite(count) && count > 0) {
-            REAL.writeLong(out, floatToRawIntBits((float) result));
-        }
-        else {
-            out.appendNull();
-        }
-    }
-
-    @AggregationFunction("regr_r2")
-    @OutputFunction(StandardTypes.REAL)
-    public static void regrR2(@AggregationState RegressionState state, BlockBuilder out)
-    {
-        double result = getRegressionR2(state);
-        if (Double.isFinite(result)) {
-            REAL.writeLong(out, floatToRawIntBits((float) result));
-        }
-        else {
-            out.appendNull();
-        }
-    }
-
-    @AggregationFunction("regr_count")
-    @OutputFunction(StandardTypes.REAL)
-    public static void regrCount(@AggregationState RegressionState state, BlockBuilder out)
-    {
-        double result = getRegressionCount(state);
-        if (Double.isFinite(result) && result > 0) {
-            REAL.writeLong(out, floatToRawIntBits((float) result));
-        }
-        else {
-            out.appendNull();
-        }
-    }
-
-    @AggregationFunction("regr_avgy")
-    @OutputFunction(StandardTypes.REAL)
-    public static void regrAvgy(@AggregationState RegressionState state, BlockBuilder out)
-    {
-        double result = getRegressionAvgy(state);
-        double count = getRegressionCount(state);
-        if (Double.isFinite(result) && Double.isFinite(count) && count > 0) {
-            REAL.writeLong(out, floatToRawIntBits((float) result));
-        }
-        else {
-            out.appendNull();
-        }
-    }
-
-    @AggregationFunction("regr_avgx")
-    @OutputFunction(StandardTypes.REAL)
-    public static void regrAvgx(@AggregationState RegressionState state, BlockBuilder out)
-    {
-        double result = getRegressionAvgx(state);
-        double count = getRegressionCount(state);
-        if (Double.isFinite(result) && Double.isFinite(count) && count > 0) {
-            REAL.writeLong(out, floatToRawIntBits((float) result));
-        }
-        else {
-            out.appendNull();
-        }
-    }
 }
diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/RealRegressionExtendedAggregation.java b/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/RealRegressionExtendedAggregation.java
new file mode 100644
index 0000000000000..2d0335ae9aca6
--- /dev/null
+++ b/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/RealRegressionExtendedAggregation.java
@@ -0,0 +1,149 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.operator.aggregation;
+
+import com.facebook.presto.common.block.BlockBuilder;
+import com.facebook.presto.common.type.StandardTypes;
+import com.facebook.presto.operator.aggregation.state.ExtendedRegressionState;
+import com.facebook.presto.spi.function.AggregationFunction;
+import com.facebook.presto.spi.function.AggregationState;
+import com.facebook.presto.spi.function.CombineFunction;
+import com.facebook.presto.spi.function.InputFunction;
+import com.facebook.presto.spi.function.OutputFunction;
+import com.facebook.presto.spi.function.SqlType;
+
+import static com.facebook.presto.common.type.RealType.REAL;
+import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionAvgx;
+import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionAvgy;
+import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionCount;
+import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionR2;
+import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionSxx;
+import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionSxy;
+import static com.facebook.presto.operator.aggregation.AggregationUtils.getRegressionSyy;
+import static java.lang.Float.floatToRawIntBits;
+import static java.lang.Float.intBitsToFloat;
+
+@AggregationFunction
+public class RealRegressionExtendedAggregation
+{
+    private RealRegressionExtendedAggregation() {}
+
+    @InputFunction
+    public static void input(@AggregationState ExtendedRegressionState state, @SqlType(StandardTypes.REAL) long dependentValue, @SqlType(StandardTypes.REAL) long independentValue)
+    {
+        DoubleRegressionExtendedAggregation.input(state, intBitsToFloat((int) dependentValue), intBitsToFloat((int) independentValue));
+    }
+
+    @CombineFunction
+    public static void combine(@AggregationState ExtendedRegressionState state, @AggregationState ExtendedRegressionState otherState)
+    {
+        DoubleRegressionExtendedAggregation.combine(state, otherState);
+    }
+
+    @AggregationFunction("regr_sxy")
+    @OutputFunction(StandardTypes.REAL)
+    public static void regrSxy(@AggregationState ExtendedRegressionState state, BlockBuilder out)
+    {
+        double result = getRegressionSxy(state);
+        double count = getRegressionCount(state);
+        if (Double.isFinite(result) && Double.isFinite(count) && count > 0) {
+            REAL.writeLong(out, floatToRawIntBits((float) result));
+        }
+        else {
+            out.appendNull();
+        }
+    }
+
+    @AggregationFunction("regr_sxx")
+    @OutputFunction(StandardTypes.REAL)
+    public static void regrSxx(@AggregationState ExtendedRegressionState state, BlockBuilder out)
+    {
+        double result = getRegressionSxx(state);
+        double count = getRegressionCount(state);
+        if (Double.isFinite(result) && Double.isFinite(count) && count > 0) {
+            REAL.writeLong(out, floatToRawIntBits((float) result));
+        }
+        else {
+            out.appendNull();
+        }
+    }
+
+    @AggregationFunction("regr_syy")
+    @OutputFunction(StandardTypes.REAL)
+    public static void regrSyy(@AggregationState ExtendedRegressionState state, BlockBuilder out)
+    {
+        double result = getRegressionSyy(state);
+        double count = getRegressionCount(state);
+        if (Double.isFinite(result) && Double.isFinite(count) && count > 0) {
+            REAL.writeLong(out, floatToRawIntBits((float) result));
+        }
+        else {
+            out.appendNull();
+        }
+    }
+
+    @AggregationFunction("regr_r2")
+    @OutputFunction(StandardTypes.REAL)
+    public static void regrR2(@AggregationState ExtendedRegressionState state, BlockBuilder out)
+    {
+        double result = getRegressionR2(state);
+        if (Double.isFinite(result)) {
+            REAL.writeLong(out, floatToRawIntBits((float) result));
+        }
+        else {
+            out.appendNull();
+        }
+    }
+
+    @AggregationFunction("regr_count")
+    @OutputFunction(StandardTypes.REAL)
+    public static void regrCount(@AggregationState ExtendedRegressionState state, BlockBuilder out)
+    {
+        double result = getRegressionCount(state);
+        if (Double.isFinite(result) && result > 0) {
+            REAL.writeLong(out, floatToRawIntBits((float) result));
+        }
+        else {
+            out.appendNull();
+        }
+    }
+
+    @AggregationFunction("regr_avgy")
+    @OutputFunction(StandardTypes.REAL)
+    public static void regrAvgy(@AggregationState ExtendedRegressionState state, BlockBuilder out)
+    {
+        double result = getRegressionAvgy(state);
+        double count = getRegressionCount(state);
+        if (Double.isFinite(result) && Double.isFinite(count) && count > 0) {
+            REAL.writeLong(out, floatToRawIntBits((float) result));
+        }
+        else {
+            out.appendNull();
+        }
+    }
+
+    @AggregationFunction("regr_avgx")
+    @OutputFunction(StandardTypes.REAL)
+    public static void regrAvgx(@AggregationState ExtendedRegressionState state, BlockBuilder out)
+    {
+        double result = getRegressionAvgx(state);
+        double count = getRegressionCount(state);
+        if (Double.isFinite(result) && Double.isFinite(count) && count > 0) {
+            REAL.writeLong(out, floatToRawIntBits((float) result));
+        }
+        else {
+            out.appendNull();
+        }
+    }
+}
diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/state/ExtendedRegressionState.java b/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/state/ExtendedRegressionState.java
new file mode 100644
index 0000000000000..64a9883174158
--- /dev/null
+++ b/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/state/ExtendedRegressionState.java
@@ -0,0 +1,22 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.operator.aggregation.state;
+
+public interface ExtendedRegressionState
+        extends RegressionState
+{
+    double getM2Y();
+
+    void setM2Y(double value);
+}
diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/state/RegressionState.java b/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/state/RegressionState.java
index 79837f90c0c11..ae3af6f46dc43 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/state/RegressionState.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/operator/aggregation/state/RegressionState.java
@@ -19,8 +19,4 @@ public interface RegressionState
     double getM2X();
 
     void setM2X(double value);
-
-    double getM2Y();
-
-    void setM2Y(double value);
 }

From 0d75bc32e4d8943c626502db3d7195bf67124fe4 Mon Sep 17 00:00:00 2001
From: Xin Zhang 
Date: Thu, 4 Sep 2025 19:03:24 +0100
Subject: [PATCH 056/113] Remove register-test-functions from
 PrestoSparkNativeQueryRunnerUtils

---
 .../facebook/presto/spark/PrestoSparkNativeQueryRunnerUtils.java | 1 -
 1 file changed, 1 deletion(-)

diff --git a/presto-native-execution/src/test/java/com/facebook/presto/spark/PrestoSparkNativeQueryRunnerUtils.java b/presto-native-execution/src/test/java/com/facebook/presto/spark/PrestoSparkNativeQueryRunnerUtils.java
index 14113b22d4fb3..21ef570381795 100644
--- a/presto-native-execution/src/test/java/com/facebook/presto/spark/PrestoSparkNativeQueryRunnerUtils.java
+++ b/presto-native-execution/src/test/java/com/facebook/presto/spark/PrestoSparkNativeQueryRunnerUtils.java
@@ -79,7 +79,6 @@ public static Map getNativeExecutionSparkConfigs()
                 .put("catalog.config-dir", "/")
                 .put("task.info-update-interval", "100ms")
                 .put("spark.initial-partition-count", "1")
-                .put("register-test-functions", "true")
                 .put("native-execution-program-arguments", "--logtostderr=1 --minloglevel=3")
                 .put("spark.partition-count-auto-tune-enabled", "false");
 

From 3257215cfe2ff4ab22329c7e2827a10a8f11b4e8 Mon Sep 17 00:00:00 2001
From: Pramod Satya 
Date: Thu, 4 Sep 2025 20:13:43 -0700
Subject: [PATCH 057/113] [native] Use subscript operator to retrieve function
 handle

---
 .../sql/relational/FunctionResolution.java     |  2 +-
 .../sidecar/TestNativeSidecarPlugin.java       | 18 ++++++++++++++++++
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/relational/FunctionResolution.java b/presto-main-base/src/main/java/com/facebook/presto/sql/relational/FunctionResolution.java
index a3fe25d7b65fa..337582257bad4 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/sql/relational/FunctionResolution.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/sql/relational/FunctionResolution.java
@@ -281,7 +281,7 @@ public boolean isEqualsFunction(FunctionHandle functionHandle)
     @Override
     public FunctionHandle subscriptFunction(Type baseType, Type indexType)
     {
-        return functionAndTypeResolver.lookupFunction(SUBSCRIPT.getFunctionName().getObjectName(), fromTypes(baseType, indexType));
+        return functionAndTypeResolver.resolveOperator(SUBSCRIPT, fromTypes(baseType, indexType));
     }
 
     @Override
diff --git a/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java b/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java
index 1c5d05d55b274..b4716f98362cf 100644
--- a/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java
+++ b/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java
@@ -14,6 +14,7 @@
 package com.facebook.presto.sidecar;
 
 import com.facebook.airlift.units.DataSize;
+import com.facebook.presto.Session;
 import com.facebook.presto.nativeworker.PrestoNativeQueryRunnerUtils;
 import com.facebook.presto.sidecar.functionNamespace.FunctionDefinitionProvider;
 import com.facebook.presto.sidecar.functionNamespace.NativeFunctionDefinitionProvider;
@@ -44,6 +45,7 @@
 import java.util.stream.Collectors;
 
 import static com.facebook.airlift.units.DataSize.Unit.MEGABYTE;
+import static com.facebook.presto.SystemSessionProperties.REMOVE_MAP_CAST;
 import static com.facebook.presto.common.Utils.checkArgument;
 import static com.facebook.presto.common.type.BigintType.BIGINT;
 import static com.facebook.presto.nativeworker.NativeQueryRunnerUtils.createLineitem;
@@ -405,6 +407,22 @@ public void testGeometryQueries()
                 "Error from native plan checker: .SpatialJoinNode no abstract type PlanNode ");
     }
 
+    @Test
+    public void testRemoveMapCast()
+    {
+        Session enableOptimization = Session.builder(getSession())
+                .setSystemProperty(REMOVE_MAP_CAST, "true")
+                .build();
+        assertQuery(enableOptimization, "select feature[key] from (values (map(array[cast(1 as integer), 2, 3, 4], array[0.3, 0.5, 0.9, 0.1]), cast(2 as bigint)), (map(array[cast(1 as integer), 2, 3, 4], array[0.3, 0.5, 0.9, 0.1]), 4)) t(feature,  key)",
+                "values 0.5, 0.1");
+        assertQuery(enableOptimization, "select element_at(feature, key) from (values (map(array[cast(1 as integer), 2, 3, 4], array[0.3, 0.5, 0.9, 0.1]), cast(2 as bigint)), (map(array[cast(1 as integer), 2, 3, 4], array[0.3, 0.5, 0.9, 0.1]), 4)) t(feature,  key)",
+                "values 0.5, 0.1");
+        assertQuery(enableOptimization, "select element_at(feature, key) from (values (map(array[cast(1 as integer), 2, 3, 4], array[0.3, 0.5, 0.9, 0.1]), cast(2 as bigint)), (map(array[cast(1 as integer), 2, 3, 4], array[0.3, 0.5, 0.9, 0.1]), 400000000000)) t(feature, key)",
+                "values 0.5, null");
+        assertQuery(enableOptimization, "select feature[key] from (values (map(array[cast(1 as varchar), '2', '3', '4'], array[0.3, 0.5, 0.9, 0.1]), cast('2' as varchar)), (map(array[cast(1 as varchar), '2', '3', '4'], array[0.3, 0.5, 0.9, 0.1]), '4')) t(feature,  key)",
+                "values 0.5, 0.1");
+    }
+
     private String generateRandomTableName()
     {
         String tableName = "tmp_presto_" + UUID.randomUUID().toString().replace("-", "");

From 551500e4f6c2f030c0be0c437e6a8b11d344a994 Mon Sep 17 00:00:00 2001
From: Mahadevuni Naveen Kumar 
Date: Fri, 5 Sep 2025 15:11:31 +0530
Subject: [PATCH 058/113] Implement Iceberg system.bucket scalar function
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-Authored-By: Yuya Ebihara 
Co-Authored-By: Pádraig O'Sullivan 
---
 .../presto/iceberg/IcebergConnector.java      |   3 +
 .../function/IcebergBucketFunction.java       | 111 ++++++++++++++++++
 .../iceberg/TestIcebergScalarFunctions.java   |  80 +++++++++++++
 .../operator/scalar/FunctionAssertions.java   |   6 +
 4 files changed, 200 insertions(+)
 create mode 100644 presto-iceberg/src/main/java/com/facebook/presto/iceberg/function/IcebergBucketFunction.java
 create mode 100644 presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestIcebergScalarFunctions.java

diff --git a/presto-iceberg/src/main/java/com/facebook/presto/iceberg/IcebergConnector.java b/presto-iceberg/src/main/java/com/facebook/presto/iceberg/IcebergConnector.java
index 6936b7fee96be..b48f916db6e6b 100644
--- a/presto-iceberg/src/main/java/com/facebook/presto/iceberg/IcebergConnector.java
+++ b/presto-iceberg/src/main/java/com/facebook/presto/iceberg/IcebergConnector.java
@@ -15,6 +15,7 @@
 
 import com.facebook.airlift.bootstrap.LifeCycleManager;
 import com.facebook.presto.hive.HiveTransactionHandle;
+import com.facebook.presto.iceberg.function.IcebergBucketFunction;
 import com.facebook.presto.iceberg.function.changelog.ApplyChangelogFunction;
 import com.facebook.presto.spi.SystemTable;
 import com.facebook.presto.spi.classloader.ThreadContextClassLoader;
@@ -225,6 +226,8 @@ public Set> getSystemFunctions()
     {
         return ImmutableSet.>builder()
                 .add(ApplyChangelogFunction.class)
+                .add(IcebergBucketFunction.class)
+                .add(IcebergBucketFunction.Bucket.class)
                 .build();
     }
 }
diff --git a/presto-iceberg/src/main/java/com/facebook/presto/iceberg/function/IcebergBucketFunction.java b/presto-iceberg/src/main/java/com/facebook/presto/iceberg/function/IcebergBucketFunction.java
new file mode 100644
index 0000000000000..429b296f79d8e
--- /dev/null
+++ b/presto-iceberg/src/main/java/com/facebook/presto/iceberg/function/IcebergBucketFunction.java
@@ -0,0 +1,111 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.iceberg.function;
+
+import com.facebook.presto.common.type.StandardTypes;
+import com.facebook.presto.spi.function.LiteralParameter;
+import com.facebook.presto.spi.function.LiteralParameters;
+import com.facebook.presto.spi.function.ScalarFunction;
+import com.facebook.presto.spi.function.SqlType;
+import io.airlift.slice.Slice;
+import org.apache.iceberg.transforms.Transforms;
+import org.apache.iceberg.types.Types;
+
+import java.math.BigDecimal;
+import java.math.MathContext;
+
+import static com.facebook.presto.common.type.DateTimeEncoding.unpackMillisUtc;
+import static com.facebook.presto.common.type.Decimals.decodeUnscaledValue;
+import static com.facebook.presto.common.type.SqlTimestamp.MICROSECONDS_PER_MILLISECOND;
+
+public final class IcebergBucketFunction
+{
+    private IcebergBucketFunction() {}
+
+    @ScalarFunction("bucket")
+    @SqlType(StandardTypes.BIGINT)
+    public static long bucketInteger(@SqlType(StandardTypes.BIGINT) long value, @SqlType(StandardTypes.INTEGER) long numberOfBuckets)
+    {
+        return Transforms.bucket((int) numberOfBuckets)
+                .bind(Types.LongType.get())
+                .apply(value);
+    }
+
+    @ScalarFunction("bucket")
+    @SqlType(StandardTypes.BIGINT)
+    public static long bucketVarchar(@SqlType(StandardTypes.VARCHAR) Slice value, @SqlType(StandardTypes.INTEGER) long numberOfBuckets)
+    {
+        return (long) Transforms.bucket((int) numberOfBuckets)
+                .bind(Types.StringType.get())
+                .apply(value.toStringUtf8());
+    }
+
+    @ScalarFunction("bucket")
+    @SqlType(StandardTypes.BIGINT)
+    public static long bucketVarbinary(@SqlType(StandardTypes.VARBINARY) Slice value, @SqlType(StandardTypes.INTEGER) long numberOfBuckets)
+    {
+        return (long) Transforms.bucket((int) numberOfBuckets)
+                .bind(Types.BinaryType.get())
+                .apply(value.toByteBuffer());
+    }
+
+    @ScalarFunction("bucket")
+    @SqlType(StandardTypes.BIGINT)
+    public static long bucketDate(@SqlType(StandardTypes.DATE) long value, @SqlType(StandardTypes.INTEGER) long numberOfBuckets)
+    {
+        return Transforms.bucket((int) numberOfBuckets)
+                .bind(Types.DateType.get())
+                .apply((int) value);
+    }
+
+    @ScalarFunction("bucket")
+    @SqlType(StandardTypes.BIGINT)
+    public static long bucketTimestamp(@SqlType(StandardTypes.TIMESTAMP) long value, @SqlType(StandardTypes.INTEGER) long numberOfBuckets)
+    {
+        return Transforms.bucket((int) numberOfBuckets)
+                .bind(Types.TimestampType.withoutZone())
+                .apply(value);
+    }
+
+    @ScalarFunction("bucket")
+    @SqlType(StandardTypes.BIGINT)
+    public static long bucketTimestampWithTimeZone(@SqlType(StandardTypes.TIMESTAMP_WITH_TIME_ZONE) long value, @SqlType(StandardTypes.INTEGER) long numberOfBuckets)
+    {
+        return Transforms.bucket((int) numberOfBuckets)
+                .bind(Types.TimestampType.withZone())
+                .apply(unpackMillisUtc(value) * MICROSECONDS_PER_MILLISECOND);
+    }
+
+    @ScalarFunction("bucket")
+    public static final class Bucket
+    {
+        @LiteralParameters({"p", "s"})
+        @SqlType(StandardTypes.BIGINT)
+        public static long bucketShortDecimal(@LiteralParameter("p") long numPrecision, @LiteralParameter("s") long numScale, @SqlType("decimal(p, s)") long value, @SqlType(StandardTypes.INTEGER) long numberOfBuckets)
+        {
+            return Transforms.bucket((int) numberOfBuckets)
+                    .bind(Types.DecimalType.of((int) numPrecision, (int) numScale))
+                    .apply(BigDecimal.valueOf(value));
+        }
+
+        @LiteralParameters({"p", "s"})
+        @SqlType(StandardTypes.BIGINT)
+        public static long bucketLongDecimal(@LiteralParameter("p") long numPrecision, @LiteralParameter("s") long numScale, @SqlType("decimal(p, s)") Slice value, @SqlType(StandardTypes.INTEGER) long numberOfBuckets)
+        {
+            return Transforms.bucket((int) numberOfBuckets)
+                    .bind(Types.DecimalType.of((int) numPrecision, (int) numScale))
+                    .apply(new BigDecimal(decodeUnscaledValue(value), (int) numScale, new MathContext((int) numPrecision)));
+        }
+    }
+}
diff --git a/presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestIcebergScalarFunctions.java b/presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestIcebergScalarFunctions.java
new file mode 100644
index 0000000000000..c6253afcb9638
--- /dev/null
+++ b/presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestIcebergScalarFunctions.java
@@ -0,0 +1,80 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.iceberg;
+
+import com.facebook.presto.common.CatalogSchemaName;
+import com.facebook.presto.iceberg.function.IcebergBucketFunction;
+import com.facebook.presto.metadata.FunctionExtractor;
+import com.facebook.presto.operator.scalar.AbstractTestFunctions;
+import com.facebook.presto.sql.analyzer.FeaturesConfig;
+import com.facebook.presto.sql.analyzer.FunctionsConfig;
+import com.facebook.presto.type.DateOperators;
+import com.facebook.presto.type.TimestampOperators;
+import com.facebook.presto.type.TimestampWithTimeZoneOperators;
+import org.testcontainers.shaded.com.google.common.collect.ImmutableList;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import java.math.BigDecimal;
+
+import static com.facebook.presto.SessionTestUtils.TEST_SESSION;
+import static com.facebook.presto.common.type.BigintType.BIGINT;
+import static com.facebook.presto.common.type.Decimals.encodeScaledValue;
+import static com.facebook.presto.iceberg.function.IcebergBucketFunction.Bucket.bucketLongDecimal;
+import static com.facebook.presto.iceberg.function.IcebergBucketFunction.Bucket.bucketShortDecimal;
+import static com.facebook.presto.iceberg.function.IcebergBucketFunction.bucketDate;
+import static com.facebook.presto.iceberg.function.IcebergBucketFunction.bucketInteger;
+import static com.facebook.presto.iceberg.function.IcebergBucketFunction.bucketTimestamp;
+import static com.facebook.presto.iceberg.function.IcebergBucketFunction.bucketTimestampWithTimeZone;
+import static com.facebook.presto.iceberg.function.IcebergBucketFunction.bucketVarbinary;
+import static com.facebook.presto.iceberg.function.IcebergBucketFunction.bucketVarchar;
+import static io.airlift.slice.Slices.utf8Slice;
+
+public class TestIcebergScalarFunctions
+        extends AbstractTestFunctions
+{
+    public TestIcebergScalarFunctions()
+    {
+        super(TEST_SESSION, new FeaturesConfig(), new FunctionsConfig(), false);
+    }
+
+    @BeforeClass
+    public void registerFunction()
+    {
+        ImmutableList.Builder> functions = ImmutableList.builder();
+        functions.add(IcebergBucketFunction.class)
+                .add(IcebergBucketFunction.Bucket.class);
+        functionAssertions.addConnectorFunctions(FunctionExtractor.extractFunctions(functions.build(),
+                new CatalogSchemaName("iceberg", "system")), "iceberg");
+    }
+
+    @Test
+    public void testBucketFunction()
+    {
+        String catalogSchema = "iceberg.system";
+        functionAssertions.assertFunction(catalogSchema + ".bucket(cast(10 as tinyint), 3)", BIGINT, bucketInteger(10, 3));
+        functionAssertions.assertFunction(catalogSchema + ".bucket(cast(1950 as smallint), 4)", BIGINT, bucketInteger(1950, 4));
+        functionAssertions.assertFunction(catalogSchema + ".bucket(cast(2375645 as int), 5)", BIGINT, bucketInteger(2375645, 5));
+        functionAssertions.assertFunction(catalogSchema + ".bucket(cast(2779099983928392323 as bigint), 6)", BIGINT, bucketInteger(2779099983928392323L, 6));
+        functionAssertions.assertFunction(catalogSchema + ".bucket(cast(456.43 as DECIMAL(5,2)), 12)", BIGINT, bucketShortDecimal(5, 2, 45643, 12));
+        functionAssertions.assertFunction(catalogSchema + ".bucket(cast('12345678901234567890.1234567890' as DECIMAL(30,10)), 12)", BIGINT, bucketLongDecimal(30, 10, encodeScaledValue(new BigDecimal("12345678901234567890.1234567890")), 12));
+
+        functionAssertions.assertFunction(catalogSchema + ".bucket(cast('nasdbsdnsdms' as varchar), 7)", BIGINT, bucketVarchar(utf8Slice("nasdbsdnsdms"), 7));
+        functionAssertions.assertFunction(catalogSchema + ".bucket(cast('nasdbsdnsdms' as varbinary), 8)", BIGINT, bucketVarbinary(utf8Slice("nasdbsdnsdms"), 8));
+
+        functionAssertions.assertFunction(catalogSchema + ".bucket(cast('2018-04-06' as date), 9)", BIGINT, bucketDate(DateOperators.castFromSlice(utf8Slice("2018-04-06")), 9));
+        functionAssertions.assertFunction(catalogSchema + ".bucket(CAST('2018-04-06 04:35:00.000' AS TIMESTAMP),10)", BIGINT, bucketTimestamp(TimestampOperators.castFromSlice(TEST_SESSION.getSqlFunctionProperties(), utf8Slice("2018-04-06 04:35:00.000")), 10));
+        functionAssertions.assertFunction(catalogSchema + ".bucket(CAST('2018-04-06 04:35:00.000 GMT' AS TIMESTAMP WITH TIME ZONE), 11)", BIGINT, bucketTimestampWithTimeZone(TimestampWithTimeZoneOperators.castFromSlice(TEST_SESSION.getSqlFunctionProperties(), utf8Slice("2018-04-06 04:35:00.000 GMT")), 11));
+    }
+}
diff --git a/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/FunctionAssertions.java b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/FunctionAssertions.java
index b2965a54329c5..d46bb9a0d0bd2 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/FunctionAssertions.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/FunctionAssertions.java
@@ -278,6 +278,12 @@ public FunctionAssertions addScalarFunctions(Class clazz)
         return this;
     }
 
+    public FunctionAssertions addConnectorFunctions(List functionInfos, String namespace)
+    {
+        metadata.registerConnectorFunctions(namespace, functionInfos);
+        return this;
+    }
+
     public void assertFunction(String projection, Type expectedType, Object expected)
     {
         if (expected instanceof Slice) {

From b6c6cbf2d591b96aff67ab5f54c80712b522c516 Mon Sep 17 00:00:00 2001
From: Pratik Joseph Dabre 
Date: Fri, 5 Sep 2025 11:27:30 -0700
Subject: [PATCH 059/113] [native] Introduce
 presto-native-sql-invoked-functions-plugin for sidecar enabled clusters

Adds a new plugin : presto-native-sql-invoked-functions-plugin that contains all inlined SQL functions except those with overridden native
implementations. This plugin is intended to be loaded only in sidecar enabled clusters.
---
 .../prestocpp-linux-build-and-unit-test.yml   |   2 +-
 pom.xml                                       |   1 +
 .../optimizations/KeyBasedSampler.java        |   4 +-
 .../presto/sql/relational/Expressions.java    |   7 -
 .../nativeworker/NativeQueryRunnerUtils.java  |   4 +-
 presto-native-sidecar-plugin/pom.xml          |  16 +++
 .../NativeSidecarPluginQueryRunnerUtils.java  |   2 +
 .../sidecar/TestNativeSidecarPlugin.java      | 126 +++++++++++++++++-
 .../pom.xml                                   |  29 ++++
 .../scalar/sql/NativeArraySqlFunctions.java   |  74 ++++++++++
 .../scalar/sql/NativeMapSqlFunctions.java     |  48 +++++++
 .../sql/NativeSimpleSamplingPercent.java      |  33 +++++
 .../sql/NativeSqlInvokedFunctionsPlugin.java  |  33 +++++
 presto-native-tests/pom.xml                   |   7 +
 presto-plan-checker-router-plugin/pom.xml     |   7 +
 .../conf/docker/common/compose-commons.sh     |  10 ++
 presto-server/src/main/provisio/presto.xml    |   6 +
 17 files changed, 392 insertions(+), 17 deletions(-)
 create mode 100644 presto-native-sql-invoked-functions-plugin/pom.xml
 create mode 100644 presto-native-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/NativeArraySqlFunctions.java
 create mode 100644 presto-native-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/NativeMapSqlFunctions.java
 create mode 100644 presto-native-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/NativeSimpleSamplingPercent.java
 create mode 100644 presto-native-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/NativeSqlInvokedFunctionsPlugin.java

diff --git a/.github/workflows/prestocpp-linux-build-and-unit-test.yml b/.github/workflows/prestocpp-linux-build-and-unit-test.yml
index b28f2292252ea..f3fc4843c3321 100644
--- a/.github/workflows/prestocpp-linux-build-and-unit-test.yml
+++ b/.github/workflows/prestocpp-linux-build-and-unit-test.yml
@@ -370,7 +370,7 @@ jobs:
           # Use different Maven options to install.
           MAVEN_OPTS: "-Xmx2G -XX:+ExitOnOutOfMemoryError"
         run: |
-          for i in $(seq 1 3); do ./mvnw clean install $MAVEN_FAST_INSTALL -pl 'presto-native-execution' -am && s=0 && break || s=$? && sleep 10; done; (exit $s)
+          for i in $(seq 1 3); do ./mvnw clean install $MAVEN_FAST_INSTALL -pl 'presto-native-sidecar-plugin' -am && s=0 && break || s=$? && sleep 10; done; (exit $s)
 
       - name: Run presto-native sidecar tests
         if: |
diff --git a/pom.xml b/pom.xml
index d5ac80313ec19..f73719dacd94a 100644
--- a/pom.xml
+++ b/pom.xml
@@ -224,6 +224,7 @@
         presto-router-example-plugin-scheduler
         presto-plan-checker-router-plugin
         presto-sql-invoked-functions-plugin
+        presto-native-sql-invoked-functions-plugin
     
 
     
diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/KeyBasedSampler.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/KeyBasedSampler.java
index f34359aeba6af..1c7856a62caea 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/KeyBasedSampler.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/KeyBasedSampler.java
@@ -14,7 +14,6 @@
 package com.facebook.presto.sql.planner.optimizations;
 
 import com.facebook.presto.Session;
-import com.facebook.presto.common.QualifiedObjectName;
 import com.facebook.presto.common.function.OperatorType;
 import com.facebook.presto.common.type.Type;
 import com.facebook.presto.common.type.Varchars;
@@ -55,7 +54,6 @@
 import static com.facebook.presto.common.type.BooleanType.BOOLEAN;
 import static com.facebook.presto.common.type.DoubleType.DOUBLE;
 import static com.facebook.presto.common.type.VarcharType.VARCHAR;
-import static com.facebook.presto.metadata.BuiltInTypeAndFunctionNamespaceManager.JAVA_BUILTIN_NAMESPACE;
 import static com.facebook.presto.metadata.CastType.CAST;
 import static com.facebook.presto.spi.StandardErrorCode.FUNCTION_NOT_FOUND;
 import static com.facebook.presto.spi.StandardWarningCode.SAMPLED_FIELDS;
@@ -150,7 +148,7 @@ private PlanNode addSamplingFilter(PlanNode tableScanNode, Optional arguments)
-    {
-        FunctionHandle functionHandle = functionAndTypeManager.lookupFunction(qualifiedObjectName, fromTypes(arguments.stream().map(RowExpression::getType).collect(toImmutableList())));
-        return call(String.valueOf(qualifiedObjectName), functionHandle, returnType, arguments);
-    }
-
     public static CallExpression call(FunctionAndTypeResolver functionAndTypeResolver, String name, Type returnType, RowExpression... arguments)
     {
         FunctionHandle functionHandle = functionAndTypeResolver.lookupFunction(name, fromTypes(Arrays.stream(arguments).map(RowExpression::getType).collect(toImmutableList())));
diff --git a/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/NativeQueryRunnerUtils.java b/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/NativeQueryRunnerUtils.java
index 3650f8cb36531..59a60b0024842 100644
--- a/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/NativeQueryRunnerUtils.java
+++ b/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/NativeQueryRunnerUtils.java
@@ -29,7 +29,7 @@ private NativeQueryRunnerUtils() {}
     public static Map getNativeWorkerHiveProperties()
     {
         return ImmutableMap.of("hive.parquet.pushdown-filter-enabled", "true",
-            "hive.orc-compression-codec", "ZSTD", "hive.storage-format", "DWRF");
+                "hive.orc-compression-codec", "ZSTD", "hive.storage-format", "DWRF");
     }
 
     public static Map getNativeWorkerIcebergProperties()
@@ -59,6 +59,8 @@ public static Map getNativeSidecarProperties()
                 .put("coordinator-sidecar-enabled", "true")
                 .put("exclude-invalid-worker-session-properties", "true")
                 .put("presto.default-namespace", "native.default")
+                // inline-sql-functions is overridden to be true in sidecar enabled native clusters.
+                .put("inline-sql-functions", "true")
                 .build();
     }
 
diff --git a/presto-native-sidecar-plugin/pom.xml b/presto-native-sidecar-plugin/pom.xml
index b2bc40e8f2bd4..d04424440469a 100644
--- a/presto-native-sidecar-plugin/pom.xml
+++ b/presto-native-sidecar-plugin/pom.xml
@@ -260,9 +260,25 @@
                 
             
         
+
         
             com.facebook.presto
             presto-built-in-worker-function-tools
+            ${project.version}
+        
+
+        
+            com.facebook.presto
+            presto-native-sql-invoked-functions-plugin
+            ${project.version}
+            test
+        
+
+        
+            com.facebook.presto
+            presto-sql-invoked-functions-plugin
+            ${project.version}
+            test
         
     
 
diff --git a/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/NativeSidecarPluginQueryRunnerUtils.java b/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/NativeSidecarPluginQueryRunnerUtils.java
index 776d4920e2f16..c8c7e1123f974 100644
--- a/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/NativeSidecarPluginQueryRunnerUtils.java
+++ b/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/NativeSidecarPluginQueryRunnerUtils.java
@@ -13,6 +13,7 @@
  */
 package com.facebook.presto.sidecar;
 
+import com.facebook.presto.scalar.sql.NativeSqlInvokedFunctionsPlugin;
 import com.facebook.presto.sidecar.functionNamespace.NativeFunctionNamespaceManagerFactory;
 import com.facebook.presto.sidecar.sessionpropertyproviders.NativeSystemSessionPropertyProviderFactory;
 import com.facebook.presto.sidecar.typemanager.NativeTypeManagerFactory;
@@ -37,5 +38,6 @@ public static void setupNativeSidecarPlugin(QueryRunner queryRunner)
                         "function-implementation-type", "CPP"));
         queryRunner.loadTypeManager(NativeTypeManagerFactory.NAME);
         queryRunner.loadPlanCheckerProviderManager("native", ImmutableMap.of());
+        queryRunner.installPlugin(new NativeSqlInvokedFunctionsPlugin());
     }
 }
diff --git a/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java b/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java
index b4716f98362cf..fe0b18c24b2fb 100644
--- a/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java
+++ b/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java
@@ -16,6 +16,8 @@
 import com.facebook.airlift.units.DataSize;
 import com.facebook.presto.Session;
 import com.facebook.presto.nativeworker.PrestoNativeQueryRunnerUtils;
+import com.facebook.presto.scalar.sql.NativeSqlInvokedFunctionsPlugin;
+import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin;
 import com.facebook.presto.sidecar.functionNamespace.FunctionDefinitionProvider;
 import com.facebook.presto.sidecar.functionNamespace.NativeFunctionDefinitionProvider;
 import com.facebook.presto.sidecar.functionNamespace.NativeFunctionNamespaceManager;
@@ -45,9 +47,12 @@
 import java.util.stream.Collectors;
 
 import static com.facebook.airlift.units.DataSize.Unit.MEGABYTE;
+import static com.facebook.presto.SystemSessionProperties.INLINE_SQL_FUNCTIONS;
+import static com.facebook.presto.SystemSessionProperties.KEY_BASED_SAMPLING_ENABLED;
 import static com.facebook.presto.SystemSessionProperties.REMOVE_MAP_CAST;
 import static com.facebook.presto.common.Utils.checkArgument;
 import static com.facebook.presto.common.type.BigintType.BIGINT;
+import static com.facebook.presto.nativeworker.NativeQueryRunnerUtils.createCustomer;
 import static com.facebook.presto.nativeworker.NativeQueryRunnerUtils.createLineitem;
 import static com.facebook.presto.nativeworker.NativeQueryRunnerUtils.createNation;
 import static com.facebook.presto.nativeworker.NativeQueryRunnerUtils.createOrders;
@@ -65,6 +70,7 @@ public class TestNativeSidecarPlugin
     private static final String REGEX_FUNCTION_NAMESPACE = "native.default.*";
     private static final String REGEX_SESSION_NAMESPACE = "Native Execution only.*";
     private static final long SIDECAR_HTTP_CLIENT_MAX_CONTENT_SIZE_MB = 128;
+    private static final int INLINED_SQL_FUNCTIONS_COUNT = 7;
 
     @Override
     protected void createTables()
@@ -75,6 +81,7 @@ protected void createTables()
         createOrders(queryRunner);
         createOrdersEx(queryRunner);
         createRegion(queryRunner);
+        createCustomer(queryRunner);
     }
 
     @Override
@@ -93,9 +100,11 @@ protected QueryRunner createQueryRunner()
     protected QueryRunner createExpectedQueryRunner()
             throws Exception
     {
-        return PrestoNativeQueryRunnerUtils.javaHiveQueryRunnerBuilder()
+        QueryRunner queryRunner = PrestoNativeQueryRunnerUtils.javaHiveQueryRunnerBuilder()
                 .setAddStorageFormatToPath(true)
                 .build();
+        queryRunner.installPlugin(new SqlInvokedFunctionsPlugin());
+        return queryRunner;
     }
 
     public static void setupNativeSidecarPlugin(QueryRunner queryRunner)
@@ -113,6 +122,7 @@ public static void setupNativeSidecarPlugin(QueryRunner queryRunner)
                         "sidecar.http-client.max-content-length", SIDECAR_HTTP_CLIENT_MAX_CONTENT_SIZE_MB + "MB"));
         queryRunner.loadTypeManager(NativeTypeManagerFactory.NAME);
         queryRunner.loadPlanCheckerProviderManager("native", ImmutableMap.of());
+        queryRunner.installPlugin(new NativeSqlInvokedFunctionsPlugin());
     }
 
     @Test
@@ -163,6 +173,7 @@ public void testSetNativeWorkerSessionProperty()
     @Test
     public void testShowFunctions()
     {
+        int inlinedSQLFunctionsCount = 0;
         @Language("SQL") String sql = "SHOW FUNCTIONS";
         MaterializedResult actualResult = computeActual(sql);
         List actualRows = actualResult.getMaterializedRows();
@@ -176,11 +187,17 @@ public void testShowFunctions()
 
             // function namespace should be present.
             String fullFunctionName = row.get(5).toString();
-            if (Pattern.matches(REGEX_FUNCTION_NAMESPACE, fullFunctionName)) {
-                continue;
+            if (!Pattern.matches(REGEX_FUNCTION_NAMESPACE, fullFunctionName)) {
+                // If no namespace match found, check if it's an inlined SQL Invoked function.
+                String language = row.get(9).toString();
+                if (language.equalsIgnoreCase("SQL")) {
+                    inlinedSQLFunctionsCount++;
+                    continue;
+                }
+                fail(format("No namespace match found for row: %s", row));
             }
-            fail(format("No namespace match found for row: %s", row));
         }
+        assertEquals(inlinedSQLFunctionsCount, INLINED_SQL_FUNCTIONS_COUNT);
     }
 
     @Test
@@ -321,7 +338,7 @@ public void testApproxPercentile()
     public void testInformationSchemaTables()
     {
         assertQuery("select lower(table_name) from information_schema.tables "
-                        + "where table_name = 'lineitem' or table_name = 'LINEITEM' ");
+                + "where table_name = 'lineitem' or table_name = 'LINEITEM' ");
     }
 
     @Test
@@ -423,6 +440,105 @@ public void testRemoveMapCast()
                 "values 0.5, 0.1");
     }
 
+    @Test
+    public void testOverriddenInlinedSqlInvokedFunctions()
+    {
+        // String functions
+        assertQuery("SELECT trail(comment, cast(nationkey as integer)) FROM nation");
+        assertQuery("SELECT name, comment, replace_first(comment, 'iron', 'gold') from nation");
+
+        // Array functions
+        assertQuery("SELECT array_intersect(ARRAY['apple', 'banana', 'cherry'], ARRAY['apple', 'mango', 'fig'])");
+        assertQuery("SELECT array_frequency(split(comment, '')) from nation");
+        assertQuery("SELECT array_duplicates(ARRAY[regionkey]), array_duplicates(ARRAY[comment]) from nation");
+        assertQuery("SELECT array_has_duplicates(ARRAY[custkey]) from orders");
+        assertQuery("SELECT array_max_by(ARRAY[comment], x -> length(x)) from orders");
+        assertQuery("SELECT array_min_by(ARRAY[ROW('USA', 1), ROW('INDIA', 2), ROW('UK', 3)], x -> x[2])");
+        assertQuery("SELECT array_sort_desc(map_keys(map_union(quantity_by_linenumber))) FROM orders_ex");
+        assertQuery("SELECT remove_nulls(ARRAY[CAST(regionkey AS VARCHAR), comment, NULL]) from nation");
+        assertQuery("SELECT array_top_n(ARRAY[CAST(nationkey AS VARCHAR)], 3) from nation");
+        assertQuerySucceeds("SELECT array_sort_desc(quantities, x -> abs(x)) FROM orders_ex");
+
+        // Map functions
+        assertQuery("SELECT map_normalize(MAP(ARRAY['a', 'b', 'c'], ARRAY[1, 4, 5]))");
+        assertQuery("SELECT map_normalize(MAP(ARRAY['a', 'b', 'c'], ARRAY[1, 0, -1]))");
+        assertQuery("SELECT name, map_normalize(MAP(ARRAY['regionkey', 'length'], ARRAY[regionkey, length(comment)])) from nation");
+        assertQuery("SELECT name, map_remove_null_values(map(ARRAY['region', 'comment', 'nullable'], " +
+                "ARRAY[CAST(regionkey AS VARCHAR), comment, NULL])) from nation");
+        assertQuery("SELECT name, map_key_exists(map(ARRAY['nation', 'comment'], ARRAY[CAST(nationkey AS VARCHAR), comment]), 'comment') from nation");
+        assertQuery("SELECT map_keys_by_top_n_values(MAP(ARRAY[orderkey], ARRAY[custkey]), 2) from orders");
+        assertQuery("SELECT map_top_n(MAP(ARRAY[CAST(nationkey AS VARCHAR)], ARRAY[comment]), 3) from nation");
+        assertQuery("SELECT map_top_n_keys(MAP(ARRAY[orderkey], ARRAY[custkey]), 3) from orders");
+        assertQuery("SELECT map_top_n_values(MAP(ARRAY[orderkey], ARRAY[custkey]), 3) from orders");
+        assertQuery("SELECT all_keys_match(MAP(ARRAY[comment], ARRAY[custkey]), k -> length(k) > 5) from orders");
+        assertQuery("SELECT any_keys_match(MAP(ARRAY[comment], ARRAY[custkey]), k -> starts_with(k, 'abc')) from orders");
+        assertQuery("SELECT any_values_match(MAP(ARRAY[orderkey], ARRAY[totalprice]), k -> abs(k) > 20) from orders");
+        assertQuery("SELECT no_values_match(MAP(ARRAY[orderkey], ARRAY[comment]), k -> length(k) > 2) from orders");
+        assertQuery("SELECT no_keys_match(MAP(ARRAY[comment], ARRAY[custkey]), k -> ends_with(k, 'a')) from orders");
+    }
+
+    @Test
+    public void testNonOverriddenInlinedSqlInvokedFunctionsWhenConfigEnabled()
+    {
+        // Array functions
+        assertQuery("SELECT array_split_into_chunks(split(comment, ''), 2) from nation");
+        assertQuery("SELECT array_least_frequent(quantities) from orders_ex");
+        assertQuery("SELECT array_least_frequent(split(comment, ''), 5) from nation");
+        assertQuerySucceeds("SELECT array_top_n(ARRAY[orderkey], 25, (x, y) -> if (x < y, cast(1 as bigint), if (x > y, cast(-1 as bigint), cast(0 as bigint)))) from orders");
+
+        // Map functions
+        assertQuerySucceeds("SELECT map_top_n_values(MAP(ARRAY[comment], ARRAY[nationkey]), 2, (x, y) -> if (x < y, cast(1 as bigint), if (x > y, cast(-1 as bigint), cast(0 as bigint)))) from nation");
+        assertQuerySucceeds("SELECT map_top_n_keys(MAP(ARRAY[regionkey], ARRAY[nationkey]), 5, (x, y) -> if (x < y, cast(1 as bigint), if (x > y, cast(-1 as bigint), cast(0 as bigint)))) from nation");
+
+        Session sessionWithKeyBasedSampling = Session.builder(getSession())
+                .setSystemProperty(KEY_BASED_SAMPLING_ENABLED, "true")
+                .build();
+
+        @Language("SQL") String query = "select count(1) FROM lineitem l left JOIN orders o ON l.orderkey = o.orderkey JOIN customer c ON o.custkey = c.custkey";
+
+        assertQuery(query, "select cast(60175 as bigint)");
+        assertQuery(sessionWithKeyBasedSampling, query, "select cast(16185 as bigint)");
+    }
+
+    @Test
+    public void testNonOverriddenInlinedSqlInvokedFunctionsWhenConfigDisabled()
+    {
+        // When inline_sql_functions is set to false, the below queries should fail as the implementations don't exist on the native worker
+        Session session = Session.builder(getSession())
+                .setSystemProperty(KEY_BASED_SAMPLING_ENABLED, "true")
+                .setSystemProperty(INLINE_SQL_FUNCTIONS, "false")
+                .build();
+
+        // Array functions
+        assertQueryFails(session,
+                "SELECT array_split_into_chunks(split(comment, ''), 2) from nation",
+                ".*Scalar function name not registered: native.default.array_split_into_chunks.*");
+        assertQueryFails(session,
+                "SELECT array_least_frequent(quantities) from orders_ex",
+                ".*Scalar function name not registered: native.default.array_least_frequent.*");
+        assertQueryFails(session,
+                "SELECT array_least_frequent(split(comment, ''), 2) from nation",
+                ".*Scalar function name not registered: native.default.array_least_frequent.*");
+        assertQueryFails(session,
+                "SELECT array_top_n(ARRAY[orderkey], 25, (x, y) -> if (x < y, cast(1 as bigint), if (x > y, cast(-1 as bigint), cast(0 as bigint)))) from orders",
+                " Scalar function native\\.default\\.array_top_n not registered with arguments.*",
+                true);
+
+        // Map functions
+        assertQueryFails(session,
+                "SELECT map_top_n_values(MAP(ARRAY[comment], ARRAY[nationkey]), 2, (x, y) -> if (x < y, cast(1 as bigint), if (x > y, cast(-1 as bigint), cast(0 as bigint)))) from nation",
+                ".*Scalar function native\\.default\\.map_top_n_values not registered with arguments.*",
+                true);
+        assertQueryFails(session,
+                "SELECT map_top_n_keys(MAP(ARRAY[regionkey], ARRAY[nationkey]), 5, (x, y) -> if (x < y, cast(1 as bigint), if (x > y, cast(-1 as bigint), cast(0 as bigint)))) from nation",
+                ".*Scalar function native\\.default\\.map_top_n_keys not registered with arguments.*",
+                true);
+
+        assertQueryFails(session,
+                "select count(1) FROM lineitem l left JOIN orders o ON l.orderkey = o.orderkey JOIN customer c ON o.custkey = c.custkey",
+                ".*Scalar function name not registered: native.default.key_sampling_percent.*");
+    }
+
     private String generateRandomTableName()
     {
         String tableName = "tmp_presto_" + UUID.randomUUID().toString().replace("-", "");
diff --git a/presto-native-sql-invoked-functions-plugin/pom.xml b/presto-native-sql-invoked-functions-plugin/pom.xml
new file mode 100644
index 0000000000000..7d837a28a8bfb
--- /dev/null
+++ b/presto-native-sql-invoked-functions-plugin/pom.xml
@@ -0,0 +1,29 @@
+
+    4.0.0
+    
+        com.facebook.presto
+        presto-root
+        0.295-SNAPSHOT
+    
+
+    presto-native-sql-invoked-functions-plugin
+    Presto Native - Sql invoked functions plugin
+    presto-plugin
+
+    
+        ${project.parent.basedir}
+    
+
+    
+        
+            com.facebook.presto
+            presto-spi
+            provided
+        
+        
+            com.google.guava
+            guava
+        
+    
+
diff --git a/presto-native-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/NativeArraySqlFunctions.java b/presto-native-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/NativeArraySqlFunctions.java
new file mode 100644
index 0000000000000..841883d99ae8f
--- /dev/null
+++ b/presto-native-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/NativeArraySqlFunctions.java
@@ -0,0 +1,74 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.scalar.sql;
+
+import com.facebook.presto.spi.function.Description;
+import com.facebook.presto.spi.function.SqlInvokedScalarFunction;
+import com.facebook.presto.spi.function.SqlParameter;
+import com.facebook.presto.spi.function.SqlParameters;
+import com.facebook.presto.spi.function.SqlType;
+import com.facebook.presto.spi.function.TypeParameter;
+
+public class NativeArraySqlFunctions
+{
+    private NativeArraySqlFunctions() {}
+
+    @SqlInvokedScalarFunction(value = "array_split_into_chunks", deterministic = true, calledOnNullInput = false)
+    @Description("Returns an array of arrays splitting input array into chunks of given length. " +
+            "If array is not evenly divisible it will split into as many possible chunks and " +
+            "return the left over elements for the last array. Returns null for null inputs, but not elements.")
+    @TypeParameter("T")
+    @SqlParameters({@SqlParameter(name = "input", type = "array(T)"), @SqlParameter(name = "sz", type = "int")})
+    @SqlType("array(array(T))")
+    public static String arraySplitIntoChunks()
+    {
+        return "RETURN IF(sz <= 0, " +
+                "fail('Invalid slice size: ' || cast(sz as varchar) || '. Size must be greater than zero.'), " +
+                "IF(cardinality(input) / sz > 10000, " +
+                "fail('Cannot split array of size: ' || cast(cardinality(input) as varchar) || ' into more than 10000 parts.'), " +
+                "transform(" +
+                "sequence(1, cardinality(input), sz), " +
+                "x -> slice(input, x, sz))))";
+    }
+
+    @SqlInvokedScalarFunction(value = "array_least_frequent", deterministic = true, calledOnNullInput = true)
+    @Description("Determines the least frequent element in the array. If there are multiple elements, the function returns the smallest element")
+    @TypeParameter("T")
+    @SqlParameter(name = "input", type = "array(T)")
+    @SqlType("array")
+    public static String arrayLeastFrequent()
+    {
+        return "RETURN IF(COALESCE(CARDINALITY(REMOVE_NULLS(input)), 0) = 0, NULL, TRANSFORM(SLICE(ARRAY_SORT(TRANSFORM(MAP_ENTRIES(ARRAY_FREQUENCY(REMOVE_NULLS(input))), x -> ROW(x[2], x[1]))), 1, 1), x -> x[2]))";
+    }
+
+    @SqlInvokedScalarFunction(value = "array_least_frequent", deterministic = true, calledOnNullInput = true)
+    @Description("Determines the n least frequent element in the array in the ascending order of the elements.")
+    @TypeParameter("T")
+    @SqlParameters({@SqlParameter(name = "input", type = "array(T)"), @SqlParameter(name = "n", type = "bigint")})
+    @SqlType("array")
+    public static String arrayNLeastFrequent()
+    {
+        return "RETURN IF(n < 0, fail('n must be greater than or equal to 0'), IF(COALESCE(CARDINALITY(REMOVE_NULLS(input)), 0) = 0, NULL, TRANSFORM(SLICE(ARRAY_SORT(TRANSFORM(MAP_ENTRIES(ARRAY_FREQUENCY(REMOVE_NULLS(input))), x -> ROW(x[2], x[1]))), 1, n), x -> x[2])))";
+    }
+
+    @SqlInvokedScalarFunction(value = "array_top_n", deterministic = true, calledOnNullInput = true)
+    @Description("Returns the top N values of the given map sorted using the provided lambda comparator.")
+    @TypeParameter("T")
+    @SqlParameters({@SqlParameter(name = "input", type = "array(T)"), @SqlParameter(name = "n", type = "int"), @SqlParameter(name = "f", type = "function(T, T, bigint)")})
+    @SqlType("array")
+    public static String arrayTopNComparator()
+    {
+        return "RETURN IF(n < 0, fail('Parameter n: ' || cast(n as varchar) || ' to ARRAY_TOP_N is negative'), SLICE(REVERSE(ARRAY_SORT(input, f)), 1, n))";
+    }
+}
diff --git a/presto-native-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/NativeMapSqlFunctions.java b/presto-native-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/NativeMapSqlFunctions.java
new file mode 100644
index 0000000000000..9eccc84d6d8c8
--- /dev/null
+++ b/presto-native-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/NativeMapSqlFunctions.java
@@ -0,0 +1,48 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.scalar.sql;
+
+import com.facebook.presto.spi.function.Description;
+import com.facebook.presto.spi.function.SqlInvokedScalarFunction;
+import com.facebook.presto.spi.function.SqlParameter;
+import com.facebook.presto.spi.function.SqlParameters;
+import com.facebook.presto.spi.function.SqlType;
+import com.facebook.presto.spi.function.TypeParameter;
+
+public class NativeMapSqlFunctions
+{
+    private NativeMapSqlFunctions() {}
+
+    @SqlInvokedScalarFunction(value = "map_top_n_keys", deterministic = true, calledOnNullInput = true)
+    @Description("Returns the top N keys of the given map sorting its keys using the provided lambda comparator.")
+    @TypeParameter("K")
+    @TypeParameter("V")
+    @SqlParameters({@SqlParameter(name = "input", type = "map(K, V)"), @SqlParameter(name = "n", type = "bigint"), @SqlParameter(name = "f", type = "function(K, K, bigint)")})
+    @SqlType("array")
+    public static String mapTopNKeysComparator()
+    {
+        return "RETURN IF(n < 0, fail('n must be greater than or equal to 0'), slice(reverse(array_sort(map_keys(input), f)), 1, n))";
+    }
+
+    @SqlInvokedScalarFunction(value = "map_top_n_values", deterministic = true, calledOnNullInput = true)
+    @Description("Returns the top N values of the given map sorted using the provided lambda comparator.")
+    @TypeParameter("K")
+    @TypeParameter("V")
+    @SqlParameters({@SqlParameter(name = "input", type = "map(K, V)"), @SqlParameter(name = "n", type = "bigint"), @SqlParameter(name = "f", type = "function(V, V, bigint)")})
+    @SqlType("array")
+    public static String mapTopNValuesComparator()
+    {
+        return "RETURN IF(n < 0, fail('n must be greater than or equal to 0'), slice(reverse(array_sort(remove_nulls(map_values(input)), f)) || filter(map_values(input), x -> x is null), 1, n))";
+    }
+}
diff --git a/presto-native-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/NativeSimpleSamplingPercent.java b/presto-native-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/NativeSimpleSamplingPercent.java
new file mode 100644
index 0000000000000..a710391760714
--- /dev/null
+++ b/presto-native-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/NativeSimpleSamplingPercent.java
@@ -0,0 +1,33 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.scalar.sql;
+
+import com.facebook.presto.spi.function.Description;
+import com.facebook.presto.spi.function.SqlInvokedScalarFunction;
+import com.facebook.presto.spi.function.SqlParameter;
+import com.facebook.presto.spi.function.SqlType;
+
+public class NativeSimpleSamplingPercent
+{
+    private NativeSimpleSamplingPercent() {}
+
+    @SqlInvokedScalarFunction(value = "key_sampling_percent", deterministic = true, calledOnNullInput = false)
+    @Description("Returns a value between 0.0 and 1.0 using the hash of the given input string")
+    @SqlParameter(name = "input", type = "varchar")
+    @SqlType("double")
+    public static String keySamplingPercent()
+    {
+        return "return (abs(from_ieee754_64(xxhash64(cast(input as varbinary)))) % 100) / 100. ";
+    }
+}
diff --git a/presto-native-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/NativeSqlInvokedFunctionsPlugin.java b/presto-native-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/NativeSqlInvokedFunctionsPlugin.java
new file mode 100644
index 0000000000000..69d7ff1e78522
--- /dev/null
+++ b/presto-native-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/NativeSqlInvokedFunctionsPlugin.java
@@ -0,0 +1,33 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.scalar.sql;
+
+import com.facebook.presto.spi.Plugin;
+import com.google.common.collect.ImmutableSet;
+
+import java.util.Set;
+
+public class NativeSqlInvokedFunctionsPlugin
+        implements Plugin
+{
+    @Override
+    public Set> getSqlInvokedFunctions()
+    {
+        return ImmutableSet.>builder()
+                .add(NativeArraySqlFunctions.class)
+                .add(NativeMapSqlFunctions.class)
+                .add(NativeSimpleSamplingPercent.class)
+                .build();
+    }
+}
diff --git a/presto-native-tests/pom.xml b/presto-native-tests/pom.xml
index e90785cbb6940..7e23670117b1d 100644
--- a/presto-native-tests/pom.xml
+++ b/presto-native-tests/pom.xml
@@ -192,6 +192,13 @@
             units
             test
         
+
+        
+            com.facebook.presto
+            presto-native-sql-invoked-functions-plugin
+            ${project.version}
+            test
+        
     
 
     
diff --git a/presto-plan-checker-router-plugin/pom.xml b/presto-plan-checker-router-plugin/pom.xml
index 60e34e4147d19..db1c4b0f1736c 100644
--- a/presto-plan-checker-router-plugin/pom.xml
+++ b/presto-plan-checker-router-plugin/pom.xml
@@ -223,6 +223,13 @@
             presto-hive-metastore
             test
         
+
+        
+            com.facebook.presto
+            presto-native-sql-invoked-functions-plugin
+            ${project.version}
+            test
+        
     
 
     
diff --git a/presto-product-tests/conf/docker/common/compose-commons.sh b/presto-product-tests/conf/docker/common/compose-commons.sh
index eae9f18ce9583..5c20783716b60 100644
--- a/presto-product-tests/conf/docker/common/compose-commons.sh
+++ b/presto-product-tests/conf/docker/common/compose-commons.sh
@@ -39,6 +39,16 @@ if [[ -z "${PRESTO_SERVER_DIR:-}" ]]; then
     source "${PRODUCT_TESTS_ROOT}/target/classes/presto.env"
     PRESTO_SERVER_DIR="${PROJECT_ROOT}/presto-server/target/presto-server-${PRESTO_VERSION}/"
 fi
+
+# The following plugin results in a function signature conflict when loaded in Java/ sidecar disabled native clusters.
+# This plugin is only meant for sidecar enabled native clusters, hence exclude it.
+PLUGIN_TO_EXCLUDE="native-sql-invoked-functions-plugin"
+
+if [[ -d "${PRESTO_SERVER_DIR}/plugin/${PLUGIN_TO_EXCLUDE}" ]]; then
+    echo "Excluding plugin: $PLUGIN_TO_EXCLUDE"
+    rm -rf "${PRESTO_SERVER_DIR}/plugin/${PLUGIN_TO_EXCLUDE}"
+fi
+
 export_canonical_path PRESTO_SERVER_DIR
 
 if [[ -z "${PRESTO_CLI_JAR:-}" ]]; then
diff --git a/presto-server/src/main/provisio/presto.xml b/presto-server/src/main/provisio/presto.xml
index b14b36a768e69..d15a041c7d1f5 100644
--- a/presto-server/src/main/provisio/presto.xml
+++ b/presto-server/src/main/provisio/presto.xml
@@ -292,4 +292,10 @@
             
         
     
+
+    
+        
+            
+        
+    
 

From 4c543a03ba96fa06dd60b7fcfb1e071d54262933 Mon Sep 17 00:00:00 2001
From: mohsaka <135669458+mohsaka@users.noreply.github.com>
Date: Fri, 25 Jul 2025 10:33:12 +0100
Subject: [PATCH 060/113] Analyze table function invocation

Changes adapted from trino/PR#11336, 12910
Original commits:
9e8d51ad45f57267f5f7fa6bf8e8c4ec56103dda
f0508a7ab420449c6e2960ecf1d0a8d7058242da
Author: kasiafi

Modifications were made to adapt to Presto including:
Removal of Node Location from TrinoException
Added new SemanticErrorCodes
Changed Void context to SqlPlannerContext in RelationPlanner.java
Add newUnqualified to Field class with Presto specification
Add getCanonicalValue to Identifier.java

Co-authored-by: Pratik Joseph Dabre 
Co-authored-by: kasiafi <30203062+kasiafi@users.noreply.github.com>
Co-authored-by: Xin Zhang 
---
 .../presto/sql/analyzer/Analysis.java         |  70 ++++
 .../facebook/presto/sql/analyzer/Field.java   |   8 +
 .../sql/analyzer/SemanticErrorCode.java       |   7 +
 .../metadata/FunctionAndTypeManager.java      |   5 +
 .../sql/analyzer/StatementAnalyzer.java       | 293 ++++++++++++++-
 .../tvf/MockConnectorColumnHandle.java        |  81 +++++
 .../tvf/MockConnectorTableHandle.java         |  98 +++++
 .../connector/tvf/TestingTableFunctions.java  | 338 ++++++++++++++++++
 .../metadata/TestTableFunctionRegistry.java   | 197 +---------
 .../sql/analyzer/AbstractAnalyzerTest.java    |  11 +
 .../presto/sql/analyzer/TestAnalyzer.java     |  71 ++++
 .../facebook/presto/sql/tree/Expression.java  |   6 +
 .../com/facebook/presto/sql/tree/Node.java    |   5 +
 .../tree/TableFunctionDescriptorArgument.java |   6 +
 .../sql/tree/TableFunctionTableArgument.java  |   6 +
 .../function/table/ArgumentSpecification.java |   6 +
 .../DescriptorArgumentSpecification.java      |   7 +
 .../table/ReturnTypeSpecification.java        |  27 ++
 .../table/ScalarArgumentSpecification.java    |   7 +
 .../table/TableArgumentSpecification.java     |   7 +
 20 files changed, 1063 insertions(+), 193 deletions(-)
 create mode 100644 presto-main-base/src/test/java/com/facebook/presto/connector/tvf/MockConnectorColumnHandle.java
 create mode 100644 presto-main-base/src/test/java/com/facebook/presto/connector/tvf/MockConnectorTableHandle.java
 create mode 100644 presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestingTableFunctions.java

diff --git a/presto-analyzer/src/main/java/com/facebook/presto/sql/analyzer/Analysis.java b/presto-analyzer/src/main/java/com/facebook/presto/sql/analyzer/Analysis.java
index 66e3f7b0e4dc2..0deb5e1ca1fb7 100644
--- a/presto-analyzer/src/main/java/com/facebook/presto/sql/analyzer/Analysis.java
+++ b/presto-analyzer/src/main/java/com/facebook/presto/sql/analyzer/Analysis.java
@@ -19,13 +19,17 @@
 import com.facebook.presto.common.type.Type;
 import com.facebook.presto.spi.ColumnHandle;
 import com.facebook.presto.spi.ColumnMetadata;
+import com.facebook.presto.spi.ConnectorId;
 import com.facebook.presto.spi.TableHandle;
 import com.facebook.presto.spi.analyzer.AccessControlInfo;
 import com.facebook.presto.spi.analyzer.AccessControlInfoForTable;
 import com.facebook.presto.spi.analyzer.AccessControlReferences;
 import com.facebook.presto.spi.analyzer.AccessControlRole;
+import com.facebook.presto.spi.connector.ConnectorTransactionHandle;
 import com.facebook.presto.spi.function.FunctionHandle;
 import com.facebook.presto.spi.function.FunctionKind;
+import com.facebook.presto.spi.function.table.Argument;
+import com.facebook.presto.spi.function.table.ConnectorTableFunctionHandle;
 import com.facebook.presto.spi.security.AccessControl;
 import com.facebook.presto.spi.security.AccessControlContext;
 import com.facebook.presto.spi.security.AllowAllAccessControl;
@@ -51,6 +55,7 @@
 import com.facebook.presto.sql.tree.Statement;
 import com.facebook.presto.sql.tree.SubqueryExpression;
 import com.facebook.presto.sql.tree.Table;
+import com.facebook.presto.sql.tree.TableFunctionInvocation;
 import com.google.common.collect.ArrayListMultimap;
 import com.google.common.collect.HashMultiset;
 import com.google.common.collect.ImmutableList;
@@ -199,6 +204,8 @@ public class Analysis
     // Keeps track of the subquery we are visiting, so we have access to base query information when processing materialized view status
     private Optional currentQuerySpecification = Optional.empty();
 
+    private final Map, TableFunctionInvocationAnalysis> tableFunctionAnalyses = new LinkedHashMap<>();
+
     public Analysis(@Nullable Statement root, Map, Expression> parameters, boolean isDescribe)
     {
         this.root = root;
@@ -1061,6 +1068,16 @@ public Map getColumnMasks(Table table)
         return columnMasks.getOrDefault(NodeRef.of(table), ImmutableMap.of());
     }
 
+    public void setTableFunctionAnalysis(TableFunctionInvocation node, TableFunctionInvocationAnalysis analysis)
+    {
+        tableFunctionAnalyses.put(NodeRef.of(node), analysis);
+    }
+
+    public TableFunctionInvocationAnalysis getTableFunctionAnalysis(TableFunctionInvocation node)
+    {
+        return tableFunctionAnalyses.get(NodeRef.of(node));
+    }
+
     @Immutable
     public static final class Insert
     {
@@ -1311,4 +1328,57 @@ public int hashCode()
             return Objects.hash(table, column, identity);
         }
     }
+
+    /**
+     * Encapsulates the result of analyzing a table function invocation.
+     * Includes the connector ID, function name, argument bindings, and the
+     * connector-specific table function handle needed for planning and execution.
+     */
+    public static class TableFunctionInvocationAnalysis
+    {
+        private final ConnectorId connectorId;
+        private final String functionName;
+        private final Map arguments;
+        private final ConnectorTableFunctionHandle connectorTableFunctionHandle;
+        private final ConnectorTransactionHandle transactionHandle;
+
+        public TableFunctionInvocationAnalysis(
+                ConnectorId connectorId,
+                String functionName,
+                Map arguments,
+                ConnectorTableFunctionHandle connectorTableFunctionHandle,
+                ConnectorTransactionHandle transactionHandle)
+        {
+            this.connectorId = requireNonNull(connectorId, "connectorId is null");
+            this.functionName = requireNonNull(functionName, "functionName is null");
+            this.arguments = ImmutableMap.copyOf(arguments);
+            this.connectorTableFunctionHandle = requireNonNull(connectorTableFunctionHandle, "connectorTableFunctionHandle is null");
+            this.transactionHandle = requireNonNull(transactionHandle, "transactionHandle is null");
+        }
+
+        public ConnectorId getConnectorId()
+        {
+            return connectorId;
+        }
+
+        public String getFunctionName()
+        {
+            return functionName;
+        }
+
+        public Map getArguments()
+        {
+            return arguments;
+        }
+
+        public ConnectorTableFunctionHandle getConnectorTableFunctionHandle()
+        {
+            return connectorTableFunctionHandle;
+        }
+
+        public ConnectorTransactionHandle getTransactionHandle()
+        {
+            return transactionHandle;
+        }
+    }
 }
diff --git a/presto-analyzer/src/main/java/com/facebook/presto/sql/analyzer/Field.java b/presto-analyzer/src/main/java/com/facebook/presto/sql/analyzer/Field.java
index 630f4670f6cc2..15c33950ce14b 100644
--- a/presto-analyzer/src/main/java/com/facebook/presto/sql/analyzer/Field.java
+++ b/presto-analyzer/src/main/java/com/facebook/presto/sql/analyzer/Field.java
@@ -86,6 +86,14 @@ public Field(Optional nodeLocation, Optional relati
         this.aliased = aliased;
     }
 
+    public static Field newUnqualified(Optional name, Type type)
+    {
+        requireNonNull(name, "name is null");
+        requireNonNull(type, "type is null");
+
+        return new Field(Optional.empty(), Optional.empty(), name, type, false, Optional.empty(), Optional.empty(), false);
+    }
+
     public Optional getNodeLocation()
     {
         return nodeLocation;
diff --git a/presto-analyzer/src/main/java/com/facebook/presto/sql/analyzer/SemanticErrorCode.java b/presto-analyzer/src/main/java/com/facebook/presto/sql/analyzer/SemanticErrorCode.java
index 656b6dbaab079..5c9be03540f67 100644
--- a/presto-analyzer/src/main/java/com/facebook/presto/sql/analyzer/SemanticErrorCode.java
+++ b/presto-analyzer/src/main/java/com/facebook/presto/sql/analyzer/SemanticErrorCode.java
@@ -58,6 +58,13 @@ public enum SemanticErrorCode
     DUPLICATE_PARAMETER_NAME,
     EXCEPTIONS_WHEN_RESOLVING_FUNCTIONS,
 
+    MISSING_RETURN_TYPE,
+    AMBIGUOUS_RETURN_TYPE,
+    MISSING_ARGUMENT,
+    INVALID_FUNCTION_ARGUMENT,
+    INVALID_ARGUMENTS,
+    NOT_IMPLEMENTED,
+
     ORDER_BY_MUST_BE_IN_SELECT,
     ORDER_BY_MUST_BE_IN_AGGREGATE,
     REFERENCE_TO_OUTPUT_ATTRIBUTE_WITHIN_ORDER_BY_GROUPING,
diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionAndTypeManager.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionAndTypeManager.java
index bb0c6cf311eb6..9f18bbed9fc00 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionAndTypeManager.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionAndTypeManager.java
@@ -476,6 +476,11 @@ public void addTypeManagerFactory(TypeManagerFactory factory)
         }
     }
 
+    public TransactionManager getTransactionManager()
+    {
+        return transactionManager;
+    }
+
     public void registerBuiltInFunctions(List functions)
     {
         builtInTypeAndFunctionNamespaceManager.registerBuiltInFunctions(functions);
diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/StatementAnalyzer.java b/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/StatementAnalyzer.java
index b16042d87537b..bb9e068da5aa7 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/StatementAnalyzer.java
+++ b/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/StatementAnalyzer.java
@@ -31,15 +31,19 @@
 import com.facebook.presto.common.type.TimestampWithTimeZoneType;
 import com.facebook.presto.common.type.Type;
 import com.facebook.presto.common.type.VarcharType;
+import com.facebook.presto.metadata.CatalogMetadata;
 import com.facebook.presto.metadata.Metadata;
 import com.facebook.presto.metadata.OperatorNotFoundException;
+import com.facebook.presto.metadata.TableFunctionMetadata;
 import com.facebook.presto.spi.ColumnHandle;
 import com.facebook.presto.spi.ColumnMetadata;
+import com.facebook.presto.spi.ConnectorId;
 import com.facebook.presto.spi.MaterializedViewDefinition;
 import com.facebook.presto.spi.MaterializedViewStatus;
 import com.facebook.presto.spi.PrestoException;
 import com.facebook.presto.spi.PrestoWarning;
 import com.facebook.presto.spi.SchemaTableName;
+import com.facebook.presto.spi.StandardErrorCode;
 import com.facebook.presto.spi.TableHandle;
 import com.facebook.presto.spi.WarningCollector;
 import com.facebook.presto.spi.analyzer.AccessControlInfo;
@@ -47,9 +51,20 @@
 import com.facebook.presto.spi.analyzer.MetadataResolver;
 import com.facebook.presto.spi.analyzer.ViewDefinition;
 import com.facebook.presto.spi.connector.ConnectorTableVersion;
+import com.facebook.presto.spi.connector.ConnectorTransactionHandle;
 import com.facebook.presto.spi.function.FunctionKind;
 import com.facebook.presto.spi.function.Signature;
 import com.facebook.presto.spi.function.SqlFunction;
+import com.facebook.presto.spi.function.table.Argument;
+import com.facebook.presto.spi.function.table.ArgumentSpecification;
+import com.facebook.presto.spi.function.table.ConnectorTableFunction;
+import com.facebook.presto.spi.function.table.Descriptor;
+import com.facebook.presto.spi.function.table.DescriptorArgumentSpecification;
+import com.facebook.presto.spi.function.table.ReturnTypeSpecification;
+import com.facebook.presto.spi.function.table.ScalarArgument;
+import com.facebook.presto.spi.function.table.ScalarArgumentSpecification;
+import com.facebook.presto.spi.function.table.TableArgumentSpecification;
+import com.facebook.presto.spi.function.table.TableFunctionAnalysis;
 import com.facebook.presto.spi.relation.DomainTranslator;
 import com.facebook.presto.spi.relation.RowExpression;
 import com.facebook.presto.spi.security.AccessControl;
@@ -127,6 +142,7 @@
 import com.facebook.presto.sql.tree.NodeRef;
 import com.facebook.presto.sql.tree.Offset;
 import com.facebook.presto.sql.tree.OrderBy;
+import com.facebook.presto.sql.tree.Parameter;
 import com.facebook.presto.sql.tree.Prepare;
 import com.facebook.presto.sql.tree.Property;
 import com.facebook.presto.sql.tree.QualifiedName;
@@ -158,6 +174,10 @@
 import com.facebook.presto.sql.tree.StartTransaction;
 import com.facebook.presto.sql.tree.Statement;
 import com.facebook.presto.sql.tree.Table;
+import com.facebook.presto.sql.tree.TableFunctionArgument;
+import com.facebook.presto.sql.tree.TableFunctionDescriptorArgument;
+import com.facebook.presto.sql.tree.TableFunctionInvocation;
+import com.facebook.presto.sql.tree.TableFunctionTableArgument;
 import com.facebook.presto.sql.tree.TableSubquery;
 import com.facebook.presto.sql.tree.TruncateTable;
 import com.facebook.presto.sql.tree.Union;
@@ -171,6 +191,7 @@
 import com.facebook.presto.sql.tree.With;
 import com.facebook.presto.sql.tree.WithQuery;
 import com.facebook.presto.sql.util.AstUtils;
+import com.facebook.presto.transaction.TransactionManager;
 import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
@@ -208,9 +229,7 @@
 import static com.facebook.presto.metadata.MetadataUtil.getConnectorIdOrThrow;
 import static com.facebook.presto.metadata.MetadataUtil.toSchemaTableName;
 import static com.facebook.presto.spi.StandardErrorCode.DATATYPE_MISMATCH;
-import static com.facebook.presto.spi.StandardErrorCode.INVALID_ARGUMENTS;
 import static com.facebook.presto.spi.StandardErrorCode.INVALID_COLUMN_MASK;
-import static com.facebook.presto.spi.StandardErrorCode.INVALID_FUNCTION_ARGUMENT;
 import static com.facebook.presto.spi.StandardErrorCode.INVALID_ROW_FILTER;
 import static com.facebook.presto.spi.StandardWarningCode.PERFORMANCE_WARNING;
 import static com.facebook.presto.spi.StandardWarningCode.REDUNDANT_ORDER_BY;
@@ -243,12 +262,16 @@
 import static com.facebook.presto.sql.analyzer.RefreshMaterializedViewPredicateAnalyzer.extractTablePredicates;
 import static com.facebook.presto.sql.analyzer.ScopeReferenceExtractor.hasReferencesToScope;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.AMBIGUOUS_ATTRIBUTE;
+import static com.facebook.presto.sql.analyzer.SemanticErrorCode.AMBIGUOUS_RETURN_TYPE;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.COLUMN_NAME_NOT_SPECIFIED;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.COLUMN_TYPE_UNKNOWN;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.DUPLICATE_COLUMN_NAME;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.DUPLICATE_PARAMETER_NAME;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.DUPLICATE_PROPERTY;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.DUPLICATE_RELATION;
+import static com.facebook.presto.sql.analyzer.SemanticErrorCode.FUNCTION_NOT_FOUND;
+import static com.facebook.presto.sql.analyzer.SemanticErrorCode.INVALID_ARGUMENTS;
+import static com.facebook.presto.sql.analyzer.SemanticErrorCode.INVALID_FUNCTION_ARGUMENT;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.INVALID_FUNCTION_NAME;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.INVALID_OFFSET_ROW_COUNT;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.INVALID_ORDINAL;
@@ -258,15 +281,18 @@
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.MATERIALIZED_VIEW_IS_RECURSIVE;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.MISMATCHED_COLUMN_ALIASES;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.MISMATCHED_SET_COLUMN_TYPES;
+import static com.facebook.presto.sql.analyzer.SemanticErrorCode.MISSING_ARGUMENT;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.MISSING_ATTRIBUTE;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.MISSING_COLUMN;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.MISSING_MATERIALIZED_VIEW;
+import static com.facebook.presto.sql.analyzer.SemanticErrorCode.MISSING_RETURN_TYPE;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.MISSING_SCHEMA;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.MISSING_TABLE;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.MUST_BE_WINDOW_FUNCTION;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.NESTED_WINDOW;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.NONDETERMINISTIC_ORDER_BY_EXPRESSION_WITH_SELECT_DISTINCT;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.NON_NUMERIC_SAMPLE_PERCENTAGE;
+import static com.facebook.presto.sql.analyzer.SemanticErrorCode.NOT_IMPLEMENTED;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.NOT_SUPPORTED;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.ORDER_BY_MUST_BE_IN_SELECT;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.TABLE_ALREADY_EXISTS;
@@ -1255,7 +1281,7 @@ else if (expressionType instanceof MapType) {
                     outputFields.add(Field.newUnqualified(expression.getLocation(), Optional.empty(), ((MapType) expressionType).getValueType()));
                 }
                 else {
-                    throw new PrestoException(INVALID_FUNCTION_ARGUMENT, "Cannot unnest type: " + expressionType);
+                    throw new PrestoException(StandardErrorCode.INVALID_FUNCTION_ARGUMENT, "Cannot unnest type: " + expressionType);
                 }
             }
             if (node.isWithOrdinality()) {
@@ -1272,6 +1298,265 @@ protected Scope visitLateral(Lateral node, Optional scope)
             return createAndAssignScope(node, scope, queryScope.getRelationType());
         }
 
+        @Override
+        protected Scope visitTableFunctionInvocation(TableFunctionInvocation node, Optional scope)
+        {
+            TableFunctionMetadata tableFunctionMetadata = metadata.getFunctionAndTypeManager()
+                    .getTableFunctionRegistry()
+                    .resolve(session, node.getName())
+                    .orElseThrow(() -> new SemanticException(
+                            FUNCTION_NOT_FOUND,
+                            node,
+                            "Table function %s not registered",
+                            node.getName()));
+
+            ConnectorTableFunction function = tableFunctionMetadata.getFunction();
+            ConnectorId connectorId = tableFunctionMetadata.getConnectorId();
+
+            QualifiedObjectName functionName = new QualifiedObjectName(connectorId.getCatalogName(), function.getSchema(), function.getName());
+
+            Map passedArguments = analyzeArguments(node, function.getArguments(), node.getArguments());
+
+            TransactionManager transactionManager = metadata.getFunctionAndTypeManager().getTransactionManager();
+            CatalogMetadata registrationCatalogMetadata = transactionManager.getOptionalCatalogMetadata(session.getRequiredTransactionId(), connectorId.getCatalogName()).orElseThrow(() -> new IllegalStateException("Missing catalog metadata"));
+            // a call to getRequiredCatalogHandle() is necessary so that the catalog is recorded by the TransactionManager
+            ConnectorTransactionHandle transactionHandle = transactionManager.getConnectorTransaction(
+                    session.getRequiredTransactionId(), registrationCatalogMetadata.getConnectorId());
+
+            TableFunctionAnalysis functionAnalysis = function.analyze(session.toConnectorSession(connectorId), transactionHandle, passedArguments);
+            analysis.setTableFunctionAnalysis(node, new Analysis.TableFunctionInvocationAnalysis(connectorId, functionName.toString(), passedArguments, functionAnalysis.getHandle(), transactionHandle));
+
+            // TODO process the copartitioning:
+            // 1. validate input table references
+            // 2. the copartitioned tables in each set must be partitioned, and have the same number of partitioning columns
+            // 3. the corresponding columns must be comparable
+            // 4. within a set, determine and record coercions of the corresponding columns to a common supertype
+            // Note that if a table is part of multiple copartitioning sets, it might require a different coercion for a column
+            // per each set. Additionally, there might be another coercion required by the Table Function logic. Also, since
+            // all partitioning columns are passed-through, we also need an un-coerced copy.
+            // See ExpressionAnalyzer.sortKeyCoercionsForFrameBoundCalculation for multiple coercions on a column.
+            if (!node.getCopartitioning().isEmpty()) {
+                throw new SemanticException(NOT_IMPLEMENTED, node, "COPARTITION clause is not yet supported for table functions");
+            }
+
+            // determine the result relation type.
+            // The result relation type of a table function consists of:
+            // 1. passed columns from input tables:
+            // - for tables with the "pass through columns" option, these are all columns of the table,
+            // - for tables without the "pass through columns" option, these are the partitioning columns of the table, if any.
+            // 2. columns created by the table function, called the proper columns.
+            ReturnTypeSpecification returnTypeSpecification = function.getReturnTypeSpecification();
+            Optional analyzedProperColumnsDescriptor = functionAnalysis.getReturnedType();
+            Descriptor properColumnsDescriptor;
+            switch (returnTypeSpecification.getReturnType()) {
+                case ReturnTypeSpecification.OnlyPassThrough.returnType:
+                    throw new SemanticException(NOT_IMPLEMENTED, node, "Returning only pass through columns is not yet supported for table functions");
+                case ReturnTypeSpecification.GenericTable.returnType:
+                    properColumnsDescriptor = analyzedProperColumnsDescriptor
+                            .orElseThrow(() -> new SemanticException(MISSING_RETURN_TYPE, node, "Cannot determine returned relation type for table function " + node.getName()));
+                    break;
+                default:
+                    // returned type is statically declared at function declaration and cannot be overridden
+                    if (analyzedProperColumnsDescriptor.isPresent()) {
+                        throw new SemanticException(AMBIGUOUS_RETURN_TYPE, node, "Returned relation type for table function %s is ambiguous", node.getName());
+                    }
+                    properColumnsDescriptor = ((ReturnTypeSpecification.DescribedTable) returnTypeSpecification).getDescriptor();
+            }
+            // currently we don't support input tables, so the output consists of proper columns only
+            List fields = properColumnsDescriptor.getFields().stream()
+                    // per spec, field names are mandatory
+                    .map(field -> Field.newUnqualified((field.getName()), field.getType().orElseThrow(() -> new IllegalStateException("missing returned type for proper field"))))
+                    .collect(toImmutableList());
+
+            return createAndAssignScope(node, scope, fields);
+        }
+
+        private Map analyzeArguments(Node node, List argumentSpecifications, List arguments)
+        {
+            Node errorLocation = node;
+            if (!arguments.isEmpty()) {
+                errorLocation = arguments.get(0);
+            }
+
+            if (argumentSpecifications.size() < arguments.size()) {
+                throw new SemanticException(INVALID_ARGUMENTS, errorLocation, "Too many arguments. Expected at most %s arguments, got %s arguments", argumentSpecifications.size(), arguments.size());
+            }
+
+            if (argumentSpecifications.isEmpty()) {
+                return ImmutableMap.of();
+            }
+
+            boolean argumentsPassedByName = !arguments.isEmpty() && arguments.stream().allMatch(argument -> argument.getName().isPresent());
+            boolean argumentsPassedByPosition = arguments.stream().allMatch(argument -> !argument.getName().isPresent());
+            if (!argumentsPassedByName && !argumentsPassedByPosition) {
+                throw new SemanticException(INVALID_ARGUMENTS, errorLocation, "All arguments must be passed by name or all must be passed positionally");
+            }
+
+            if (argumentsPassedByName) {
+                return mapTableFunctionsArgsByName(argumentSpecifications, arguments, errorLocation);
+            }
+            else {
+                return mapTableFunctionArgsByPosition(argumentSpecifications, arguments, errorLocation);
+            }
+        }
+
+        private Map mapTableFunctionsArgsByName(List argumentSpecifications, List arguments, Node errorLocation)
+        {
+            ImmutableMap.Builder passedArguments = ImmutableMap.builder();
+            Map argumentSpecificationsByName = new HashMap<>();
+            for (ArgumentSpecification argumentSpecification : argumentSpecifications) {
+                if (argumentSpecificationsByName.put(argumentSpecification.getName(), argumentSpecification) != null) {
+                    // this should never happen, because the argument names are validated at function registration time
+                    throw new IllegalStateException("Duplicate argument specification for name: " + argumentSpecification.getName());
+                }
+            }
+            Set uniqueArgumentNames = new HashSet<>();
+            for (TableFunctionArgument argument : arguments) {
+                String argumentName = argument.getName().orElseThrow(() -> new IllegalStateException("Missing table function argument name")).getCanonicalValue();
+                if (!uniqueArgumentNames.add(argumentName)) {
+                    throw new SemanticException(INVALID_FUNCTION_ARGUMENT, argument, "Duplicate argument name: %s", argumentName);
+                }
+                ArgumentSpecification argumentSpecification = argumentSpecificationsByName.remove(argumentName);
+                if (argumentSpecification == null) {
+                    throw new SemanticException(INVALID_FUNCTION_ARGUMENT, argument, "Unexpected argument name: %s", argumentName);
+                }
+                passedArguments.put(argumentSpecification.getName(), analyzeArgument(argumentSpecification, argument));
+            }
+            // apply defaults for not specified arguments
+            for (Map.Entry entry : argumentSpecificationsByName.entrySet()) {
+                ArgumentSpecification argumentSpecification = entry.getValue();
+                passedArguments.put(argumentSpecification.getName(), analyzeDefault(argumentSpecification, errorLocation));
+            }
+            return passedArguments.build();
+        }
+
+        private Map mapTableFunctionArgsByPosition(List argumentSpecifications, List arguments, Node errorLocation)
+        {
+            ImmutableMap.Builder passedArguments = ImmutableMap.builder();
+            for (int i = 0; i < arguments.size(); i++) {
+                TableFunctionArgument argument = arguments.get(i);
+                ArgumentSpecification argumentSpecification = argumentSpecifications.get(i); // TODO args passed positionally - can one only pass some prefix of args?
+                passedArguments.put(argumentSpecification.getName(), analyzeArgument(argumentSpecification, argument));
+            }
+            // apply defaults for not specified arguments
+            for (int i = arguments.size(); i < argumentSpecifications.size(); i++) {
+                ArgumentSpecification argumentSpecification = argumentSpecifications.get(i);
+                passedArguments.put(argumentSpecification.getName(), analyzeDefault(argumentSpecification, errorLocation));
+            }
+            return passedArguments.build();
+        }
+
+        private Argument analyzeArgument(ArgumentSpecification argumentSpecification, TableFunctionArgument argument)
+        {
+            String actualType = getArgumentTypeString(argument);
+            switch (argumentSpecification.getArgumentType()){
+                case TableArgumentSpecification.argumentType:
+                    return analyzeTableArgument(argument, argumentSpecification, actualType);
+                case DescriptorArgumentSpecification.argumentType:
+                    return analyzeDescriptorArgument(argument, argumentSpecification, actualType);
+                case ScalarArgumentSpecification.argumentType:
+                    return analyzeScalarArgument(argument, argumentSpecification, actualType);
+                default:
+                    throw new IllegalStateException("Unexpected argument specification: " + argumentSpecification.getClass().getSimpleName());
+            }
+        }
+
+        private Argument analyzeDefault(ArgumentSpecification argumentSpecification, Node errorLocation)
+        {
+            if (argumentSpecification.isRequired()) {
+                throw new SemanticException(MISSING_ARGUMENT, errorLocation, "Missing argument: " + argumentSpecification.getName());
+            }
+
+            checkArgument(!(argumentSpecification instanceof TableArgumentSpecification), "invalid table argument specification: default set");
+
+            if (argumentSpecification instanceof DescriptorArgumentSpecification) {
+                throw new SemanticException(NOT_IMPLEMENTED, errorLocation, "Descriptor arguments are not yet supported for table functions");
+            }
+            if (argumentSpecification instanceof ScalarArgumentSpecification) {
+                return ScalarArgument.builder()
+                        .type(((ScalarArgumentSpecification) argumentSpecification).getType())
+                        .value(argumentSpecification.getDefaultValue())
+                        .build();
+            }
+
+            throw new IllegalStateException("Unexpected argument specification: " + argumentSpecification.getClass().getSimpleName());
+        }
+
+        private String getArgumentTypeString(TableFunctionArgument argument)
+        {
+            try {
+                return argument.getValue().getArgumentTypeString();
+            }
+            catch (IllegalArgumentException e) {
+                throw new SemanticException(INVALID_FUNCTION_ARGUMENT, argument, "Unexpected table function argument type: ", argument.getClass().getSimpleName());
+            }
+        }
+
+        private Argument analyzeScalarArgument(TableFunctionArgument argument, ArgumentSpecification argumentSpecification, String actualType)
+        {
+            Type type = ((ScalarArgumentSpecification) argumentSpecification).getType();
+            if (!(argument.getValue() instanceof Expression)) {
+                throw new SemanticException(INVALID_FUNCTION_ARGUMENT, argument, "Invalid argument %s. Expected expression, got %s", argumentSpecification.getName(), actualType);
+            }
+            Expression expression = (Expression) argument.getValue();
+            // 'descriptor' as a function name is not allowed in this context
+            if (expression instanceof FunctionCall && ((FunctionCall) expression).getName().hasSuffix(QualifiedName.of("descriptor"))) { // function name is always compared case-insensitive
+                throw new SemanticException(INVALID_FUNCTION_ARGUMENT, argument, "'descriptor' function is not allowed as a table function argument");
+            }
+            // inline parameters
+            Expression inlined = ExpressionTreeRewriter.rewriteWith(new ExpressionRewriter()
+            {
+                @Override
+                public Expression rewriteParameter(Parameter node, Void context, ExpressionTreeRewriter treeRewriter)
+                {
+                    if (analysis.isDescribe()) {
+                        // We cannot handle DESCRIBE when a table function argument involves a parameter.
+                        // In DESCRIBE, the parameter values are not known. We cannot pass a dummy value for a parameter.
+                        // The value of a table function argument can affect the returned relation type. The returned
+                        // relation type can affect the assumed types for other parameters in the query.
+                        throw new SemanticException(NOT_SUPPORTED, node, "DESCRIBE is not supported if a table function uses parameters");
+                    }
+                    return analysis.getParameters().get(NodeRef.of(node));
+                }
+            }, expression);
+            // currently, only constant arguments are supported
+            Object constantValue = ExpressionInterpreter.evaluateConstantExpression(inlined, type, metadata, session, analysis.getParameters());
+            return ScalarArgument.builder()
+                    .type(type)
+                    .value(constantValue)
+                    .build();
+        }
+
+        private Argument analyzeTableArgument(TableFunctionArgument argument, ArgumentSpecification argumentSpecification, String actualType)
+        {
+            if (!(argument.getValue() instanceof TableFunctionTableArgument)) {
+                if (argument.getValue() instanceof FunctionCall) {
+                    // probably an attempt to pass a table function call, which is not supported, and was parsed as a function call
+                    throw new SemanticException(NOT_IMPLEMENTED, argument, "Invalid table argument %s. Table functions are not allowed as table function arguments", argumentSpecification.getName());
+                }
+                throw new SemanticException(INVALID_FUNCTION_ARGUMENT, argument, "Invalid argument %s. Expected table, got %s", argumentSpecification.getName(), actualType);
+            }
+            // TODO analyze the argument
+            // 1. process the Relation
+            // 2. partitioning and ordering must only apply to tables with set semantics
+            // 3. validate partitioning and ordering using `validateAndGetInputField()`
+            // 4. validate the prune when empty property vs argument specification (forbidden for row semantics; override? -> check spec)
+            // 5. return Argument
+            throw new SemanticException(NOT_IMPLEMENTED, argument, "Table arguments are not yet supported for table functions");
+        }
+
+        private Argument analyzeDescriptorArgument(TableFunctionArgument argument, ArgumentSpecification argumentSpecification, String actualType)
+        {
+            if (!(argument.getValue() instanceof TableFunctionDescriptorArgument)) {
+                if (argument.getValue() instanceof FunctionCall && ((FunctionCall) argument.getValue()).getName().hasSuffix(QualifiedName.of("descriptor"))) { // function name is always compared case-insensitive
+                    // malformed descriptor which parsed as a function call
+                    throw new SemanticException(INVALID_FUNCTION_ARGUMENT, argument, "Invalid descriptor argument %s. Descriptors should be formatted as 'DESCRIPTOR(name [type], ...)'", (Object) argumentSpecification.getName());
+                }
+                throw new SemanticException(INVALID_FUNCTION_ARGUMENT, argument, "Invalid argument %s. Expected descriptor, got %s", argumentSpecification.getName(), actualType);
+            }
+            throw new SemanticException(NOT_IMPLEMENTED, argument, "Descriptor arguments are not yet supported for table functions");
+        }
+
         @Override
         protected Scope visitTable(Table table, Optional scope)
         {
@@ -1458,7 +1743,7 @@ private Optional processTableVersion(Table table, QualifiedObjectNa
             analysis.recordSubqueries(table, expressionAnalysis);
             Type stateExprType = expressionAnalysis.getType(stateExpr);
             if (stateExprType == UNKNOWN) {
-                throw new PrestoException(INVALID_ARGUMENTS, format("Table version AS OF/BEFORE expression cannot be NULL for %s", name.toString()));
+                throw new PrestoException(StandardErrorCode.INVALID_ARGUMENTS, format("Table version AS OF/BEFORE expression cannot be NULL for %s", name.toString()));
             }
             Object evalStateExpr = evaluateConstantExpression(stateExpr, stateExprType, metadata, session, analysis.getParameters());
             if (tableVersionType == TIMESTAMP) {
diff --git a/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/MockConnectorColumnHandle.java b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/MockConnectorColumnHandle.java
new file mode 100644
index 0000000000000..84fbd45bfa9b9
--- /dev/null
+++ b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/MockConnectorColumnHandle.java
@@ -0,0 +1,81 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.connector.tvf;
+
+import com.facebook.presto.common.type.Type;
+import com.facebook.presto.spi.ColumnHandle;
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+import java.util.Objects;
+
+import static com.google.common.base.MoreObjects.toStringHelper;
+import static java.util.Objects.requireNonNull;
+
+public class MockConnectorColumnHandle
+        implements ColumnHandle
+{
+    private final String name;
+    private final Type type;
+
+    @JsonCreator
+    public MockConnectorColumnHandle(
+            @JsonProperty("name") String name,
+            @JsonProperty("type") Type type)
+    {
+        this.name = requireNonNull(name, "name is null");
+        this.type = requireNonNull(type, "type is null");
+    }
+
+    @JsonProperty
+    public String getName()
+    {
+        return name;
+    }
+
+    @JsonProperty
+    public Type getType()
+    {
+        return type;
+    }
+
+    @Override
+    public String toString()
+    {
+        return toStringHelper(this)
+                .add("name", name)
+                .add("type", type)
+                .toString();
+    }
+
+    @Override
+    public boolean equals(Object o)
+    {
+        if (this == o) {
+            return true;
+        }
+        if ((o == null) || (getClass() != o.getClass())) {
+            return false;
+        }
+        MockConnectorColumnHandle other = (MockConnectorColumnHandle) o;
+        return Objects.equals(name, other.name) &&
+                Objects.equals(type, other.type);
+    }
+
+    @Override
+    public int hashCode()
+    {
+        return Objects.hash(name, type);
+    }
+}
diff --git a/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/MockConnectorTableHandle.java b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/MockConnectorTableHandle.java
new file mode 100644
index 0000000000000..b29826b684dd7
--- /dev/null
+++ b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/MockConnectorTableHandle.java
@@ -0,0 +1,98 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.connector.tvf;
+
+import com.facebook.presto.common.predicate.TupleDomain;
+import com.facebook.presto.spi.ColumnHandle;
+import com.facebook.presto.spi.ConnectorTableHandle;
+import com.facebook.presto.spi.SchemaTableName;
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.google.common.collect.ImmutableList;
+
+import java.util.List;
+import java.util.Objects;
+import java.util.Optional;
+
+import static java.util.Objects.requireNonNull;
+
+public class MockConnectorTableHandle
+        implements ConnectorTableHandle
+{
+    private final SchemaTableName tableName;
+    private final TupleDomain constraint;
+    private final Optional> columns;
+
+    public MockConnectorTableHandle(SchemaTableName tableName)
+    {
+        this(tableName, TupleDomain.all(), Optional.empty());
+    }
+
+    @JsonCreator
+    public MockConnectorTableHandle(
+            @JsonProperty SchemaTableName tableName,
+            @JsonProperty("constraint") TupleDomain constraint,
+            @JsonProperty("columns") Optional> columns)
+    {
+        this.tableName = requireNonNull(tableName, "tableName is null");
+        this.constraint = requireNonNull(constraint, "constraint is null");
+        requireNonNull(columns, "columns is null");
+        this.columns = columns.map(ImmutableList::copyOf);
+    }
+
+    @JsonProperty
+    public SchemaTableName getTableName()
+    {
+        return tableName;
+    }
+
+    @JsonProperty
+    public TupleDomain getConstraint()
+    {
+        return constraint;
+    }
+
+    @JsonProperty
+    public Optional> getColumns()
+    {
+        return columns;
+    }
+
+    @Override
+    public boolean equals(Object o)
+    {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        MockConnectorTableHandle other = (MockConnectorTableHandle) o;
+        return Objects.equals(tableName, other.tableName) &&
+                Objects.equals(constraint, other.constraint) &&
+                Objects.equals(columns, other.columns);
+    }
+
+    @Override
+    public int hashCode()
+    {
+        return Objects.hash(tableName, constraint, columns);
+    }
+
+    @Override
+    public String toString()
+    {
+        return tableName.toString();
+    }
+}
diff --git a/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestingTableFunctions.java b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestingTableFunctions.java
new file mode 100644
index 0000000000000..1986401eac75a
--- /dev/null
+++ b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestingTableFunctions.java
@@ -0,0 +1,338 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.connector.tvf;
+
+import com.facebook.presto.common.predicate.TupleDomain;
+import com.facebook.presto.spi.ConnectorSession;
+import com.facebook.presto.spi.SchemaTableName;
+import com.facebook.presto.spi.connector.ConnectorTransactionHandle;
+import com.facebook.presto.spi.function.SchemaFunctionName;
+import com.facebook.presto.spi.function.table.AbstractConnectorTableFunction;
+import com.facebook.presto.spi.function.table.Argument;
+import com.facebook.presto.spi.function.table.ConnectorTableFunctionHandle;
+import com.facebook.presto.spi.function.table.Descriptor;
+import com.facebook.presto.spi.function.table.DescriptorArgumentSpecification;
+import com.facebook.presto.spi.function.table.ReturnTypeSpecification;
+import com.facebook.presto.spi.function.table.ScalarArgument;
+import com.facebook.presto.spi.function.table.ScalarArgumentSpecification;
+import com.facebook.presto.spi.function.table.TableArgumentSpecification;
+import com.facebook.presto.spi.function.table.TableFunctionAnalysis;
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonInclude;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.google.common.collect.ImmutableList;
+import io.airlift.slice.Slice;
+
+import java.util.Arrays;
+import java.util.Map;
+import java.util.Optional;
+
+import static com.facebook.presto.common.type.BigintType.BIGINT;
+import static com.facebook.presto.common.type.BooleanType.BOOLEAN;
+import static com.facebook.presto.common.type.IntegerType.INTEGER;
+import static com.facebook.presto.common.type.VarcharType.VARCHAR;
+import static com.facebook.presto.spi.function.table.ReturnTypeSpecification.GenericTable.GENERIC_TABLE;
+import static io.airlift.slice.Slices.utf8Slice;
+import static java.util.Objects.requireNonNull;
+
+public class TestingTableFunctions
+{
+    private static final String SCHEMA_NAME = "system";
+    private static final String TABLE_NAME = "table";
+    private static final String COLUMN_NAME = "column";
+    private static final ConnectorTableFunctionHandle HANDLE = new TestingTableFunctionPushdownHandle();
+    private static final TableFunctionAnalysis ANALYSIS = TableFunctionAnalysis.builder()
+            .handle(HANDLE)
+            .returnedType(new Descriptor(ImmutableList.of(new Descriptor.Field(COLUMN_NAME, Optional.of(BOOLEAN)))))
+            .build();
+
+    public static class TestConnectorTableFunction
+            extends AbstractConnectorTableFunction
+    {
+        private static final String TEST_FUNCTION = "test_function";
+
+        public TestConnectorTableFunction()
+        {
+            super(SCHEMA_NAME, TEST_FUNCTION, ImmutableList.of(), ReturnTypeSpecification.GenericTable.GENERIC_TABLE);
+        }
+
+        @Override
+        public TableFunctionAnalysis analyze(ConnectorSession session, ConnectorTransactionHandle transaction, Map arguments)
+        {
+            return TableFunctionAnalysis.builder()
+                    .handle(new TestingTableFunctionHandle(new SchemaFunctionName(SCHEMA_NAME, TEST_FUNCTION)))
+                    .returnedType(new Descriptor(ImmutableList.of(new Descriptor.Field("c1", Optional.of(BOOLEAN)))))
+                    .build();
+        }
+    }
+
+    public static class TestConnectorTableFunction2
+            extends AbstractConnectorTableFunction
+    {
+        private static final String TEST_FUNCTION_2 = "test_function2";
+
+        public TestConnectorTableFunction2()
+        {
+            super(SCHEMA_NAME, TEST_FUNCTION_2, ImmutableList.of(), ReturnTypeSpecification.OnlyPassThrough.ONLY_PASS_THROUGH);
+        }
+
+        @Override
+        public TableFunctionAnalysis analyze(ConnectorSession session, ConnectorTransactionHandle transaction, Map arguments)
+        {
+            return null;
+        }
+    }
+
+    public static class NullArgumentsTableFunction
+            extends AbstractConnectorTableFunction
+    {
+        private static final String NULL_ARGUMENTS_FUNCTION = "null_arguments_function";
+
+        public NullArgumentsTableFunction()
+        {
+            super(SCHEMA_NAME, NULL_ARGUMENTS_FUNCTION, null, ReturnTypeSpecification.OnlyPassThrough.ONLY_PASS_THROUGH);
+        }
+
+        @Override
+        public TableFunctionAnalysis analyze(ConnectorSession session, ConnectorTransactionHandle transaction, Map arguments)
+        {
+            return null;
+        }
+    }
+
+    public static class DuplicateArgumentsTableFunction
+            extends AbstractConnectorTableFunction
+    {
+        private static final String DUPLICATE_ARGUMENTS_FUNCTION = "duplicate_arguments_function";
+        public DuplicateArgumentsTableFunction()
+        {
+            super(
+                    SCHEMA_NAME,
+                    DUPLICATE_ARGUMENTS_FUNCTION,
+                    ImmutableList.of(
+                            ScalarArgumentSpecification.builder().name("a").type(INTEGER).build(),
+                            ScalarArgumentSpecification.builder().name("a").type(INTEGER).build()),
+                    ReturnTypeSpecification.OnlyPassThrough.ONLY_PASS_THROUGH);
+        }
+
+        @Override
+        public TableFunctionAnalysis analyze(ConnectorSession session, ConnectorTransactionHandle transaction, Map arguments)
+        {
+            return null;
+        }
+    }
+
+    public static class MultipleRSTableFunction
+            extends AbstractConnectorTableFunction
+    {
+        private static final String MULTIPLE_SOURCES_FUNCTION = "multiple_sources_function";
+        public MultipleRSTableFunction()
+        {
+            super(
+                    SCHEMA_NAME,
+                    MULTIPLE_SOURCES_FUNCTION,
+                    ImmutableList.of(TableArgumentSpecification.builder().name("t").rowSemantics().build(),
+                            TableArgumentSpecification.builder().name("t2").rowSemantics().build()),
+                    ReturnTypeSpecification.OnlyPassThrough.ONLY_PASS_THROUGH);
+        }
+
+        @Override
+        public TableFunctionAnalysis analyze(ConnectorSession session, ConnectorTransactionHandle transaction, Map arguments)
+        {
+            return null;
+        }
+    }
+
+    /**
+     * A table function returning a table with single empty column of type BOOLEAN.
+     * The argument `COLUMN` is the column name.
+     * The argument `IGNORED` is ignored.
+     * Both arguments are optional.
+     */
+    public static class SimpleTableFunction
+            extends AbstractConnectorTableFunction
+    {
+        private static final String SCHEMA_NAME = "system";
+        private static final String FUNCTION_NAME = "simple_table_function";
+        private static final String TABLE_NAME = "simple_table";
+
+        public SimpleTableFunction()
+        {
+            super(
+                    SCHEMA_NAME,
+                    FUNCTION_NAME,
+                    Arrays.asList(
+                            ScalarArgumentSpecification.builder()
+                                    .name("COLUMN")
+                                    .type(VARCHAR)
+                                    .defaultValue(utf8Slice("col"))
+                                    .build(),
+                            ScalarArgumentSpecification.builder()
+                                    .name("IGNORED")
+                                    .type(BIGINT)
+                                    .defaultValue(0L)
+                                    .build()),
+                    GENERIC_TABLE);
+        }
+
+        @Override
+        public TableFunctionAnalysis analyze(ConnectorSession session, ConnectorTransactionHandle transaction, Map arguments)
+        {
+            ScalarArgument argument = (ScalarArgument) arguments.get("COLUMN");
+            String columnName = ((Slice) argument.getValue()).toStringUtf8();
+
+            return TableFunctionAnalysis.builder()
+                    .handle(new SimpleTableFunctionHandle(getSchema(), TABLE_NAME, columnName))
+                    .returnedType(new Descriptor(ImmutableList.of(new Descriptor.Field(columnName, Optional.of(BOOLEAN)))))
+                    .build();
+        }
+
+        public static class SimpleTableFunctionHandle
+                implements ConnectorTableFunctionHandle
+        {
+            private final MockConnectorTableHandle tableHandle;
+
+            public SimpleTableFunctionHandle(String schema, String table, String column)
+            {
+                this.tableHandle = new MockConnectorTableHandle(
+                        new SchemaTableName(schema, table),
+                        TupleDomain.all(),
+                        Optional.of(ImmutableList.of(new MockConnectorColumnHandle(column, BOOLEAN))));
+            }
+
+            public MockConnectorTableHandle getTableHandle()
+            {
+                return tableHandle;
+            }
+        }
+    }
+
+    public static class TwoScalarArgumentsFunction
+            extends AbstractConnectorTableFunction
+    {
+        public TwoScalarArgumentsFunction()
+        {
+            super(
+                    SCHEMA_NAME,
+                    "two_arguments_function",
+                    ImmutableList.of(
+                            ScalarArgumentSpecification.builder()
+                                    .name("TEXT")
+                                    .type(VARCHAR)
+                                    .build(),
+                            ScalarArgumentSpecification.builder()
+                                    .name("NUMBER")
+                                    .type(BIGINT)
+                                    .defaultValue(null)
+                                    .build()),
+                    GENERIC_TABLE);
+        }
+
+        @Override
+        public TableFunctionAnalysis analyze(ConnectorSession session, ConnectorTransactionHandle transaction, Map arguments)
+        {
+            return ANALYSIS;
+        }
+    }
+
+    public static class TableArgumentFunction
+            extends AbstractConnectorTableFunction
+    {
+        public static final String FUNCTION_NAME = "table_argument_function";
+
+        public TableArgumentFunction()
+        {
+            super(
+                    SCHEMA_NAME,
+                    FUNCTION_NAME,
+                    ImmutableList.of(
+                            TableArgumentSpecification.builder()
+                                    .name("INPUT")
+                                    .keepWhenEmpty()
+                                    .build()),
+                    GENERIC_TABLE);
+        }
+
+        @Override
+        public TableFunctionAnalysis analyze(ConnectorSession session, ConnectorTransactionHandle transaction, Map arguments)
+        {
+            return TableFunctionAnalysis.builder()
+                    .handle(new TestingTableFunctionHandle(new SchemaFunctionName(SCHEMA_NAME, FUNCTION_NAME)))
+                    .returnedType(new Descriptor(ImmutableList.of(new Descriptor.Field(COLUMN_NAME, Optional.of(BOOLEAN)))))
+                    .requiredColumns("INPUT", ImmutableList.of(0))
+                    .build();
+        }
+    }
+
+    public static class DescriptorArgumentFunction
+            extends AbstractConnectorTableFunction
+    {
+        public DescriptorArgumentFunction()
+        {
+            super(
+                    SCHEMA_NAME,
+                    "descriptor_argument_function",
+                    ImmutableList.of(
+                            DescriptorArgumentSpecification.builder()
+                                    .name("SCHEMA")
+                                    .defaultValue(null)
+                                    .build()),
+                    GENERIC_TABLE);
+        }
+
+        @Override
+        public TableFunctionAnalysis analyze(ConnectorSession session, ConnectorTransactionHandle transaction, Map arguments)
+        {
+            return ANALYSIS;
+        }
+    }
+
+    public static class TestingTableFunctionPushdownHandle
+            implements ConnectorTableFunctionHandle
+    {
+        private final MockConnectorTableHandle tableHandle;
+
+        public TestingTableFunctionPushdownHandle()
+        {
+            this.tableHandle = new MockConnectorTableHandle(
+                    new SchemaTableName(SCHEMA_NAME, TABLE_NAME),
+                    TupleDomain.all(),
+                    Optional.of(ImmutableList.of(new MockConnectorColumnHandle(COLUMN_NAME, BOOLEAN))));
+        }
+
+        public MockConnectorTableHandle getTableHandle()
+        {
+            return tableHandle;
+        }
+    }
+
+    @JsonInclude(JsonInclude.Include.ALWAYS)
+    public static class TestingTableFunctionHandle
+            implements ConnectorTableFunctionHandle
+    {
+        private final SchemaFunctionName schemaFunctionName;
+
+        @JsonCreator
+        public TestingTableFunctionHandle(@JsonProperty("schemaFunctionName") SchemaFunctionName schemaFunctionName)
+        {
+            this.schemaFunctionName = requireNonNull(schemaFunctionName, "schemaFunctionName is null");
+        }
+
+        @JsonProperty
+        public SchemaFunctionName getSchemaFunctionName()
+        {
+            return schemaFunctionName;
+        }
+    }
+}
diff --git a/presto-main-base/src/test/java/com/facebook/presto/metadata/TestTableFunctionRegistry.java b/presto-main-base/src/test/java/com/facebook/presto/metadata/TestTableFunctionRegistry.java
index 439d7b35a7c9e..e7670a5af9b5d 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/metadata/TestTableFunctionRegistry.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/metadata/TestTableFunctionRegistry.java
@@ -15,25 +15,18 @@
 
 import com.facebook.presto.Session;
 import com.facebook.presto.spi.ConnectorId;
-import com.facebook.presto.spi.ConnectorSession;
-import com.facebook.presto.spi.connector.ConnectorTransactionHandle;
-import com.facebook.presto.spi.function.table.Argument;
-import com.facebook.presto.spi.function.table.ArgumentSpecification;
-import com.facebook.presto.spi.function.table.ConnectorTableFunction;
-import com.facebook.presto.spi.function.table.ReturnTypeSpecification;
-import com.facebook.presto.spi.function.table.ScalarArgumentSpecification;
-import com.facebook.presto.spi.function.table.TableArgumentSpecification;
-import com.facebook.presto.spi.function.table.TableFunctionAnalysis;
 import com.facebook.presto.spi.security.Identity;
 import com.facebook.presto.sql.tree.QualifiedName;
 import com.google.common.collect.ImmutableList;
 import org.testng.annotations.Test;
 
-import java.util.List;
-import java.util.Map;
 import java.util.Optional;
 
-import static com.facebook.presto.common.type.IntegerType.INTEGER;
+import static com.facebook.presto.connector.tvf.TestingTableFunctions.DuplicateArgumentsTableFunction;
+import static com.facebook.presto.connector.tvf.TestingTableFunctions.MultipleRSTableFunction;
+import static com.facebook.presto.connector.tvf.TestingTableFunctions.NullArgumentsTableFunction;
+import static com.facebook.presto.connector.tvf.TestingTableFunctions.TestConnectorTableFunction;
+import static com.facebook.presto.connector.tvf.TestingTableFunctions.TestConnectorTableFunction2;
 import static com.facebook.presto.testing.TestingSession.testSessionBuilder;
 import static com.facebook.presto.testing.assertions.Assert.assertEquals;
 import static org.testng.Assert.assertFalse;
@@ -44,7 +37,7 @@ public class TestTableFunctionRegistry
 {
     private static final String CATALOG = "test_catalog";
     private static final String USER = "user";
-    private static final String SCHEMA = "test_schema";
+    private static final String SCHEMA = "system";
     private static final Session SESSION = testSessionBuilder()
             .setCatalog(CATALOG)
             .setSchema(SCHEMA)
@@ -86,188 +79,14 @@ public void testTableFunctionRegistry()
 
         // Verify that null arguments table functions cannot be added.
         ex = expectThrows(NullPointerException.class, () -> testFunctionRegistry.addTableFunctions(id, ImmutableList.of(new NullArgumentsTableFunction())));
-        assertTrue(ex.getMessage().contains("table function arguments is null"), ex.getMessage());
+        assertTrue(ex.getMessage().contains("arguments is null"), ex.getMessage());
 
         // Verify that duplicate arguments table functions cannot be added.
         ex = expectThrows(IllegalArgumentException.class, () -> testFunctionRegistry.addTableFunctions(id, ImmutableList.of(new DuplicateArgumentsTableFunction())));
         assertTrue(ex.getMessage().contains("duplicate argument name: a"), ex.getMessage());
 
         // Verify that two row semantic table function arguments functions cannot be added.
-        ex = expectThrows(IllegalArgumentException.class, () -> testFunctionRegistry.addTableFunctions(id, ImmutableList.of(new MultipleRTSTableFunction())));
+        ex = expectThrows(IllegalArgumentException.class, () -> testFunctionRegistry.addTableFunctions(id, ImmutableList.of(new MultipleRSTableFunction())));
         assertTrue(ex.getMessage().contains("more than one table argument with row semantics"), ex.getMessage());
     }
-
-    private static class TestConnectorTableFunction
-            implements ConnectorTableFunction
-    {
-        @Override
-        public String getSchema()
-        {
-            return SCHEMA;
-        }
-
-        @Override
-        public String getName()
-        {
-            return TEST_FUNCTION;
-        }
-
-        @Override
-        public List getArguments()
-        {
-            return ImmutableList.of();
-        }
-
-        @Override
-        public ReturnTypeSpecification getReturnTypeSpecification()
-        {
-            return ReturnTypeSpecification.OnlyPassThrough.ONLY_PASS_THROUGH;
-        }
-
-        @Override
-        public TableFunctionAnalysis analyze(ConnectorSession session, ConnectorTransactionHandle transaction, Map arguments)
-        {
-            return null;
-        }
-    }
-
-    private static class TestConnectorTableFunction2
-            implements ConnectorTableFunction
-    {
-        @Override
-        public String getSchema()
-        {
-            return SCHEMA;
-        }
-
-        @Override
-        public String getName()
-        {
-            return TEST_FUNCTION_2;
-        }
-
-        @Override
-        public List getArguments()
-        {
-            return ImmutableList.of();
-        }
-
-        @Override
-        public ReturnTypeSpecification getReturnTypeSpecification()
-        {
-            return ReturnTypeSpecification.OnlyPassThrough.ONLY_PASS_THROUGH;
-        }
-
-        @Override
-        public TableFunctionAnalysis analyze(ConnectorSession session, ConnectorTransactionHandle transaction, Map arguments)
-        {
-            return null;
-        }
-    }
-
-    private static class NullArgumentsTableFunction
-            implements ConnectorTableFunction
-    {
-        @Override
-        public String getSchema()
-        {
-            return SCHEMA;
-        }
-
-        @Override
-        public String getName()
-        {
-            return TEST_FUNCTION;
-        }
-
-        @Override
-        public List getArguments()
-        {
-            return null;
-        }
-
-        @Override
-        public ReturnTypeSpecification getReturnTypeSpecification()
-        {
-            return ReturnTypeSpecification.OnlyPassThrough.ONLY_PASS_THROUGH;
-        }
-
-        @Override
-        public TableFunctionAnalysis analyze(ConnectorSession session, ConnectorTransactionHandle transaction, Map arguments)
-        {
-            return null;
-        }
-    }
-
-    private static class DuplicateArgumentsTableFunction
-            implements ConnectorTableFunction
-    {
-        @Override
-        public String getSchema()
-        {
-            return SCHEMA;
-        }
-
-        @Override
-        public String getName()
-        {
-            return TEST_FUNCTION;
-        }
-
-        @Override
-        public List getArguments()
-        {
-            ScalarArgumentSpecification s = ScalarArgumentSpecification.builder().name("a").type(INTEGER).build();
-            ScalarArgumentSpecification s2 = ScalarArgumentSpecification.builder().name("a").type(INTEGER).build();
-            return ImmutableList.of(s, s2);
-        }
-
-        @Override
-        public ReturnTypeSpecification getReturnTypeSpecification()
-        {
-            return ReturnTypeSpecification.OnlyPassThrough.ONLY_PASS_THROUGH;
-        }
-
-        @Override
-        public TableFunctionAnalysis analyze(ConnectorSession session, ConnectorTransactionHandle transaction, Map arguments)
-        {
-            return null;
-        }
-    }
-
-    private static class MultipleRTSTableFunction
-            implements ConnectorTableFunction
-    {
-        @Override
-        public String getSchema()
-        {
-            return SCHEMA;
-        }
-
-        @Override
-        public String getName()
-        {
-            return TEST_FUNCTION;
-        }
-
-        @Override
-        public List getArguments()
-        {
-            TableArgumentSpecification t = TableArgumentSpecification.builder().name("t").rowSemantics().build();
-            TableArgumentSpecification t2 = TableArgumentSpecification.builder().name("t2").rowSemantics().build();
-            return ImmutableList.of(t, t2);
-        }
-
-        @Override
-        public ReturnTypeSpecification getReturnTypeSpecification()
-        {
-            return ReturnTypeSpecification.OnlyPassThrough.ONLY_PASS_THROUGH;
-        }
-
-        @Override
-        public TableFunctionAnalysis analyze(ConnectorSession session, ConnectorTransactionHandle transaction, Map arguments)
-        {
-            return null;
-        }
-    }
 }
diff --git a/presto-main-base/src/test/java/com/facebook/presto/sql/analyzer/AbstractAnalyzerTest.java b/presto-main-base/src/test/java/com/facebook/presto/sql/analyzer/AbstractAnalyzerTest.java
index 3702dc249a00a..5819eadadcc6a 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/sql/analyzer/AbstractAnalyzerTest.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/sql/analyzer/AbstractAnalyzerTest.java
@@ -22,6 +22,10 @@
 import com.facebook.presto.common.type.StandardTypes;
 import com.facebook.presto.connector.informationSchema.InformationSchemaConnector;
 import com.facebook.presto.connector.system.SystemConnector;
+import com.facebook.presto.connector.tvf.TestingTableFunctions.DescriptorArgumentFunction;
+import com.facebook.presto.connector.tvf.TestingTableFunctions.SimpleTableFunction;
+import com.facebook.presto.connector.tvf.TestingTableFunctions.TableArgumentFunction;
+import com.facebook.presto.connector.tvf.TestingTableFunctions.TwoScalarArgumentsFunction;
 import com.facebook.presto.execution.warnings.WarningCollectorConfig;
 import com.facebook.presto.functionNamespace.SqlInvokedFunctionNamespaceManagerConfig;
 import com.facebook.presto.functionNamespace.execution.NoopSqlFunctionExecutor;
@@ -150,6 +154,13 @@ public void setup()
 
         metadata.getFunctionAndTypeManager().createFunction(SQL_FUNCTION_SQUARE, true);
 
+        metadata.getFunctionAndTypeManager().getTableFunctionRegistry().addTableFunctions(TPCH_CONNECTOR_ID,
+                ImmutableList.of(
+                        new SimpleTableFunction(),
+                        new TwoScalarArgumentsFunction(),
+                        new TableArgumentFunction(),
+                        new DescriptorArgumentFunction()));
+
         Catalog tpchTestCatalog = createTestingCatalog(TPCH_CATALOG, TPCH_CONNECTOR_ID);
         catalogManager.registerCatalog(tpchTestCatalog);
         metadata.getAnalyzePropertyManager().addProperties(TPCH_CONNECTOR_ID, tpchTestCatalog.getConnector(TPCH_CONNECTOR_ID).getAnalyzeProperties());
diff --git a/presto-main-base/src/test/java/com/facebook/presto/sql/analyzer/TestAnalyzer.java b/presto-main-base/src/test/java/com/facebook/presto/sql/analyzer/TestAnalyzer.java
index 48f442cf750ca..816bcd2a0aa5c 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/sql/analyzer/TestAnalyzer.java
+++ b/presto-main-base/src/test/java/com/facebook/presto/sql/analyzer/TestAnalyzer.java
@@ -44,6 +44,10 @@
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.DUPLICATE_COLUMN_NAME;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.DUPLICATE_PROPERTY;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.DUPLICATE_RELATION;
+import static com.facebook.presto.sql.analyzer.SemanticErrorCode.EXPRESSION_NOT_CONSTANT;
+import static com.facebook.presto.sql.analyzer.SemanticErrorCode.FUNCTION_NOT_FOUND;
+import static com.facebook.presto.sql.analyzer.SemanticErrorCode.INVALID_ARGUMENTS;
+import static com.facebook.presto.sql.analyzer.SemanticErrorCode.INVALID_FUNCTION_ARGUMENT;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.INVALID_FUNCTION_NAME;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.INVALID_LITERAL;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.INVALID_OFFSET_ROW_COUNT;
@@ -55,6 +59,7 @@
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.INVALID_WINDOW_FRAME;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.MISMATCHED_COLUMN_ALIASES;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.MISMATCHED_SET_COLUMN_TYPES;
+import static com.facebook.presto.sql.analyzer.SemanticErrorCode.MISSING_ARGUMENT;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.MISSING_ATTRIBUTE;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.MISSING_CATALOG;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.MISSING_COLUMN;
@@ -1942,4 +1947,70 @@ public void testInvalidTemporaryFunctionName()
         assertFails(INVALID_FUNCTION_NAME, "CREATE TEMPORARY FUNCTION sum() RETURNS INT RETURN 1");
         assertFails(INVALID_FUNCTION_NAME, "CREATE TEMPORARY FUNCTION dev.test.foo() RETURNS INT RETURN 1");
     }
+
+    @Test
+    public void testTableFunctionNotFound()
+    {
+        assertFails(FUNCTION_NOT_FOUND,
+                "line 1:21: Table function non_existent_table_function not registered",
+                "SELECT * FROM TABLE(non_existent_table_function())");
+    }
+
+    @Test
+    public void testTableFunctionArguments()
+    {
+        assertFails(INVALID_ARGUMENTS, "line 1:51: Too many arguments. Expected at most 2 arguments, got 3 arguments", "SELECT * FROM TABLE(system.two_arguments_function(1, 2, 3))");
+
+        analyze("SELECT * FROM TABLE(system.two_arguments_function('foo'))");
+        analyze("SELECT * FROM TABLE(system.two_arguments_function(text => 'foo'))");
+        analyze("SELECT * FROM TABLE(system.two_arguments_function('foo', 1))");
+        analyze("SELECT * FROM TABLE(system.two_arguments_function(text => 'foo', number => 1))");
+
+        assertFails(INVALID_ARGUMENTS,
+                "line 1:51: All arguments must be passed by name or all must be passed positionally",
+                "SELECT * FROM TABLE(system.two_arguments_function('foo', number => 1))");
+
+        assertFails(INVALID_ARGUMENTS,
+                "line 1:51: All arguments must be passed by name or all must be passed positionally",
+                "SELECT * FROM TABLE(system.two_arguments_function(text => 'foo', 1))");
+
+        assertFails(INVALID_FUNCTION_ARGUMENT,
+                "line 1:66: Duplicate argument name: TEXT",
+                "SELECT * FROM TABLE(system.two_arguments_function(text => 'foo', text => 'bar'))");
+
+        // argument names are resolved in the canonical form
+        assertFails(INVALID_FUNCTION_ARGUMENT,
+                "line 1:66: Duplicate argument name: TEXT",
+                "SELECT * FROM TABLE(system.two_arguments_function(text => 'foo', TeXt => 'bar'))");
+
+        assertFails(INVALID_FUNCTION_ARGUMENT,
+                "line 1:66: Unexpected argument name: BAR",
+                "SELECT * FROM TABLE(system.two_arguments_function(text => 'foo', bar => 'bar'))");
+
+        assertFails(MISSING_ARGUMENT,
+                "line 1:51: Missing argument: TEXT",
+                "SELECT * FROM TABLE(system.two_arguments_function(number => 1))");
+    }
+
+    @Test
+    public void testScalarArgument()
+    {
+        analyze("SELECT * FROM TABLE(system.two_arguments_function('foo', 1))");
+
+        assertFails(INVALID_FUNCTION_ARGUMENT,
+                "line 1:64: Invalid argument NUMBER. Expected expression, got descriptor",
+                "SELECT * FROM TABLE(system.two_arguments_function(text => 'a', number => DESCRIPTOR(x integer, y boolean)))");
+
+        assertFails(INVALID_FUNCTION_ARGUMENT,
+                "line 1:64: 'descriptor' function is not allowed as a table function argument",
+                "SELECT * FROM TABLE(system.two_arguments_function(text => 'a', number => DESCRIPTOR(1 + 2)))");
+
+        assertFails(INVALID_FUNCTION_ARGUMENT,
+                "line 1:64: Invalid argument NUMBER. Expected expression, got table",
+                "SELECT * FROM TABLE(system.two_arguments_function(text => 'a', number => TABLE(t1)))");
+
+        assertFails(EXPRESSION_NOT_CONSTANT,
+                "line 1:74: Constant expression cannot contain a subquery",
+                "SELECT * FROM TABLE(system.two_arguments_function(text => 'a', number => (SELECT 1)))");
+    }
 }
diff --git a/presto-parser/src/main/java/com/facebook/presto/sql/tree/Expression.java b/presto-parser/src/main/java/com/facebook/presto/sql/tree/Expression.java
index 693852f8a8c39..48761444c2b29 100644
--- a/presto-parser/src/main/java/com/facebook/presto/sql/tree/Expression.java
+++ b/presto-parser/src/main/java/com/facebook/presto/sql/tree/Expression.java
@@ -39,4 +39,10 @@ public final String toString()
     {
         return ExpressionFormatter.formatExpression(this, Optional.empty()); // This will not replace parameters, but we don't have access to them here
     }
+
+    @Override
+    public String getArgumentTypeString()
+    {
+        return "expression";
+    }
 }
diff --git a/presto-parser/src/main/java/com/facebook/presto/sql/tree/Node.java b/presto-parser/src/main/java/com/facebook/presto/sql/tree/Node.java
index 2317c5258c1c9..b1c1354b09d8c 100644
--- a/presto-parser/src/main/java/com/facebook/presto/sql/tree/Node.java
+++ b/presto-parser/src/main/java/com/facebook/presto/sql/tree/Node.java
@@ -40,6 +40,11 @@ public Optional getLocation()
         return location;
     }
 
+    public String getArgumentTypeString()
+    {
+        throw new IllegalArgumentException("Unexpected table function argument type: " + this.getClass().getSimpleName());
+    }
+
     public abstract List getChildren();
 
     // Force subclasses to have a proper equals and hashcode implementation
diff --git a/presto-parser/src/main/java/com/facebook/presto/sql/tree/TableFunctionDescriptorArgument.java b/presto-parser/src/main/java/com/facebook/presto/sql/tree/TableFunctionDescriptorArgument.java
index 7c6150ee320d8..e762c99b0fa70 100644
--- a/presto-parser/src/main/java/com/facebook/presto/sql/tree/TableFunctionDescriptorArgument.java
+++ b/presto-parser/src/main/java/com/facebook/presto/sql/tree/TableFunctionDescriptorArgument.java
@@ -83,4 +83,10 @@ public String toString()
     {
         return descriptor.map(Descriptor::toString).orElse("CAST (NULL AS DESCRIPTOR)");
     }
+
+    @Override
+    public String getArgumentTypeString()
+    {
+        return "descriptor";
+    }
 }
diff --git a/presto-parser/src/main/java/com/facebook/presto/sql/tree/TableFunctionTableArgument.java b/presto-parser/src/main/java/com/facebook/presto/sql/tree/TableFunctionTableArgument.java
index b23da149cabae..ef0eb8bc20d96 100644
--- a/presto-parser/src/main/java/com/facebook/presto/sql/tree/TableFunctionTableArgument.java
+++ b/presto-parser/src/main/java/com/facebook/presto/sql/tree/TableFunctionTableArgument.java
@@ -120,4 +120,10 @@ public String toString()
 
         return builder.toString();
     }
+
+    @Override
+    public String getArgumentTypeString()
+    {
+        return "table";
+    }
 }
diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/function/table/ArgumentSpecification.java b/presto-spi/src/main/java/com/facebook/presto/spi/function/table/ArgumentSpecification.java
index 5443475c98ce0..73c822095863f 100644
--- a/presto-spi/src/main/java/com/facebook/presto/spi/function/table/ArgumentSpecification.java
+++ b/presto-spi/src/main/java/com/facebook/presto/spi/function/table/ArgumentSpecification.java
@@ -30,6 +30,7 @@
  */
 public abstract class ArgumentSpecification
 {
+    public static final String argumentType = "Abstract";
     private final String name;
     private final boolean required;
 
@@ -58,4 +59,9 @@ public Object getDefaultValue()
     {
         return defaultValue;
     }
+
+    public String getArgumentType()
+    {
+        return argumentType;
+    }
 }
diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/function/table/DescriptorArgumentSpecification.java b/presto-spi/src/main/java/com/facebook/presto/spi/function/table/DescriptorArgumentSpecification.java
index 637d25fac73be..11ed93f02cd00 100644
--- a/presto-spi/src/main/java/com/facebook/presto/spi/function/table/DescriptorArgumentSpecification.java
+++ b/presto-spi/src/main/java/com/facebook/presto/spi/function/table/DescriptorArgumentSpecification.java
@@ -16,6 +16,7 @@
 public class DescriptorArgumentSpecification
         extends ArgumentSpecification
 {
+    public static final String argumentType = "DescriptorArgumentSpecification";
     private DescriptorArgumentSpecification(String name, boolean required, Descriptor defaultValue)
     {
         super(name, required, defaultValue);
@@ -52,4 +53,10 @@ public DescriptorArgumentSpecification build()
             return new DescriptorArgumentSpecification(name, required, defaultValue);
         }
     }
+
+    @Override
+    public String getArgumentType()
+    {
+        return argumentType;
+    }
 }
diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/function/table/ReturnTypeSpecification.java b/presto-spi/src/main/java/com/facebook/presto/spi/function/table/ReturnTypeSpecification.java
index f93e9d2b601e3..2c77503cadada 100644
--- a/presto-spi/src/main/java/com/facebook/presto/spi/function/table/ReturnTypeSpecification.java
+++ b/presto-spi/src/main/java/com/facebook/presto/spi/function/table/ReturnTypeSpecification.java
@@ -27,6 +27,7 @@
  */
 public abstract class ReturnTypeSpecification
 {
+    public static final String returnType = "Abstract";
     /**
      * The proper columns of the table function are not known at function declaration time.
      * They must be determined at query analysis time based on the actual call arguments.
@@ -34,9 +35,16 @@ public abstract class ReturnTypeSpecification
     public static class GenericTable
             extends ReturnTypeSpecification
     {
+        public static final String returnType = "GenericTable";
         public static final GenericTable GENERIC_TABLE = new GenericTable();
 
         private GenericTable() {}
+
+        @Override
+        public String getReturnType()
+        {
+            return returnType;
+        }
     }
 
     /**
@@ -45,9 +53,16 @@ private GenericTable() {}
     public static class OnlyPassThrough
             extends ReturnTypeSpecification
     {
+        public static final String returnType = "OnlyPassThrough";
         public static final OnlyPassThrough ONLY_PASS_THROUGH = new OnlyPassThrough();
 
         private OnlyPassThrough() {}
+
+        @Override
+        public String getReturnType()
+        {
+            return returnType;
+        }
     }
 
     /**
@@ -57,6 +72,7 @@ private OnlyPassThrough() {}
     public static class DescribedTable
             extends ReturnTypeSpecification
     {
+        public static final String returnType = "DescribedTable";
         private final Descriptor descriptor;
 
         public DescribedTable(Descriptor descriptor)
@@ -70,5 +86,16 @@ public Descriptor getDescriptor()
         {
             return descriptor;
         }
+
+        @Override
+        public String getReturnType()
+        {
+            return returnType;
+        }
+    }
+
+    public String getReturnType()
+    {
+        return returnType;
     }
 }
diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/function/table/ScalarArgumentSpecification.java b/presto-spi/src/main/java/com/facebook/presto/spi/function/table/ScalarArgumentSpecification.java
index 6016b4fc9b5ce..94f98bafe949a 100644
--- a/presto-spi/src/main/java/com/facebook/presto/spi/function/table/ScalarArgumentSpecification.java
+++ b/presto-spi/src/main/java/com/facebook/presto/spi/function/table/ScalarArgumentSpecification.java
@@ -23,6 +23,7 @@
 public class ScalarArgumentSpecification
         extends ArgumentSpecification
 {
+    public static final String argumentType = "ScalarArgumentSpecification";
     private final Type type;
 
     private ScalarArgumentSpecification(String name, Type type, boolean required, Object defaultValue)
@@ -77,4 +78,10 @@ public ScalarArgumentSpecification build()
             return new ScalarArgumentSpecification(name, type, required, defaultValue);
         }
     }
+
+    @Override
+    public String getArgumentType()
+    {
+        return argumentType;
+    }
 }
diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/function/table/TableArgumentSpecification.java b/presto-spi/src/main/java/com/facebook/presto/spi/function/table/TableArgumentSpecification.java
index e1c1e638c6a9e..44c74258152dd 100644
--- a/presto-spi/src/main/java/com/facebook/presto/spi/function/table/TableArgumentSpecification.java
+++ b/presto-spi/src/main/java/com/facebook/presto/spi/function/table/TableArgumentSpecification.java
@@ -19,6 +19,7 @@
 public class TableArgumentSpecification
         extends ArgumentSpecification
 {
+    public static final String argumentType = "TableArgumentSpecification";
     private final boolean rowSemantics;
     private final boolean pruneWhenEmpty;
     private final boolean passThroughColumns;
@@ -107,4 +108,10 @@ public TableArgumentSpecification build()
             return new TableArgumentSpecification(name, rowSemantics, pruneWhenEmpty, passThroughColumns);
         }
     }
+
+    @Override
+    public String getArgumentType()
+    {
+        return argumentType;
+    }
 }

From 7b9613fc1135d562ef513a1f71ca34f310c3e4ee Mon Sep 17 00:00:00 2001
From: Amit Dutta 
Date: Sun, 7 Sep 2025 15:54:10 -0700
Subject: [PATCH 061/113] [native] Advance velox

---
 .../presto_cpp/main/connectors/SystemConnector.cpp       | 1 +
 .../presto_cpp/main/connectors/SystemConnector.h         | 5 +++++
 .../main/connectors/arrow_flight/ArrowFlightConnector.h  | 9 ++++++++-
 presto-native-execution/velox                            | 2 +-
 4 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/presto-native-execution/presto_cpp/main/connectors/SystemConnector.cpp b/presto-native-execution/presto_cpp/main/connectors/SystemConnector.cpp
index 2fc0b7338e351..d75dc5435c60b 100644
--- a/presto-native-execution/presto_cpp/main/connectors/SystemConnector.cpp
+++ b/presto-native-execution/presto_cpp/main/connectors/SystemConnector.cpp
@@ -74,6 +74,7 @@ SystemTableHandle::SystemTableHandle(
     std::string schemaName,
     std::string tableName)
     : ConnectorTableHandle(std::move(connectorId)),
+      name_(fmt::format("{}.{}", schemaName, tableName)),
       schemaName_(std::move(schemaName)),
       tableName_(std::move(tableName)) {
   VELOX_USER_CHECK_EQ(
diff --git a/presto-native-execution/presto_cpp/main/connectors/SystemConnector.h b/presto-native-execution/presto_cpp/main/connectors/SystemConnector.h
index 5e77af7a43c66..f491a4b3b98a9 100644
--- a/presto-native-execution/presto_cpp/main/connectors/SystemConnector.h
+++ b/presto-native-execution/presto_cpp/main/connectors/SystemConnector.h
@@ -42,6 +42,10 @@ class SystemTableHandle : public velox::connector::ConnectorTableHandle {
 
   std::string toString() const override;
 
+  const std::string& name() const override {
+    return name_;
+  }
+
   const std::string& schemaName() const {
     return schemaName_;
   }
@@ -53,6 +57,7 @@ class SystemTableHandle : public velox::connector::ConnectorTableHandle {
   const velox::RowTypePtr taskSchema() const;
 
  private:
+  const std::string name_;
   const std::string schemaName_;
   const std::string tableName_;
 };
diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/ArrowFlightConnector.h b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/ArrowFlightConnector.h
index 6123b270e2349..d1c431575c677 100644
--- a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/ArrowFlightConnector.h
+++ b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/ArrowFlightConnector.h
@@ -31,7 +31,14 @@ namespace facebook::presto {
 class ArrowFlightTableHandle : public velox::connector::ConnectorTableHandle {
  public:
   explicit ArrowFlightTableHandle(const std::string& connectorId)
-      : ConnectorTableHandle(connectorId) {}
+      : ConnectorTableHandle(connectorId), name_("arrow_flight") {}
+
+  const std::string& name() const override {
+    return name_;
+  }
+
+ private:
+  const std::string name_;
 };
 
 struct ArrowFlightSplit : public velox::connector::ConnectorSplit {
diff --git a/presto-native-execution/velox b/presto-native-execution/velox
index eac4b4bc61445..7992424ce4d78 160000
--- a/presto-native-execution/velox
+++ b/presto-native-execution/velox
@@ -1 +1 @@
-Subproject commit eac4b4bc61445e58436dfaa6fa907353308db1e5
+Subproject commit 7992424ce4d7886681270019b1729bb44a91e7ab

From ba13f42f298a93411e2618aa719f643ee268b539 Mon Sep 17 00:00:00 2001
From: Nishitha-Bhaskaran 
Date: Fri, 29 Aug 2025 15:08:40 +0530
Subject: [PATCH 062/113] Upgrade objenesis version to 3.4

---
 presto-bigquery/pom.xml | 2 +-
 presto-hive/pom.xml     | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/presto-bigquery/pom.xml b/presto-bigquery/pom.xml
index 07ad6c02c7b80..611d8703d3226 100644
--- a/presto-bigquery/pom.xml
+++ b/presto-bigquery/pom.xml
@@ -342,7 +342,7 @@
         
             org.objenesis
             objenesis
-            2.6
+            3.4
             test
             
                 
diff --git a/presto-hive/pom.xml b/presto-hive/pom.xml
index 4b855fb5df084..c5bc57b835548 100644
--- a/presto-hive/pom.xml
+++ b/presto-hive/pom.xml
@@ -485,7 +485,7 @@
         
             org.objenesis
             objenesis
-            2.6
+            3.4
             test
             
                 

From a642598233b2294e2b45c366c2e438c06470f35a Mon Sep 17 00:00:00 2001
From: sumi-mathew 
Date: Mon, 8 Sep 2025 20:51:41 +0530
Subject: [PATCH 063/113] Upgrade mongo server version

---
 presto-mongodb/pom.xml                                    | 2 +-
 .../com/facebook/presto/mongodb/SyncMemoryBackend.java    | 8 ++++----
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/presto-mongodb/pom.xml b/presto-mongodb/pom.xml
index 8644f91fb7809..8b4d7b916f61e 100644
--- a/presto-mongodb/pom.xml
+++ b/presto-mongodb/pom.xml
@@ -15,7 +15,7 @@
     
         ${project.parent.basedir}
         3.12.14
-        1.5.0
+        1.47.0
         true
     
 
diff --git a/presto-mongodb/src/test/java/com/facebook/presto/mongodb/SyncMemoryBackend.java b/presto-mongodb/src/test/java/com/facebook/presto/mongodb/SyncMemoryBackend.java
index 828230a075424..c8c7abdcc392c 100644
--- a/presto-mongodb/src/test/java/com/facebook/presto/mongodb/SyncMemoryBackend.java
+++ b/presto-mongodb/src/test/java/com/facebook/presto/mongodb/SyncMemoryBackend.java
@@ -13,7 +13,7 @@
  */
 package com.facebook.presto.mongodb;
 
-import de.bwaldvogel.mongo.MongoBackend;
+import de.bwaldvogel.mongo.backend.CursorRegistry;
 import de.bwaldvogel.mongo.backend.memory.MemoryBackend;
 import de.bwaldvogel.mongo.backend.memory.MemoryDatabase;
 import de.bwaldvogel.mongo.exception.MongoServerException;
@@ -25,16 +25,16 @@ public class SyncMemoryBackend
     public MemoryDatabase openOrCreateDatabase(String databaseName)
             throws MongoServerException
     {
-        return new SyncMemoryDatabase(this, databaseName);
+        return new SyncMemoryDatabase(databaseName, this.getCursorRegistry());
     }
 
     private static class SyncMemoryDatabase
             extends MemoryDatabase
     {
-        public SyncMemoryDatabase(MongoBackend backend, String databaseName)
+        public SyncMemoryDatabase(String databaseName, CursorRegistry cursorRegistry)
                 throws MongoServerException
         {
-            super(backend, databaseName);
+            super(databaseName, cursorRegistry);
         }
     }
 }

From c9df88d3cfb5c9c7c003bb93a318e543eed209bb Mon Sep 17 00:00:00 2001
From: Li Zhou 
Date: Mon, 8 Sep 2025 18:23:10 +0100
Subject: [PATCH 064/113] Update 0.294 release notes about executable jars
 (#25982)

## Description
Update 0.294 release notes about executable jars, related to PR
https://github.com/prestodb/prestodb.github.io/pull/303

## Motivation and Context
https://github.com/prestodb/prestodb.github.io/pull/303

## Impact
0.294 release notes

## Test Plan
N/A

## Contributor checklist

- [ ] Please make sure your submission complies with our [contributing
guide](https://github.com/prestodb/presto/blob/master/CONTRIBUTING.md),
in particular [code
style](https://github.com/prestodb/presto/blob/master/CONTRIBUTING.md#code-style)
and [commit
standards](https://github.com/prestodb/presto/blob/master/CONTRIBUTING.md#commit-standards).
- [ ] PR description addresses the issue accurately and concisely. If
the change is non-trivial, a GitHub Issue is referenced.
- [ ] Documented new properties (with its default value), SQL syntax,
functions, or other functionality.
- [ ] If release notes are required, they follow the [release notes
guidelines](https://github.com/prestodb/presto/wiki/Release-Notes-Guidelines).
- [ ] Adequate tests were added if applicable.
- [ ] CI passed.

## Release Notes
Please follow [release notes
guidelines](https://github.com/prestodb/presto/wiki/Release-Notes-Guidelines)
and fill in the release notes below.

```
== NO RELEASE NOTE ==
```

## Summary by Sourcery

Documentation:
- Add details on building and running executable JARs in the 0.294
release notes

---------

Co-authored-by: sourcery-ai[bot] <58596630+sourcery-ai[bot]@users.noreply.github.com>
Co-authored-by: Steve Burnett 
---
 presto-docs/src/main/sphinx/release/release-0.294.rst | 1 +
 1 file changed, 1 insertion(+)

diff --git a/presto-docs/src/main/sphinx/release/release-0.294.rst b/presto-docs/src/main/sphinx/release/release-0.294.rst
index 5fdc161bef1d0..73b7d871ff7de 100644
--- a/presto-docs/src/main/sphinx/release/release-0.294.rst
+++ b/presto-docs/src/main/sphinx/release/release-0.294.rst
@@ -12,6 +12,7 @@ Release 0.294
 * Add mixed case support for schema and table names. `#24551 `_
 * Add case-sensitive support for column names. It can be enabled for JDBC based connector by setting ``case-sensitive-name-matching=true`` at the catalog level. `#24983 `_
 * Update ``presto-plan-checker-router-plugin router`` plugin to use ``EXPLAIN (TYPE VALIDATE)`` in place of ``EXPLAIN (TYPE DISTRIBUTED)``, enabling faster routing of queries to either native or Java clusters. `#25545 `_
+* From release 0.294, due to Maven Central publishing limitations, executable jar files including ``presto-cli``, ``presto-benchmark-driver``, and ``presto-test-server-launcher`` are no longer published in the Maven Central repository. These jars can now be found on the `Presto GitHub release page `_.
 
 **Details**
 ===========

From 6d36242049e4ed124e813760004fcc57a183bdcd Mon Sep 17 00:00:00 2001
From: Jialiang Tan 
Date: Sun, 7 Sep 2025 11:37:11 -0700
Subject: [PATCH 065/113] Make POS count failed task info and perform dedup on
 them

---
 .../PrestoSparkQueryExecutionFactory.java     |  2 +-
 .../AbstractPrestoSparkQueryExecution.java    | 66 +++++++++++++++++--
 2 files changed, 63 insertions(+), 5 deletions(-)

diff --git a/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkQueryExecutionFactory.java b/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkQueryExecutionFactory.java
index 042b56ce62e1a..8eef5d1db658b 100644
--- a/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkQueryExecutionFactory.java
+++ b/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkQueryExecutionFactory.java
@@ -679,7 +679,7 @@ else if (preparedQuery.isExplainTypeValidate()) {
                 planAndMore = queryPlanner.createQueryPlan(session, preparedQuery, warningCollector, variableAllocator, planNodeIdAllocator, sparkContext, sql);
                 JavaSparkContext javaSparkContext = new JavaSparkContext(sparkContext);
                 CollectionAccumulator taskInfoCollector = new CollectionAccumulator<>();
-                taskInfoCollector.register(sparkContext, Option.empty(), false);
+                taskInfoCollector.register(sparkContext, Option.empty(), true);
                 CollectionAccumulator shuffleStatsCollector = new CollectionAccumulator<>();
                 shuffleStatsCollector.register(sparkContext, Option.empty(), false);
                 TempStorage tempStorage = tempStorageManager.getTempStorage(storageBasedBroadcastJoinStorage);
diff --git a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/AbstractPrestoSparkQueryExecution.java b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/AbstractPrestoSparkQueryExecution.java
index 2e03dd64a54d1..b76da7f2ddc2e 100644
--- a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/AbstractPrestoSparkQueryExecution.java
+++ b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/AbstractPrestoSparkQueryExecution.java
@@ -29,7 +29,9 @@
 import com.facebook.presto.execution.QueryState;
 import com.facebook.presto.execution.QueryStateTimer;
 import com.facebook.presto.execution.StageInfo;
+import com.facebook.presto.execution.TaskId;
 import com.facebook.presto.execution.TaskInfo;
+import com.facebook.presto.execution.TaskState;
 import com.facebook.presto.execution.scheduler.ExecutionWriterTarget;
 import com.facebook.presto.execution.scheduler.StreamingPlanSection;
 import com.facebook.presto.execution.scheduler.StreamingSubPlan;
@@ -116,6 +118,7 @@
 import java.util.TreeMap;
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
 import static com.facebook.airlift.units.DataSize.Unit.BYTE;
@@ -572,23 +575,78 @@ protected void validateStorageCapabilities(TempStorage tempStorage)
         }
     }
 
+    /**
+     * Updates the taskInfoMap to ensure it stores the most relevant {@link TaskInfo} for each
+     * logical task, identified by task ID (excluding attempt number).
+     * 

+ * This method ensures that, for each logical task, the map retains the latest successful + * attempt if available, or otherwise the most recent attempt based on attempt number. Warnings + * are logged in cases of unexpected duplicate or multiple successful attempts. + * + * @param taskInfoMap the map from logical task ID (taskId excluding attempt number) to + * {@link TaskInfo} + * @param taskInfo the {@link TaskInfo} to consider for updating the map + */ + private void updateTaskInfoMap(HashMap taskInfoMap, TaskInfo taskInfo) + { + TaskId newTaskId = taskInfo.getTaskId(); + String taskIdWithoutAttemptId = new StringBuilder() + .append(newTaskId.getStageExecutionId().toString()) + .append(".") + .append(newTaskId.getId()) + .toString(); + if (!taskInfoMap.containsKey(taskIdWithoutAttemptId)) { + taskInfoMap.put(taskIdWithoutAttemptId, taskInfo); + return; + } + + TaskInfo storedTaskInfo = taskInfoMap.get(taskIdWithoutAttemptId); + TaskId storedTaskId = storedTaskInfo.getTaskId(); + TaskState storedTaskState = storedTaskInfo.getTaskStatus().getState(); + TaskState newTaskState = taskInfo.getTaskStatus().getState(); + if (storedTaskState == TaskState.FINISHED) { + if (newTaskState == TaskState.FINISHED) { + log.warn("Multiple attempts of the same task have succeeded %s vs %s", + storedTaskId.toString(), newTaskId.toString()); + } + // Successful one has been stored. Nothing needs to be done. + return; + } + + int storedAttemptNumber = storedTaskId.getAttemptNumber(); + int newAttemptNumber = newTaskId.getAttemptNumber(); + if (newTaskState == TaskState.FINISHED || storedAttemptNumber < newAttemptNumber) { + taskInfoMap.put(taskIdWithoutAttemptId, taskInfo); + } + if (storedAttemptNumber == newAttemptNumber) { + log.warn("Received multiple identical TaskId %s vs %s", + storedTaskId.toString(), newTaskId.toString()); + } + } + protected void queryCompletedEvent(Optional failureInfo, OptionalLong updateCount) { List serializedTaskInfos = taskInfoCollector.value(); - ImmutableList.Builder taskInfos = ImmutableList.builder(); + HashMap taskInfoMap = new HashMap<>(); long totalSerializedTaskInfoSizeInBytes = 0; for (SerializedTaskInfo serializedTaskInfo : serializedTaskInfos) { byte[] bytes = serializedTaskInfo.getBytesAndClear(); totalSerializedTaskInfoSizeInBytes += bytes.length; TaskInfo taskInfo = deserializeZstdCompressed(taskInfoCodec, bytes); - taskInfos.add(taskInfo); + updateTaskInfoMap(taskInfoMap, taskInfo); } taskInfoCollector.reset(); - log.info("Total serialized task info size: %s", DataSize.succinctBytes(totalSerializedTaskInfoSizeInBytes)); + log.info("Total serialized task info count %s size: %s. Total deduped task info count %s", + serializedTaskInfos.size(), + DataSize.succinctBytes(totalSerializedTaskInfoSizeInBytes), + taskInfoMap.size()); Optional stageInfoOptional = getFinalFragmentedPlan().map(finalFragmentedPlan -> - PrestoSparkQueryExecutionFactory.createStageInfo(session.getQueryId(), finalFragmentedPlan, taskInfos.build())); + PrestoSparkQueryExecutionFactory.createStageInfo( + session.getQueryId(), + finalFragmentedPlan, + taskInfoMap.values().stream().collect(Collectors.toList()))); QueryState queryState = failureInfo.isPresent() ? FAILED : FINISHED; QueryInfo queryInfo = PrestoSparkQueryExecutionFactory.createQueryInfo( From 5b391d836dd52a46b6b05bc33cba3b848938c4d6 Mon Sep 17 00:00:00 2001 From: "Jalpreet Singh Nanda (:imjalpreet)" Date: Sat, 6 Sep 2025 02:03:17 +0530 Subject: [PATCH 066/113] Fix NullPointerException when catalogSessionProperties is not configured --- .../file/TestFileSessionPropertyManager.java | 58 +++++++++++++++++++ .../presto/session/SessionMatchSpec.java | 9 ++- 2 files changed, 64 insertions(+), 3 deletions(-) diff --git a/presto-file-session-property-manager/src/test/java/com/facebook/presto/session/file/TestFileSessionPropertyManager.java b/presto-file-session-property-manager/src/test/java/com/facebook/presto/session/file/TestFileSessionPropertyManager.java index 77f5f79ffdb1d..f071f2547b353 100644 --- a/presto-file-session-property-manager/src/test/java/com/facebook/presto/session/file/TestFileSessionPropertyManager.java +++ b/presto-file-session-property-manager/src/test/java/com/facebook/presto/session/file/TestFileSessionPropertyManager.java @@ -19,12 +19,14 @@ import com.facebook.presto.session.SessionMatchSpec; import com.facebook.presto.spi.session.SessionPropertyConfigurationManager; import com.google.common.collect.ImmutableMap; +import org.testng.annotations.Test; import java.io.IOException; import java.nio.file.Files; import java.nio.file.Path; import java.util.Arrays; import java.util.Map; +import java.util.Optional; import static com.facebook.presto.session.file.FileSessionPropertyManager.CODEC; import static org.testng.Assert.assertEquals; @@ -59,4 +61,60 @@ protected void assertProperties(Map defaultProperties, Map> catalogProperties = ImmutableMap.of("CATALOG", ImmutableMap.of("PROPERTY", "VALUE")); + SessionMatchSpec spec = new SessionMatchSpec( + Optional.empty(), + Optional.empty(), + Optional.empty(), + Optional.empty(), + Optional.empty(), + Optional.empty(), + Optional.empty(), + null, + catalogProperties); + + assertProperties(ImmutableMap.of(), ImmutableMap.of(), catalogProperties, spec); + } + + @Test + public void testNullCatalogSessionProperties() + throws IOException + { + Map properties = ImmutableMap.of("PROPERTY1", "VALUE1", "PROPERTY2", "VALUE2"); + SessionMatchSpec spec = new SessionMatchSpec( + Optional.empty(), + Optional.empty(), + Optional.empty(), + Optional.empty(), + Optional.empty(), + Optional.empty(), + Optional.empty(), + properties, + null); + + assertProperties(properties, spec); + } + + @Test(expectedExceptions = IllegalArgumentException.class, expectedExceptionsMessageRegExp = "Either sessionProperties or catalogSessionProperties must be provided") + public void testNullBothSessionProperties() + throws IOException + { + SessionMatchSpec spec = new SessionMatchSpec( + Optional.empty(), + Optional.empty(), + Optional.empty(), + Optional.empty(), + Optional.empty(), + Optional.empty(), + Optional.empty(), + null, + null); + + assertProperties(ImmutableMap.of(), spec); + } } diff --git a/presto-session-property-managers-common/src/main/java/com/facebook/presto/session/SessionMatchSpec.java b/presto-session-property-managers-common/src/main/java/com/facebook/presto/session/SessionMatchSpec.java index 8d7a88826d90d..6a70aaae42d69 100644 --- a/presto-session-property-managers-common/src/main/java/com/facebook/presto/session/SessionMatchSpec.java +++ b/presto-session-property-managers-common/src/main/java/com/facebook/presto/session/SessionMatchSpec.java @@ -69,9 +69,12 @@ public SessionMatchSpec( this.resourceGroupRegex = requireNonNull(resourceGroupRegex, "resourceGroupRegex is null"); this.clientInfoRegex = requireNonNull(clientInfoRegex, "clientInfoRegex is null"); this.overrideSessionProperties = requireNonNull(overrideSessionProperties, "overrideSessionProperties is null"); - requireNonNull(sessionProperties, "sessionProperties is null"); - this.sessionProperties = ImmutableMap.copyOf(sessionProperties); - this.catalogSessionProperties = ImmutableMap.copyOf(catalogSessionProperties); + + checkArgument(sessionProperties != null || catalogSessionProperties != null, + "Either sessionProperties or catalogSessionProperties must be provided"); + + this.sessionProperties = ImmutableMap.copyOf(Optional.ofNullable(sessionProperties).orElse(ImmutableMap.of())); + this.catalogSessionProperties = ImmutableMap.copyOf(Optional.ofNullable(catalogSessionProperties).orElse(ImmutableMap.of())); } public Map match(Map object, SessionConfigurationContext context) From 814cdd63d28c292da2a899701fedceb467e36683 Mon Sep 17 00:00:00 2001 From: wangd Date: Mon, 8 Sep 2025 08:56:03 +0800 Subject: [PATCH 067/113] [Native]Fix typo in NativeSidecarPluginQueryRunner configuration --- presto-native-execution/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/presto-native-execution/README.md b/presto-native-execution/README.md index 59f4fe00424ee..5617ab2510676 100644 --- a/presto-native-execution/README.md +++ b/presto-native-execution/README.md @@ -222,7 +222,7 @@ Run NativeSidecarPluginQueryRunner: * VM options : `-ea -Xmx5G -XX:+ExitOnOutOfMemoryError -Duser.timezone=America/Bahia_Banderas -Dhive.security=legacy`. * Working directory: `$MODULE_DIR$` * Environment variables: `PRESTO_SERVER=/Users//git/presto/presto-native-execution/cmake-build-debug/presto_cpp/main/presto_server;DATA_DIR=/Users//Desktop/data;WORKER_COUNT=0` - * Use classpath of module: choose `presto-native-execution` module. + * Use classpath of module: choose `presto-native-sidecar-plugin` module. Run CLion: * File->Close Project if any is open. From 8828a3bfc7ca4a1c3c5178af063eee9f73213fba Mon Sep 17 00:00:00 2001 From: "Jalpreet Singh Nanda (:imjalpreet)" Date: Mon, 8 Sep 2025 19:05:00 +0530 Subject: [PATCH 068/113] Fix `current_time` handling for legacy timestamp semantics MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Correct current_time behavior in legacy timestamp semantics when the session time zone’s current offset differs from its offset on 1970-01-01. --- .../presto/operator/scalar/DateTimeFunctions.java | 13 ++++++------- .../operator/scalar/TestDateTimeFunctions.java | 15 --------------- .../scalar/TestDateTimeFunctionsBase.java | 13 +++++++++++++ .../scalar/TestDateTimeFunctionsLegacy.java | 15 --------------- 4 files changed, 19 insertions(+), 37 deletions(-) diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/DateTimeFunctions.java b/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/DateTimeFunctions.java index 43f4de46c1e9b..aa2bc0ce44fa0 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/DateTimeFunctions.java +++ b/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/DateTimeFunctions.java @@ -117,13 +117,12 @@ public static long currentTime(SqlFunctionProperties properties) // and we need to have UTC millis for packDateTimeWithZone long millis = UTC_CHRONOLOGY.millisOfDay().get(properties.getSessionStartTime()); - if (!properties.isLegacyTimestamp()) { - // However, those UTC millis are pointing to the correct UTC timestamp - // Our TIME WITH TIME ZONE representation does use UTC 1970-01-01 representation - // So we have to hack here in order to get valid representation - // of TIME WITH TIME ZONE - millis -= valueToSessionTimeZoneOffsetDiff(properties.getSessionStartTime(), getDateTimeZone(properties.getTimeZoneKey())); - } + // However, those UTC millis are pointing to the correct UTC timestamp + // Our TIME WITH TIME ZONE representation does use UTC 1970-01-01 representation + // So we have to hack here in order to get valid representation + // of TIME WITH TIME ZONE + millis -= valueToSessionTimeZoneOffsetDiff(properties.getSessionStartTime(), getDateTimeZone(properties.getTimeZoneKey())); + try { return packDateTimeWithZone(millis, properties.getTimeZoneKey()); } diff --git a/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctions.java b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctions.java index 1a939a01b3461..79d1ebd87ad1b 100644 --- a/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctions.java +++ b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctions.java @@ -20,7 +20,6 @@ import org.joda.time.DateTime; import org.testng.annotations.Test; -import static com.facebook.presto.common.type.TimeWithTimeZoneType.TIME_WITH_TIME_ZONE; import static com.facebook.presto.common.type.TimestampWithTimeZoneType.TIMESTAMP_WITH_TIME_ZONE; import static com.facebook.presto.common.type.VarcharType.createVarcharType; @@ -57,20 +56,6 @@ public void testLocalTime() } } - @Test - public void testCurrentTime() - { - Session localSession = Session.builder(session) - // we use Asia/Kathmandu here to test the difference in semantic change of current_time - // between legacy and non-legacy timestamp - .setTimeZoneKey(KATHMANDU_ZONE_KEY) - .setStartTime(new DateTime(2017, 3, 1, 15, 45, 0, 0, KATHMANDU_ZONE).getMillis()) - .build(); - try (FunctionAssertions localAssertion = new FunctionAssertions(localSession)) { - localAssertion.assertFunctionString("CURRENT_TIME", TIME_WITH_TIME_ZONE, "15:45:00.000 Asia/Kathmandu"); - } - } - @Test public void testLocalTimestamp() { diff --git a/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctionsBase.java b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctionsBase.java index 7fa1263fcd5e1..e3a6bc59bfe82 100644 --- a/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctionsBase.java +++ b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctionsBase.java @@ -186,6 +186,19 @@ private static long epochDaysInZone(TimeZoneKey timeZoneKey, long instant) return LocalDate.from(Instant.ofEpochMilli(instant).atZone(ZoneId.of(timeZoneKey.getId()))).toEpochDay(); } + @Test + public void testCurrentTime() + { + Session localSession = Session.builder(session) + // we use Asia/Kathmandu here, as it has different zone offset on 2017-03-01 and on 1970-01-01 + .setTimeZoneKey(KATHMANDU_ZONE_KEY) + .setStartTime(new DateTime(2017, 3, 1, 15, 45, 0, 0, KATHMANDU_ZONE).getMillis()) + .build(); + try (FunctionAssertions localAssertion = new FunctionAssertions(localSession)) { + localAssertion.assertFunctionString("CURRENT_TIME", TIME_WITH_TIME_ZONE, "15:45:00.000 Asia/Kathmandu"); + } + } + @Test public void testFromUnixTime() { diff --git a/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctionsLegacy.java b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctionsLegacy.java index 97527e6b00d10..609b1a52c14a3 100644 --- a/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctionsLegacy.java +++ b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctionsLegacy.java @@ -20,7 +20,6 @@ import org.joda.time.DateTime; import org.testng.annotations.Test; -import static com.facebook.presto.common.type.TimeWithTimeZoneType.TIME_WITH_TIME_ZONE; import static com.facebook.presto.common.type.TimestampWithTimeZoneType.TIMESTAMP_WITH_TIME_ZONE; import static com.facebook.presto.common.type.VarcharType.VARCHAR; import static com.facebook.presto.common.type.VarcharType.createVarcharType; @@ -56,20 +55,6 @@ public void testLocalTime() } } - @Test - public void testCurrentTime() - { - Session localSession = Session.builder(session) - // we use Asia/Kathmandu here to test the difference in semantic change of current_time - // between legacy and non-legacy timestamp - .setTimeZoneKey(KATHMANDU_ZONE_KEY) - .setStartTime(new DateTime(2017, 3, 1, 15, 45, 0, 0, KATHMANDU_ZONE).getMillis()) - .build(); - try (FunctionAssertions localAssertion = new FunctionAssertions(localSession)) { - localAssertion.assertFunctionString("CURRENT_TIME", TIME_WITH_TIME_ZONE, "15:30:00.000 Asia/Kathmandu"); - } - } - @Test public void testLocalTimestamp() { From 3425c7ef6e8c7b0390d22da7c386c1221ee495ca Mon Sep 17 00:00:00 2001 From: "Jalpreet Singh Nanda (:imjalpreet)" Date: Tue, 9 Sep 2025 03:31:11 +0530 Subject: [PATCH 069/113] Fix `localtime` handling for legacy timestamp semantics MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Correct `localtime` behavior in legacy timestamp semantics when the session time zone’s current offset differs from its offset on 1970-01-01. --- .../presto/operator/scalar/DateTimeFunctions.java | 3 ++- .../operator/scalar/TestDateTimeFunctions.java | 12 ------------ .../operator/scalar/TestDateTimeFunctionsBase.java | 11 +++++++++++ .../operator/scalar/TestDateTimeFunctionsLegacy.java | 12 ------------ 4 files changed, 13 insertions(+), 25 deletions(-) diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/DateTimeFunctions.java b/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/DateTimeFunctions.java index aa2bc0ce44fa0..8db043aab9da2 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/DateTimeFunctions.java +++ b/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/DateTimeFunctions.java @@ -143,7 +143,8 @@ public static long currentTime(SqlFunctionProperties properties) public static long localTime(SqlFunctionProperties properties) { if (properties.isLegacyTimestamp()) { - return UTC_CHRONOLOGY.millisOfDay().get(properties.getSessionStartTime()); + long millis = UTC_CHRONOLOGY.millisOfDay().get(properties.getSessionStartTime()); + return millis - valueToSessionTimeZoneOffsetDiff(properties.getSessionStartTime(), getDateTimeZone(properties.getTimeZoneKey())); } ISOChronology localChronology = getChronology(properties.getTimeZoneKey()); return localChronology.millisOfDay().get(properties.getSessionStartTime()); diff --git a/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctions.java b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctions.java index 79d1ebd87ad1b..cf45e5643f9cb 100644 --- a/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctions.java +++ b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctions.java @@ -15,7 +15,6 @@ package com.facebook.presto.operator.scalar; import com.facebook.presto.Session; -import com.facebook.presto.common.type.TimeType; import com.facebook.presto.common.type.TimestampType; import org.joda.time.DateTime; import org.testng.annotations.Test; @@ -45,17 +44,6 @@ public void testFormatDateCannotImplicitlyAddTimeZoneToTimestampLiteral() "format_datetime for TIMESTAMP type, cannot use 'Z' nor 'z' in format, as this type does not contain TZ information"); } - @Test - public void testLocalTime() - { - Session localSession = Session.builder(session) - .setStartTime(new DateTime(2017, 3, 1, 14, 30, 0, 0, DATE_TIME_ZONE).getMillis()) - .build(); - try (FunctionAssertions localAssertion = new FunctionAssertions(localSession)) { - localAssertion.assertFunctionString("LOCALTIME", TimeType.TIME, "14:30:00.000"); - } - } - @Test public void testLocalTimestamp() { diff --git a/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctionsBase.java b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctionsBase.java index e3a6bc59bfe82..733108e7d8dcc 100644 --- a/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctionsBase.java +++ b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctionsBase.java @@ -186,6 +186,17 @@ private static long epochDaysInZone(TimeZoneKey timeZoneKey, long instant) return LocalDate.from(Instant.ofEpochMilli(instant).atZone(ZoneId.of(timeZoneKey.getId()))).toEpochDay(); } + @Test + public void testLocalTime() + { + Session localSession = Session.builder(session) + .setStartTime(new DateTime(2017, 3, 1, 14, 30, 0, 0, DATE_TIME_ZONE).getMillis()) + .build(); + try (FunctionAssertions localAssertion = new FunctionAssertions(localSession)) { + localAssertion.assertFunctionString("LOCALTIME", TimeType.TIME, "14:30:00.000"); + } + } + @Test public void testCurrentTime() { diff --git a/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctionsLegacy.java b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctionsLegacy.java index 609b1a52c14a3..2293dd0ceb865 100644 --- a/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctionsLegacy.java +++ b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestDateTimeFunctionsLegacy.java @@ -15,7 +15,6 @@ package com.facebook.presto.operator.scalar; import com.facebook.presto.Session; -import com.facebook.presto.common.type.TimeType; import com.facebook.presto.common.type.TimestampType; import org.joda.time.DateTime; import org.testng.annotations.Test; @@ -44,17 +43,6 @@ public void testFormatDateCanImplicitlyAddTimeZoneToTimestampLiteral() assertFunction("format_datetime(" + TIMESTAMP_LITERAL + ", 'YYYY/MM/dd HH:mm ZZZZ')", VARCHAR, "2001/08/22 03:04 " + DATE_TIME_ZONE.getID()); } - @Test - public void testLocalTime() - { - Session localSession = Session.builder(session) - .setStartTime(new DateTime(2017, 3, 1, 14, 30, 0, 0, DATE_TIME_ZONE).getMillis()) - .build(); - try (FunctionAssertions localAssertion = new FunctionAssertions(localSession)) { - localAssertion.assertFunctionString("LOCALTIME", TimeType.TIME, "13:30:00.000"); - } - } - @Test public void testLocalTimestamp() { From 7a209483caadf1a90f3ac43c4e2cbd662489e9ce Mon Sep 17 00:00:00 2001 From: James Gill Date: Tue, 26 Aug 2025 15:55:51 -0400 Subject: [PATCH 070/113] feat: Add SpatialJoinNode to presto_protocol To send SpatialJoinNodes to Velox, we need to serialize and deserialize them via presto_protocol. This change requires https://github.com/facebookincubator/velox/pull/14339 for spatial joins to not cause an error. After this PR and the above lands, Spatial Joins should be enabled implemented as Nested Loop Joins. Not efficient, but it should be correct. --- .../main/types/PrestoToVeloxQueryPlan.cpp | 29 +++++ .../main/types/PrestoToVeloxQueryPlan.h | 5 + .../core/presto_protocol_core.cpp | 119 ++++++++++++++++++ .../core/presto_protocol_core.h | 21 ++++ .../core/presto_protocol_core.yml | 2 + .../presto_protocol/presto_protocol.yml | 2 + .../sidecar/TestNativeSidecarPlugin.java | 5 +- 7 files changed, 180 insertions(+), 3 deletions(-) diff --git a/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.cpp b/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.cpp index ccd4b1d0f6294..6c6efc582ac04 100644 --- a/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.cpp +++ b/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.cpp @@ -1186,6 +1186,17 @@ core::JoinType toJoinType(protocol::JoinType type) { VELOX_UNSUPPORTED("Unknown join type"); } + +core::JoinType toJoinType(protocol::SpatialJoinType type) { + switch (type) { + case protocol::SpatialJoinType::INNER: + return core::JoinType::kInner; + case protocol::SpatialJoinType::LEFT: + return core::JoinType::kLeft; + } + + VELOX_UNSUPPORTED("Unknown spatial join type"); +} } // namespace core::PlanNodePtr VeloxQueryPlanConverterBase::toVeloxQueryPlan( @@ -1264,6 +1275,20 @@ core::PlanNodePtr VeloxQueryPlanConverterBase::toVeloxQueryPlan( ROW(std::move(outputNames), std::move(outputTypes))); } +core::PlanNodePtr VeloxQueryPlanConverterBase::toVeloxQueryPlan( + const std::shared_ptr& node, + const std::shared_ptr& tableWriteInfo, + const protocol::TaskId& taskId) { + auto joinType = toJoinType(node->type); + + return std::make_shared( + node->id, + joinType, + exprConverter_.toVeloxExpr(node->filter), + toVeloxQueryPlan(node->left, tableWriteInfo, taskId), + toVeloxQueryPlan(node->right, tableWriteInfo, taskId), + toRowType(node->outputVariables, typeParser_));} + std::shared_ptr VeloxQueryPlanConverterBase::toVeloxQueryPlan( const std::shared_ptr& node, @@ -1842,6 +1867,10 @@ core::PlanNodePtr VeloxQueryPlanConverterBase::toVeloxQueryPlan( std::dynamic_pointer_cast(node)) { return toVeloxQueryPlan(join, tableWriteInfo, taskId); } + if (auto spatialJoin = + std::dynamic_pointer_cast(node)) { + return toVeloxQueryPlan(spatialJoin, tableWriteInfo, taskId); + } if (auto remoteSource = std::dynamic_pointer_cast(node)) { return toVeloxQueryPlan(remoteSource, tableWriteInfo, taskId); diff --git a/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.h b/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.h index 3d13090b3801c..15c87c94e0984 100644 --- a/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.h +++ b/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.h @@ -110,6 +110,11 @@ class VeloxQueryPlanConverterBase { const std::shared_ptr& tableWriteInfo, const protocol::TaskId& taskId); + velox::core::PlanNodePtr toVeloxQueryPlan( + const std::shared_ptr& node, + const std::shared_ptr& tableWriteInfo, + const protocol::TaskId& taskId); + std::shared_ptr toVeloxQueryPlan( const std::shared_ptr& node, const std::shared_ptr& tableWriteInfo, diff --git a/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.cpp b/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.cpp index e804c9a35925e..6f0b26946a186 100644 --- a/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.cpp +++ b/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.cpp @@ -728,6 +728,10 @@ void to_json(json& j, const std::shared_ptr& p) { j = *std::static_pointer_cast(p); return; } + if (type == ".SpatialJoinNode") { + j = *std::static_pointer_cast(p); + return; + } if (type == ".TableScanNode") { j = *std::static_pointer_cast(p); return; @@ -896,6 +900,12 @@ void from_json(const json& j, std::shared_ptr& p) { p = std::static_pointer_cast(k); return; } + if (type == ".SpatialJoinNode") { + std::shared_ptr k = std::make_shared(); + j.get_to(*k); + p = std::static_pointer_cast(k); + return; + } if (type == ".TableScanNode") { std::shared_ptr k = std::make_shared(); j.get_to(*k); @@ -9485,6 +9495,115 @@ void from_json(const json& j, SortedRangeSet& p) { namespace facebook::presto::protocol { // Loosly copied this here from NLOHMANN_JSON_SERIALIZE_ENUM() +// NOLINTNEXTLINE: cppcoreguidelines-avoid-c-arrays +static const std::pair SpatialJoinType_enum_table[] = + { // NOLINT: cert-err58-cpp + {SpatialJoinType::INNER, "INNER"}, + {SpatialJoinType::LEFT, "LEFT"}}; +void to_json(json& j, const SpatialJoinType& e) { + static_assert( + std::is_enum::value, "SpatialJoinType must be an enum!"); + const auto* it = std::find_if( + std::begin(SpatialJoinType_enum_table), + std::end(SpatialJoinType_enum_table), + [e](const std::pair& ej_pair) -> bool { + return ej_pair.first == e; + }); + j = ((it != std::end(SpatialJoinType_enum_table)) + ? it + : std::begin(SpatialJoinType_enum_table)) + ->second; +} +void from_json(const json& j, SpatialJoinType& e) { + static_assert( + std::is_enum::value, "SpatialJoinType must be an enum!"); + const auto* it = std::find_if( + std::begin(SpatialJoinType_enum_table), + std::end(SpatialJoinType_enum_table), + [&j](const std::pair& ej_pair) -> bool { + return ej_pair.second == j; + }); + e = ((it != std::end(SpatialJoinType_enum_table)) + ? it + : std::begin(SpatialJoinType_enum_table)) + ->first; +} +} // namespace facebook::presto::protocol +namespace facebook::presto::protocol { +SpatialJoinNode::SpatialJoinNode() noexcept { + _type = ".SpatialJoinNode"; +} + +void to_json(json& j, const SpatialJoinNode& p) { + j = json::object(); + j["@type"] = ".SpatialJoinNode"; + to_json_key(j, "id", p.id, "SpatialJoinNode", "PlanNodeId", "id"); + to_json_key(j, "type", p.type, "SpatialJoinNode", "SpatialJoinType", "type"); + to_json_key(j, "left", p.left, "SpatialJoinNode", "PlanNode", "left"); + to_json_key(j, "right", p.right, "SpatialJoinNode", "PlanNode", "right"); + to_json_key( + j, + "outputVariables", + p.outputVariables, + "SpatialJoinNode", + "List", + "outputVariables"); + to_json_key( + j, "filter", p.filter, "SpatialJoinNode", "RowExpression", "filter"); + to_json_key( + j, + "leftPartitionVariable", + p.leftPartitionVariable, + "SpatialJoinNode", + "VariableReferenceExpression", + "leftPartitionVariable"); + to_json_key( + j, + "rightPartitionVariable", + p.rightPartitionVariable, + "SpatialJoinNode", + "VariableReferenceExpression", + "rightPartitionVariable"); + to_json_key(j, "kdbTree", p.kdbTree, "SpatialJoinNode", "String", "kdbTree"); +} + +void from_json(const json& j, SpatialJoinNode& p) { + p._type = j["@type"]; + from_json_key(j, "id", p.id, "SpatialJoinNode", "PlanNodeId", "id"); + from_json_key( + j, "type", p.type, "SpatialJoinNode", "SpatialJoinType", "type"); + from_json_key(j, "left", p.left, "SpatialJoinNode", "PlanNode", "left"); + from_json_key(j, "right", p.right, "SpatialJoinNode", "PlanNode", "right"); + from_json_key( + j, + "outputVariables", + p.outputVariables, + "SpatialJoinNode", + "List", + "outputVariables"); + from_json_key( + j, "filter", p.filter, "SpatialJoinNode", "RowExpression", "filter"); + from_json_key( + j, + "leftPartitionVariable", + p.leftPartitionVariable, + "SpatialJoinNode", + "VariableReferenceExpression", + "leftPartitionVariable"); + from_json_key( + j, + "rightPartitionVariable", + p.rightPartitionVariable, + "SpatialJoinNode", + "VariableReferenceExpression", + "rightPartitionVariable"); + from_json_key( + j, "kdbTree", p.kdbTree, "SpatialJoinNode", "String", "kdbTree"); +} +} // namespace facebook::presto::protocol +namespace facebook::presto::protocol { +// Loosly copied this here from NLOHMANN_JSON_SERIALIZE_ENUM() + // NOLINTNEXTLINE: cppcoreguidelines-avoid-c-arrays static const std::pair Form_enum_table[] = { // NOLINT: cert-err58-cpp diff --git a/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.h b/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.h index 15656aa6e1780..6abc39d461005 100644 --- a/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.h +++ b/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.h @@ -2188,6 +2188,27 @@ void to_json(json& j, const SortedRangeSet& p); void from_json(const json& j, SortedRangeSet& p); } // namespace facebook::presto::protocol namespace facebook::presto::protocol { +enum class SpatialJoinType { INNER, LEFT }; +extern void to_json(json& j, const SpatialJoinType& e); +extern void from_json(const json& j, SpatialJoinType& e); +} // namespace facebook::presto::protocol +namespace facebook::presto::protocol { +struct SpatialJoinNode : public PlanNode { + SpatialJoinType type = {}; + std::shared_ptr left = {}; + std::shared_ptr right = {}; + List outputVariables = {}; + std::shared_ptr filter = {}; + std::shared_ptr leftPartitionVariable = {}; + std::shared_ptr rightPartitionVariable = {}; + std::shared_ptr kdbTree = {}; + + SpatialJoinNode() noexcept; +}; +void to_json(json& j, const SpatialJoinNode& p); +void from_json(const json& j, SpatialJoinNode& p); +} // namespace facebook::presto::protocol +namespace facebook::presto::protocol { enum class Form { IF, NULL_IF, diff --git a/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.yml b/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.yml index bd710882c51d1..aec2731384959 100644 --- a/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.yml +++ b/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.yml @@ -157,6 +157,7 @@ AbstractClasses: - { name: RemoteSourceNode, key: com.facebook.presto.sql.planner.plan.RemoteSourceNode } - { name: SampleNode, key: com.facebook.presto.sql.planner.plan.SampleNode } - { name: SemiJoinNode, key: .SemiJoinNode } + - { name: SpatialJoinNode, key: .SpatialJoinNode } - { name: TableScanNode, key: .TableScanNode } - { name: TableWriterNode, key: .TableWriterNode } - { name: TableWriterMergeNode, key: com.facebook.presto.sql.planner.plan.TableWriterMergeNode } @@ -317,6 +318,7 @@ JavaClasses: - presto-spi/src/main/java/com/facebook/presto/spi/plan/JoinNode.java - presto-spi/src/main/java/com/facebook/presto/spi/plan/SemiJoinNode.java - presto-spi/src/main/java/com/facebook/presto/spi/plan/MergeJoinNode.java + - presto-spi/src/main/java/com/facebook/presto/spi/plan/SpatialJoinNode.java - presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/IndexJoinNode.java - presto-spi/src/main/java/com/facebook/presto/spi/plan/IndexSourceNode.java - presto-spi/src/main/java/com/facebook/presto/spi/plan/TopNNode.java diff --git a/presto-native-execution/presto_cpp/presto_protocol/presto_protocol.yml b/presto-native-execution/presto_cpp/presto_protocol/presto_protocol.yml index 5b407cb8df8b9..219f9bf4bc23a 100644 --- a/presto-native-execution/presto_cpp/presto_protocol/presto_protocol.yml +++ b/presto-native-execution/presto_cpp/presto_protocol/presto_protocol.yml @@ -155,6 +155,7 @@ AbstractClasses: - { name: RemoteSourceNode, key: com.facebook.presto.sql.planner.plan.RemoteSourceNode } - { name: SampleNode, key: com.facebook.presto.sql.planner.plan.SampleNode } - { name: SemiJoinNode, key: .SemiJoinNode } + - { name: SpatialJoinNode, key: .SpatialJoinNode } - { name: TableScanNode, key: .TableScanNode } - { name: TableWriterNode, key: .TableWriterNode } - { name: TableWriterMergeNode, key: com.facebook.presto.sql.planner.plan.TableWriterMergeNode } @@ -360,6 +361,7 @@ JavaClasses: - presto-spi/src/main/java/com/facebook/presto/spi/plan/JoinNode.java - presto-spi/src/main/java/com/facebook/presto/spi/plan/SemiJoinNode.java - presto-spi/src/main/java/com/facebook/presto/spi/plan/MergeJoinNode.java + - presto-spi/src/main/java/com/facebook/presto/spi/plan/SpatialJoinNode.java - presto-spi/src/main/java/com/facebook/presto/spi/plan/TopNNode.java - presto-hive/src/main/java/com/facebook/presto/hive/HivePartitioningHandle.java - presto-main/src/main/java/com/facebook/presto/split/EmptySplit.java diff --git a/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java b/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java index fe0b18c24b2fb..3e96d8c50a0bd 100644 --- a/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java +++ b/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java @@ -412,7 +412,7 @@ public void testGeometryQueries() assertQuery("SELECT " + "ST_DISTANCE(ST_POINT(a.nationkey, a.regionkey), ST_POINT(b.nationkey, b.regionkey)) " + "FROM nation a JOIN nation b ON a.nationkey < b.nationkey"); - assertQueryFails( + assertQuery( "WITH regions(name, geom) AS (VALUES" + " ('A', ST_GeometryFromText('POLYGON ((0 0, 0 5, 5 5, 5 0, 0 0))'))," + " ('B', ST_GeometryFromText('POLYGON ((5 0, 5 5, 10 5, 10 0, 5 0))')))," + @@ -420,8 +420,7 @@ public void testGeometryQueries() " ('P1', ST_Point(1, 1))," + " ('P2', ST_Point(6, 1))," + " ('P3', ST_Point(8, 4)))" + - "SELECT p.id, r.name FROM points p LEFT JOIN regions r ON ST_Within(p.geom, r.geom)", - "Error from native plan checker: .SpatialJoinNode no abstract type PlanNode "); + "SELECT p.id, r.name FROM points p LEFT JOIN regions r ON ST_Within(p.geom, r.geom)"); } @Test From b49b416ec0ab4b94cd19535a1670f1feec081e64 Mon Sep 17 00:00:00 2001 From: wangd Date: Mon, 8 Sep 2025 17:01:05 +0800 Subject: [PATCH 071/113] Make all in-tree connectors utilize new SPI method in ConnectorMetadata --- .../presto/accumulo/AccumuloMetadata.java | 4 ++-- .../com/facebook/presto/atop/AtopMetadata.java | 5 +++-- .../com/facebook/plugin/arrow/ArrowMetadata.java | 8 ++++++-- .../facebook/presto/plugin/jdbc/JdbcMetadata.java | 8 ++++++-- .../presto/plugin/bigquery/BigQueryMetadata.java | 4 ++-- .../plugin/blackhole/BlackHoleMetadata.java | 4 ++-- .../presto/cassandra/CassandraMetadata.java | 8 ++++++-- .../presto/cassandra/TestCassandraConnector.java | 5 ++--- .../plugin/clickhouse/ClickHouseMetadata.java | 8 ++++++-- .../com/facebook/presto/delta/DeltaMetadata.java | 4 ++-- .../com/facebook/presto/druid/DruidMetadata.java | 8 ++++++-- .../elasticsearch/ElasticsearchMetadata.java | 8 ++++++-- .../facebook/presto/example/ExampleMetadata.java | 8 ++++++-- .../presto/google/sheets/SheetsMetadata.java | 10 ++++++---- .../com/facebook/presto/hive/HiveMetadata.java | 15 +++++++++------ .../presto/hive/AbstractTestHiveClient.java | 6 +++--- .../presto/hive/AbstractTestHiveFileSystem.java | 4 ++-- .../presto/hive/HiveFileSystemTestUtils.java | 13 ++++++------- .../com/facebook/presto/hudi/HudiMetadata.java | 8 ++++++-- .../presto/connector/jmx/JmxMetadata.java | 8 ++++++-- .../com/facebook/presto/kafka/KafkaMetadata.java | 8 ++++++-- .../com/facebook/presto/kudu/KuduMetadata.java | 5 +++-- .../presto/lark/sheets/LarkSheetsMetadata.java | 5 +++-- .../presto/localfile/LocalFileMetadata.java | 8 ++++++-- .../presto/plugin/memory/MemoryMetadata.java | 4 ++-- .../presto/plugin/memory/TestMemoryMetadata.java | 6 +++--- .../facebook/presto/mongodb/MongoMetadata.java | 11 +++++++---- .../com/facebook/presto/pinot/PinotMetadata.java | 4 ++-- .../plugin/prometheus/PrometheusMetadata.java | 8 ++++++-- .../com/facebook/presto/redis/RedisMetadata.java | 7 +++---- .../presto/connector/thrift/ThriftMetadata.java | 5 ++--- .../com/facebook/presto/tpcds/TpcdsMetadata.java | 7 +++---- .../com/facebook/presto/tpch/TpchMetadata.java | 7 +++---- .../facebook/presto/tpch/TestTpchMetadata.java | 5 +---- 34 files changed, 143 insertions(+), 93 deletions(-) diff --git a/presto-accumulo/src/main/java/com/facebook/presto/accumulo/AccumuloMetadata.java b/presto-accumulo/src/main/java/com/facebook/presto/accumulo/AccumuloMetadata.java index d585bf8464f88..c1ea2bc36dcb5 100644 --- a/presto-accumulo/src/main/java/com/facebook/presto/accumulo/AccumuloMetadata.java +++ b/presto-accumulo/src/main/java/com/facebook/presto/accumulo/AccumuloMetadata.java @@ -259,7 +259,7 @@ public ConnectorTableHandle getTableHandle(ConnectorSession session, SchemaTable } @Override - public List getTableLayouts( + public ConnectorTableLayoutResult getTableLayoutForConstraint( ConnectorSession session, ConnectorTableHandle table, Constraint constraint, @@ -267,7 +267,7 @@ public List getTableLayouts( { AccumuloTableHandle tableHandle = (AccumuloTableHandle) table; ConnectorTableLayout layout = new ConnectorTableLayout(new AccumuloTableLayoutHandle(tableHandle, constraint.getSummary())); - return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary())); + return new ConnectorTableLayoutResult(layout, constraint.getSummary()); } @Override diff --git a/presto-atop/src/main/java/com/facebook/presto/atop/AtopMetadata.java b/presto-atop/src/main/java/com/facebook/presto/atop/AtopMetadata.java index 04f2e84478935..16710c146f604 100644 --- a/presto-atop/src/main/java/com/facebook/presto/atop/AtopMetadata.java +++ b/presto-atop/src/main/java/com/facebook/presto/atop/AtopMetadata.java @@ -85,7 +85,8 @@ public ConnectorTableHandle getTableHandle(ConnectorSession session, SchemaTable } @Override - public List getTableLayouts(ConnectorSession session, + public ConnectorTableLayoutResult getTableLayoutForConstraint( + ConnectorSession session, ConnectorTableHandle table, Constraint constraint, Optional> desiredColumns) @@ -104,7 +105,7 @@ public List getTableLayouts(ConnectorSession session } AtopTableLayoutHandle layoutHandle = new AtopTableLayoutHandle(tableHandle, startTimeDomain, endTimeDomain); ConnectorTableLayout tableLayout = getTableLayout(session, layoutHandle); - return ImmutableList.of(new ConnectorTableLayoutResult(tableLayout, constraint.getSummary())); + return new ConnectorTableLayoutResult(tableLayout, constraint.getSummary()); } @Override diff --git a/presto-base-arrow-flight/src/main/java/com/facebook/plugin/arrow/ArrowMetadata.java b/presto-base-arrow-flight/src/main/java/com/facebook/plugin/arrow/ArrowMetadata.java index 4018e8f25fbaf..15dda40404686 100644 --- a/presto-base-arrow-flight/src/main/java/com/facebook/plugin/arrow/ArrowMetadata.java +++ b/presto-base-arrow-flight/src/main/java/com/facebook/plugin/arrow/ArrowMetadata.java @@ -112,7 +112,11 @@ public Map getColumnHandles(ConnectorSession session, Conn } @Override - public List getTableLayouts(ConnectorSession session, ConnectorTableHandle table, Constraint constraint, Optional> desiredColumns) + public ConnectorTableLayoutResult getTableLayoutForConstraint( + ConnectorSession session, + ConnectorTableHandle table, + Constraint constraint, + Optional> desiredColumns) { checkArgument(table instanceof ArrowTableHandle, "Invalid table handle: Expected an instance of ArrowTableHandle but received %s", @@ -129,7 +133,7 @@ public List getTableLayouts(ConnectorSession session } ConnectorTableLayout layout = new ConnectorTableLayout(new ArrowTableLayoutHandle(tableHandle, columns, constraint.getSummary())); - return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary())); + return new ConnectorTableLayoutResult(layout, constraint.getSummary()); } @Override diff --git a/presto-base-jdbc/src/main/java/com/facebook/presto/plugin/jdbc/JdbcMetadata.java b/presto-base-jdbc/src/main/java/com/facebook/presto/plugin/jdbc/JdbcMetadata.java index b56e81b2912a1..9c37bd8a3d731 100644 --- a/presto-base-jdbc/src/main/java/com/facebook/presto/plugin/jdbc/JdbcMetadata.java +++ b/presto-base-jdbc/src/main/java/com/facebook/presto/plugin/jdbc/JdbcMetadata.java @@ -85,11 +85,15 @@ public JdbcTableHandle getTableHandle(ConnectorSession session, SchemaTableName } @Override - public List getTableLayouts(ConnectorSession session, ConnectorTableHandle table, Constraint constraint, Optional> desiredColumns) + public ConnectorTableLayoutResult getTableLayoutForConstraint( + ConnectorSession session, + ConnectorTableHandle table, + Constraint constraint, + Optional> desiredColumns) { JdbcTableHandle tableHandle = (JdbcTableHandle) table; ConnectorTableLayout layout = new ConnectorTableLayout(new JdbcTableLayoutHandle(session.getSqlFunctionProperties(), tableHandle, constraint.getSummary(), Optional.empty())); - return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary())); + return new ConnectorTableLayoutResult(layout, constraint.getSummary()); } @Override diff --git a/presto-bigquery/src/main/java/com/facebook/presto/plugin/bigquery/BigQueryMetadata.java b/presto-bigquery/src/main/java/com/facebook/presto/plugin/bigquery/BigQueryMetadata.java index ad0bf4775a42b..45dfb474efb81 100644 --- a/presto-bigquery/src/main/java/com/facebook/presto/plugin/bigquery/BigQueryMetadata.java +++ b/presto-bigquery/src/main/java/com/facebook/presto/plugin/bigquery/BigQueryMetadata.java @@ -127,7 +127,7 @@ public ConnectorTableHandle getTableHandle(ConnectorSession session, SchemaTable } @Override - public List getTableLayouts( + public ConnectorTableLayoutResult getTableLayoutForConstraint( ConnectorSession session, ConnectorTableHandle table, Constraint constraint, @@ -139,7 +139,7 @@ public List getTableLayouts( bigQueryTableHandle = bigQueryTableHandle.withProjectedColumns(ImmutableList.copyOf(desiredColumns.get())); } BigQueryTableLayoutHandle bigQueryTableLayoutHandle = new BigQueryTableLayoutHandle(bigQueryTableHandle); - return ImmutableList.of(new ConnectorTableLayoutResult(new ConnectorTableLayout(bigQueryTableLayoutHandle), constraint.getSummary())); + return new ConnectorTableLayoutResult(new ConnectorTableLayout(bigQueryTableLayoutHandle), constraint.getSummary()); } @Override diff --git a/presto-blackhole/src/main/java/com/facebook/presto/plugin/blackhole/BlackHoleMetadata.java b/presto-blackhole/src/main/java/com/facebook/presto/plugin/blackhole/BlackHoleMetadata.java index 97ea968c5bbdc..afe0f1efd8dc4 100644 --- a/presto-blackhole/src/main/java/com/facebook/presto/plugin/blackhole/BlackHoleMetadata.java +++ b/presto-blackhole/src/main/java/com/facebook/presto/plugin/blackhole/BlackHoleMetadata.java @@ -253,7 +253,7 @@ public Optional finishInsert(ConnectorSession session, } @Override - public List getTableLayouts( + public ConnectorTableLayoutResult getTableLayoutForConstraint( ConnectorSession session, ConnectorTableHandle handle, Constraint constraint, @@ -266,7 +266,7 @@ public List getTableLayouts( blackHoleHandle.getRowsPerPage(), blackHoleHandle.getFieldsLength(), blackHoleHandle.getPageProcessingDelay()); - return ImmutableList.of(new ConnectorTableLayoutResult(getTableLayout(session, layoutHandle), constraint.getSummary())); + return new ConnectorTableLayoutResult(getTableLayout(session, layoutHandle), constraint.getSummary()); } @Override diff --git a/presto-cassandra/src/main/java/com/facebook/presto/cassandra/CassandraMetadata.java b/presto-cassandra/src/main/java/com/facebook/presto/cassandra/CassandraMetadata.java index f609fb22a47ff..e393803401b95 100644 --- a/presto-cassandra/src/main/java/com/facebook/presto/cassandra/CassandraMetadata.java +++ b/presto-cassandra/src/main/java/com/facebook/presto/cassandra/CassandraMetadata.java @@ -201,7 +201,11 @@ public ColumnMetadata getColumnMetadata(ConnectorSession session, ConnectorTable } @Override - public List getTableLayouts(ConnectorSession session, ConnectorTableHandle table, Constraint constraint, Optional> desiredColumns) + public ConnectorTableLayoutResult getTableLayoutForConstraint( + ConnectorSession session, + ConnectorTableHandle table, + Constraint constraint, + Optional> desiredColumns) { CassandraTableHandle handle = (CassandraTableHandle) table; CassandraPartitionResult partitionResult = partitionManager.getPartitions(handle, constraint.getSummary()); @@ -224,7 +228,7 @@ public List getTableLayouts(ConnectorSession session handle, partitionResult.getPartitions(), clusteringKeyPredicates)); - return ImmutableList.of(new ConnectorTableLayoutResult(layout, unenforcedConstraint)); + return new ConnectorTableLayoutResult(layout, unenforcedConstraint); } @Override diff --git a/presto-cassandra/src/test/java/com/facebook/presto/cassandra/TestCassandraConnector.java b/presto-cassandra/src/test/java/com/facebook/presto/cassandra/TestCassandraConnector.java index 268e368918a65..4f5856a09e4e6 100644 --- a/presto-cassandra/src/test/java/com/facebook/presto/cassandra/TestCassandraConnector.java +++ b/presto-cassandra/src/test/java/com/facebook/presto/cassandra/TestCassandraConnector.java @@ -67,7 +67,6 @@ import static com.facebook.presto.spi.connector.ConnectorSplitManager.SplitSchedulingStrategy.UNGROUPED_SCHEDULING; import static com.facebook.presto.spi.connector.NotPartitionedPartitionHandle.NOT_PARTITIONED; import static com.google.common.base.Preconditions.checkArgument; -import static com.google.common.collect.Iterables.getOnlyElement; import static java.util.Locale.ENGLISH; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertNull; @@ -185,8 +184,8 @@ public void testGetRecords() ConnectorTransactionHandle transaction = CassandraTransactionHandle.INSTANCE; - List layouts = metadata.getTableLayouts(SESSION, tableHandle, Constraint.alwaysTrue(), Optional.empty()); - ConnectorTableLayoutHandle layout = getOnlyElement(layouts).getTableLayout().getHandle(); + ConnectorTableLayoutResult layoutResult = metadata.getTableLayoutForConstraint(SESSION, tableHandle, Constraint.alwaysTrue(), Optional.empty()); + ConnectorTableLayoutHandle layout = layoutResult.getTableLayout().getHandle(); List splits = getAllSplits(splitManager.getSplits(transaction, SESSION, layout, new SplitSchedulingContext(UNGROUPED_SCHEDULING, false, WarningCollector.NOOP))); long rowNumber = 0; diff --git a/presto-clickhouse/src/main/java/com/facebook/presto/plugin/clickhouse/ClickHouseMetadata.java b/presto-clickhouse/src/main/java/com/facebook/presto/plugin/clickhouse/ClickHouseMetadata.java index 6e8fdf91e40ea..749da9652ded4 100755 --- a/presto-clickhouse/src/main/java/com/facebook/presto/plugin/clickhouse/ClickHouseMetadata.java +++ b/presto-clickhouse/src/main/java/com/facebook/presto/plugin/clickhouse/ClickHouseMetadata.java @@ -84,11 +84,15 @@ public ClickHouseTableHandle getTableHandle(ConnectorSession session, SchemaTabl } @Override - public List getTableLayouts(ConnectorSession session, ConnectorTableHandle table, Constraint constraint, Optional> desiredColumns) + public ConnectorTableLayoutResult getTableLayoutForConstraint( + ConnectorSession session, + ConnectorTableHandle table, + Constraint constraint, + Optional> desiredColumns) { ClickHouseTableHandle tableHandle = (ClickHouseTableHandle) table; ConnectorTableLayout layout = new ConnectorTableLayout(new ClickHouseTableLayoutHandle(tableHandle, constraint.getSummary(), Optional.empty(), Optional.empty(), Optional.empty())); - return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary())); + return new ConnectorTableLayoutResult(layout, constraint.getSummary()); } @Override diff --git a/presto-delta/src/main/java/com/facebook/presto/delta/DeltaMetadata.java b/presto-delta/src/main/java/com/facebook/presto/delta/DeltaMetadata.java index d0b45969dd390..22e67e9d47f6e 100644 --- a/presto-delta/src/main/java/com/facebook/presto/delta/DeltaMetadata.java +++ b/presto-delta/src/main/java/com/facebook/presto/delta/DeltaMetadata.java @@ -231,7 +231,7 @@ public DeltaTableHandle getTableHandle(ConnectorSession session, SchemaTableName } @Override - public List getTableLayouts( + public ConnectorTableLayoutResult getTableLayoutForConstraint( ConnectorSession session, ConnectorTableHandle table, Constraint constraint, @@ -259,7 +259,7 @@ public List getTableLayouts( ImmutableList.of(), Optional.empty()); - return ImmutableList.of(new ConnectorTableLayoutResult(newLayout, unenforcedPredicate)); + return new ConnectorTableLayoutResult(newLayout, unenforcedPredicate); } @Override diff --git a/presto-druid/src/main/java/com/facebook/presto/druid/DruidMetadata.java b/presto-druid/src/main/java/com/facebook/presto/druid/DruidMetadata.java index bcc6df536e167..ce0d7f40e3ec7 100644 --- a/presto-druid/src/main/java/com/facebook/presto/druid/DruidMetadata.java +++ b/presto-druid/src/main/java/com/facebook/presto/druid/DruidMetadata.java @@ -85,11 +85,15 @@ public ConnectorTableHandle getTableHandle(ConnectorSession session, SchemaTable } @Override - public List getTableLayouts(ConnectorSession session, ConnectorTableHandle table, Constraint constraint, Optional> desiredColumns) + public ConnectorTableLayoutResult getTableLayoutForConstraint( + ConnectorSession session, + ConnectorTableHandle table, + Constraint constraint, + Optional> desiredColumns) { DruidTableHandle handle = (DruidTableHandle) table; ConnectorTableLayout layout = new ConnectorTableLayout(new DruidTableLayoutHandle(handle, constraint.getSummary())); - return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary())); + return new ConnectorTableLayoutResult(layout, constraint.getSummary()); } @Override diff --git a/presto-elasticsearch/src/main/java/com/facebook/presto/elasticsearch/ElasticsearchMetadata.java b/presto-elasticsearch/src/main/java/com/facebook/presto/elasticsearch/ElasticsearchMetadata.java index aa8b62bbad357..e6b9b1c0e3ca7 100644 --- a/presto-elasticsearch/src/main/java/com/facebook/presto/elasticsearch/ElasticsearchMetadata.java +++ b/presto-elasticsearch/src/main/java/com/facebook/presto/elasticsearch/ElasticsearchMetadata.java @@ -158,11 +158,15 @@ public ElasticsearchTableHandle getTableHandle(ConnectorSession session, SchemaT } @Override - public List getTableLayouts(ConnectorSession session, ConnectorTableHandle table, Constraint constraint, Optional> desiredColumns) + public ConnectorTableLayoutResult getTableLayoutForConstraint( + ConnectorSession session, + ConnectorTableHandle table, + Constraint constraint, + Optional> desiredColumns) { ElasticsearchTableHandle handle = (ElasticsearchTableHandle) table; ConnectorTableLayout layout = new ConnectorTableLayout(new ElasticsearchTableLayoutHandle(handle, constraint.getSummary())); - return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary())); + return new ConnectorTableLayoutResult(layout, constraint.getSummary()); } @Override diff --git a/presto-example-http/src/main/java/com/facebook/presto/example/ExampleMetadata.java b/presto-example-http/src/main/java/com/facebook/presto/example/ExampleMetadata.java index 9653a09277cac..9b6f4f2c504b9 100644 --- a/presto-example-http/src/main/java/com/facebook/presto/example/ExampleMetadata.java +++ b/presto-example-http/src/main/java/com/facebook/presto/example/ExampleMetadata.java @@ -81,11 +81,15 @@ public ExampleTableHandle getTableHandle(ConnectorSession session, SchemaTableNa } @Override - public List getTableLayouts(ConnectorSession session, ConnectorTableHandle table, Constraint constraint, Optional> desiredColumns) + public ConnectorTableLayoutResult getTableLayoutForConstraint( + ConnectorSession session, + ConnectorTableHandle table, + Constraint constraint, + Optional> desiredColumns) { ExampleTableHandle tableHandle = (ExampleTableHandle) table; ConnectorTableLayout layout = new ConnectorTableLayout(new ExampleTableLayoutHandle(tableHandle)); - return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary())); + return new ConnectorTableLayoutResult(layout, constraint.getSummary()); } @Override diff --git a/presto-google-sheets/src/main/java/com/facebook/presto/google/sheets/SheetsMetadata.java b/presto-google-sheets/src/main/java/com/facebook/presto/google/sheets/SheetsMetadata.java index c833915fb6eda..c077c77c0e323 100644 --- a/presto-google-sheets/src/main/java/com/facebook/presto/google/sheets/SheetsMetadata.java +++ b/presto-google-sheets/src/main/java/com/facebook/presto/google/sheets/SheetsMetadata.java @@ -81,13 +81,15 @@ public SheetsTableHandle getTableHandle(ConnectorSession session, SchemaTableNam } @Override - public List getTableLayouts( - ConnectorSession session, ConnectorTableHandle table, - Constraint constraint, Optional> desiredColumns) + public ConnectorTableLayoutResult getTableLayoutForConstraint( + ConnectorSession session, + ConnectorTableHandle table, + Constraint constraint, + Optional> desiredColumns) { SheetsTableHandle tableHandle = (SheetsTableHandle) table; ConnectorTableLayout layout = new ConnectorTableLayout(new SheetsTableLayoutHandle(tableHandle)); - return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary())); + return new ConnectorTableLayoutResult(layout, constraint.getSummary()); } @Override diff --git a/presto-hive/src/main/java/com/facebook/presto/hive/HiveMetadata.java b/presto-hive/src/main/java/com/facebook/presto/hive/HiveMetadata.java index e0d07e7ce4a58..1efa0d67f3aaa 100644 --- a/presto-hive/src/main/java/com/facebook/presto/hive/HiveMetadata.java +++ b/presto-hive/src/main/java/com/facebook/presto/hive/HiveMetadata.java @@ -109,7 +109,6 @@ import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import com.google.common.collect.ImmutableSortedMap; -import com.google.common.collect.Iterables; import com.google.common.collect.Lists; import com.google.common.collect.Maps; import com.google.common.collect.Sets; @@ -2648,8 +2647,8 @@ private List getOrComputePartitions(HiveTableLayoutHandle layoutH else { TupleDomain partitionColumnPredicate = layoutHandle.getPartitionColumnPredicate(); Predicate> predicate = convertToPredicate(partitionColumnPredicate); - List tableLayoutResults = getTableLayouts(session, tableHandle, new Constraint<>(partitionColumnPredicate, predicate), Optional.empty()); - return ((HiveTableLayoutHandle) Iterables.getOnlyElement(tableLayoutResults).getTableLayout().getHandle()).getPartitions().get(); + ConnectorTableLayoutResult tableLayoutResult = getTableLayoutForConstraint(session, tableHandle, new Constraint<>(partitionColumnPredicate, predicate), Optional.empty()); + return ((HiveTableLayoutHandle) tableLayoutResult.getTableLayout().getHandle()).getPartitions().get(); } } @@ -2726,7 +2725,11 @@ private String createTableLayoutString( } @Override - public List getTableLayouts(ConnectorSession session, ConnectorTableHandle tableHandle, Constraint constraint, Optional> desiredColumns) + public ConnectorTableLayoutResult getTableLayoutForConstraint( + ConnectorSession session, + ConnectorTableHandle tableHandle, + Constraint constraint, + Optional> desiredColumns) { HiveTableHandle handle = (HiveTableHandle) tableHandle; HivePartitionResult hivePartitionResult; @@ -2754,7 +2757,7 @@ public List getTableLayouts(ConnectorSession session String layoutString = createTableLayoutString(session, handle.getSchemaTableName(), hivePartitionResult.getBucketHandle(), hivePartitionResult.getBucketFilter(), TRUE_CONSTANT, domainPredicate); Optional> requestedColumns = desiredColumns.map(columns -> columns.stream().map(column -> (HiveColumnHandle) column).collect(toImmutableSet())); - return ImmutableList.of(new ConnectorTableLayoutResult( + return new ConnectorTableLayoutResult( getTableLayout( session, new HiveTableLayoutHandle.Builder() @@ -2777,7 +2780,7 @@ public List getTableLayouts(ConnectorSession session .setAppendRowNumberEnabled(false) .setHiveTableHandle(handle) .build()), - hivePartitionResult.getUnenforcedConstraint())); + hivePartitionResult.getUnenforcedConstraint()); } private static Subfield toSubfield(ColumnHandle columnHandle) diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/AbstractTestHiveClient.java b/presto-hive/src/test/java/com/facebook/presto/hive/AbstractTestHiveClient.java index 17b2e70cff0ec..38bd3ece02ad3 100644 --- a/presto-hive/src/test/java/com/facebook/presto/hive/AbstractTestHiveClient.java +++ b/presto-hive/src/test/java/com/facebook/presto/hive/AbstractTestHiveClient.java @@ -2328,7 +2328,7 @@ private void doTestBucketedTableEvolution(HiveStorageFormat storageFormat, Schem Optional.empty()).getLayout().getHandle(); } else { - layoutHandle = getOnlyElement(metadata.getTableLayouts(session, tableHandle, new Constraint<>(TupleDomain.fromFixedValues(ImmutableMap.of(bucketColumnHandle(), singleBucket))), Optional.empty())).getTableLayout().getHandle(); + layoutHandle = metadata.getTableLayoutForConstraint(session, tableHandle, new Constraint<>(TupleDomain.fromFixedValues(ImmutableMap.of(bucketColumnHandle(), singleBucket))), Optional.empty()).getTableLayout().getHandle(); } result = readTable( @@ -2686,8 +2686,8 @@ protected ConnectorTableLayout getTableLayout(ConnectorSession session, Connecto Optional.empty()).getLayout(); } - List tableLayoutResults = metadata.getTableLayouts(session, tableHandle, constraint, Optional.empty()); - return getOnlyElement(tableLayoutResults).getTableLayout(); + ConnectorTableLayoutResult tableLayoutResult = metadata.getTableLayoutForConstraint(session, tableHandle, constraint, Optional.empty()); + return tableLayoutResult.getTableLayout(); } @Test diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/AbstractTestHiveFileSystem.java b/presto-hive/src/test/java/com/facebook/presto/hive/AbstractTestHiveFileSystem.java index 5ec04cc1aa4c3..866d3f797d7fc 100644 --- a/presto-hive/src/test/java/com/facebook/presto/hive/AbstractTestHiveFileSystem.java +++ b/presto-hive/src/test/java/com/facebook/presto/hive/AbstractTestHiveFileSystem.java @@ -470,8 +470,8 @@ private void createTable(MetastoreContext metastoreContext, SchemaTableName tabl assertEquals(filterNonHiddenColumnMetadata(tableMetadata.getColumns()), columns); // verify the data - List tableLayoutResults = metadata.getTableLayouts(session, hiveTableHandle, Constraint.alwaysTrue(), Optional.empty()); - HiveTableLayoutHandle layoutHandle = (HiveTableLayoutHandle) getOnlyElement(tableLayoutResults).getTableLayout().getHandle(); + ConnectorTableLayoutResult tableLayoutResult = metadata.getTableLayoutForConstraint(session, hiveTableHandle, Constraint.alwaysTrue(), Optional.empty()); + HiveTableLayoutHandle layoutHandle = (HiveTableLayoutHandle) tableLayoutResult.getTableLayout().getHandle(); assertEquals(layoutHandle.getPartitions().get().size(), 1); ConnectorSplitSource splitSource = splitManager.getSplits(transaction.getTransactionHandle(), session, layoutHandle, SPLIT_SCHEDULING_CONTEXT); ConnectorSplit split = getOnlyElement(getAllSplits(splitSource)); diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/HiveFileSystemTestUtils.java b/presto-hive/src/test/java/com/facebook/presto/hive/HiveFileSystemTestUtils.java index 0d509285ed16a..41f5000b446be 100644 --- a/presto-hive/src/test/java/com/facebook/presto/hive/HiveFileSystemTestUtils.java +++ b/presto-hive/src/test/java/com/facebook/presto/hive/HiveFileSystemTestUtils.java @@ -50,7 +50,6 @@ import static com.facebook.presto.testing.MaterializedResult.materializeSourceDataStream; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.collect.ImmutableList.toImmutableList; -import static com.google.common.collect.Iterables.getOnlyElement; public class HiveFileSystemTestUtils { @@ -74,8 +73,8 @@ public static MaterializedResult readTable(SchemaTableName tableName, ConnectorTableHandle table = getTableHandle(metadata, tableName, session); List columnHandles = ImmutableList.copyOf(metadata.getColumnHandles(session, table).values()); - List tableLayoutResults = metadata.getTableLayouts(session, table, Constraint.alwaysTrue(), Optional.empty()); - HiveTableLayoutHandle layoutHandle = (HiveTableLayoutHandle) getOnlyElement(tableLayoutResults).getTableLayout().getHandle(); + ConnectorTableLayoutResult tableLayoutResult = metadata.getTableLayoutForConstraint(session, table, Constraint.alwaysTrue(), Optional.empty()); + HiveTableLayoutHandle layoutHandle = (HiveTableLayoutHandle) tableLayoutResult.getTableLayout().getHandle(); TableHandle tableHandle = new TableHandle(new ConnectorId(tableName.getSchemaName()), table, transaction.getTransactionHandle(), Optional.of(layoutHandle)); metadata.beginQuery(session); @@ -134,8 +133,8 @@ public static MaterializedResult filterTable(SchemaTableName tableName, session = newSession(config); ConnectorTableHandle table = getTableHandle(metadata, tableName, session); - List tableLayoutResults = metadata.getTableLayouts(session, table, Constraint.alwaysTrue(), Optional.empty()); - HiveTableLayoutHandle layoutHandle = (HiveTableLayoutHandle) getOnlyElement(tableLayoutResults).getTableLayout().getHandle(); + ConnectorTableLayoutResult tableLayoutResult = metadata.getTableLayoutForConstraint(session, table, Constraint.alwaysTrue(), Optional.empty()); + HiveTableLayoutHandle layoutHandle = (HiveTableLayoutHandle) tableLayoutResult.getTableLayout().getHandle(); TableHandle tableHandle = new TableHandle(new ConnectorId(tableName.getSchemaName()), table, transaction.getTransactionHandle(), Optional.of(layoutHandle)); metadata.beginQuery(session); @@ -190,8 +189,8 @@ public static int getSplitsCount(SchemaTableName tableName, session = newSession(config); ConnectorTableHandle table = getTableHandle(metadata, tableName, session); - List tableLayoutResults = metadata.getTableLayouts(session, table, Constraint.alwaysTrue(), Optional.empty()); - HiveTableLayoutHandle layoutHandle = (HiveTableLayoutHandle) getOnlyElement(tableLayoutResults).getTableLayout().getHandle(); + ConnectorTableLayoutResult tableLayoutResult = metadata.getTableLayoutForConstraint(session, table, Constraint.alwaysTrue(), Optional.empty()); + HiveTableLayoutHandle layoutHandle = (HiveTableLayoutHandle) tableLayoutResult.getTableLayout().getHandle(); TableHandle tableHandle = new TableHandle(new ConnectorId(tableName.getSchemaName()), table, transaction.getTransactionHandle(), Optional.of(layoutHandle)); metadata.beginQuery(session); diff --git a/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiMetadata.java b/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiMetadata.java index 34d2676d86015..24a04a82c44e1 100644 --- a/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiMetadata.java +++ b/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiMetadata.java @@ -115,7 +115,11 @@ public Optional getSystemTable(ConnectorSession session, SchemaTabl } @Override - public List getTableLayouts(ConnectorSession session, ConnectorTableHandle tableHandle, Constraint constraint, Optional> desiredColumns) + public ConnectorTableLayoutResult getTableLayoutForConstraint( + ConnectorSession session, + ConnectorTableHandle tableHandle, + Constraint constraint, + Optional> desiredColumns) { HudiTableHandle handle = (HudiTableHandle) tableHandle; Table table = getTable(session, tableHandle); @@ -127,7 +131,7 @@ public List getTableLayouts(ConnectorSession session partitionColumns, table.getParameters(), constraint.getSummary())); - return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary())); + return new ConnectorTableLayoutResult(layout, constraint.getSummary()); } @Override diff --git a/presto-jmx/src/main/java/com/facebook/presto/connector/jmx/JmxMetadata.java b/presto-jmx/src/main/java/com/facebook/presto/connector/jmx/JmxMetadata.java index f1b958959b246..a9de60e9f0699 100644 --- a/presto-jmx/src/main/java/com/facebook/presto/connector/jmx/JmxMetadata.java +++ b/presto-jmx/src/main/java/com/facebook/presto/connector/jmx/JmxMetadata.java @@ -263,11 +263,15 @@ public Map> listTableColumns(ConnectorSess } @Override - public List getTableLayouts(ConnectorSession session, ConnectorTableHandle table, Constraint constraint, Optional> desiredColumns) + public ConnectorTableLayoutResult getTableLayoutForConstraint( + ConnectorSession session, + ConnectorTableHandle table, + Constraint constraint, + Optional> desiredColumns) { JmxTableHandle handle = (JmxTableHandle) table; ConnectorTableLayout layout = new ConnectorTableLayout(new JmxTableLayoutHandle(handle, constraint.getSummary())); - return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary())); + return new ConnectorTableLayoutResult(layout, constraint.getSummary()); } @Override diff --git a/presto-kafka/src/main/java/com/facebook/presto/kafka/KafkaMetadata.java b/presto-kafka/src/main/java/com/facebook/presto/kafka/KafkaMetadata.java index 8792b566309bc..d2320f2ebdbc6 100644 --- a/presto-kafka/src/main/java/com/facebook/presto/kafka/KafkaMetadata.java +++ b/presto-kafka/src/main/java/com/facebook/presto/kafka/KafkaMetadata.java @@ -200,7 +200,11 @@ public ColumnMetadata getColumnMetadata(ConnectorSession session, ConnectorTable } @Override - public List getTableLayouts(ConnectorSession session, ConnectorTableHandle table, Constraint constraint, Optional> desiredColumns) + public ConnectorTableLayoutResult getTableLayoutForConstraint( + ConnectorSession session, + ConnectorTableHandle table, + Constraint constraint, + Optional> desiredColumns) { KafkaTableHandle handle = convertTableHandle(table); long startTimestamp = 0; @@ -224,7 +228,7 @@ public List getTableLayouts(ConnectorSession session } ConnectorTableLayout layout = new ConnectorTableLayout(new KafkaTableLayoutHandle(handle, startTimestamp, endTimestamp)); - return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary())); + return new ConnectorTableLayoutResult(layout, constraint.getSummary()); } @Override diff --git a/presto-kudu/src/main/java/com/facebook/presto/kudu/KuduMetadata.java b/presto-kudu/src/main/java/com/facebook/presto/kudu/KuduMetadata.java index 77eccc43f7a09..603b82c56716d 100755 --- a/presto-kudu/src/main/java/com/facebook/presto/kudu/KuduMetadata.java +++ b/presto-kudu/src/main/java/com/facebook/presto/kudu/KuduMetadata.java @@ -203,7 +203,8 @@ public KuduTableHandle getTableHandle(ConnectorSession session, SchemaTableName } @Override - public List getTableLayouts(ConnectorSession session, + public ConnectorTableLayoutResult getTableLayoutForConstraint( + ConnectorSession session, ConnectorTableHandle tableHandle, Constraint constraint, Optional> desiredColumns) @@ -211,7 +212,7 @@ public List getTableLayouts(ConnectorSession session KuduTableHandle handle = (KuduTableHandle) tableHandle; ConnectorTableLayout layout = new ConnectorTableLayout( new KuduTableLayoutHandle(handle, constraint.getSummary(), desiredColumns)); - return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary())); + return new ConnectorTableLayoutResult(layout, constraint.getSummary()); } @Override diff --git a/presto-lark-sheets/src/main/java/com/facebook/presto/lark/sheets/LarkSheetsMetadata.java b/presto-lark-sheets/src/main/java/com/facebook/presto/lark/sheets/LarkSheetsMetadata.java index 9549708c317d9..0e1ba5d8c0d6c 100644 --- a/presto-lark-sheets/src/main/java/com/facebook/presto/lark/sheets/LarkSheetsMetadata.java +++ b/presto-lark-sheets/src/main/java/com/facebook/presto/lark/sheets/LarkSheetsMetadata.java @@ -127,14 +127,15 @@ public Optional getSystemTable(ConnectorSession session, SchemaTabl } @Override - public List getTableLayouts(ConnectorSession session, + public ConnectorTableLayoutResult getTableLayoutForConstraint( + ConnectorSession session, ConnectorTableHandle table, Constraint constraint, Optional> desiredColumns) { LarkSheetsTableHandle tableHandle = (LarkSheetsTableHandle) table; ConnectorTableLayout layout = new ConnectorTableLayout(new LarkSheetsTableLayoutHandle(tableHandle)); - return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary())); + return new ConnectorTableLayoutResult(layout, constraint.getSummary()); } @Override diff --git a/presto-local-file/src/main/java/com/facebook/presto/localfile/LocalFileMetadata.java b/presto-local-file/src/main/java/com/facebook/presto/localfile/LocalFileMetadata.java index 3db1028497578..c4224fbc7c761 100644 --- a/presto-local-file/src/main/java/com/facebook/presto/localfile/LocalFileMetadata.java +++ b/presto-local-file/src/main/java/com/facebook/presto/localfile/LocalFileMetadata.java @@ -84,11 +84,15 @@ public List listTables(ConnectorSession session, String schemaN } @Override - public List getTableLayouts(ConnectorSession session, ConnectorTableHandle table, Constraint constraint, Optional> desiredColumns) + public ConnectorTableLayoutResult getTableLayoutForConstraint( + ConnectorSession session, + ConnectorTableHandle table, + Constraint constraint, + Optional> desiredColumns) { LocalFileTableHandle tableHandle = (LocalFileTableHandle) table; ConnectorTableLayout layout = new ConnectorTableLayout(new LocalFileTableLayoutHandle(tableHandle, constraint.getSummary())); - return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary())); + return new ConnectorTableLayoutResult(layout, constraint.getSummary()); } @Override diff --git a/presto-memory/src/main/java/com/facebook/presto/plugin/memory/MemoryMetadata.java b/presto-memory/src/main/java/com/facebook/presto/plugin/memory/MemoryMetadata.java index 82232eb153f70..b9741caf478c9 100644 --- a/presto-memory/src/main/java/com/facebook/presto/plugin/memory/MemoryMetadata.java +++ b/presto-memory/src/main/java/com/facebook/presto/plugin/memory/MemoryMetadata.java @@ -356,7 +356,7 @@ private void updateRowsOnHosts(MemoryTableHandle table, Collection fragme } @Override - public synchronized List getTableLayouts( + public synchronized ConnectorTableLayoutResult getTableLayoutForConstraint( ConnectorSession session, ConnectorTableHandle handle, Constraint constraint, @@ -375,7 +375,7 @@ public synchronized List getTableLayouts( tableDataFragments.get(memoryTableHandle.getTableId()).values()); MemoryTableLayoutHandle layoutHandle = new MemoryTableLayoutHandle(memoryTableHandle, expectedFragments); - return ImmutableList.of(new ConnectorTableLayoutResult(getTableLayout(session, layoutHandle), constraint.getSummary())); + return new ConnectorTableLayoutResult(getTableLayout(session, layoutHandle), constraint.getSummary()); } @Override diff --git a/presto-memory/src/test/java/com/facebook/presto/plugin/memory/TestMemoryMetadata.java b/presto-memory/src/test/java/com/facebook/presto/plugin/memory/TestMemoryMetadata.java index 529b46cf580e4..3eb2834008551 100644 --- a/presto-memory/src/test/java/com/facebook/presto/plugin/memory/TestMemoryMetadata.java +++ b/presto-memory/src/test/java/com/facebook/presto/plugin/memory/TestMemoryMetadata.java @@ -150,9 +150,9 @@ public void testReadTableBeforeCreationCompleted() assertTrue(tableNames.size() == 1, "Expected exactly one table"); ConnectorTableHandle tableHandle = metadata.getTableHandle(SESSION, tableName); - List tableLayouts = metadata.getTableLayouts(SESSION, tableHandle, Constraint.alwaysTrue(), Optional.empty()); - assertTrue(tableLayouts.size() == 1, "Expected exactly one layout."); - ConnectorTableLayout tableLayout = tableLayouts.get(0).getTableLayout(); + ConnectorTableLayoutResult tableLayoutResult = metadata.getTableLayoutForConstraint(SESSION, tableHandle, Constraint.alwaysTrue(), Optional.empty()); + assertTrue(tableLayoutResult != null, "Table layout is null."); + ConnectorTableLayout tableLayout = tableLayoutResult.getTableLayout(); ConnectorTableLayoutHandle tableLayoutHandle = tableLayout.getHandle(); assertTrue(tableLayoutHandle instanceof MemoryTableLayoutHandle); assertTrue(((MemoryTableLayoutHandle) tableLayoutHandle).getDataFragments().isEmpty(), "Data fragments should be empty"); diff --git a/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoMetadata.java b/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoMetadata.java index 6133db8f0ac53..cc348ff4ad08a 100644 --- a/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoMetadata.java +++ b/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoMetadata.java @@ -153,7 +153,11 @@ public ColumnMetadata getColumnMetadata(ConnectorSession session, ConnectorTable } @Override - public List getTableLayouts(ConnectorSession session, ConnectorTableHandle table, Constraint constraint, Optional> desiredColumns) + public ConnectorTableLayoutResult getTableLayoutForConstraint( + ConnectorSession session, + ConnectorTableHandle table, + Constraint constraint, + Optional> desiredColumns) { MongoTableHandle tableHandle = (MongoTableHandle) table; @@ -183,7 +187,7 @@ public List getTableLayouts(ConnectorSession session Optional.empty(), localProperties.build()); - return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary())); + return new ConnectorTableLayoutResult(layout, constraint.getSummary()); } @Override @@ -192,8 +196,7 @@ public ConnectorTableLayout getTableLayout(ConnectorSession session, ConnectorTa MongoTableLayoutHandle layout = (MongoTableLayoutHandle) handle; // tables in this connector have a single layout - return getTableLayouts(session, layout.getTable(), Constraint.alwaysTrue(), Optional.empty()) - .get(0) + return getTableLayoutForConstraint(session, layout.getTable(), Constraint.alwaysTrue(), Optional.empty()) .getTableLayout(); } diff --git a/presto-pinot-toolkit/src/main/java/com/facebook/presto/pinot/PinotMetadata.java b/presto-pinot-toolkit/src/main/java/com/facebook/presto/pinot/PinotMetadata.java index 6dd350c8e7491..09af2042494d4 100644 --- a/presto-pinot-toolkit/src/main/java/com/facebook/presto/pinot/PinotMetadata.java +++ b/presto-pinot-toolkit/src/main/java/com/facebook/presto/pinot/PinotMetadata.java @@ -80,7 +80,7 @@ public PinotTableHandle getTableHandle(ConnectorSession session, SchemaTableName } @Override - public List getTableLayouts( + public ConnectorTableLayoutResult getTableLayoutForConstraint( ConnectorSession session, ConnectorTableHandle table, Constraint constraint, @@ -89,7 +89,7 @@ public List getTableLayouts( // Constraint's don't need to be pushed down since they are already taken care off by the pushdown logic PinotTableHandle pinotTableHandle = (PinotTableHandle) table; ConnectorTableLayout layout = new ConnectorTableLayout(new PinotTableLayoutHandle(pinotTableHandle)); - return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary())); + return new ConnectorTableLayoutResult(layout, constraint.getSummary()); } @Override diff --git a/presto-prometheus/src/main/java/com/facebook/presto/plugin/prometheus/PrometheusMetadata.java b/presto-prometheus/src/main/java/com/facebook/presto/plugin/prometheus/PrometheusMetadata.java index bb223c7836421..07b68eee5cb1f 100644 --- a/presto-prometheus/src/main/java/com/facebook/presto/plugin/prometheus/PrometheusMetadata.java +++ b/presto-prometheus/src/main/java/com/facebook/presto/plugin/prometheus/PrometheusMetadata.java @@ -77,11 +77,15 @@ public ConnectorTableHandle getTableHandle(ConnectorSession session, SchemaTable } @Override - public List getTableLayouts(ConnectorSession session, ConnectorTableHandle table, Constraint constraint, Optional> desiredColumns) + public ConnectorTableLayoutResult getTableLayoutForConstraint( + ConnectorSession session, + ConnectorTableHandle table, + Constraint constraint, + Optional> desiredColumns) { PrometheusTableHandle tableHandle = (PrometheusTableHandle) table; ConnectorTableLayout layout = new ConnectorTableLayout(new PrometheusTableLayoutHandle(tableHandle, constraint.getSummary())); - return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary())); + return new ConnectorTableLayoutResult(layout, constraint.getSummary()); } @Override diff --git a/presto-redis/src/main/java/com/facebook/presto/redis/RedisMetadata.java b/presto-redis/src/main/java/com/facebook/presto/redis/RedisMetadata.java index acb4fef35a7e8..8295dd7d0d90c 100644 --- a/presto-redis/src/main/java/com/facebook/presto/redis/RedisMetadata.java +++ b/presto-redis/src/main/java/com/facebook/presto/redis/RedisMetadata.java @@ -124,7 +124,7 @@ public ConnectorTableMetadata getTableMetadata(ConnectorSession session, Connect } @Override - public List getTableLayouts( + public ConnectorTableLayoutResult getTableLayoutForConstraint( ConnectorSession session, ConnectorTableHandle table, Constraint constraint, @@ -134,7 +134,7 @@ public List getTableLayouts( ConnectorTableLayout layout = new ConnectorTableLayout(new RedisTableLayoutHandle(tableHandle)); - return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary())); + return new ConnectorTableLayoutResult(layout, constraint.getSummary()); } @Override @@ -143,8 +143,7 @@ public ConnectorTableLayout getTableLayout(ConnectorSession session, ConnectorTa RedisTableLayoutHandle layout = convertLayout(handle); // tables in this connector have a single layout - return getTableLayouts(session, layout.getTable(), Constraint.alwaysTrue(), Optional.empty()) - .get(0) + return getTableLayoutForConstraint(session, layout.getTable(), Constraint.alwaysTrue(), Optional.empty()) .getTableLayout(); } diff --git a/presto-thrift-connector/src/main/java/com/facebook/presto/connector/thrift/ThriftMetadata.java b/presto-thrift-connector/src/main/java/com/facebook/presto/connector/thrift/ThriftMetadata.java index d0291feb4035e..63d499bcdc752 100644 --- a/presto-thrift-connector/src/main/java/com/facebook/presto/connector/thrift/ThriftMetadata.java +++ b/presto-thrift-connector/src/main/java/com/facebook/presto/connector/thrift/ThriftMetadata.java @@ -42,7 +42,6 @@ import com.google.common.cache.CacheBuilder; import com.google.common.cache.CacheLoader; import com.google.common.cache.LoadingCache; -import com.google.common.collect.ImmutableList; import jakarta.inject.Inject; import java.util.List; @@ -110,7 +109,7 @@ public ConnectorTableHandle getTableHandle(ConnectorSession session, SchemaTable } @Override - public List getTableLayouts( + public ConnectorTableLayoutResult getTableLayoutForConstraint( ConnectorSession session, ConnectorTableHandle table, Constraint constraint, @@ -122,7 +121,7 @@ public List getTableLayouts( tableHandle.getTableName(), desiredColumns, constraint.getSummary()); - return ImmutableList.of(new ConnectorTableLayoutResult(new ConnectorTableLayout(layoutHandle), constraint.getSummary())); + return new ConnectorTableLayoutResult(new ConnectorTableLayout(layoutHandle), constraint.getSummary()); } @Override diff --git a/presto-tpcds/src/main/java/com/facebook/presto/tpcds/TpcdsMetadata.java b/presto-tpcds/src/main/java/com/facebook/presto/tpcds/TpcdsMetadata.java index ab57475bb199a..7d9e964e6a541 100644 --- a/presto-tpcds/src/main/java/com/facebook/presto/tpcds/TpcdsMetadata.java +++ b/presto-tpcds/src/main/java/com/facebook/presto/tpcds/TpcdsMetadata.java @@ -100,7 +100,7 @@ public ConnectorTableHandle getTableHandle(ConnectorSession session, SchemaTable } @Override - public List getTableLayouts( + public ConnectorTableLayoutResult getTableLayoutForConstraint( ConnectorSession session, ConnectorTableHandle table, Constraint constraint, @@ -116,7 +116,7 @@ public List getTableLayouts( Optional.empty(), ImmutableList.of()); - return ImmutableList.of(new ConnectorTableLayoutResult(layout, constraint.getSummary())); + return new ConnectorTableLayoutResult(layout, constraint.getSummary()); } @Override @@ -124,8 +124,7 @@ public ConnectorTableLayout getTableLayout(ConnectorSession session, ConnectorTa { TpcdsTableLayoutHandle layout = (TpcdsTableLayoutHandle) handle; - return getTableLayouts(session, layout.getTable(), Constraint.alwaysTrue(), Optional.empty()) - .get(0) + return getTableLayoutForConstraint(session, layout.getTable(), Constraint.alwaysTrue(), Optional.empty()) .getTableLayout(); } diff --git a/presto-tpch/src/main/java/com/facebook/presto/tpch/TpchMetadata.java b/presto-tpch/src/main/java/com/facebook/presto/tpch/TpchMetadata.java index 6e27f7df994b6..255e80cbdafe2 100644 --- a/presto-tpch/src/main/java/com/facebook/presto/tpch/TpchMetadata.java +++ b/presto-tpch/src/main/java/com/facebook/presto/tpch/TpchMetadata.java @@ -186,7 +186,7 @@ public ConnectorTableHandle getTableHandleForStatisticsCollection(ConnectorSessi } @Override - public List getTableLayouts( + public ConnectorTableLayoutResult getTableLayoutForConstraint( ConnectorSession session, ConnectorTableHandle table, Constraint constraint, @@ -252,7 +252,7 @@ else if (tableHandle.getTableName().equals(TpchTable.LINE_ITEM.getTableName())) Optional.empty(), localProperties); - return ImmutableList.of(new ConnectorTableLayoutResult(layout, unenforcedConstraint)); + return new ConnectorTableLayoutResult(layout, unenforcedConstraint); } private Set filterValues(Set nullableValues, TpchColumn column, Constraint constraint) @@ -269,8 +269,7 @@ public ConnectorTableLayout getTableLayout(ConnectorSession session, ConnectorTa TpchTableLayoutHandle layout = (TpchTableLayoutHandle) handle; // tables in this connector have a single layout - return getTableLayouts(session, layout.getTable(), Constraint.alwaysTrue(), Optional.empty()) - .get(0) + return getTableLayoutForConstraint(session, layout.getTable(), Constraint.alwaysTrue(), Optional.empty()) .getTableLayout(); } diff --git a/presto-tpch/src/test/java/com/facebook/presto/tpch/TestTpchMetadata.java b/presto-tpch/src/test/java/com/facebook/presto/tpch/TestTpchMetadata.java index d6875ae59f612..b7a2737d46029 100644 --- a/presto-tpch/src/test/java/com/facebook/presto/tpch/TestTpchMetadata.java +++ b/presto-tpch/src/test/java/com/facebook/presto/tpch/TestTpchMetadata.java @@ -46,7 +46,6 @@ import static com.facebook.presto.spi.Constraint.alwaysTrue; import static com.facebook.presto.tpch.TpchMetadata.getPrestoType; import static com.facebook.presto.tpch.util.PredicateUtils.filterOutColumnFromPredicate; -import static com.google.common.collect.MoreCollectors.onlyElement; import static io.airlift.slice.Slices.utf8Slice; import static io.airlift.tpch.CustomerColumn.MARKET_SEGMENT; import static io.airlift.tpch.CustomerColumn.NAME; @@ -389,9 +388,7 @@ private static TupleDomain fixedValueTupleDomain(TpchMetadata tpch private static ConnectorTableLayoutResult getTableOnlyLayout(TpchMetadata tpchMetadata, ConnectorSession session, ConnectorTableHandle tableHandle, Constraint constraint) { - return tpchMetadata.getTableLayouts(session, tableHandle, constraint, Optional.empty()) - .stream() - .collect(onlyElement()); + return tpchMetadata.getTableLayoutForConstraint(session, tableHandle, constraint, Optional.empty()); } private ColumnStatistics noColumnStatistics() From 29b16570633670287c9954a1ffc0c2bd4f3b123a Mon Sep 17 00:00:00 2001 From: Christian Zentgraf Date: Wed, 20 Aug 2025 09:00:25 -0700 Subject: [PATCH 072/113] [native] Pass Presto build options to Velox build options directly Some Presto build options directly set Velox build option to enable certain features. Previously, configuration values that do not match Velox defaults might not set the Velox config appropriatly. For example, turning off parquet support would not actually disable it at the build level because the Velox config would be left to the default value which is ON. --- presto-native-execution/CMakeLists.txt | 60 +++++++++++--------------- 1 file changed, 24 insertions(+), 36 deletions(-) diff --git a/presto-native-execution/CMakeLists.txt b/presto-native-execution/CMakeLists.txt index d09c9fad6c98b..aea12fd9a8456 100644 --- a/presto-native-execution/CMakeLists.txt +++ b/presto-native-execution/CMakeLists.txt @@ -83,58 +83,46 @@ option(PRESTO_ENABLE_SPATIAL "Enable spatial support" ON) # and turn on int128. add_compile_definitions(FOLLY_HAVE_INT128_T=1 FOLLY_CFG_NO_COROUTINES) -if(PRESTO_ENABLE_S3) - set(VELOX_ENABLE_S3 - ON - CACHE BOOL "Build S3 support") -endif() +set(VELOX_ENABLE_S3 + ${PRESTO_ENABLE_S3} + CACHE BOOL "Build S3 support") -if(PRESTO_ENABLE_HDFS) - set(VELOX_ENABLE_HDFS - ON - CACHE BOOL "Build HDFS support") -endif() +set(VELOX_ENABLE_HDFS + ${PRESTO_ENABLE_HDFS} + CACHE BOOL "Build HDFS support") -if(PRESTO_ENABLE_GCS) - set(VELOX_ENABLE_GCS - ON - CACHE BOOL "Build GCS support") -endif() +set(VELOX_ENABLE_GCS + ${PRESTO_ENABLE_GCS} + CACHE BOOL "Build GCS support") -if(PRESTO_ENABLE_ABFS) - set(VELOX_ENABLE_ABFS - ON - CACHE BOOL "Build ABFS support") -endif() +set(VELOX_ENABLE_ABFS + ${PRESTO_ENABLE_ABFS} + CACHE BOOL "Build ABFS support") -if(PRESTO_ENABLE_PARQUET) - set(VELOX_ENABLE_PARQUET - ON - CACHE BOOL "Enable Parquet support") -endif() +set(VELOX_ENABLE_PARQUET + ${PRESTO_ENABLE_PARQUET} + CACHE BOOL "Enable Parquet support") +set(VELOX_ENABLE_REMOTE_FUNCTIONS + ${PRESTO_ENABLE_REMOTE_FUNCTIONS} + CACHE BOOL "Enable remote function support in Velox") if(PRESTO_ENABLE_REMOTE_FUNCTIONS) - set(VELOX_ENABLE_REMOTE_FUNCTIONS - ON - CACHE BOOL "Enable remote function support in Velox") add_compile_definitions(PRESTO_ENABLE_REMOTE_FUNCTIONS) endif() +set(VELOX_ENABLE_CUDF + ${PRESTO_ENABLE_CUDF} + CACHE BOOL "Enable cuDF support") if(PRESTO_ENABLE_CUDF) - set(VELOX_ENABLE_CUDF - ON - CACHE BOOL "Enable cuDF support") add_compile_definitions(PRESTO_ENABLE_CUDF) enable_language(CUDA) # Determine CUDA_ARCHITECTURES automatically. cmake_policy(SET CMP0104 NEW) endif() -if(PRESTO_ENABLE_SPATIAL) - set(VELOX_ENABLE_GEO - ON - CACHE BOOL "Enable Velox Geometry (aka spatial) support") -endif() +set(VELOX_ENABLE_GEO + ${PRESTO_ENABLE_SPATIAL} + CACHE BOOL "Enable Velox Geometry (aka spatial) support") set(VELOX_BUILD_TESTING OFF From eab926355537063fd5838015a0d55b89e4a51142 Mon Sep 17 00:00:00 2001 From: PRASHANT GOLASH <40184733+prashantgolash@users.noreply.github.com> Date: Tue, 9 Sep 2025 14:28:46 -0700 Subject: [PATCH 073/113] [Coordinator throttling] Scheduling Policies for Admission Control based on worker load (#25689) **Admission control scheduling policy** **Logic** Gather worker overload data from the added end point in PR - [https://github.com/prestodb/presto/pull/25687](https://l.facebook.com/l.php?u=https%3A%2F%2Fgithub.com%2Fprestodb%2Fpresto%2Fpull%2F25687&h=AT20QdO6Ld9h4KN4cMS_xFUtKTtNuKLc-glQIqSWbPgMh0KxW_3-UYhiBxO2gl-I8P_i8wG5gixSBWcGmglCWFShRKtOUlWnYZ5RqgBNHhmkGYrK2XJfKIjqXNj4Ltd4oWghm0eSCgnKL0UGvLr9mXLyJyA) Based on configured policies (cnt of overloaded workers or pct of overloaded workers) and cluster overload, queue the queries **Background** RFC PR: [https://github.com/prestodb/rfcs/pull/42](https://l.facebook.com/l.php?u=https%3A%2F%2Fgithub.com%2Fprestodb%2Frfcs%2Fpull%2F42&h=AT20QdO6Ld9h4KN4cMS_xFUtKTtNuKLc-glQIqSWbPgMh0KxW_3-UYhiBxO2gl-I8P_i8wG5gixSBWcGmglCWFShRKtOUlWnYZ5RqgBNHhmkGYrK2XJfKIjqXNj4Ltd4oWghm0eSCgnKL0UGvLr9mXLyJyA) **Metrics on queuing due to this feature:** Added following metrics - ClusterOverloadDuration - ClusterOverloadCount **Feature flag:** Right now feature is disabled. We can use coordinator configs to enable / add thresholds ## Summary by Sourcery Add cluster overload-based admission control scheduling policy and supporting infrastructure to throttle queries based on worker load. New Features: - Introduce ClusterResourceChecker for periodic cluster overload detection and query throttling based on worker load metrics - Implement CpuMemoryOverloadPolicy with count- and percentage-based overload thresholds and expose overload detection count and duration metrics - Add AdmissionControlBypassConfig to allow certain queries (e.g., DDL) to bypass admission control Enhancements: - Refactor DiscoveryNodeManager and SPI to use RemoteNodeStats for asynchronous fetching of node load metrics - Integrate cluster overload checks into InternalResourceGroup scheduling logic and update eligibility propagation on overload state changes - Introduce ClusterOverloadPolicyModule and DI bindings for policy selection and checker Tests: - Add unit tests for ClusterResourceChecker, CpuMemoryOverloadPolicy, and AdmissionControlBypassConfig - Update existing resource group and node manager tests to support clustering resource checker mocks Chores: - Add new configuration properties for cluster-overload throttling and internal-communication stats polling intervals --- .../execution/ClusterOverloadConfig.java | 111 +++++++ .../resourceGroups/InternalResourceGroup.java | 30 +- .../InternalResourceGroupManager.java | 37 ++- .../ClusterOverloadPolicy.java | 40 +++ .../ClusterOverloadPolicyFactory.java | 78 +++++ .../ClusterOverloadPolicyModule.java | 44 +++ .../ClusterOverloadStateListener.java | 37 +++ .../ClusterResourceChecker.java | 248 ++++++++++++++ .../CpuMemoryOverloadPolicy.java | 169 ++++++++++ .../presto/metadata/InMemoryNodeManager.java | 8 + .../presto/metadata/InternalNodeManager.java | 4 + .../presto/metadata/RemoteNodeStats.java | 56 ++++ .../metadata/ThriftRemoteNodeStats.java | 110 +++++++ .../server/InternalCommunicationConfig.java | 28 ++ .../execution/TestClusterOverloadConfig.java | 57 ++++ .../BenchmarkResourceGroup.java | 31 +- .../TestInternalResourceGroupManager.java | 8 +- .../resourceGroups/TestResourceGroups.java | 64 ++-- .../TestClusterResourceChecker.java | 305 ++++++++++++++++++ .../TestCpuMemoryOverloadPolicy.java | 254 +++++++++++++++ .../TestInternalCommunicationConfig.java | 10 +- .../presto/server/TestQueryStateInfo.java | 31 +- .../presto/metadata/DiscoveryNodeManager.java | 38 ++- .../presto/metadata/HttpRemoteNodeStats.java | 121 +++++++ .../presto/server/ServerMainModule.java | 6 + ...dingResourceGroupConfigurationManager.java | 35 +- .../presto/spark/PrestoSparkModule.java | 4 + .../node/PrestoSparkInternalNodeManager.java | 8 + 28 files changed, 1922 insertions(+), 50 deletions(-) create mode 100644 presto-main-base/src/main/java/com/facebook/presto/execution/ClusterOverloadConfig.java create mode 100644 presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/ClusterOverloadPolicy.java create mode 100644 presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/ClusterOverloadPolicyFactory.java create mode 100644 presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/ClusterOverloadPolicyModule.java create mode 100644 presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/ClusterOverloadStateListener.java create mode 100644 presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/ClusterResourceChecker.java create mode 100644 presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/CpuMemoryOverloadPolicy.java create mode 100644 presto-main-base/src/main/java/com/facebook/presto/metadata/RemoteNodeStats.java create mode 100644 presto-main-base/src/main/java/com/facebook/presto/metadata/ThriftRemoteNodeStats.java create mode 100644 presto-main-base/src/test/java/com/facebook/presto/execution/TestClusterOverloadConfig.java create mode 100644 presto-main-base/src/test/java/com/facebook/presto/execution/scheduler/clusterOverload/TestClusterResourceChecker.java create mode 100644 presto-main-base/src/test/java/com/facebook/presto/execution/scheduler/clusterOverload/TestCpuMemoryOverloadPolicy.java create mode 100644 presto-main/src/main/java/com/facebook/presto/metadata/HttpRemoteNodeStats.java diff --git a/presto-main-base/src/main/java/com/facebook/presto/execution/ClusterOverloadConfig.java b/presto-main-base/src/main/java/com/facebook/presto/execution/ClusterOverloadConfig.java new file mode 100644 index 0000000000000..a8ddf72707447 --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/execution/ClusterOverloadConfig.java @@ -0,0 +1,111 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.execution; + +import com.facebook.airlift.configuration.Config; + +public class ClusterOverloadConfig +{ + public static final String OVERLOAD_POLICY_CNT_BASED = "overload_worker_cnt_based_throttling"; + public static final String OVERLOAD_POLICY_PCT_BASED = "overload_worker_pct_based_throttling"; + private boolean clusterOverloadThrottlingEnabled; + private double allowedOverloadWorkersPct = 0.01; + private int allowedOverloadWorkersCnt; + private String overloadPolicyType = OVERLOAD_POLICY_CNT_BASED; + private int overloadCheckCacheTtlInSecs = 5; + + /** + * Gets the time-to-live for the cached cluster overload state. + * This determines how frequently the system will re-evaluate whether the cluster is overloaded. + * + * @return the cache TTL duration + */ + public int getOverloadCheckCacheTtlInSecs() + { + return overloadCheckCacheTtlInSecs; + } + + /** + * Gets the time-to-live for the cached cluster overload state. + * This determines how frequently the system will re-evaluate whether the cluster is overloaded. + * + * @return the cache TTL duration + */ + public int getOverloadCheckCacheTtlMillis() + { + return overloadCheckCacheTtlInSecs * 1000; + } + + /** + * Sets the time-to-live for the cached cluster overload state. + * + * @param overloadCheckCacheTtlInSecs the cache TTL duration + * @return this for chaining + */ + @Config("cluster.overload-check-cache-ttl-secs") + public ClusterOverloadConfig setOverloadCheckCacheTtlInSecs(int overloadCheckCacheTtlInSecs) + { + this.overloadCheckCacheTtlInSecs = overloadCheckCacheTtlInSecs; + return this; + } + + @Config("cluster-overload.enable-throttling") + public ClusterOverloadConfig setClusterOverloadThrottlingEnabled(boolean clusterOverloadThrottlingEnabled) + { + this.clusterOverloadThrottlingEnabled = clusterOverloadThrottlingEnabled; + return this; + } + + public boolean isClusterOverloadThrottlingEnabled() + { + return this.clusterOverloadThrottlingEnabled; + } + + @Config("cluster-overload.allowed-overload-workers-pct") + public ClusterOverloadConfig setAllowedOverloadWorkersPct(Double allowedOverloadWorkersPct) + { + this.allowedOverloadWorkersPct = allowedOverloadWorkersPct; + return this; + } + + public double getAllowedOverloadWorkersPct() + { + return this.allowedOverloadWorkersPct; + } + + @Config("cluster-overload.allowed-overload-workers-cnt") + public ClusterOverloadConfig setAllowedOverloadWorkersCnt(int allowedOverloadWorkersCnt) + { + this.allowedOverloadWorkersCnt = allowedOverloadWorkersCnt; + return this; + } + + public double getAllowedOverloadWorkersCnt() + { + return this.allowedOverloadWorkersCnt; + } + + @Config("cluster-overload.overload-policy-type") + public ClusterOverloadConfig setOverloadPolicyType(String overloadPolicyType) + { + // validate + this.overloadPolicyType = overloadPolicyType; + return this; + } + + public String getOverloadPolicyType() + { + return this.overloadPolicyType; + } +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/execution/resourceGroups/InternalResourceGroup.java b/presto-main-base/src/main/java/com/facebook/presto/execution/resourceGroups/InternalResourceGroup.java index e1504e7b50f21..82274c04b3550 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/execution/resourceGroups/InternalResourceGroup.java +++ b/presto-main-base/src/main/java/com/facebook/presto/execution/resourceGroups/InternalResourceGroup.java @@ -18,6 +18,7 @@ import com.facebook.airlift.units.Duration; import com.facebook.presto.execution.ManagedQueryExecution; import com.facebook.presto.execution.resourceGroups.WeightedFairQueue.Usage; +import com.facebook.presto.execution.scheduler.clusterOverload.ClusterResourceChecker; import com.facebook.presto.metadata.InternalNodeManager; import com.facebook.presto.server.QueryStateInfo; import com.facebook.presto.server.ResourceGroupInfo; @@ -96,6 +97,7 @@ public class InternalResourceGroup private final Function> additionalRuntimeInfo; private final Predicate shouldWaitForResourceManagerUpdate; private final InternalNodeManager nodeManager; + private final ClusterResourceChecker clusterResourceChecker; // Configuration // ============= @@ -166,12 +168,14 @@ protected InternalResourceGroup( boolean staticResourceGroup, Function> additionalRuntimeInfo, Predicate shouldWaitForResourceManagerUpdate, - InternalNodeManager nodeManager) + InternalNodeManager nodeManager, + ClusterResourceChecker clusterResourceChecker) { this.parent = requireNonNull(parent, "parent is null"); this.jmxExportListener = requireNonNull(jmxExportListener, "jmxExportListener is null"); this.executor = requireNonNull(executor, "executor is null"); this.nodeManager = requireNonNull(nodeManager, "node manager is null"); + this.clusterResourceChecker = requireNonNull(clusterResourceChecker, "clusterResourceChecker is null"); requireNonNull(name, "name is null"); if (parent.isPresent()) { id = new ResourceGroupId(parent.get().id, name); @@ -671,7 +675,8 @@ public InternalResourceGroup getOrCreateSubGroup(String name, boolean staticSegm staticResourceGroup && staticSegment, additionalRuntimeInfo, shouldWaitForResourceManagerUpdate, - nodeManager); + nodeManager, + clusterResourceChecker); // Sub group must use query priority to ensure ordering if (schedulingPolicy == QUERY_PRIORITY) { subGroup.setSchedulingPolicy(QUERY_PRIORITY); @@ -770,7 +775,7 @@ private void enqueueQuery(ManagedQueryExecution query) } // This method must be called whenever the group's eligibility to run more queries may have changed. - private void updateEligibility() + protected void updateEligibility() { checkState(Thread.holdsLock(root), "Must hold lock to update eligibility"); synchronized (root) { @@ -1019,6 +1024,11 @@ private boolean canRunMore() { checkState(Thread.holdsLock(root), "Must hold lock"); synchronized (root) { + // Check if more queries can be run on the cluster based on cluster overload + if (clusterResourceChecker.isClusterCurrentlyOverloaded()) { + return false; + } + if (cpuUsageMillis >= hardCpuLimitMillis) { return false; } @@ -1135,7 +1145,8 @@ public RootInternalResourceGroup( Executor executor, Function> additionalRuntimeInfo, Predicate shouldWaitForResourceManagerUpdate, - InternalNodeManager nodeManager) + InternalNodeManager nodeManager, + ClusterResourceChecker clusterResourceChecker) { super(Optional.empty(), name, @@ -1144,7 +1155,16 @@ public RootInternalResourceGroup( true, additionalRuntimeInfo, shouldWaitForResourceManagerUpdate, - nodeManager); + nodeManager, + clusterResourceChecker); + } + + public synchronized void updateEligibilityRecursively(InternalResourceGroup group) + { + group.updateEligibility(); + for (InternalResourceGroup subGroup : group.subGroups()) { + updateEligibilityRecursively(subGroup); + } } public synchronized void processQueuedQueries() diff --git a/presto-main-base/src/main/java/com/facebook/presto/execution/resourceGroups/InternalResourceGroupManager.java b/presto-main-base/src/main/java/com/facebook/presto/execution/resourceGroups/InternalResourceGroupManager.java index 6dbf4115d0442..f760cfa19b031 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/execution/resourceGroups/InternalResourceGroupManager.java +++ b/presto-main-base/src/main/java/com/facebook/presto/execution/resourceGroups/InternalResourceGroupManager.java @@ -19,6 +19,8 @@ import com.facebook.presto.execution.ManagedQueryExecution; import com.facebook.presto.execution.QueryManagerConfig; import com.facebook.presto.execution.resourceGroups.InternalResourceGroup.RootInternalResourceGroup; +import com.facebook.presto.execution.scheduler.clusterOverload.ClusterOverloadStateListener; +import com.facebook.presto.execution.scheduler.clusterOverload.ClusterResourceChecker; import com.facebook.presto.metadata.InternalNodeManager; import com.facebook.presto.resourcemanager.ResourceGroupService; import com.facebook.presto.server.ResourceGroupInfo; @@ -81,7 +83,7 @@ @ThreadSafe public final class InternalResourceGroupManager - implements ResourceGroupManager + implements ResourceGroupManager, ClusterOverloadStateListener { private static final Logger log = Logger.get(InternalResourceGroupManager.class); private static final File RESOURCE_GROUPS_CONFIGURATION = new File("etc/resource-groups.properties"); @@ -112,6 +114,7 @@ public final class InternalResourceGroupManager private final QueryManagerConfig queryManagerConfig; private final InternalNodeManager nodeManager; private AtomicBoolean isConfigurationManagerLoaded; + private final ClusterResourceChecker clusterResourceChecker; @Inject public InternalResourceGroupManager( @@ -121,7 +124,8 @@ public InternalResourceGroupManager( MBeanExporter exporter, ResourceGroupService resourceGroupService, ServerConfig serverConfig, - InternalNodeManager nodeManager) + InternalNodeManager nodeManager, + ClusterResourceChecker clusterResourceChecker) { this.queryManagerConfig = requireNonNull(queryManagerConfig, "queryManagerConfig is null"); this.exporter = requireNonNull(exporter, "exporter is null"); @@ -137,6 +141,7 @@ public InternalResourceGroupManager( this.resourceGroupRuntimeExecutor = new PeriodicTaskExecutor(resourceGroupRuntimeInfoRefreshInterval.toMillis(), refreshExecutor, this::refreshResourceGroupRuntimeInfo); configurationManagerFactories.putIfAbsent(LegacyResourceGroupConfigurationManager.NAME, new LegacyResourceGroupConfigurationManager.Factory()); this.isConfigurationManagerLoaded = new AtomicBoolean(false); + this.clusterResourceChecker = clusterResourceChecker; } @Override @@ -254,6 +259,8 @@ public ResourceGroupConfigurationManager getConfigurationManager() @PreDestroy public void destroy() { + // Unregister from cluster overload state changes + clusterResourceChecker.removeListener(this); refreshExecutor.shutdownNow(); resourceGroupRuntimeExecutor.stop(); } @@ -275,6 +282,9 @@ public void start() if (isResourceManagerEnabled) { resourceGroupRuntimeExecutor.start(); } + + // Register as listener for cluster overload state changes + clusterResourceChecker.addListener(this); } } @@ -396,7 +406,7 @@ private synchronized void createGroupIfNecessary(SelectionContext context, Ex else { RootInternalResourceGroup root; if (!isResourceManagerEnabled) { - root = new RootInternalResourceGroup(id.getSegments().get(0), this::exportGroup, executor, ignored -> Optional.empty(), rg -> false, nodeManager); + root = new RootInternalResourceGroup(id.getSegments().get(0), this::exportGroup, executor, ignored -> Optional.empty(), rg -> false, nodeManager, clusterResourceChecker); } else { root = new RootInternalResourceGroup( @@ -409,7 +419,8 @@ private synchronized void createGroupIfNecessary(SelectionContext context, Ex resourceGroupRuntimeInfosSnapshot::get, lastUpdatedResourceGroupRuntimeInfo::get, concurrencyThreshold), - nodeManager); + nodeManager, + clusterResourceChecker); } group = root; rootGroups.add(root); @@ -463,6 +474,24 @@ public int getQueriesQueuedOnInternal() return queriesQueuedInternal; } + @Override + public void onClusterEnteredOverloadedState() + { + // Resource groups will handle overload state through their existing admission control logic + // No additional action needed here as queries will be queued automatically + } + + @Override + public void onClusterExitedOverloadedState() + { + log.info("Cluster exited overloaded state, updating eligibility for all resource groups"); + for (RootInternalResourceGroup rootGroup : rootGroups) { + synchronized (rootGroup) { + rootGroup.updateEligibilityRecursively(rootGroup); + } + } + } + @Managed public long getLastSchedulingCycleRuntimeDelayMs() { diff --git a/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/ClusterOverloadPolicy.java b/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/ClusterOverloadPolicy.java new file mode 100644 index 0000000000000..175c18afe5405 --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/ClusterOverloadPolicy.java @@ -0,0 +1,40 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.facebook.presto.execution.scheduler.clusterOverload; + +import com.facebook.presto.metadata.InternalNodeManager; + +/** + * Interface for policies that determine if cluster is overloaded. + * Implementations can check various metrics from NodeStats to determine + * if a worker is overloaded and queries should be throttled. + */ +public interface ClusterOverloadPolicy +{ + /** + * Checks if cluster is overloaded. + * + * @param nodeManager The node manager to get node information + * @return true if cluster is overloaded, false otherwise + */ + boolean isClusterOverloaded(InternalNodeManager nodeManager); + + /** + * Gets the name of the policy. + * + * @return The name of the policy + */ + String getName(); +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/ClusterOverloadPolicyFactory.java b/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/ClusterOverloadPolicyFactory.java new file mode 100644 index 0000000000000..2d6b9b82a8962 --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/ClusterOverloadPolicyFactory.java @@ -0,0 +1,78 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.execution.scheduler.clusterOverload; + +import com.google.common.collect.ImmutableMap; +import jakarta.inject.Inject; + +import java.util.Map; +import java.util.Optional; + +import static java.util.Objects.requireNonNull; + +/** + * Factory for creating ClusterOverloadPolicy instances. + * This allows for extensible policy creation based on configuration. + */ +public class ClusterOverloadPolicyFactory +{ + private final Map policies; + + @Inject + public ClusterOverloadPolicyFactory(ClusterOverloadPolicy clusterOverloadPolicy) + { + requireNonNull(clusterOverloadPolicy, "clusterOverloadPolicy is null"); + + // Register available policies + ImmutableMap.Builder policiesBuilder = ImmutableMap.builder(); + + // Add the default overload policy - use the injected instance + policiesBuilder.put(clusterOverloadPolicy.getName(), clusterOverloadPolicy); + + // Add more policies here as they are implemented + this.policies = policiesBuilder.build(); + } + + /** + * Get a policy by name. + * + * @param name The name of the policy to get + * @return The policy, or empty if no policy with that name exists + */ + public Optional getPolicy(String name) + { + return Optional.ofNullable(policies.get(name)); + } + + /** + * Get the default policy. + * + * @return The default policy + */ + public ClusterOverloadPolicy getDefaultPolicy() + { + // Default to CPU/Memory policy + return policies.get("cpu-memory-overload"); + } + + /** + * Get all available policies. + * + * @return Map of policy name to policy + */ + public Map getPolicies() + { + return policies; + } +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/ClusterOverloadPolicyModule.java b/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/ClusterOverloadPolicyModule.java new file mode 100644 index 0000000000000..207a0fd6dd198 --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/ClusterOverloadPolicyModule.java @@ -0,0 +1,44 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.execution.scheduler.clusterOverload; + +import com.facebook.airlift.configuration.AbstractConfigurationAwareModule; +import com.facebook.presto.execution.ClusterOverloadConfig; +import com.google.inject.Binder; + +import static com.facebook.airlift.configuration.ConfigBinder.configBinder; +import static com.google.inject.Scopes.SINGLETON; + +/** + * Provides bindings for the node overload policy and cluster resource checker. + */ +public class ClusterOverloadPolicyModule + extends AbstractConfigurationAwareModule +{ + @Override + protected void setup(Binder binder) + { + // Bind the default node overload policy + binder.bind(ClusterOverloadPolicy.class).to(CpuMemoryOverloadPolicy.class).in(SINGLETON); + + // Bind the node overload policy factory + binder.bind(ClusterOverloadPolicyFactory.class).in(SINGLETON); + + // Bind the cluster resource checker + binder.bind(ClusterResourceChecker.class).in(SINGLETON); + + // Bind the cluster overload config + configBinder(binder).bindConfig(ClusterOverloadConfig.class); + } +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/ClusterOverloadStateListener.java b/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/ClusterOverloadStateListener.java new file mode 100644 index 0000000000000..86d0845545e01 --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/ClusterOverloadStateListener.java @@ -0,0 +1,37 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.execution.scheduler.clusterOverload; + +/** + * Listener interface for receiving notifications about cluster overload state changes. + * This interface allows components to react to changes in cluster capacity and + * resource availability, particularly when the cluster transitions from an overloaded + * state back to a normal state. + */ +public interface ClusterOverloadStateListener +{ + /** + * Called when the cluster enters an overloaded state. + * This indicates that the cluster resources are under stress and + * new query admissions should be restricted. + */ + void onClusterEnteredOverloadedState(); + + /** + * Called when the cluster exits the overloaded state. + * This indicates that cluster resources have recovered and + * normal query processing can resume. + */ + void onClusterExitedOverloadedState(); +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/ClusterResourceChecker.java b/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/ClusterResourceChecker.java new file mode 100644 index 0000000000000..683fd703e970d --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/ClusterResourceChecker.java @@ -0,0 +1,248 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.execution.scheduler.clusterOverload; + +import com.facebook.airlift.log.Logger; +import com.facebook.airlift.stats.CounterStat; +import com.facebook.airlift.stats.TimeStat; +import com.facebook.presto.execution.ClusterOverloadConfig; +import com.facebook.presto.metadata.InternalNodeManager; +import com.google.errorprone.annotations.ThreadSafe; +import com.google.inject.Singleton; +import jakarta.annotation.PostConstruct; +import jakarta.annotation.PreDestroy; +import jakarta.inject.Inject; +import org.weakref.jmx.Managed; +import org.weakref.jmx.Nested; + +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicLong; + +import static com.facebook.airlift.concurrent.Threads.threadsNamed; +import static java.util.Objects.requireNonNull; +import static java.util.concurrent.Executors.newSingleThreadScheduledExecutor; + +/** + * Provides methods to check if more queries can be run on the cluster + * based on various resource constraints. + */ +@Singleton +@ThreadSafe +public class ClusterResourceChecker +{ + private static final Logger log = Logger.get(ClusterResourceChecker.class); + + private final ClusterOverloadPolicy clusterOverloadPolicy; + private final ClusterOverloadConfig config; + private final AtomicBoolean cachedOverloadState = new AtomicBoolean(false); + private final AtomicLong lastCheckTimeMillis = new AtomicLong(0); + private final CounterStat overloadDetectionCount = new CounterStat(); + private final TimeStat timeSinceLastCheck = new TimeStat(); + private final AtomicLong overloadStartTimeMillis = new AtomicLong(0); + private final CopyOnWriteArrayList listeners = new CopyOnWriteArrayList<>(); + + private final ScheduledExecutorService overloadCheckerExecutor; + private final InternalNodeManager nodeManager; + + @Inject + public ClusterResourceChecker(ClusterOverloadPolicy clusterOverloadPolicy, ClusterOverloadConfig config, InternalNodeManager nodeManager) + { + this.clusterOverloadPolicy = requireNonNull(clusterOverloadPolicy, "clusterOverloadPolicy is null"); + this.config = requireNonNull(config, "config is null"); + this.nodeManager = requireNonNull(nodeManager, "nodeManager is null"); + this.overloadCheckerExecutor = newSingleThreadScheduledExecutor(threadsNamed("cluster-overload-checker-%s")); + } + + @PostConstruct + public void start() + { + if (config.isClusterOverloadThrottlingEnabled()) { + long checkIntervalMillis = Math.max(1000, config.getOverloadCheckCacheTtlInSecs() * 1000L); + overloadCheckerExecutor.scheduleWithFixedDelay(() -> { + try { + performPeriodicOverloadCheck(); + } + catch (Exception e) { + log.error(e, "Error polling cluster overload state"); + } + }, checkIntervalMillis, checkIntervalMillis, TimeUnit.MILLISECONDS); + log.info("Started periodic cluster overload checker with interval: %d milliseconds", checkIntervalMillis); + // Perform initial check + performPeriodicOverloadCheck(); + } + } + + @PreDestroy + public void stop() + { + overloadCheckerExecutor.shutdownNow(); + } + + /** + * Registers a listener to be notified when the cluster exits the overloaded state. + * + * @param listener the listener to register + */ + public void addListener(ClusterOverloadStateListener listener) + { + requireNonNull(listener, "listener is null"); + listeners.add(listener); + } + + /** + * Removes a previously registered listener. + * + * @param listener the listener to remove + */ + public void removeListener(ClusterOverloadStateListener listener) + { + listeners.remove(listener); + } + + /** + * Performs a periodic check of cluster overload state. + * This method is called by the periodic task when throttling is enabled. + * Updates JMX metrics and notifies listeners when cluster exits overloaded state. + */ + private void performPeriodicOverloadCheck() + { + try { + long currentTimeMillis = System.currentTimeMillis(); + long lastCheckTime = lastCheckTimeMillis.get(); + + if (lastCheckTime > 0) { + timeSinceLastCheck.add(currentTimeMillis - lastCheckTime, TimeUnit.MILLISECONDS); + } + + boolean isOverloaded = clusterOverloadPolicy.isClusterOverloaded(nodeManager); + synchronized (this) { + boolean wasOverloaded = cachedOverloadState.getAndSet(isOverloaded); + lastCheckTimeMillis.set(currentTimeMillis); + + if (isOverloaded && !wasOverloaded) { + overloadDetectionCount.update(1); + overloadStartTimeMillis.set(currentTimeMillis); + log.info("Cluster entered overloaded state via periodic check"); + } + else if (!isOverloaded && wasOverloaded) { + long overloadDuration = currentTimeMillis - overloadStartTimeMillis.get(); + log.info("Cluster exited overloaded state after %d ms via periodic check", overloadDuration); + overloadStartTimeMillis.set(0); + // Notify listeners that cluster exited overload state + notifyClusterExitedOverloadedState(); + } + } + + log.debug("Periodic overload check completed: %s", isOverloaded ? "OVERLOADED" : "NOT OVERLOADED"); + } + catch (Exception e) { + log.error(e, "Error during periodic cluster overload check"); + } + } + + /** + * Returns the current overload state of the cluster. + * @return true if cluster is overloaded, false otherwise + */ + public boolean isClusterCurrentlyOverloaded() + { + if (!config.isClusterOverloadThrottlingEnabled()) { + return false; + } + + return cachedOverloadState.get(); + } + + /** + * Notifies all registered listeners that the cluster has exited the overloaded state. + */ + private void notifyClusterExitedOverloadedState() + { + for (ClusterOverloadStateListener listener : listeners) { + listener.onClusterExitedOverloadedState(); + } + } + + /** + * Returns whether cluster overload throttling is enabled. + * When disabled, the cluster overload check will be bypassed. + * + * @return true if throttling is enabled, false otherwise + */ + @Managed + public boolean isClusterOverloadThrottlingEnabled() + { + return config.isClusterOverloadThrottlingEnabled(); + } + + /** + * Returns whether the cluster is currently in an overloaded state. + * This is exposed as a JMX metric for monitoring. + * + * @return true if the cluster is overloaded, false otherwise + */ + @Managed + public boolean isClusterOverloaded() + { + return cachedOverloadState.get(); + } + + /** + * Returns the number of times the cluster has entered an overloaded state. + * + * @return counter of overload detections + */ + @Managed + @Nested + public CounterStat getOverloadDetectionCount() + { + return overloadDetectionCount; + } + + /** + * Returns statistics about the time between overload checks. + * + * @return time statistics for overload checks + */ + @Managed + @Nested + public TimeStat getTimeSinceLastCheck() + { + return timeSinceLastCheck; + } + + /** + * Returns the duration in milliseconds that the cluster has been in an overloaded state. + * Returns 0 if the cluster is not currently overloaded. + * + * Note: This method reads two atomic fields but doesn't need synchronization because: + * 1. Single writer (periodic task) ensures consistent updates + * 2. Atomic fields provide memory visibility guarantees + * 3. Slight inconsistency in edge cases is acceptable for monitoring metrics + * + * @return duration in milliseconds of current overload state + */ + @Managed + public long getOverloadDurationMillis() + { + long startTime = overloadStartTimeMillis.get(); + if (startTime == 0 || !cachedOverloadState.get()) { + return 0; + } + return System.currentTimeMillis() - startTime; + } +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/CpuMemoryOverloadPolicy.java b/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/CpuMemoryOverloadPolicy.java new file mode 100644 index 0000000000000..14d5387454178 --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/clusterOverload/CpuMemoryOverloadPolicy.java @@ -0,0 +1,169 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.facebook.presto.execution.scheduler.clusterOverload; + +import com.facebook.airlift.log.Logger; +import com.facebook.presto.execution.ClusterOverloadConfig; +import com.facebook.presto.metadata.InternalNode; +import com.facebook.presto.metadata.InternalNodeManager; +import com.facebook.presto.spi.NodeLoadMetrics; +import jakarta.inject.Inject; + +import java.util.HashSet; +import java.util.Optional; +import java.util.Set; + +import static com.facebook.presto.spi.NodeState.ACTIVE; +import static java.lang.String.format; +import static java.util.Objects.requireNonNull; + +/** + * A policy that checks if cluster is overloaded based on CPU or memory metrics. + * Supports two modes of operation: + * - Percentage-based: Checks if the percentage of overloaded workers exceeds a threshold + * - Count-based: Checks if the absolute count of overloaded workers exceeds a threshold + */ +public class CpuMemoryOverloadPolicy + implements ClusterOverloadPolicy +{ + private static final Logger log = Logger.get(CpuMemoryOverloadPolicy.class); + + private final double allowedOverloadWorkersPct; + private final double allowedOverloadWorkersCnt; + private final String policyType; + + @Inject + public CpuMemoryOverloadPolicy(ClusterOverloadConfig config) + { + this.allowedOverloadWorkersPct = config.getAllowedOverloadWorkersPct(); + this.allowedOverloadWorkersCnt = config.getAllowedOverloadWorkersCnt(); + this.policyType = requireNonNull(config.getOverloadPolicyType(), "policyType is null"); + } + + @Override + public boolean isClusterOverloaded(InternalNodeManager nodeManager) + { + Set activeNodes = nodeManager.getNodes(ACTIVE); + if (activeNodes.isEmpty()) { + return false; + } + + OverloadStats stats = collectOverloadStats(activeNodes, nodeManager); + return evaluateOverload(stats, activeNodes.size()); + } + + private OverloadStats collectOverloadStats(Set activeNodes, InternalNodeManager nodeManager) + { + Set overloadedNodeIds = new HashSet<>(); + int overloadedWorkersCnt = 0; + + for (InternalNode node : activeNodes) { + Optional metricsOptional = nodeManager.getNodeLoadMetrics(node.getNodeIdentifier()); + String nodeId = node.getNodeIdentifier(); + + if (!metricsOptional.isPresent()) { + continue; + } + + NodeLoadMetrics metrics = metricsOptional.get(); + + // Check for CPU overload + if (metrics.getCpuOverload()) { + overloadedNodeIds.add(String.format("%s (CPU overloaded)", nodeId)); + overloadedWorkersCnt++; + continue; + } + + // Check for memory overload + if (metrics.getMemoryOverload()) { + overloadedNodeIds.add(String.format("%s (Memory overloaded)", nodeId)); + overloadedWorkersCnt++; + } + } + + return new OverloadStats(overloadedNodeIds, overloadedWorkersCnt); + } + + private boolean evaluateOverload(OverloadStats stats, int totalNodes) + { + boolean isOverloaded; + + if (ClusterOverloadConfig.OVERLOAD_POLICY_PCT_BASED.equals(policyType)) { + double overloadedWorkersPct = (double) stats.getOverloadedWorkersCnt() / totalNodes; + isOverloaded = overloadedWorkersPct > allowedOverloadWorkersPct; + + if (isOverloaded) { + logOverload( + String.format("%s%% of workers are overloaded (threshold: %s%%)", + format("%.2f", overloadedWorkersPct * 100), + format("%.2f", allowedOverloadWorkersPct * 100)), + stats.getOverloadedNodeIds()); + } + } + else if (ClusterOverloadConfig.OVERLOAD_POLICY_CNT_BASED.equals(policyType)) { + isOverloaded = stats.getOverloadedWorkersCnt() > allowedOverloadWorkersCnt; + + if (isOverloaded) { + logOverload( + String.format("%s workers are overloaded (threshold: %s workers)", + stats.getOverloadedWorkersCnt(), allowedOverloadWorkersCnt), + stats.getOverloadedNodeIds()); + } + } + else { + throw new IllegalStateException("Unknown cluster overload policy type: " + policyType); + } + + return isOverloaded; + } + + private void logOverload(String message, Set overloadedNodeIds) + { + log.warn("Cluster is overloaded: " + message); + if (!overloadedNodeIds.isEmpty()) { + log.warn("Overloaded nodes: %s", String.join(", ", overloadedNodeIds)); + } + } + + @Override + public String getName() + { + return "cpu-memory-overload-" + + (ClusterOverloadConfig.OVERLOAD_POLICY_PCT_BASED.equals(policyType) ? "pct" : "cnt"); + } + + // Helper class to encapsulate overload statistics + private static class OverloadStats + { + private final Set overloadedNodeIds; + private final int overloadedWorkersCnt; + + public OverloadStats(Set overloadedNodeIds, int overloadedWorkersCnt) + { + this.overloadedNodeIds = overloadedNodeIds; + this.overloadedWorkersCnt = overloadedWorkersCnt; + } + + public Set getOverloadedNodeIds() + { + return overloadedNodeIds; + } + + public int getOverloadedWorkersCnt() + { + return overloadedWorkersCnt; + } + } +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/InMemoryNodeManager.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/InMemoryNodeManager.java index 69ed24cb1c5bd..f0a3ca28578fa 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/metadata/InMemoryNodeManager.java +++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/InMemoryNodeManager.java @@ -15,6 +15,7 @@ import com.facebook.presto.client.NodeVersion; import com.facebook.presto.spi.ConnectorId; +import com.facebook.presto.spi.NodeLoadMetrics; import com.facebook.presto.spi.NodeState; import com.google.common.collect.HashMultimap; import com.google.common.collect.ImmutableList; @@ -27,6 +28,7 @@ import java.net.URI; import java.util.ArrayList; import java.util.List; +import java.util.Optional; import java.util.Set; import java.util.function.Consumer; import java.util.stream.Stream; @@ -179,4 +181,10 @@ public synchronized void removeNodeChangeListener(Consumer listener) { listeners.remove(requireNonNull(listener, "listener is null")); } + + @Override + public Optional getNodeLoadMetrics(String nodeIdentifier) + { + return Optional.empty(); + } } diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/InternalNodeManager.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/InternalNodeManager.java index 81fa60c8165d7..c080dc40fd42d 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/metadata/InternalNodeManager.java +++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/InternalNodeManager.java @@ -14,8 +14,10 @@ package com.facebook.presto.metadata; import com.facebook.presto.spi.ConnectorId; +import com.facebook.presto.spi.NodeLoadMetrics; import com.facebook.presto.spi.NodeState; +import java.util.Optional; import java.util.Set; import java.util.function.Consumer; @@ -46,4 +48,6 @@ public interface InternalNodeManager void addNodeChangeListener(Consumer listener); void removeNodeChangeListener(Consumer listener); + + Optional getNodeLoadMetrics(String nodeIdentifier); } diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/RemoteNodeStats.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/RemoteNodeStats.java new file mode 100644 index 0000000000000..a9645cd524491 --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/RemoteNodeStats.java @@ -0,0 +1,56 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.metadata; + +import com.facebook.presto.spi.NodeStats; + +import java.util.Optional; + +/** + * Interface for retrieving statistics from remote nodes in a Presto cluster. + *

+ * This interface provides a mechanism to asynchronously fetch and cache node statistics + * from remote Presto worker nodes. Implementations handle the communication protocol + * details (HTTP, Thrift, etc.) and provide a unified way to access node health and + * performance metrics. + *

+ * The interface supports lazy loading and caching of node statistics to minimize + * network overhead while ensuring that cluster management components have access + * to current node state information for scheduling and health monitoring decisions. + */ +public interface RemoteNodeStats +{ + /** + * Returns the cached node statistics if available. + *

+ * This method returns the most recently fetched statistics for the remote node. + * If no statistics have been fetched yet or if the last fetch failed, this + * method returns an empty Optional. + * + * @return an Optional containing the node statistics if available, empty otherwise + */ + Optional getNodeStats(); + + /** + * Triggers an asynchronous refresh of the node statistics. + *

+ * This method initiates a background request to fetch the latest statistics + * from the remote node. The operation is non-blocking and the results will + * be available through subsequent calls to {@link #getNodeStats()}. + *

+ * Implementations should handle network failures gracefully and avoid + * overwhelming the remote node with excessive requests. + */ + void asyncRefresh(); +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/ThriftRemoteNodeStats.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/ThriftRemoteNodeStats.java new file mode 100644 index 0000000000000..185df3be1c143 --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/ThriftRemoteNodeStats.java @@ -0,0 +1,110 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.metadata; + +import com.facebook.airlift.log.Logger; +import com.facebook.airlift.units.Duration; +import com.facebook.drift.client.DriftClient; +import com.facebook.presto.server.thrift.ThriftServerInfoClient; +import com.facebook.presto.spi.NodeState; +import com.facebook.presto.spi.NodeStats; +import com.google.common.util.concurrent.FutureCallback; +import com.google.common.util.concurrent.Futures; +import com.google.common.util.concurrent.ListenableFuture; +import com.google.errorprone.annotations.ThreadSafe; +import jakarta.annotation.Nullable; + +import java.net.URI; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; + +import static com.facebook.airlift.units.Duration.nanosSince; +import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.util.concurrent.MoreExecutors.directExecutor; +import static java.util.Objects.requireNonNull; +import static java.util.concurrent.TimeUnit.SECONDS; + +@ThreadSafe +public class ThriftRemoteNodeStats + implements RemoteNodeStats +{ + private static final Logger log = Logger.get(ThriftRemoteNodeStats.class); + + private final ThriftServerInfoClient thriftClient; + private final long refreshIntervalMillis; + private final AtomicReference> nodeStats = new AtomicReference<>(Optional.empty()); + private final AtomicBoolean requestInflight = new AtomicBoolean(); + private final AtomicLong lastUpdateNanos = new AtomicLong(); + private final URI stateInfoUri; + private final AtomicLong lastWarningLogged = new AtomicLong(); + + public ThriftRemoteNodeStats(DriftClient thriftClient, URI stateInfoUri, long refreshIntervalMillis) + { + requireNonNull(stateInfoUri, "stateInfoUri is null"); + checkArgument(stateInfoUri.getScheme().equals("thrift"), "unexpected scheme %s", stateInfoUri.getScheme()); + + this.stateInfoUri = stateInfoUri; + this.refreshIntervalMillis = refreshIntervalMillis; + this.thriftClient = requireNonNull(thriftClient, "thriftClient is null").get(Optional.of(stateInfoUri.getAuthority())); + } + + @Override + public Optional getNodeStats() + { + return nodeStats.get(); + } + + @Override + public void asyncRefresh() + { + Duration sinceUpdate = nanosSince(lastUpdateNanos.get()); + if (nanosSince(lastWarningLogged.get()).toMillis() > 1_000 && + sinceUpdate.toMillis() > 10_000 && + requestInflight.get()) { + log.warn("Node state update request to %s has not returned in %s", stateInfoUri, sinceUpdate.toString(SECONDS)); + lastWarningLogged.set(System.nanoTime()); + } + + if (sinceUpdate.toMillis() > refreshIntervalMillis && requestInflight.compareAndSet(false, true)) { + ListenableFuture responseFuture = thriftClient.getServerState(); + + Futures.addCallback(responseFuture, new FutureCallback() + { + @Override + public void onSuccess(@Nullable Integer result) + { + lastUpdateNanos.set(System.nanoTime()); + requestInflight.compareAndSet(true, false); + if (result != null) { + NodeStats nodeStats1 = new NodeStats(NodeState.valueOf(result), null); + nodeStats.set(Optional.of(nodeStats1)); + } + else { + log.warn("Node statistics endpoint %s returned null response, using cached statistics", stateInfoUri); + } + } + + @Override + public void onFailure(Throwable t) + { + log.error("Error fetching node stats from %s: %s", stateInfoUri, t.getMessage()); + lastUpdateNanos.set(System.nanoTime()); + requestInflight.compareAndSet(true, false); + } + }, directExecutor()); + } + } +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/server/InternalCommunicationConfig.java b/presto-main-base/src/main/java/com/facebook/presto/server/InternalCommunicationConfig.java index 06990e3057559..47812dc99f384 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/server/InternalCommunicationConfig.java +++ b/presto-main-base/src/main/java/com/facebook/presto/server/InternalCommunicationConfig.java @@ -49,6 +49,8 @@ public class InternalCommunicationConfig private CommunicationProtocol serverInfoCommunicationProtocol = CommunicationProtocol.HTTP; private boolean memoizeDeadNodesEnabled; private String sharedSecret; + private long nodeStatsRefreshIntervalMillis = 1_000; + private long nodeDiscoveryPollingIntervalMillis = 5_000; private boolean internalJwtEnabled; @@ -312,6 +314,32 @@ public InternalCommunicationConfig setSharedSecret(String sharedSecret) return this; } + public long getNodeStatsRefreshIntervalMillis() + { + return nodeStatsRefreshIntervalMillis; + } + + @Config("internal-communication.node-stats-refresh-interval-millis") + @ConfigDescription("Interval in milliseconds for refreshing node statistics") + public InternalCommunicationConfig setNodeStatsRefreshIntervalMillis(long nodeStatsRefreshIntervalMillis) + { + this.nodeStatsRefreshIntervalMillis = nodeStatsRefreshIntervalMillis; + return this; + } + + public long getNodeDiscoveryPollingIntervalMillis() + { + return nodeDiscoveryPollingIntervalMillis; + } + + @Config("internal-communication.node-discovery-polling-interval-millis") + @ConfigDescription("Interval in milliseconds for polling node discovery and refreshing node states") + public InternalCommunicationConfig setNodeDiscoveryPollingIntervalMillis(long nodeDiscoveryPollingIntervalMillis) + { + this.nodeDiscoveryPollingIntervalMillis = nodeDiscoveryPollingIntervalMillis; + return this; + } + public boolean isInternalJwtEnabled() { return internalJwtEnabled; diff --git a/presto-main-base/src/test/java/com/facebook/presto/execution/TestClusterOverloadConfig.java b/presto-main-base/src/test/java/com/facebook/presto/execution/TestClusterOverloadConfig.java new file mode 100644 index 0000000000000..b28aee410fb20 --- /dev/null +++ b/presto-main-base/src/test/java/com/facebook/presto/execution/TestClusterOverloadConfig.java @@ -0,0 +1,57 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.execution; + +import com.facebook.airlift.configuration.testing.ConfigAssertions; +import com.google.common.collect.ImmutableMap; +import org.testng.annotations.Test; + +import java.util.Map; + +import static com.facebook.presto.execution.ClusterOverloadConfig.OVERLOAD_POLICY_CNT_BASED; + +public class TestClusterOverloadConfig +{ + @Test + public void testDefaults() + { + ConfigAssertions.assertRecordedDefaults(ConfigAssertions.recordDefaults(ClusterOverloadConfig.class) + .setClusterOverloadThrottlingEnabled(false) + .setAllowedOverloadWorkersPct(0.01) + .setAllowedOverloadWorkersCnt(0) + .setOverloadPolicyType(OVERLOAD_POLICY_CNT_BASED) + .setOverloadCheckCacheTtlInSecs(5)); + } + + @Test + public void testExplicitPropertyMappings() + { + Map properties = new ImmutableMap.Builder() + .put("cluster-overload.enable-throttling", "true") + .put("cluster-overload.allowed-overload-workers-pct", "0.05") + .put("cluster-overload.allowed-overload-workers-cnt", "5") + .put("cluster-overload.overload-policy-type", "overload_worker_pct_based_throttling") + .put("cluster.overload-check-cache-ttl-secs", "10") + .build(); + + ClusterOverloadConfig expected = new ClusterOverloadConfig() + .setClusterOverloadThrottlingEnabled(true) + .setAllowedOverloadWorkersPct(0.05) + .setAllowedOverloadWorkersCnt(5) + .setOverloadPolicyType("overload_worker_pct_based_throttling") + .setOverloadCheckCacheTtlInSecs(10); + + ConfigAssertions.assertFullMapping(properties, expected); + } +} diff --git a/presto-main-base/src/test/java/com/facebook/presto/execution/resourceGroups/BenchmarkResourceGroup.java b/presto-main-base/src/test/java/com/facebook/presto/execution/resourceGroups/BenchmarkResourceGroup.java index dfa9235e2f4f7..1ab8ba0bc30e7 100644 --- a/presto-main-base/src/test/java/com/facebook/presto/execution/resourceGroups/BenchmarkResourceGroup.java +++ b/presto-main-base/src/test/java/com/facebook/presto/execution/resourceGroups/BenchmarkResourceGroup.java @@ -14,9 +14,13 @@ package com.facebook.presto.execution.resourceGroups; import com.facebook.airlift.units.DataSize; +import com.facebook.presto.execution.ClusterOverloadConfig; import com.facebook.presto.execution.MockManagedQueryExecution; import com.facebook.presto.execution.resourceGroups.InternalResourceGroup.RootInternalResourceGroup; +import com.facebook.presto.execution.scheduler.clusterOverload.ClusterOverloadPolicy; +import com.facebook.presto.execution.scheduler.clusterOverload.ClusterResourceChecker; import com.facebook.presto.metadata.InMemoryNodeManager; +import com.facebook.presto.metadata.InternalNodeManager; import org.openjdk.jmh.annotations.Benchmark; import org.openjdk.jmh.annotations.BenchmarkMode; import org.openjdk.jmh.annotations.Fork; @@ -73,7 +77,7 @@ public static class BenchmarkData @Setup public void setup() { - root = new RootInternalResourceGroup("root", (group, export) -> {}, executor, ignored -> Optional.empty(), rg -> false, new InMemoryNodeManager()); + root = new RootInternalResourceGroup("root", (group, export) -> {}, executor, ignored -> Optional.empty(), rg -> false, new InMemoryNodeManager(), createClusterResourceChecker()); root.setSoftMemoryLimit(new DataSize(1, MEGABYTE)); root.setMaxQueuedQueries(queries); root.setHardConcurrencyLimit(queries); @@ -89,6 +93,31 @@ public void setup() } } + private ClusterResourceChecker createClusterResourceChecker() + { + // Create a mock cluster overload policy that never reports overload + ClusterOverloadPolicy mockPolicy = new ClusterOverloadPolicy() + { + @Override + public boolean isClusterOverloaded(InternalNodeManager nodeManager) + { + return false; // Never overloaded for benchmarks + } + + @Override + public String getName() + { + return "benchmark-policy"; + } + }; + + // Create a config with throttling disabled for benchmarks + ClusterOverloadConfig config = new ClusterOverloadConfig() + .setClusterOverloadThrottlingEnabled(false); + + return new ClusterResourceChecker(mockPolicy, config, new InMemoryNodeManager()); + } + @TearDown public void tearDown() { diff --git a/presto-main-base/src/test/java/com/facebook/presto/execution/resourceGroups/TestInternalResourceGroupManager.java b/presto-main-base/src/test/java/com/facebook/presto/execution/resourceGroups/TestInternalResourceGroupManager.java index bc4cf99103531..257d957f8e578 100644 --- a/presto-main-base/src/test/java/com/facebook/presto/execution/resourceGroups/TestInternalResourceGroupManager.java +++ b/presto-main-base/src/test/java/com/facebook/presto/execution/resourceGroups/TestInternalResourceGroupManager.java @@ -15,8 +15,11 @@ package com.facebook.presto.execution.resourceGroups; import com.facebook.airlift.node.NodeInfo; +import com.facebook.presto.execution.ClusterOverloadConfig; import com.facebook.presto.execution.MockManagedQueryExecution; import com.facebook.presto.execution.QueryManagerConfig; +import com.facebook.presto.execution.scheduler.clusterOverload.ClusterResourceChecker; +import com.facebook.presto.execution.scheduler.clusterOverload.CpuMemoryOverloadPolicy; import com.facebook.presto.metadata.InMemoryNodeManager; import com.facebook.presto.server.ServerConfig; import com.facebook.presto.spi.PrestoException; @@ -32,8 +35,7 @@ public class TestInternalResourceGroupManager @Test(expectedExceptions = PrestoException.class, expectedExceptionsMessageRegExp = ".*Presto server is still initializing.*") public void testQueryFailsWithInitializingConfigurationManager() { - InternalResourceGroupManager> internalResourceGroupManager = new InternalResourceGroupManager<>((poolId, listener) -> {}, - new QueryManagerConfig(), new NodeInfo("test"), new MBeanExporter(new TestingMBeanServer()), () -> null, new ServerConfig(), new InMemoryNodeManager()); + InternalResourceGroupManager> internalResourceGroupManager = new InternalResourceGroupManager<>((poolId, listener) -> {}, new QueryManagerConfig(), new NodeInfo("test"), new MBeanExporter(new TestingMBeanServer()), () -> null, new ServerConfig(), new InMemoryNodeManager(), new ClusterResourceChecker(new CpuMemoryOverloadPolicy(new ClusterOverloadConfig()), new ClusterOverloadConfig(), new InMemoryNodeManager())); internalResourceGroupManager.submit(new MockManagedQueryExecution(0), new SelectionContext<>(new ResourceGroupId("global"), ImmutableMap.of()), command -> {}); } @@ -42,7 +44,7 @@ public void testQuerySucceedsWhenConfigurationManagerLoaded() throws Exception { InternalResourceGroupManager> internalResourceGroupManager = new InternalResourceGroupManager<>((poolId, listener) -> {}, - new QueryManagerConfig(), new NodeInfo("test"), new MBeanExporter(new TestingMBeanServer()), () -> null, new ServerConfig(), new InMemoryNodeManager()); + new QueryManagerConfig(), new NodeInfo("test"), new MBeanExporter(new TestingMBeanServer()), () -> null, new ServerConfig(), new InMemoryNodeManager(), new ClusterResourceChecker(new CpuMemoryOverloadPolicy(new ClusterOverloadConfig()), new ClusterOverloadConfig(), new InMemoryNodeManager())); internalResourceGroupManager.loadConfigurationManager(); internalResourceGroupManager.submit(new MockManagedQueryExecution(0), new SelectionContext<>(new ResourceGroupId("global"), ImmutableMap.of()), command -> {}); } diff --git a/presto-main-base/src/test/java/com/facebook/presto/execution/resourceGroups/TestResourceGroups.java b/presto-main-base/src/test/java/com/facebook/presto/execution/resourceGroups/TestResourceGroups.java index 2b0eb9ac7d299..5890a210c44e1 100644 --- a/presto-main-base/src/test/java/com/facebook/presto/execution/resourceGroups/TestResourceGroups.java +++ b/presto-main-base/src/test/java/com/facebook/presto/execution/resourceGroups/TestResourceGroups.java @@ -15,8 +15,11 @@ import com.facebook.airlift.units.DataSize; import com.facebook.airlift.units.Duration; +import com.facebook.presto.execution.ClusterOverloadConfig; import com.facebook.presto.execution.MockManagedQueryExecution; import com.facebook.presto.execution.resourceGroups.InternalResourceGroup.RootInternalResourceGroup; +import com.facebook.presto.execution.scheduler.clusterOverload.ClusterOverloadPolicy; +import com.facebook.presto.execution.scheduler.clusterOverload.ClusterResourceChecker; import com.facebook.presto.metadata.InMemoryNodeManager; import com.facebook.presto.metadata.InternalNode; import com.facebook.presto.metadata.InternalNodeManager; @@ -69,7 +72,7 @@ public class TestResourceGroups @Test(timeOut = 10_000) public void testQueueFull() { - RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager()); + RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker()); root.setSoftMemoryLimit(new DataSize(1, MEGABYTE)); root.setMaxQueuedQueries(1); root.setHardConcurrencyLimit(1); @@ -91,7 +94,7 @@ public void testQueueFull() @Test(timeOut = 10_000) public void testFairEligibility() { - RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager()); + RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker()); root.setSoftMemoryLimit(new DataSize(1, MEGABYTE)); root.setMaxQueuedQueries(4); root.setHardConcurrencyLimit(1); @@ -151,7 +154,7 @@ public void testFairEligibility() @Test public void testSetSchedulingPolicy() { - RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager()); + RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker()); root.setSoftMemoryLimit(new DataSize(1, MEGABYTE)); root.setMaxQueuedQueries(4); root.setHardConcurrencyLimit(1); @@ -197,7 +200,7 @@ public void testSetSchedulingPolicy() @Test(timeOut = 10_000) public void testFairQueuing() { - RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager()); + RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker()); root.setSoftMemoryLimit(new DataSize(1, MEGABYTE)); root.setMaxQueuedQueries(4); root.setHardConcurrencyLimit(1); @@ -243,7 +246,7 @@ public void testFairQueuing() @Test(timeOut = 10_000) public void testMemoryLimit() { - RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager()); + RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker()); root.setSoftMemoryLimit(new DataSize(1, BYTE)); root.setMaxQueuedQueries(4); root.setHardConcurrencyLimit(3); @@ -271,7 +274,7 @@ public void testMemoryLimit() @Test public void testSubgroupMemoryLimit() { - RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager()); + RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker()); root.setSoftMemoryLimit(new DataSize(10, BYTE)); root.setMaxQueuedQueries(4); root.setHardConcurrencyLimit(3); @@ -304,7 +307,7 @@ public void testSubgroupMemoryLimit() @Test(timeOut = 10_000) public void testSoftCpuLimit() { - RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager()); + RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker()); root.setSoftMemoryLimit(new DataSize(1, BYTE)); root.setSoftCpuLimit(new Duration(1, SECONDS)); root.setHardCpuLimit(new Duration(2, SECONDS)); @@ -341,7 +344,7 @@ public void testSoftCpuLimit() @Test(timeOut = 10_000) public void testPerWorkerQueryLimit() { - RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager()); + RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker()); root.setWorkersPerQueryLimit(5); root.setMaxQueuedQueries(2); root.setHardConcurrencyLimit(2); @@ -374,7 +377,7 @@ public void testPerWorkerQueryLimit() @Test(timeOut = 10_000) public void testPerWorkerQueryLimitMultipleGroups() { - RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager()); + RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker()); root.setWorkersPerQueryLimit(5); root.setMaxQueuedQueries(5); root.setHardConcurrencyLimit(2); @@ -417,7 +420,7 @@ public void testPerWorkerQueryLimitMultipleGroups() @Test(timeOut = 10_000) public void testHardCpuLimit() { - RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager()); + RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker()); root.setSoftMemoryLimit(new DataSize(1, BYTE)); root.setHardCpuLimit(new Duration(1, SECONDS)); root.setCpuQuotaGenerationMillisPerSecond(2000); @@ -444,7 +447,7 @@ public void testHardCpuLimit() @Test(timeOut = 10_000) public void testPriorityScheduling() { - RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager()); + RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker()); root.setSoftMemoryLimit(new DataSize(1, MEGABYTE)); root.setMaxQueuedQueries(100); // Start with zero capacity, so that nothing starts running until we've added all the queries @@ -494,7 +497,7 @@ public void testPriorityScheduling() @Test(timeOut = 20_000) public void testWeightedScheduling() { - RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager()); + RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker()); root.setSoftMemoryLimit(new DataSize(1, MEGABYTE)); root.setMaxQueuedQueries(4); // Start with zero capacity, so that nothing starts running until we've added all the queries @@ -543,7 +546,7 @@ public void testWeightedScheduling() @Test(timeOut = 30_000) public void testWeightedFairScheduling() { - RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager()); + RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker()); root.setSoftMemoryLimit(new DataSize(1, MEGABYTE)); root.setMaxQueuedQueries(50); // Start with zero capacity, so that nothing starts running until we've added all the queries @@ -586,7 +589,7 @@ public void testWeightedFairScheduling() @Test(timeOut = 10_000) public void testWeightedFairSchedulingEqualWeights() { - RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager()); + RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker()); root.setSoftMemoryLimit(new DataSize(1, MEGABYTE)); root.setMaxQueuedQueries(50); // Start with zero capacity, so that nothing starts running until we've added all the queries @@ -645,7 +648,7 @@ public void testWeightedFairSchedulingEqualWeights() @Test(timeOut = 20_000) public void testWeightedFairSchedulingNoStarvation() { - RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager()); + RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker()); root.setSoftMemoryLimit(new DataSize(1, MEGABYTE)); root.setMaxQueuedQueries(50); // Start with zero capacity, so that nothing starts running until we've added all the queries @@ -686,7 +689,7 @@ public void testWeightedFairSchedulingNoStarvation() @Test public void testGetInfo() { - RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager()); + RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker()); root.setSoftMemoryLimit(new DataSize(1, MEGABYTE)); root.setMaxQueuedQueries(40); // Start with zero capacity, so that nothing starts running until we've added all the queries @@ -776,7 +779,7 @@ public void testGetInfo() @Test public void testGetResourceGroupStateInfo() { - RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager()); + RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker()); root.setSoftMemoryLimit(new DataSize(1, GIGABYTE)); root.setMaxQueuedQueries(40); root.setHardConcurrencyLimit(10); @@ -844,7 +847,7 @@ public void testGetResourceGroupStateInfo() @Test public void testGetStaticResourceGroupInfo() { - RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager()); + RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker()); root.setSoftMemoryLimit(new DataSize(1, GIGABYTE)); root.setMaxQueuedQueries(100); root.setHardConcurrencyLimit(10); @@ -921,7 +924,7 @@ private Optional getResourceGroupInfoForId(InternalResourceGr @Test public void testGetBlockedQueuedQueries() { - RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager()); + RootInternalResourceGroup root = new RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, createNodeManager(), createClusterResourceChecker()); root.setSoftMemoryLimit(new DataSize(1, MEGABYTE)); root.setMaxQueuedQueries(40); // Start with zero capacity, so that nothing starts running until we've added all the queries @@ -1072,4 +1075,27 @@ private InternalNodeManager createNodeManager() false)); return internalNodeManager; } + + private ClusterResourceChecker createClusterResourceChecker() + { + ClusterOverloadPolicy mockPolicy = new ClusterOverloadPolicy() + { + @Override + public boolean isClusterOverloaded(InternalNodeManager nodeManager) + { + return false; + } + + @Override + public String getName() + { + return "test-policy"; + } + }; + + ClusterOverloadConfig config = new ClusterOverloadConfig() + .setClusterOverloadThrottlingEnabled(false); + + return new ClusterResourceChecker(mockPolicy, config, createNodeManager()); + } } diff --git a/presto-main-base/src/test/java/com/facebook/presto/execution/scheduler/clusterOverload/TestClusterResourceChecker.java b/presto-main-base/src/test/java/com/facebook/presto/execution/scheduler/clusterOverload/TestClusterResourceChecker.java new file mode 100644 index 0000000000000..1a5d9c4af8371 --- /dev/null +++ b/presto-main-base/src/test/java/com/facebook/presto/execution/scheduler/clusterOverload/TestClusterResourceChecker.java @@ -0,0 +1,305 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.execution.scheduler.clusterOverload; + +import com.facebook.presto.execution.ClusterOverloadConfig; +import com.facebook.presto.metadata.AllNodes; +import com.facebook.presto.metadata.InternalNode; +import com.facebook.presto.metadata.InternalNodeManager; +import com.facebook.presto.spi.ConnectorId; +import com.facebook.presto.spi.NodeLoadMetrics; +import com.facebook.presto.spi.NodeState; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + +import java.util.Collections; +import java.util.Optional; +import java.util.Set; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Consumer; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertTrue; + +@Test(singleThreaded = true) +public class TestClusterResourceChecker +{ + private static final int CACHE_TTL_SECS = 1; + private static final int SLEEP_BUFFER_MILLIS = 200; + + private ClusterOverloadConfig config; + private TestingClusterOverloadPolicy clusterOverloadPolicy; + private ClusterResourceChecker clusterResourceChecker; + private TestingInternalNodeManager nodeManager; + + @BeforeMethod + public void setUp() + { + config = new ClusterOverloadConfig() + .setOverloadCheckCacheTtlInSecs(CACHE_TTL_SECS) + .setClusterOverloadThrottlingEnabled(true); + + clusterOverloadPolicy = new TestingClusterOverloadPolicy(); + nodeManager = new TestingInternalNodeManager(); + + clusterResourceChecker = new ClusterResourceChecker(clusterOverloadPolicy, config, nodeManager); + } + + public void testInitialState() + { + assertFalse(clusterResourceChecker.isClusterOverloaded()); + assertEquals(clusterResourceChecker.getOverloadDetectionCount().getTotalCount(), 0); + assertEquals(clusterResourceChecker.getOverloadDurationMillis(), 0); + assertTrue(clusterResourceChecker.isClusterOverloadThrottlingEnabled()); + } + + @Test + public void testIsClusterCurrentlyOverloaded() + { + // Start the periodic task + clusterResourceChecker.start(); + + // Initially not overloaded + clusterOverloadPolicy.setOverloaded(false); + assertFalse(clusterResourceChecker.isClusterCurrentlyOverloaded()); + + // Wait for periodic check to update state + clusterOverloadPolicy.setOverloaded(true); + sleep((CACHE_TTL_SECS * 1000) + SLEEP_BUFFER_MILLIS); + assertTrue(clusterResourceChecker.isClusterCurrentlyOverloaded()); + assertTrue(clusterResourceChecker.isClusterOverloaded()); + assertEquals(clusterResourceChecker.getOverloadDetectionCount().getTotalCount(), 1); + + // Set back to not overloaded + clusterOverloadPolicy.setOverloaded(false); + sleep((CACHE_TTL_SECS * 1000) + SLEEP_BUFFER_MILLIS); + assertFalse(clusterResourceChecker.isClusterCurrentlyOverloaded()); + assertFalse(clusterResourceChecker.isClusterOverloaded()); + + // Stop the periodic task + clusterResourceChecker.stop(); + } + + @Test + public void testOverloadDurationMetric() + { + // Start the periodic task + clusterResourceChecker.start(); + + // Set to overloaded and wait for periodic check + clusterOverloadPolicy.setOverloaded(true); + sleep((CACHE_TTL_SECS * 1000) + SLEEP_BUFFER_MILLIS); + assertTrue(clusterResourceChecker.isClusterCurrentlyOverloaded()); + assertTrue(clusterResourceChecker.isClusterOverloaded()); + + // Duration should be greater than 0 + sleep(100); + assertTrue(clusterResourceChecker.getOverloadDurationMillis() > 0); + + // Set back to not overloaded + clusterOverloadPolicy.setOverloaded(false); + sleep((CACHE_TTL_SECS * 1000) + SLEEP_BUFFER_MILLIS); + assertFalse(clusterResourceChecker.isClusterCurrentlyOverloaded()); + + // Duration should be 0 again + assertEquals(clusterResourceChecker.getOverloadDurationMillis(), 0); + + // Stop the periodic task + clusterResourceChecker.stop(); + } + + @Test + public void testMultipleOverloadTransitions() + { + // Start the periodic task + clusterResourceChecker.start(); + + // First transition to overloaded + clusterOverloadPolicy.setOverloaded(true); + sleep((CACHE_TTL_SECS * 1000) + SLEEP_BUFFER_MILLIS); + assertTrue(clusterResourceChecker.isClusterCurrentlyOverloaded()); + assertEquals(clusterResourceChecker.getOverloadDetectionCount().getTotalCount(), 1); + + // Wait and transition back to not overloaded + clusterOverloadPolicy.setOverloaded(false); + sleep((CACHE_TTL_SECS * 1000) + SLEEP_BUFFER_MILLIS); + assertFalse(clusterResourceChecker.isClusterCurrentlyOverloaded()); + + // Second transition to overloaded + clusterOverloadPolicy.setOverloaded(true); + sleep((CACHE_TTL_SECS * 1000) + SLEEP_BUFFER_MILLIS); + assertTrue(clusterResourceChecker.isClusterCurrentlyOverloaded()); + assertEquals(clusterResourceChecker.getOverloadDetectionCount().getTotalCount(), 2); + + // Stop the periodic task + clusterResourceChecker.stop(); + } + + @Test + public void testClusterOverloadThrottlingEnabled() + { + // Default is enabled (set in setUp) + assertTrue(clusterResourceChecker.isClusterOverloadThrottlingEnabled()); + + // Create a new config with throttling disabled + ClusterOverloadConfig disabledConfig = new ClusterOverloadConfig() + .setOverloadCheckCacheTtlInSecs(CACHE_TTL_SECS) + .setClusterOverloadThrottlingEnabled(false); + + // Create a new checker with throttling disabled + ClusterResourceChecker disabledChecker = new ClusterResourceChecker(clusterOverloadPolicy, disabledConfig, nodeManager); + assertFalse(disabledChecker.isClusterOverloadThrottlingEnabled()); + + // Even when cluster is overloaded, isClusterCurrentlyOverloaded should return false if throttling is disabled + clusterOverloadPolicy.setOverloaded(true); + assertFalse(disabledChecker.isClusterCurrentlyOverloaded()); + } + + private void sleep(long millis) + { + try { + Thread.sleep(millis); + } + catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + } + + private static class TestingClusterOverloadPolicy + implements ClusterOverloadPolicy + { + private boolean overloaded; + private final AtomicInteger checkCount = new AtomicInteger(); + + @Override + public boolean isClusterOverloaded(InternalNodeManager nodeManager) + { + checkCount.incrementAndGet(); + return overloaded; + } + + @Override + public String getName() + { + return "test-policy"; + } + + public void setOverloaded(boolean overloaded) + { + this.overloaded = overloaded; + } + + public int getCheckCount() + { + return checkCount.get(); + } + } + + private static class TestingInternalNodeManager + implements InternalNodeManager + { + @Override + public Set getNodes(NodeState state) + { + return Collections.emptySet(); + } + + @Override + public Set getActiveConnectorNodes(ConnectorId connectorId) + { + return Collections.emptySet(); + } + + @Override + public Set getAllConnectorNodes(ConnectorId connectorId) + { + return Collections.emptySet(); + } + + @Override + public InternalNode getCurrentNode() + { + return null; + } + + @Override + public Set getCoordinators() + { + return Collections.emptySet(); + } + + @Override + public Set getShuttingDownCoordinator() + { + return Collections.emptySet(); + } + + @Override + public Set getResourceManagers() + { + return Collections.emptySet(); + } + + @Override + public Set getCatalogServers() + { + return Collections.emptySet(); + } + + @Override + public Set getCoordinatorSidecars() + { + return Collections.emptySet(); + } + + @Override + public AllNodes getAllNodes() + { + return new AllNodes( + Collections.emptySet(), + Collections.emptySet(), + Collections.emptySet(), + Collections.emptySet(), + Collections.emptySet(), + Collections.emptySet(), + Collections.emptySet()); + } + + @Override + public void refreshNodes() + { + // No-op for testing + } + + @Override + public void addNodeChangeListener(Consumer listener) + { + // No-op for testing + } + + @Override + public void removeNodeChangeListener(Consumer listener) + { + // No-op for testing + } + + @Override + public Optional getNodeLoadMetrics(String nodeIdentifier) + { + return Optional.empty(); + } + } +} diff --git a/presto-main-base/src/test/java/com/facebook/presto/execution/scheduler/clusterOverload/TestCpuMemoryOverloadPolicy.java b/presto-main-base/src/test/java/com/facebook/presto/execution/scheduler/clusterOverload/TestCpuMemoryOverloadPolicy.java new file mode 100644 index 0000000000000..f3efbad3f549f --- /dev/null +++ b/presto-main-base/src/test/java/com/facebook/presto/execution/scheduler/clusterOverload/TestCpuMemoryOverloadPolicy.java @@ -0,0 +1,254 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.execution.scheduler.clusterOverload; + +import com.facebook.presto.client.NodeVersion; +import com.facebook.presto.execution.ClusterOverloadConfig; +import com.facebook.presto.metadata.AllNodes; +import com.facebook.presto.metadata.InternalNode; +import com.facebook.presto.metadata.InternalNodeManager; +import com.facebook.presto.spi.ConnectorId; +import com.facebook.presto.spi.NodeLoadMetrics; +import com.facebook.presto.spi.NodeState; +import com.google.common.collect.ImmutableSet; +import org.testng.annotations.Test; + +import java.net.URI; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import java.util.Optional; +import java.util.OptionalInt; +import java.util.Set; +import java.util.function.Consumer; + +import static com.facebook.presto.execution.ClusterOverloadConfig.OVERLOAD_POLICY_CNT_BASED; +import static com.facebook.presto.execution.ClusterOverloadConfig.OVERLOAD_POLICY_PCT_BASED; +import static com.facebook.presto.metadata.InternalNode.NodeStatus.ALIVE; +import static com.facebook.presto.spi.NodePoolType.DEFAULT; +import static com.facebook.presto.spi.NodeState.ACTIVE; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertTrue; + +public class TestCpuMemoryOverloadPolicy +{ + private static final NodeVersion TEST_VERSION = new NodeVersion("test"); + private static final URI TEST_URI = URI.create("http://test.example.com"); + + @Test + public void testIsClusterOverloadedCountBasedNoOverload() + { + ClusterOverloadConfig config = new ClusterOverloadConfig() + .setAllowedOverloadWorkersCnt(1) + .setOverloadPolicyType(OVERLOAD_POLICY_CNT_BASED); + CpuMemoryOverloadPolicy policy = new CpuMemoryOverloadPolicy(config); + + // One node is overloaded, but allowed count is 1, so not overloaded + InternalNodeManager nodeManager = createNodeManager(ImmutableSet.of(createNode("node1", true, false), createNode("node2", false, false), createNode("node3", false, false))); + assertFalse(policy.isClusterOverloaded(nodeManager)); + } + + @Test + public void testIsClusterOverloadedCountBasedOverload() + { + ClusterOverloadConfig config = new ClusterOverloadConfig() + .setAllowedOverloadWorkersCnt(1) + .setOverloadPolicyType(OVERLOAD_POLICY_CNT_BASED); + CpuMemoryOverloadPolicy policy = new CpuMemoryOverloadPolicy(config); + + // Two nodes are overloaded, but allowed count is 1, so overloaded + InternalNodeManager nodeManager = createNodeManager(ImmutableSet.of(createNode("node1", true, false), createNode("node2", false, true), createNode("node3", false, false))); + assertTrue(policy.isClusterOverloaded(nodeManager)); + } + + @Test + public void testIsClusterOverloadedPctBasedNoOverload() + { + ClusterOverloadConfig config = new ClusterOverloadConfig().setAllowedOverloadWorkersPct(0.4).setOverloadPolicyType(OVERLOAD_POLICY_PCT_BASED); + CpuMemoryOverloadPolicy policy = new CpuMemoryOverloadPolicy(config); + + // 1 out of 3 nodes (33%) are overloaded, allowed is 40%, so not overloaded + InternalNodeManager nodeManager = createNodeManager(ImmutableSet.of(createNode("node1", true, false), createNode("node2", false, false), createNode("node3", false, false))); + assertFalse(policy.isClusterOverloaded(nodeManager)); + } + + @Test + public void testIsClusterOverloadedPctBasedOverload() + { + ClusterOverloadConfig config = new ClusterOverloadConfig().setAllowedOverloadWorkersPct(0.3).setOverloadPolicyType(OVERLOAD_POLICY_PCT_BASED); + CpuMemoryOverloadPolicy policy = new CpuMemoryOverloadPolicy(config); + + // 2 out of 5 nodes (40%) are overloaded, allowed is 30%, so overloaded + InternalNodeManager nodeManager = createNodeManager(ImmutableSet.of(createNode("node1", true, false), createNode("node2", false, true), createNode("node3", false, false), createNode("node4", false, false), createNode("node5", false, false))); + assertTrue(policy.isClusterOverloaded(nodeManager)); + } + + @Test + public void testIsClusterOverloadedBothMetricsOverloaded() + { + ClusterOverloadConfig config = new ClusterOverloadConfig() + .setAllowedOverloadWorkersCnt(0) + .setOverloadPolicyType(OVERLOAD_POLICY_CNT_BASED); + CpuMemoryOverloadPolicy policy = new CpuMemoryOverloadPolicy(config); + + // Node has both CPU and memory overloaded, should only count as one overloaded node + InternalNodeManager nodeManager = createNodeManager(ImmutableSet.of(createNode("node1", true, true))); + assertTrue(policy.isClusterOverloaded(nodeManager)); + } + + @Test + public void testIsClusterOverloadedNoNodes() + { + ClusterOverloadConfig config = new ClusterOverloadConfig() + .setAllowedOverloadWorkersCnt(0) + .setOverloadPolicyType(OVERLOAD_POLICY_CNT_BASED); + CpuMemoryOverloadPolicy policy = new CpuMemoryOverloadPolicy(config); + + // No nodes, should not be overloaded + InternalNodeManager nodeManager = createNodeManager(ImmutableSet.of()); + assertFalse(policy.isClusterOverloaded(nodeManager)); + } + + @Test + public void testGetNameCountBased() + { + ClusterOverloadConfig config = new ClusterOverloadConfig() + .setOverloadPolicyType(OVERLOAD_POLICY_CNT_BASED); + CpuMemoryOverloadPolicy policy = new CpuMemoryOverloadPolicy(config); + + assertEquals(policy.getName(), "cpu-memory-overload-cnt"); + } + + @Test + public void testGetNamePctBased() + { + ClusterOverloadConfig config = new ClusterOverloadConfig() + .setOverloadPolicyType(OVERLOAD_POLICY_PCT_BASED); + CpuMemoryOverloadPolicy policy = new CpuMemoryOverloadPolicy(config); + assertEquals(policy.getName(), "cpu-memory-overload-pct"); + } + + // Store metrics separately since they're no longer part of InternalNode + private static final Map NODE_METRICS = new HashMap<>(); + + private static InternalNode createNode(String nodeId, boolean cpuOverload, boolean memoryOverload) + { + // Store metrics in the map for later retrieval + NODE_METRICS.put(nodeId, new NodeLoadMetrics(0.0, 0.0, 0, cpuOverload, memoryOverload)); + return new InternalNode( + nodeId, + TEST_URI, + OptionalInt.empty(), + TEST_VERSION, + false, + false, + false, + false, + ALIVE, + OptionalInt.empty(), + DEFAULT); + } + + private static InternalNodeManager createNodeManager(Set nodes) + { + return new InternalNodeManager() + { + @Override + public Set getNodes(NodeState state) + { + if (state == ACTIVE) { + return nodes; + } + return ImmutableSet.of(); + } + + @Override + public Set getActiveConnectorNodes(ConnectorId connectorId) + { + return ImmutableSet.of(); + } + + @Override + public Set getAllConnectorNodes(ConnectorId connectorId) + { + return Collections.emptySet(); + } + + @Override + public InternalNode getCurrentNode() + { + return null; + } + + @Override + public Set getCoordinators() + { + return ImmutableSet.of(); + } + + @Override + public Set getShuttingDownCoordinator() + { + return ImmutableSet.of(); + } + + @Override + public Set getResourceManagers() + { + return ImmutableSet.of(); + } + + @Override + public Set getCatalogServers() + { + return ImmutableSet.of(); + } + + @Override + public Set getCoordinatorSidecars() + { + return ImmutableSet.of(); + } + + @Override + public AllNodes getAllNodes() + { + return new AllNodes(nodes, ImmutableSet.of(), ImmutableSet.of(), ImmutableSet.of(), ImmutableSet.of(), ImmutableSet.of(), ImmutableSet.of()); + } + + @Override + public void refreshNodes() + { + } + + @Override + public void addNodeChangeListener(Consumer listener) + { + } + + @Override + public void removeNodeChangeListener(Consumer listener) + { + } + + @Override + public Optional getNodeLoadMetrics(String nodeIdentifier) + { + NodeLoadMetrics metrics = NODE_METRICS.get(nodeIdentifier); + return metrics != null ? Optional.of(metrics) : Optional.empty(); + } + }; + } +} diff --git a/presto-main-base/src/test/java/com/facebook/presto/server/TestInternalCommunicationConfig.java b/presto-main-base/src/test/java/com/facebook/presto/server/TestInternalCommunicationConfig.java index ba6faf2a97c53..de88e6b3fc877 100644 --- a/presto-main-base/src/test/java/com/facebook/presto/server/TestInternalCommunicationConfig.java +++ b/presto-main-base/src/test/java/com/facebook/presto/server/TestInternalCommunicationConfig.java @@ -52,7 +52,9 @@ public void testDefaults() .setSharedSecret(null) .setTaskUpdateRequestThriftSerdeEnabled(false) .setTaskInfoResponseThriftSerdeEnabled(false) - .setInternalJwtEnabled(false)); + .setInternalJwtEnabled(false) + .setNodeStatsRefreshIntervalMillis(1_000) + .setNodeDiscoveryPollingIntervalMillis(5_000)); } @Test @@ -80,6 +82,8 @@ public void testExplicitPropertyMappings() .put("internal-communication.jwt.enabled", "true") .put("experimental.internal-communication.task-update-request-thrift-serde-enabled", "true") .put("experimental.internal-communication.task-info-response-thrift-serde-enabled", "true") + .put("internal-communication.node-stats-refresh-interval-millis", "2000") + .put("internal-communication.node-discovery-polling-interval-millis", "3000") .build(); InternalCommunicationConfig expected = new InternalCommunicationConfig() @@ -103,7 +107,9 @@ public void testExplicitPropertyMappings() .setSharedSecret("secret") .setInternalJwtEnabled(true) .setTaskUpdateRequestThriftSerdeEnabled(true) - .setTaskInfoResponseThriftSerdeEnabled(true); + .setTaskInfoResponseThriftSerdeEnabled(true) + .setNodeStatsRefreshIntervalMillis(2000) + .setNodeDiscoveryPollingIntervalMillis(3000); assertFullMapping(properties, expected); } diff --git a/presto-main-base/src/test/java/com/facebook/presto/server/TestQueryStateInfo.java b/presto-main-base/src/test/java/com/facebook/presto/server/TestQueryStateInfo.java index 09230829b8e23..c4bd91d43640f 100644 --- a/presto-main-base/src/test/java/com/facebook/presto/server/TestQueryStateInfo.java +++ b/presto-main-base/src/test/java/com/facebook/presto/server/TestQueryStateInfo.java @@ -17,11 +17,15 @@ import com.facebook.airlift.units.Duration; import com.facebook.presto.common.RuntimeStats; import com.facebook.presto.cost.StatsAndCosts; +import com.facebook.presto.execution.ClusterOverloadConfig; import com.facebook.presto.execution.QueryInfo; import com.facebook.presto.execution.QueryState; import com.facebook.presto.execution.QueryStats; import com.facebook.presto.execution.resourceGroups.InternalResourceGroup; +import com.facebook.presto.execution.scheduler.clusterOverload.ClusterOverloadPolicy; +import com.facebook.presto.execution.scheduler.clusterOverload.ClusterResourceChecker; import com.facebook.presto.metadata.InMemoryNodeManager; +import com.facebook.presto.metadata.InternalNodeManager; import com.facebook.presto.spi.PrestoWarning; import com.facebook.presto.spi.QueryId; import com.facebook.presto.spi.WarningCode; @@ -57,7 +61,7 @@ public class TestQueryStateInfo @Test public void testQueryStateInfo() { - InternalResourceGroup.RootInternalResourceGroup root = new InternalResourceGroup.RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, new InMemoryNodeManager()); + InternalResourceGroup.RootInternalResourceGroup root = new InternalResourceGroup.RootInternalResourceGroup("root", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, new InMemoryNodeManager(), createClusterResourceChecker()); root.setSoftMemoryLimit(new DataSize(1, MEGABYTE)); root.setMaxQueuedQueries(40); root.setHardConcurrencyLimit(0); @@ -317,4 +321,29 @@ private QueryInfo createQueryInfo(String queryId, ResourceGroupId resourceGroupI ImmutableMap.of(), Optional.empty()); } + + private ClusterResourceChecker createClusterResourceChecker() + { + // Create a mock cluster overload policy that never reports overload + ClusterOverloadPolicy mockPolicy = new ClusterOverloadPolicy() + { + @Override + public boolean isClusterOverloaded(InternalNodeManager nodeManager) + { + return false; // Never overloaded for tests + } + + @Override + public String getName() + { + return "test-policy"; + } + }; + + // Create a config with throttling disabled for tests + ClusterOverloadConfig config = new ClusterOverloadConfig() + .setClusterOverloadThrottlingEnabled(false); + + return new ClusterResourceChecker(mockPolicy, config, new InMemoryNodeManager()); + } } diff --git a/presto-main/src/main/java/com/facebook/presto/metadata/DiscoveryNodeManager.java b/presto-main/src/main/java/com/facebook/presto/metadata/DiscoveryNodeManager.java index 02fdf0efe62cf..99d9a8d92dfd3 100644 --- a/presto-main/src/main/java/com/facebook/presto/metadata/DiscoveryNodeManager.java +++ b/presto-main/src/main/java/com/facebook/presto/metadata/DiscoveryNodeManager.java @@ -27,8 +27,10 @@ import com.facebook.presto.server.InternalCommunicationConfig.CommunicationProtocol; import com.facebook.presto.server.thrift.ThriftServerInfoClient; import com.facebook.presto.spi.ConnectorId; +import com.facebook.presto.spi.NodeLoadMetrics; import com.facebook.presto.spi.NodePoolType; import com.facebook.presto.spi.NodeState; +import com.facebook.presto.spi.NodeStats; import com.facebook.presto.statusservice.NodeStatusService; import com.google.common.base.Splitter; import com.google.common.collect.HashMultimap; @@ -93,7 +95,7 @@ public final class DiscoveryNodeManager private final FailureDetector failureDetector; private final Optional nodeStatusService; private final NodeVersion expectedNodeVersion; - private final ConcurrentHashMap nodeStates = new ConcurrentHashMap<>(); + private final ConcurrentHashMap nodeStats = new ConcurrentHashMap<>(); private final HttpClient httpClient; private final DriftClient driftClient; private final ScheduledExecutorService nodeStateUpdateExecutor; @@ -102,6 +104,7 @@ public final class DiscoveryNodeManager private final InternalNode currentNode; private final CommunicationProtocol protocol; private final boolean isMemoizeDeadNodesEnabled; + private final InternalCommunicationConfig internalCommunicationConfig; @GuardedBy("this") private SetMultimap activeNodesByConnectorId; @@ -153,6 +156,7 @@ public DiscoveryNodeManager( this.nodeStateUpdateExecutor = newSingleThreadScheduledExecutor(threadsNamed("node-state-poller-%s")); this.nodeStateEventExecutor = newCachedThreadPool(threadsNamed("node-state-events-%s")); this.httpsRequired = internalCommunicationConfig.isHttpsRequired(); + this.internalCommunicationConfig = requireNonNull(internalCommunicationConfig, "internalCommunicationConfig is null"); this.currentNode = findCurrentNode( serviceSelector.selectAllServices(), @@ -211,6 +215,7 @@ private static NodePoolType getPoolType(ServiceDescriptor service) @PostConstruct public void startPollingNodeStates() { + long pollingIntervalMillis = internalCommunicationConfig.getNodeDiscoveryPollingIntervalMillis(); nodeStateUpdateExecutor.scheduleWithFixedDelay(() -> { try { pollWorkers(); @@ -218,7 +223,7 @@ public void startPollingNodeStates() catch (Exception e) { log.error(e, "Error polling state of nodes"); } - }, 5, 5, TimeUnit.SECONDS); + }, pollingIntervalMillis, pollingIntervalMillis, TimeUnit.MILLISECONDS); pollWorkers(); } @@ -236,20 +241,20 @@ private void pollWorkers() // Remove nodes that don't exist anymore // Make a copy to materialize the set difference - Set deadNodes = difference(nodeStates.keySet(), aliveNodeIds).immutableCopy(); - nodeStates.keySet().removeAll(deadNodes); + Set deadNodes = difference(nodeStats.keySet(), aliveNodeIds).immutableCopy(); + nodeStats.keySet().removeAll(deadNodes); // Add new nodes for (InternalNode node : aliveNodes) { switch (protocol) { case HTTP: - nodeStates.putIfAbsent(node.getNodeIdentifier(), - new HttpRemoteNodeState(httpClient, uriBuilderFrom(node.getInternalUri()).appendPath("/v1/info/state").build())); + nodeStats.putIfAbsent(node.getNodeIdentifier(), + new HttpRemoteNodeStats(httpClient, uriBuilderFrom(node.getInternalUri()).appendPath("/v1/info/stats").build(), internalCommunicationConfig.getNodeStatsRefreshIntervalMillis())); break; case THRIFT: if (node.getThriftPort().isPresent()) { - nodeStates.put(node.getNodeIdentifier(), - new ThriftRemoteNodeState(driftClient, uriBuilderFrom(node.getInternalUri()).scheme("thrift").port(node.getThriftPort().getAsInt()).build())); + nodeStats.put(node.getNodeIdentifier(), + new ThriftRemoteNodeStats(driftClient, uriBuilderFrom(node.getInternalUri()).scheme("thrift").port(node.getThriftPort().getAsInt()).build(), internalCommunicationConfig.getNodeStatsRefreshIntervalMillis())); } else { // thrift port has not yet been populated; ignore the node for now @@ -259,7 +264,7 @@ private void pollWorkers() } // Schedule refresh - nodeStates.values().forEach(RemoteNodeState::asyncRefresh); + nodeStats.values().forEach(RemoteNodeStats::asyncRefresh); // update indexes refreshNodesInternal(); @@ -438,10 +443,19 @@ private NodeState getNodeState(InternalNode node) private boolean isNodeShuttingDown(String nodeId) { - Optional remoteNodeState = nodeStates.containsKey(nodeId) - ? nodeStates.get(nodeId).getNodeState() + Optional remoteNodeStats = nodeStats.containsKey(nodeId) + ? nodeStats.get(nodeId).getNodeStats() : Optional.empty(); - return remoteNodeState.isPresent() && remoteNodeState.get() == SHUTTING_DOWN; + return remoteNodeStats.isPresent() && remoteNodeStats.get().getNodeState() == SHUTTING_DOWN; + } + + @Override + public Optional getNodeLoadMetrics(String nodeId) + { + Optional remoteNodeStats = nodeStats.containsKey(nodeId) + ? nodeStats.get(nodeId).getNodeStats() + : Optional.empty(); + return remoteNodeStats.flatMap(NodeStats::getLoadMetrics); } @Override diff --git a/presto-main/src/main/java/com/facebook/presto/metadata/HttpRemoteNodeStats.java b/presto-main/src/main/java/com/facebook/presto/metadata/HttpRemoteNodeStats.java new file mode 100644 index 0000000000000..5f99a378c4d64 --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/metadata/HttpRemoteNodeStats.java @@ -0,0 +1,121 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.metadata; + +import com.facebook.airlift.http.client.FullJsonResponseHandler.JsonResponse; +import com.facebook.airlift.http.client.HttpClient; +import com.facebook.airlift.http.client.HttpClient.HttpResponseFuture; +import com.facebook.airlift.http.client.Request; +import com.facebook.airlift.log.Logger; +import com.facebook.airlift.units.Duration; +import com.facebook.presto.spi.NodeStats; +import com.google.common.util.concurrent.FutureCallback; +import com.google.common.util.concurrent.Futures; +import com.google.errorprone.annotations.ThreadSafe; +import jakarta.annotation.Nullable; + +import java.net.URI; +import java.util.Optional; +import java.util.concurrent.Future; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; + +import static com.facebook.airlift.http.client.FullJsonResponseHandler.createFullJsonResponseHandler; +import static com.facebook.airlift.http.client.HttpStatus.OK; +import static com.facebook.airlift.http.client.Request.Builder.prepareGet; +import static com.facebook.airlift.json.JsonCodec.jsonCodec; +import static com.facebook.airlift.units.Duration.nanosSince; +import static com.google.common.net.MediaType.JSON_UTF_8; +import static com.google.common.util.concurrent.MoreExecutors.directExecutor; +import static jakarta.ws.rs.core.HttpHeaders.CONTENT_TYPE; +import static java.util.Objects.requireNonNull; +import static java.util.concurrent.TimeUnit.SECONDS; + +@ThreadSafe +public class HttpRemoteNodeStats + implements RemoteNodeStats +{ + private static final Logger log = Logger.get(HttpRemoteNodeStats.class); + + private final HttpClient httpClient; + private final URI stateInfoUri; + private final long refreshIntervalMillis; + private final AtomicReference> nodeStats = new AtomicReference<>(Optional.empty()); + private final AtomicReference> future = new AtomicReference<>(); + private final AtomicLong lastUpdateNanos = new AtomicLong(); + private final AtomicLong lastWarningLogged = new AtomicLong(); + + public HttpRemoteNodeStats(HttpClient httpClient, URI stateInfoUri, long refreshIntervalMillis) + { + this.httpClient = requireNonNull(httpClient, "httpClient is null"); + this.stateInfoUri = requireNonNull(stateInfoUri, "stateInfoUri is null"); + this.refreshIntervalMillis = refreshIntervalMillis; + } + + @Override + public Optional getNodeStats() + { + return nodeStats.get(); + } + + @Override + public synchronized void asyncRefresh() + { + Duration sinceUpdate = nanosSince(lastUpdateNanos.get()); + if (nanosSince(lastWarningLogged.get()).toMillis() > 1_000 && + sinceUpdate.toMillis() > 10_000 && + future.get() != null) { + log.warn("Node state update request to %s has not returned in %s", stateInfoUri, sinceUpdate.toString(SECONDS)); + lastWarningLogged.set(System.nanoTime()); + } + if (sinceUpdate.toMillis() > refreshIntervalMillis && future.get() == null) { + Request request = prepareGet() + .setUri(stateInfoUri) + .setHeader(CONTENT_TYPE, JSON_UTF_8.toString()) + .build(); + HttpResponseFuture> responseFuture = httpClient.executeAsync(request, createFullJsonResponseHandler(jsonCodec(NodeStats.class))); + future.compareAndSet(null, responseFuture); + + Futures.addCallback(responseFuture, new FutureCallback>() + { + @Override + public void onSuccess(@Nullable JsonResponse result) + { + lastUpdateNanos.set(System.nanoTime()); + future.compareAndSet(responseFuture, null); + if (result != null) { + if (result.hasValue()) { + nodeStats.set(Optional.ofNullable(result.getValue())); + } + if (result.getStatusCode() != OK.code()) { + log.warn("Error fetching node stats from %s returned status %d", stateInfoUri, result.getStatusCode()); + return; + } + } + else { + log.warn("Node statistics endpoint %s returned null response, using cached statistics", stateInfoUri); + } + } + + @Override + public void onFailure(Throwable t) + { + log.error("Error fetching node stats from %s: %s", stateInfoUri, t.getMessage()); + lastUpdateNanos.set(System.nanoTime()); + future.compareAndSet(responseFuture, null); + } + }, directExecutor()); + } + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/server/ServerMainModule.java b/presto-main/src/main/java/com/facebook/presto/server/ServerMainModule.java index c564764143ce7..f5d2543048a46 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/ServerMainModule.java +++ b/presto-main/src/main/java/com/facebook/presto/server/ServerMainModule.java @@ -84,6 +84,8 @@ import com.facebook.presto.execution.scheduler.NodeSchedulerConfig; import com.facebook.presto.execution.scheduler.NodeSchedulerExporter; import com.facebook.presto.execution.scheduler.TableWriteInfo; +import com.facebook.presto.execution.scheduler.clusterOverload.ClusterOverloadPolicyModule; +import com.facebook.presto.execution.scheduler.clusterOverload.ClusterResourceChecker; import com.facebook.presto.execution.scheduler.nodeSelection.NodeSelectionStats; import com.facebook.presto.execution.scheduler.nodeSelection.SimpleTtlNodeSelectorConfig; import com.facebook.presto.functionNamespace.JsonBasedUdfFunctionMetadata; @@ -709,6 +711,10 @@ public ListeningExecutorService createResourceManagerExecutor(ResourceManagerCon // system connector binder.install(new SystemConnectorModule()); + // ClusterOverload policy module + binder.install(new ClusterOverloadPolicyModule()); + newExporter(binder).export(ClusterResourceChecker.class).withGeneratedName(); + // splits jsonCodecBinder(binder).bindJsonCodec(TaskUpdateRequest.class); jsonCodecBinder(binder).bindJsonCodec(ConnectorSplit.class); diff --git a/presto-resource-group-managers/src/test/java/com/facebook/presto/resourceGroups/reloading/TestReloadingResourceGroupConfigurationManager.java b/presto-resource-group-managers/src/test/java/com/facebook/presto/resourceGroups/reloading/TestReloadingResourceGroupConfigurationManager.java index 8f02abab278e5..065fbfd740daf 100644 --- a/presto-resource-group-managers/src/test/java/com/facebook/presto/resourceGroups/reloading/TestReloadingResourceGroupConfigurationManager.java +++ b/presto-resource-group-managers/src/test/java/com/facebook/presto/resourceGroups/reloading/TestReloadingResourceGroupConfigurationManager.java @@ -15,8 +15,12 @@ import com.facebook.airlift.units.DataSize; import com.facebook.airlift.units.Duration; +import com.facebook.presto.execution.ClusterOverloadConfig; import com.facebook.presto.execution.resourceGroups.InternalResourceGroup; +import com.facebook.presto.execution.scheduler.clusterOverload.ClusterOverloadPolicy; +import com.facebook.presto.execution.scheduler.clusterOverload.ClusterResourceChecker; import com.facebook.presto.metadata.InMemoryNodeManager; +import com.facebook.presto.metadata.InternalNodeManager; import com.facebook.presto.resourceGroups.VariableMap; import com.facebook.presto.resourceGroups.db.DbManagerSpecProvider; import com.facebook.presto.resourceGroups.db.DbResourceGroupConfig; @@ -74,7 +78,7 @@ public void testConfiguration() DbManagerSpecProvider dbManagerSpecProvider = new DbManagerSpecProvider(daoProvider.get(), ENVIRONMENT, new ReloadingResourceGroupConfig()); ReloadingResourceGroupConfigurationManager manager = new ReloadingResourceGroupConfigurationManager((poolId, listener) -> {}, new ReloadingResourceGroupConfig(), dbManagerSpecProvider); AtomicBoolean exported = new AtomicBoolean(); - InternalResourceGroup global = new InternalResourceGroup.RootInternalResourceGroup("global", (group, export) -> exported.set(export), directExecutor(), ignored -> Optional.empty(), rg -> false, new InMemoryNodeManager()); + InternalResourceGroup global = new InternalResourceGroup.RootInternalResourceGroup("global", (group, export) -> exported.set(export), directExecutor(), ignored -> Optional.empty(), rg -> false, new InMemoryNodeManager(), createClusterResourceChecker()); manager.configure(global, new SelectionContext<>(global.getId(), new VariableMap(ImmutableMap.of("USER", "user")))); assertEqualsResourceGroup(global, "1MB", 1000, 100, 100, WEIGHTED, DEFAULT_WEIGHT, true, new Duration(1, HOURS), new Duration(1, DAYS), new ResourceGroupQueryLimits(Optional.of(new Duration(1, HOURS)), Optional.of(new DataSize(1, MEGABYTE)), Optional.of(new Duration(1, HOURS)))); exported.set(false); @@ -97,7 +101,7 @@ public void testMissing() dao.insertSelector(2, 1, null, null, null, null, null, null); DbManagerSpecProvider dbManagerSpecProvider = new DbManagerSpecProvider(daoProvider.get(), ENVIRONMENT, new ReloadingResourceGroupConfig()); ReloadingResourceGroupConfigurationManager manager = new ReloadingResourceGroupConfigurationManager((poolId, listener) -> {}, new ReloadingResourceGroupConfig(), dbManagerSpecProvider); - InternalResourceGroup missing = new InternalResourceGroup.RootInternalResourceGroup("missing", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, new InMemoryNodeManager()); + InternalResourceGroup missing = new InternalResourceGroup.RootInternalResourceGroup("missing", (group, export) -> {}, directExecutor(), ignored -> Optional.empty(), rg -> false, new InMemoryNodeManager(), createClusterResourceChecker()); manager.configure(missing, new SelectionContext<>(missing.getId(), new VariableMap(ImmutableMap.of("USER", "user")))); } @@ -118,7 +122,7 @@ public void testReconfig() ReloadingResourceGroupConfigurationManager manager = new ReloadingResourceGroupConfigurationManager((poolId, listener) -> {}, new ReloadingResourceGroupConfig(), dbManagerSpecProvider); manager.start(); AtomicBoolean exported = new AtomicBoolean(); - InternalResourceGroup global = new InternalResourceGroup.RootInternalResourceGroup("global", (group, export) -> exported.set(export), directExecutor(), ignored -> Optional.empty(), rg -> false, new InMemoryNodeManager()); + InternalResourceGroup global = new InternalResourceGroup.RootInternalResourceGroup("global", (group, export) -> exported.set(export), directExecutor(), ignored -> Optional.empty(), rg -> false, new InMemoryNodeManager(), createClusterResourceChecker()); manager.configure(global, new SelectionContext<>(global.getId(), new VariableMap(ImmutableMap.of("USER", "user")))); InternalResourceGroup globalSub = global.getOrCreateSubGroup("sub", true); manager.configure(globalSub, new SelectionContext<>(globalSub.getId(), new VariableMap(ImmutableMap.of("USER", "user")))); @@ -247,4 +251,29 @@ private static void assertEqualsResourceGroup( assertEquals(group.getHardCpuLimit(), hardCpuLimit); assertEquals(group.getPerQueryLimits(), perQueryLimits); } + + private static ClusterResourceChecker createClusterResourceChecker() + { + // Create a mock cluster overload policy that never reports overload + ClusterOverloadPolicy mockPolicy = new ClusterOverloadPolicy() + { + @Override + public boolean isClusterOverloaded(InternalNodeManager nodeManager) + { + return false; // Never overloaded for tests + } + + @Override + public String getName() + { + return "test-policy"; + } + }; + + // Create a config with throttling disabled for tests + ClusterOverloadConfig config = new ClusterOverloadConfig() + .setClusterOverloadThrottlingEnabled(false); + + return new ClusterResourceChecker(mockPolicy, config, new InMemoryNodeManager()); + } } diff --git a/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkModule.java b/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkModule.java index 520d1d8c5bf7d..5cd272b85b277 100644 --- a/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkModule.java +++ b/presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkModule.java @@ -61,6 +61,7 @@ import com.facebook.presto.execution.resourceGroups.InternalResourceGroupManager; import com.facebook.presto.execution.resourceGroups.ResourceGroupManager; import com.facebook.presto.execution.scheduler.NodeSchedulerConfig; +import com.facebook.presto.execution.scheduler.clusterOverload.ClusterOverloadPolicyModule; import com.facebook.presto.execution.scheduler.nodeSelection.SimpleTtlNodeSelectorConfig; import com.facebook.presto.execution.warnings.WarningCollectorConfig; import com.facebook.presto.index.IndexManager; @@ -330,6 +331,9 @@ protected void setup(Binder binder) // handle resolver binder.install(new HandleJsonModule()); + // ClusterOverload policy module + binder.install(new ClusterOverloadPolicyModule()); + // plugin manager configBinder(binder).bindConfig(PluginManagerConfig.class); binder.bind(PluginManager.class).in(Scopes.SINGLETON); diff --git a/presto-spark-base/src/main/java/com/facebook/presto/spark/node/PrestoSparkInternalNodeManager.java b/presto-spark-base/src/main/java/com/facebook/presto/spark/node/PrestoSparkInternalNodeManager.java index 1046a820ad367..7df3792f4b5f7 100644 --- a/presto-spark-base/src/main/java/com/facebook/presto/spark/node/PrestoSparkInternalNodeManager.java +++ b/presto-spark-base/src/main/java/com/facebook/presto/spark/node/PrestoSparkInternalNodeManager.java @@ -18,10 +18,12 @@ import com.facebook.presto.metadata.InternalNode; import com.facebook.presto.metadata.InternalNodeManager; import com.facebook.presto.spi.ConnectorId; +import com.facebook.presto.spi.NodeLoadMetrics; import com.facebook.presto.spi.NodeState; import com.google.common.collect.ImmutableSet; import java.net.URI; +import java.util.Optional; import java.util.Set; import java.util.function.Consumer; @@ -125,4 +127,10 @@ public void removeNodeChangeListener(Consumer listener) { throw new UnsupportedOperationException(); } + + @Override + public Optional getNodeLoadMetrics(String nodeIdentifier) + { + return Optional.empty(); + } } From 197a26b61193688adc7f8ca26e74791bee47ca67 Mon Sep 17 00:00:00 2001 From: abhinavmuk04 Date: Tue, 9 Sep 2025 17:19:28 -0400 Subject: [PATCH 074/113] visitEnumLiteral fix --- .../presto/sql/TestExpressionInterpreter.java | 15 +++++++++++++++ .../facebook/presto/sql/ExpressionFormatter.java | 7 +++++++ 2 files changed, 22 insertions(+) diff --git a/presto-main-base/src/test/java/com/facebook/presto/sql/TestExpressionInterpreter.java b/presto-main-base/src/test/java/com/facebook/presto/sql/TestExpressionInterpreter.java index f07caf36cf5e8..c4d9d050dbc3d 100644 --- a/presto-main-base/src/test/java/com/facebook/presto/sql/TestExpressionInterpreter.java +++ b/presto-main-base/src/test/java/com/facebook/presto/sql/TestExpressionInterpreter.java @@ -51,6 +51,7 @@ import com.facebook.presto.sql.planner.Symbol; import com.facebook.presto.sql.planner.TypeProvider; import com.facebook.presto.sql.relational.FunctionResolution; +import com.facebook.presto.sql.tree.EnumLiteral; import com.facebook.presto.sql.tree.Expression; import com.facebook.presto.sql.tree.ExpressionRewriter; import com.facebook.presto.sql.tree.ExpressionTreeRewriter; @@ -899,6 +900,20 @@ public void testReservedWithDoubleQuotes() assertOptimizedEquals("\"time\"", "\"time\""); } + @Test + public void testEnumLiteralFormattingWithTypeAndValue() + { + java.util.function.BiFunction createEnumLiteral = (type, value) -> new EnumLiteral(Optional.empty(), type, value); + assertEquals(ExpressionFormatter.formatExpression(createEnumLiteral.apply("color", "RED"), Optional.empty()), "color: RED"); + assertEquals(ExpressionFormatter.formatExpression(createEnumLiteral.apply("level", 1), Optional.empty()), "level: 1"); + assertEquals(ExpressionFormatter.formatExpression(createEnumLiteral.apply("StatusType", "Active"), Optional.empty()), "StatusType: Active"); + assertEquals(ExpressionFormatter.formatExpression(createEnumLiteral.apply("priority", "HIGH PRIORITY"), Optional.empty()), "priority: HIGH PRIORITY"); + assertEquals(ExpressionFormatter.formatExpression(createEnumLiteral.apply("lang", "枚举"), Optional.empty()), "lang: 枚举"); + assertEquals(ExpressionFormatter.formatExpression(createEnumLiteral.apply("special", "DOLLAR$"), Optional.empty()), "special: DOLLAR$"); + assertEquals(ExpressionFormatter.formatExpression(createEnumLiteral.apply("enum_type", "VALUE_1"), Optional.empty()), "enum_type: VALUE_1"); + assertEquals(ExpressionFormatter.formatExpression(createEnumLiteral.apply("flag", true), Optional.empty()), "flag: true"); + } + @Test public void testSearchCase() { diff --git a/presto-parser/src/main/java/com/facebook/presto/sql/ExpressionFormatter.java b/presto-parser/src/main/java/com/facebook/presto/sql/ExpressionFormatter.java index bc03fb68154ea..bae8559c7f127 100644 --- a/presto-parser/src/main/java/com/facebook/presto/sql/ExpressionFormatter.java +++ b/presto-parser/src/main/java/com/facebook/presto/sql/ExpressionFormatter.java @@ -33,6 +33,7 @@ import com.facebook.presto.sql.tree.DecimalLiteral; import com.facebook.presto.sql.tree.DereferenceExpression; import com.facebook.presto.sql.tree.DoubleLiteral; +import com.facebook.presto.sql.tree.EnumLiteral; import com.facebook.presto.sql.tree.ExistsPredicate; import com.facebook.presto.sql.tree.Expression; import com.facebook.presto.sql.tree.Extract; @@ -198,6 +199,12 @@ protected String visitBooleanLiteral(BooleanLiteral node, Void context) return String.valueOf(node.getValue()); } + @Override + protected String visitEnumLiteral(EnumLiteral node, Void context) + { + return node.getType() + ": " + node.getValue(); + } + @Override protected String visitStringLiteral(StringLiteral node, Void context) { From 07a1a884e62ea2499410499ea0b67871b34ec1ca Mon Sep 17 00:00:00 2001 From: Shahim Sharafudeen Date: Wed, 10 Sep 2025 11:56:03 +0530 Subject: [PATCH 075/113] Upgrade netty to 4.1.126.Final to address CVE-2025-58056 and CVE-2025-58057 --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index f73719dacd94a..c2de59f18bcdf 100644 --- a/pom.xml +++ b/pom.xml @@ -82,7 +82,7 @@ 1.26.2 4.29.0 12.0.18 - 4.1.124.Final + 4.1.126.Final 1.2.8 2.0 2.12.1 From e366be17187e5396ae0c96dceffa21b586410ae4 Mon Sep 17 00:00:00 2001 From: HeidiHan0000 <44453261+HeidiHan0000@users.noreply.github.com> Date: Wed, 10 Sep 2025 13:40:35 -0700 Subject: [PATCH 076/113] [presto] Add enum type flag to Prestissimo worker config (#25989) Summary: Added "enum-types-enabled" boolean flag to control the parsing/usage of Bigitnenum and VarcharEnum functions in Prestissimo workers. This is to help with Capella types roll out, in case we encounter issues with correctness, etc, and if we cannot block them from gateway. Differential Revision: D82000691 ``` == NO RELEASE NOTE == ``` --- .../presto_cpp/main/common/Configs.cpp | 5 ++ .../presto_cpp/main/common/Configs.h | 8 +++ .../presto_cpp/main/types/TypeParser.cpp | 6 +- .../main/types/tests/CMakeLists.txt | 2 +- .../main/types/tests/TypeParserTest.cpp | 63 +++++++++++++++++++ 5 files changed, 82 insertions(+), 2 deletions(-) create mode 100644 presto-native-execution/presto_cpp/main/types/tests/TypeParserTest.cpp diff --git a/presto-native-execution/presto_cpp/main/common/Configs.cpp b/presto-native-execution/presto_cpp/main/common/Configs.cpp index 85e5f356ae815..76613ee440df4 100644 --- a/presto-native-execution/presto_cpp/main/common/Configs.cpp +++ b/presto-native-execution/presto_cpp/main/common/Configs.cpp @@ -262,6 +262,7 @@ SystemConfig::SystemConfig() { NUM_PROP(kMaxLocalExchangePartitionBufferSize, 65536), BOOL_PROP(kTextWriterEnabled, true), BOOL_PROP(kCharNToVarcharImplicitCast, false), + BOOL_PROP(kEnumTypesEnabled, true), }; } @@ -931,6 +932,10 @@ bool SystemConfig::charNToVarcharImplicitCast() const { return optionalProperty(kCharNToVarcharImplicitCast).value(); } +bool SystemConfig::enumTypesEnabled() const { + return optionalProperty(kEnumTypesEnabled).value(); +} + NodeConfig::NodeConfig() { registeredProps_ = std::unordered_map>{ diff --git a/presto-native-execution/presto_cpp/main/common/Configs.h b/presto-native-execution/presto_cpp/main/common/Configs.h index fc49549ca1695..c891aee3d2306 100644 --- a/presto-native-execution/presto_cpp/main/common/Configs.h +++ b/presto-native-execution/presto_cpp/main/common/Configs.h @@ -768,6 +768,12 @@ class SystemConfig : public ConfigBase { static constexpr std::string_view kCharNToVarcharImplicitCast{ "char-n-to-varchar-implicit-cast"}; + /// Enable BigintEnum and VarcharEnum types to be parsed and used in Velox. + /// When set to false, BigintEnum or VarcharEnum types will throw an + // unsupported error during type parsing. + static constexpr std::string_view kEnumTypesEnabled{ + "enum-types-enabled"}; + SystemConfig(); virtual ~SystemConfig() = default; @@ -1060,6 +1066,8 @@ class SystemConfig : public ConfigBase { bool textWriterEnabled() const; bool charNToVarcharImplicitCast() const; + + bool enumTypesEnabled() const; }; /// Provides access to node properties defined in node.properties file. diff --git a/presto-native-execution/presto_cpp/main/types/TypeParser.cpp b/presto-native-execution/presto_cpp/main/types/TypeParser.cpp index 4a19066c38fb0..15d7301c347d4 100644 --- a/presto-native-execution/presto_cpp/main/types/TypeParser.cpp +++ b/presto-native-execution/presto_cpp/main/types/TypeParser.cpp @@ -27,7 +27,11 @@ velox::TypePtr TypeParser::parse(const std::string& text) const { return velox::VARCHAR(); } } - + if (!SystemConfig::instance()->enumTypesEnabled()) { + if (text.find("BigintEnum") != std::string::npos || text.find("VarcharEnum") != std::string::npos) { + VELOX_UNSUPPORTED("Unsupported type: {}", text); + } + } auto it = cache_.find(text); if (it != cache_.end()) { return it->second; diff --git a/presto-native-execution/presto_cpp/main/types/tests/CMakeLists.txt b/presto-native-execution/presto_cpp/main/types/tests/CMakeLists.txt index 48b22e9a63328..8f98f3f76686f 100644 --- a/presto-native-execution/presto_cpp/main/types/tests/CMakeLists.txt +++ b/presto-native-execution/presto_cpp/main/types/tests/CMakeLists.txt @@ -35,7 +35,7 @@ target_link_libraries( velox_hive_partition_function) add_executable(presto_expressions_test RowExpressionTest.cpp ValuesPipeTest.cpp - PlanConverterTest.cpp) + PlanConverterTest.cpp TypeParserTest.cpp) add_test( NAME presto_expressions_test diff --git a/presto-native-execution/presto_cpp/main/types/tests/TypeParserTest.cpp b/presto-native-execution/presto_cpp/main/types/tests/TypeParserTest.cpp new file mode 100644 index 0000000000000..ee9036e5943bc --- /dev/null +++ b/presto-native-execution/presto_cpp/main/types/tests/TypeParserTest.cpp @@ -0,0 +1,63 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +#include "presto_cpp/main/common/Configs.h" +#include "presto_cpp/main/common/tests/MutableConfigs.h" +#include "presto_cpp/main/types/TypeParser.h" +#include "velox/common/file/FileSystems.h" +#include "velox/common/base/tests/GTestUtils.h" +#include "velox/functions/prestosql/types/BigintEnumRegistration.h" +#include "velox/functions/prestosql/types/VarcharEnumRegistration.h" + + +using namespace facebook::presto; +using namespace facebook::velox; + +class TypeParserTest : public ::testing::Test { + void SetUp() override { + filesystems::registerLocalFileSystem(); + test::setupMutableSystemConfig(); + registerBigintEnumType(); + registerVarcharEnumType(); + } +}; + +// Test basical functionality of TypeParser. +// More detailed tests for Presto TypeParser are in velox/functions/prestosql/types/parser/tests/TypeParserTest. +TEST_F(TypeParserTest, parseEnumTypes) { + TypeParser typeParser = TypeParser(); + + ASSERT_EQ( + typeParser.parse( + "test.enum.mood:BigintEnum(test.enum.mood{\"CURIOUS\":2, \"HAPPY\":0})")->toString(), + "test.enum.mood:BigintEnum({\"CURIOUS\": 2, \"HAPPY\": 0})"); + ASSERT_EQ( + typeParser.parse( + "test.enum.mood:VarcharEnum(test.enum.mood{\"CURIOUS\":\"ONXW2ZKWMFWHKZI=\", \"HAPPY\":\"ONXW2ZJAOZQWY5LF\" , \"SAD\":\"KNHU2RJAKZAUYVKF\"})")->toString(), + "test.enum.mood:VarcharEnum({\"CURIOUS\": \"someValue\", \"HAPPY\": \"some value\", \"SAD\": \"SOME VALUE\"})"); + + // When set to false, TypeParser will throw an unsupported error when it receives an enum type. + SystemConfig::instance()->setValue(std::string(SystemConfig::kEnumTypesEnabled), "false"); + + VELOX_ASSERT_THROW( + typeParser.parse( + "test.enum.mood:BigintEnum(test.enum.mood{\"CURIOUS\":2, \"HAPPY\":0})"), + "Unsupported type: test.enum.mood:BigintEnum(test.enum.mood{\"CURIOUS\":2, \"HAPPY\":0})"); + VELOX_ASSERT_THROW( + typeParser.parse( + "test.enum.mood:VarcharEnum(test.enum.mood{\"CURIOUS\":\"ONXW2ZKWMFWHKZI=\", \"HAPPY\":\"ONXW2ZJAOZQWY5LF\" , \"SAD\":\"KNHU2RJAKZAUYVKF\"})"), + "Unsupported type: test.enum.mood:VarcharEnum(test.enum.mood{\"CURIOUS\":\"ONXW2ZKWMFWHKZI=\", \"HAPPY\":\"ONXW2ZJAOZQWY5LF\" , \"SAD\":\"KNHU2RJAKZAUYVKF\"})"); +} From eaf1d283ba2fbfe91d383ab118031390265c4026 Mon Sep 17 00:00:00 2001 From: Pramod Satya Date: Mon, 28 Jul 2025 18:01:17 -0700 Subject: [PATCH 077/113] Allow duplicate function signatures when matching with generic candidates --- .../metadata/FunctionSignatureMatcher.java | 32 +++++++++++++++++-- .../sidecar/TestNativeSidecarPlugin.java | 14 ++++++++ 2 files changed, 44 insertions(+), 2 deletions(-) diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionSignatureMatcher.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionSignatureMatcher.java index 39cb60c65460a..46d34dc13d61b 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionSignatureMatcher.java +++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/FunctionSignatureMatcher.java @@ -71,7 +71,7 @@ public Optional match(Collection candidates, L .filter(function -> !function.getSignature().getTypeVariableConstraints().isEmpty()) .collect(Collectors.toList()); - match = matchFunctionExact(genericCandidates, parameterTypes); + match = matchFunctionGeneric(genericCandidates, parameterTypes); if (match.isPresent()) { return match; } @@ -91,6 +91,28 @@ private Optional matchFunctionExact(List candidates, Lis return matchFunction(candidates, actualParameters, false); } + private Optional matchFunctionGeneric(List candidates, List actualParameters) + { + List applicableFunctions = identifyApplicableFunctions(candidates, actualParameters, false); + if (applicableFunctions.isEmpty()) { + return Optional.empty(); + } + + if (applicableFunctions.size() == 1) { + return Optional.of(getOnlyElement(applicableFunctions).getBoundSignature()); + } + + List deduplicatedSignatures = applicableFunctions.stream() + .map(applicableFunction -> applicableFunction.boundSignature) + .distinct() + .collect(toImmutableList()); + if (deduplicatedSignatures.size() == 1) { + return Optional.of(getOnlyElement(deduplicatedSignatures)); + } + + throw new PrestoException(AMBIGUOUS_FUNCTION_CALL, getErrorMessage(applicableFunctions)); + } + private Optional matchFunctionWithCoercion(Collection candidates, List actualParameters) { return matchFunction(candidates, actualParameters, true); @@ -112,6 +134,11 @@ private Optional matchFunction(Collection cand return Optional.of(getOnlyElement(applicableFunctions).getBoundSignature()); } + throw new PrestoException(AMBIGUOUS_FUNCTION_CALL, getErrorMessage(applicableFunctions)); + } + + private String getErrorMessage(List applicableFunctions) + { StringBuilder errorMessageBuilder = new StringBuilder(); errorMessageBuilder.append("Could not choose a best candidate operator. Explicit type casts must be added.\n"); errorMessageBuilder.append("Candidates are:\n"); @@ -120,7 +147,8 @@ private Optional matchFunction(Collection cand errorMessageBuilder.append(function.getBoundSignature().toString()); errorMessageBuilder.append("\n"); } - throw new PrestoException(AMBIGUOUS_FUNCTION_CALL, errorMessageBuilder.toString()); + + return errorMessageBuilder.toString(); } private List identifyApplicableFunctions(Collection candidates, List actualParameters, boolean allowCoercion) diff --git a/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java b/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java index 3e96d8c50a0bd..40ade2247847d 100644 --- a/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java +++ b/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java @@ -334,6 +334,20 @@ public void testApproxPercentile() } } + @Test + public void testMapSubset() + { + assertQuery("select m[1], m[3] from (select map_subset(map(array[1,2,3,4], array['a', 'b', 'c', 'd']), array[1,3,10]) m)", "select 'a', 'c'"); + assertQuery("select m['p'], m['r'] from (select map_subset(map(array['p', 'q', 'r', 's'], array['a', 'b', 'c', 'd']), array['p', 'r', 'z']) m)", "select 'a', 'c'"); + assertQuery("select m[true], m[false] from (select map_subset(map(array[false, true], array['a', 'z']), array[true, false]) m)", "select 'z', 'a'"); + assertQuery("select m[DATE '2015-01-01'], m[DATE '2015-01-13'] from (select map_subset(" + + "map(array[DATE '2015-01-01', DATE '2015-02-13', DATE '2015-01-13', DATE '2015-05-15'], array['a', 'b', 'c', 'd']), " + + "array[DATE '2015-01-01', DATE '2015-01-13', DATE '2015-06-15']) m)", "select 'a', 'c'"); + assertQuery("select m[TIMESTAMP '2021-01-02 09:04:05.321'] from (select map_subset(" + + "map(array[TIMESTAMP '2021-01-02 09:04:05.321', TIMESTAMP '2022-12-22 10:07:08.456'], array['a', 'b']), " + + "array[TIMESTAMP '2021-01-02 09:04:05.321', TIMESTAMP '2022-12-22 10:07:09.246']) m)", "select 'a'"); + } + @Test public void testInformationSchemaTables() { From 987c02f809c96b669bb1143f2d7c4de754ddf435 Mon Sep 17 00:00:00 2001 From: Shrinidhi Joshi Date: Fri, 5 Sep 2025 20:39:56 -0700 Subject: [PATCH 078/113] [pos][native] Setup cpp worker memory settings based on sparkConf (for smart retries) In this PR we make the change where we base our cpp worker memory calculations on the config that is updated to higher values by spark runtime for retries. Currently this config is `spark.memory.offHeap.size`, as internally at Meta, this property is updated to higher values on retries. Note: We could possibly consider parameterizing this so that other use-cases can benefit, but as of today only Meta uses the presto-on-spark native codepath. So we leave this enhancement for future. Also note that even for first attempt we use same logic, as it will just spit out existing configured values. This simplifies the approach, compared to the presto-on-spark/java smart retries which only calculate memory settings based on attemptNumber to ensure logic is used only for retries --- .../nativeprocess/NativeExecutionProcess.java | 113 +++++++++- .../execution/TestNativeExecutionProcess.java | 213 ++++++++++++++++++ 2 files changed, 316 insertions(+), 10 deletions(-) diff --git a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/nativeprocess/NativeExecutionProcess.java b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/nativeprocess/NativeExecutionProcess.java index d0c7937c0fe52..c2c0ded8416a5 100644 --- a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/nativeprocess/NativeExecutionProcess.java +++ b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/nativeprocess/NativeExecutionProcess.java @@ -15,6 +15,7 @@ import com.facebook.airlift.json.JsonCodec; import com.facebook.airlift.log.Logger; +import com.facebook.airlift.units.DataSize; import com.facebook.airlift.units.Duration; import com.facebook.presto.Session; import com.facebook.presto.client.ServerInfo; @@ -23,6 +24,7 @@ import com.facebook.presto.spark.execution.http.server.RequestErrorTracker; import com.facebook.presto.spark.execution.http.server.smile.BaseResponse; import com.facebook.presto.spark.execution.property.NativeExecutionSystemConfig; +import com.facebook.presto.spark.execution.property.PrestoSparkWorkerProperty; import com.facebook.presto.spark.execution.property.WorkerProperty; import com.facebook.presto.spi.PrestoException; import com.google.common.annotations.VisibleForTesting; @@ -32,6 +34,7 @@ import com.google.common.util.concurrent.SettableFuture; import okhttp3.HttpUrl; import okhttp3.OkHttpClient; +import org.apache.spark.SparkConf; import org.apache.spark.SparkEnv$; import org.apache.spark.SparkFiles; @@ -61,11 +64,12 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; +import static com.facebook.airlift.units.DataSize.Unit.BYTE; +import static com.facebook.airlift.units.DataSize.Unit.GIGABYTE; import static com.facebook.presto.spi.StandardErrorCode.GENERIC_INTERNAL_ERROR; import static com.facebook.presto.spi.StandardErrorCode.NATIVE_EXECUTION_BINARY_NOT_EXIST; import static com.facebook.presto.spi.StandardErrorCode.NATIVE_EXECUTION_PROCESS_LAUNCH_ERROR; import static com.facebook.presto.spi.StandardErrorCode.NATIVE_EXECUTION_TASK_ERROR; -import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.util.concurrent.Futures.addCallback; import static com.google.common.util.concurrent.MoreExecutors.directExecutor; import static java.lang.String.format; @@ -82,6 +86,7 @@ public class NativeExecutionProcess private static final String WORKER_CONFIG_FILE = "/config.properties"; private static final String WORKER_NODE_CONFIG_FILE = "/node.properties"; private static final String WORKER_CONNECTOR_CONFIG_FILE = "/catalog/"; + private static final String NATIVE_PROCESS_MEMORY_SPARK_CONF_NAME = "spark.memory.offHeap.size"; private static final int SIGSYS = 31; private final String executablePath; @@ -131,6 +136,8 @@ public NativeExecutionProcess( scheduledExecutorService, "getting native process status"); this.workerProperty = requireNonNull(workerProperty, "workerProperty is null"); + // Update any runtime configs to be used by presto native worker + updateWorkerProperties(); } /** @@ -325,25 +332,111 @@ private static int getAvailableTcpPort(String nodeInternalAddress) } } - private String getNativeExecutionCatalogName(Session session) + private void populateConfigurationFiles(String configBasePath) + throws IOException { - checkArgument(session.getCatalog().isPresent(), "Catalog isn't set in the session."); - return session.getCatalog().get(); + workerProperty.populateAllProperties( + Paths.get(configBasePath, WORKER_CONFIG_FILE), + Paths.get(configBasePath, WORKER_NODE_CONFIG_FILE), + Paths.get(configBasePath, WORKER_CONNECTOR_CONFIG_FILE)); // Directory path for catalogs } - private void populateConfigurationFiles(String configBasePath) - throws IOException + private void updateWorkerProperties() { + // Update memory properties + updateWorkerMemoryProperties(); + // The reason we have to pick and assign the port per worker is in our prod environment, // there is no port isolation among all the containers running on the same host, so we have // to pick unique port per worker to avoid port collision. This config will be passed down to // the native execution process eventually for process initialization. workerProperty.getSystemConfig() .update(NativeExecutionSystemConfig.HTTP_SERVER_HTTP_PORT, String.valueOf(port)); - workerProperty.populateAllProperties( - Paths.get(configBasePath, WORKER_CONFIG_FILE), - Paths.get(configBasePath, WORKER_NODE_CONFIG_FILE), - Paths.get(configBasePath, WORKER_CONNECTOR_CONFIG_FILE)); // Directory path for catalogs + } + + protected SparkConf getSparkConf() + { + return SparkEnv$.MODULE$.get() == null ? null : SparkEnv$.MODULE$.get().conf(); + } + + protected PrestoSparkWorkerProperty getWorkerProperty() + { + return (PrestoSparkWorkerProperty) workerProperty; + } + + /** + * Computes values for system-memory-gb and query-memory-gb to start the native worker + * with. + * This logic is mainly useful when spark has provisioned larger containers to run + * previously OOMing tasks. Spark will provision larger container but without below + * logic the cpp process will not be able to use it. + * + * Also, we write the logic in a way that same logic applies during first attempt v/s + * subsequent OOMed larger container retry attempts + * + * The logic is simple and is as below + * - New system-memory-gb = spark.memory.offHeap.size + * - Then to calculate the new value of query-memory-gb we assume that + * the new query-memory to system-memory ratio should be same as old values. + * So we set newQueryMemory = newSystemMemory = (oldQueryMemory/oldSystemMemory) + * + * TODO: In future make this algorithm more configurable. i.e. we might want a min/max + * cap on the systemMemoryGb-queryMemoryGb buffer. Currently we just assume ratio + * is good enough + */ + protected void updateWorkerMemoryProperties() + { + // If sparkConf.NATIVE_PROCESS_MEMORY_SPARK_CONF_NAME is not set + // skip making any updates + SparkConf conf = getSparkConf(); + if (conf == null) { + log.info("Not adjusting native process memory as conf is null"); + return; + } + if (!conf.contains(NATIVE_PROCESS_MEMORY_SPARK_CONF_NAME)) { + log.info("Not adjusting native process memory as %s is not set", NATIVE_PROCESS_MEMORY_SPARK_CONF_NAME); + return; + } + DataSize offHeapMemoryBytes = DataSize.succinctDataSize( + conf.getSizeAsBytes(NATIVE_PROCESS_MEMORY_SPARK_CONF_NAME), BYTE); + DataSize currentSystemMemory = DataSize.valueOf(workerProperty.getSystemConfig().getAllProperties() + .get(NativeExecutionSystemConfig.SYSTEM_MEMORY_GB) + GIGABYTE.getUnitString()); + DataSize currentQueryMemory = DataSize.valueOf(workerProperty.getSystemConfig().getAllProperties() + .get(NativeExecutionSystemConfig.QUERY_MEMORY_GB) + GIGABYTE.getUnitString()); + if (offHeapMemoryBytes.toBytes() == 0 + || currentSystemMemory.toBytes() == 0 + || offHeapMemoryBytes.toBytes() < currentSystemMemory.toBytes()) { + log.info("Not adjusting native process memory as" + + " offHeapMemoryBytes=%s,currentSystemMemory=%s are invalid", offHeapMemoryBytes, currentSystemMemory.toBytes()); + return; + } + + log.info("Setting Native Worker system-memory-gb to offHeap: %s", offHeapMemoryBytes); + DataSize newSystemMemory = offHeapMemoryBytes.convertTo(GIGABYTE); + + double queryMemoryFraction = currentQueryMemory.toBytes() * 1.0 / currentSystemMemory.toBytes(); + DataSize newQueryMemoryBytes = DataSize.succinctDataSize( + queryMemoryFraction * newSystemMemory.toBytes(), BYTE); + log.info("Dynamically Tuning Presto Native Memory Configs. " + + "Configured SparkOffHeap: %s; " + + "[oldSystemMemory: %s, newSystemMemory: %s], queryMemoryFraction: %s, " + + "[oldQueryMemory: %s, newQueryMemory: %s]", + offHeapMemoryBytes, + currentSystemMemory, + newSystemMemory, + queryMemoryFraction, + currentQueryMemory, + newQueryMemoryBytes); + + workerProperty.getSystemConfig() + .update(NativeExecutionSystemConfig.SYSTEM_MEMORY_GB, + String.valueOf((int) newSystemMemory.getValue(GIGABYTE))); + workerProperty.getSystemConfig() + .update(NativeExecutionSystemConfig.QUERY_MEMORY_GB, + String.valueOf((int) newQueryMemoryBytes.getValue(GIGABYTE))); + workerProperty.getSystemConfig() + .update(NativeExecutionSystemConfig.QUERY_MAX_MEMORY_PER_NODE, + newQueryMemoryBytes.convertTo(GIGABYTE).toString()); } private void doGetServerInfo(SettableFuture future) diff --git a/presto-spark-base/src/test/java/com/facebook/presto/spark/execution/TestNativeExecutionProcess.java b/presto-spark-base/src/test/java/com/facebook/presto/spark/execution/TestNativeExecutionProcess.java index 4819dc6fd3e53..c1f2ee4f40a57 100644 --- a/presto-spark-base/src/test/java/com/facebook/presto/spark/execution/TestNativeExecutionProcess.java +++ b/presto-spark-base/src/test/java/com/facebook/presto/spark/execution/TestNativeExecutionProcess.java @@ -14,6 +14,7 @@ package com.facebook.presto.spark.execution; import com.facebook.airlift.json.JsonCodec; +import com.facebook.airlift.log.Logger; import com.facebook.airlift.units.Duration; import com.facebook.presto.Session; import com.facebook.presto.client.ServerInfo; @@ -28,14 +29,20 @@ import com.facebook.presto.spark.execution.property.PrestoSparkWorkerProperty; import com.facebook.presto.sql.analyzer.FeaturesConfig; import com.google.common.collect.ImmutableMap; +import okhttp3.OkHttpClient; +import org.apache.spark.SparkConf; import org.testng.annotations.Test; +import java.io.IOException; +import java.util.HashMap; +import java.util.Map; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import static com.facebook.presto.testing.TestingSession.testSessionBuilder; import static java.util.concurrent.Executors.newScheduledThreadPool; import static java.util.concurrent.Executors.newSingleThreadExecutor; +import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertFalse; import static org.testng.Assert.assertNotSame; import static org.testng.Assert.assertTrue; @@ -83,6 +90,212 @@ public void testNativeProcessShutdown() assertFalse(process.isAlive()); } + @Test + public void testUpdateWorkerMemoryPropertiesWithoutSparkEnv() + { + // Test when no SparkConf is available (SparkEnv not initialized) + TestingNativeExecutionProcess process = createTestingNativeExecutionProcess("10", "8", "8GB", null); + process.updateWorkerMemoryProperties(); + // Verify that values remain unchanged + Map properties = process.getWorkerProperty().getSystemConfig().getAllProperties(); + assertEquals(properties.get(NativeExecutionSystemConfig.SYSTEM_MEMORY_GB), "10"); + assertEquals(properties.get(NativeExecutionSystemConfig.QUERY_MEMORY_GB), "8"); + assertEquals(properties.get(NativeExecutionSystemConfig.QUERY_MAX_MEMORY_PER_NODE), "8GB"); + } + + @Test + public void testUpdateWorkerMemoryPropertiesWithOffHeapMemory() + { + // Test when spark.memory.offHeap.size is set to a value larger than current system memory + SparkConf sparkConf = new SparkConf(); + sparkConf.set("spark.memory.offHeap.size", "20g"); // 20GB + + TestingNativeExecutionProcess process = createTestingNativeExecutionProcess("10", "8", "8GB", sparkConf); + process.updateWorkerMemoryProperties(); + // Verify the updated values + Map properties = process.getWorkerProperty().getSystemConfig().getAllProperties(); + + // Expected values: + // newSystemMemory = 20GB + // queryMemoryFraction = 8/10 = 0.8 + // newQueryMemory = 20 * 0.8 = 16GB + assertEquals(properties.get(NativeExecutionSystemConfig.SYSTEM_MEMORY_GB), "20"); + assertEquals(properties.get(NativeExecutionSystemConfig.QUERY_MEMORY_GB), "16"); + assertEquals(properties.get(NativeExecutionSystemConfig.QUERY_MAX_MEMORY_PER_NODE), "16GB"); + } + + @Test + public void testUpdateWorkerMemoryPropertiesWithoutOffHeapSetting() + { + // Test when spark.memory.offHeap.size is not set + SparkConf sparkConf = new SparkConf(); + // Don't set spark.memory.offHeap.size + + TestingNativeExecutionProcess process = createTestingNativeExecutionProcess("10", "8", "8GB", sparkConf); + process.updateWorkerMemoryProperties(); + + // Verify that values remain unchanged + Map properties = process.getWorkerProperty().getSystemConfig().getAllProperties(); + assertEquals(properties.get(NativeExecutionSystemConfig.SYSTEM_MEMORY_GB), "10"); + assertEquals(properties.get(NativeExecutionSystemConfig.QUERY_MEMORY_GB), "8"); + assertEquals(properties.get(NativeExecutionSystemConfig.QUERY_MAX_MEMORY_PER_NODE), "8GB"); + } + + @Test + public void testUpdateWorkerMemoryPropertiesWithZeroOffHeapMemory() + { + // Test when spark.memory.offHeap.size is set to 0 + SparkConf sparkConf = new SparkConf(); + sparkConf.set("spark.memory.offHeap.size", "0b"); + + TestingNativeExecutionProcess process = createTestingNativeExecutionProcess("10", "8", "8GB", sparkConf); + process.updateWorkerMemoryProperties(); + + // Verify that values remain unchanged when offHeapMemory is 0 + Map properties = process.getWorkerProperty().getSystemConfig().getAllProperties(); + assertEquals(properties.get(NativeExecutionSystemConfig.SYSTEM_MEMORY_GB), "10"); + assertEquals(properties.get(NativeExecutionSystemConfig.QUERY_MEMORY_GB), "8"); + assertEquals(properties.get(NativeExecutionSystemConfig.QUERY_MAX_MEMORY_PER_NODE), "8GB"); + } + + @Test + public void testUpdateWorkerMemoryPropertiesWithSmallerOffHeapMemory() + { + // Test when spark.memory.offHeap.size is smaller than current system memory + SparkConf sparkConf = new SparkConf(); + sparkConf.set("spark.memory.offHeap.size", "5g"); // 5GB (smaller than current 10GB) + + TestingNativeExecutionProcess process = createTestingNativeExecutionProcess("10", "8", "8GB", sparkConf); + process.updateWorkerMemoryProperties(); + + // Verify that values remain unchanged when offHeapMemory is smaller than current system memory + Map properties = process.getWorkerProperty().getSystemConfig().getAllProperties(); + assertEquals(properties.get(NativeExecutionSystemConfig.SYSTEM_MEMORY_GB), "10"); + assertEquals(properties.get(NativeExecutionSystemConfig.QUERY_MEMORY_GB), "8"); + assertEquals(properties.get(NativeExecutionSystemConfig.QUERY_MAX_MEMORY_PER_NODE), "8GB"); + } + + @Test + public void testUpdateWorkerMemoryPropertiesWithDifferentQueryMemoryFraction() + { + // Test with different query memory to system memory ratio + SparkConf sparkConf = new SparkConf(); + sparkConf.set("spark.memory.offHeap.size", "30g"); // 30GB + + TestingNativeExecutionProcess process = createTestingNativeExecutionProcess("12", "6", "6GB", sparkConf); + process.updateWorkerMemoryProperties(); + + // Verify the updated values + Map properties = process.getWorkerProperty().getSystemConfig().getAllProperties(); + + // Expected values: + // newSystemMemory = 30GB + // queryMemoryFraction = 6/12 = 0.5 + // newQueryMemory = 30 * 0.5 = 15GB + assertEquals(properties.get(NativeExecutionSystemConfig.SYSTEM_MEMORY_GB), "30"); + assertEquals(properties.get(NativeExecutionSystemConfig.QUERY_MEMORY_GB), "15"); + assertEquals(properties.get(NativeExecutionSystemConfig.QUERY_MAX_MEMORY_PER_NODE), "15GB"); + } + + @Test + public void testUpdateWorkerMemoryPropertiesWithZeroCurrentSystemMemory() + { + // Test edge case when current system memory is 0 + SparkConf sparkConf = new SparkConf(); + sparkConf.set("spark.memory.offHeap.size", "20g"); // 20GB + + TestingNativeExecutionProcess process = createTestingNativeExecutionProcess("0", "8", "8GB", sparkConf); + process.updateWorkerMemoryProperties(); + + // Verify that values remain unchanged when current system memory is 0 + Map properties = process.getWorkerProperty().getSystemConfig().getAllProperties(); + assertEquals(properties.get(NativeExecutionSystemConfig.SYSTEM_MEMORY_GB), "0"); + assertEquals(properties.get(NativeExecutionSystemConfig.QUERY_MEMORY_GB), "8"); + assertEquals(properties.get(NativeExecutionSystemConfig.QUERY_MAX_MEMORY_PER_NODE), "8GB"); + } + + private TestingNativeExecutionProcess createTestingNativeExecutionProcess( + String systemMemoryGb, + String queryMemoryGb, + String queryMaxMemoryPerNode, + SparkConf sparkConf) + { + Map systemConfigs = new HashMap<>(); + systemConfigs.put(NativeExecutionSystemConfig.SYSTEM_MEMORY_GB, systemMemoryGb); + systemConfigs.put(NativeExecutionSystemConfig.QUERY_MEMORY_GB, queryMemoryGb); + systemConfigs.put(NativeExecutionSystemConfig.QUERY_MAX_MEMORY_PER_NODE, queryMaxMemoryPerNode); + + NativeExecutionSystemConfig systemConfig = new NativeExecutionSystemConfig(systemConfigs); + PrestoSparkWorkerProperty workerProperty = new PrestoSparkWorkerProperty( + new NativeExecutionCatalogProperties(ImmutableMap.of()), + new NativeExecutionNodeConfig(), + systemConfig); + + Session session = testSessionBuilder().build(); + + try { + return new TestingNativeExecutionProcess( + "/bin/echo", + "", + session, + new TestPrestoSparkHttpClient.TestingOkHttpClient(newScheduledThreadPool(4), + new TestPrestoSparkHttpClient.TestingResponseManager("test")), + newSingleThreadExecutor(), + newScheduledThreadPool(4), + SERVER_INFO_JSON_CODEC, + new Duration(10, TimeUnit.SECONDS), + workerProperty, + sparkConf); + } + catch (IOException e) { + throw new RuntimeException(e); + } + } + + /** + * Testing subclass that allows injection of custom SparkConf for testing + */ + private static class TestingNativeExecutionProcess + extends NativeExecutionProcess + { + private static final Logger log = Logger.get(TestingNativeExecutionProcess.class); + private final SparkConf testSparkConf; + public TestingNativeExecutionProcess( + String executablePath, + String programArguments, + Session session, + OkHttpClient httpClient, + java.util.concurrent.Executor executor, + ScheduledExecutorService scheduledExecutorService, + JsonCodec serverInfoCodec, + Duration maxErrorDuration, + PrestoSparkWorkerProperty workerProperty, + SparkConf testSparkConf) + throws IOException + { + super(executablePath, programArguments, session, httpClient, executor, + scheduledExecutorService, serverInfoCodec, maxErrorDuration, workerProperty); + this.testSparkConf = testSparkConf; + } + + @Override + protected SparkConf getSparkConf() + { + return testSparkConf; + } + + public PrestoSparkWorkerProperty getWorkerProperty() + { + return super.getWorkerProperty(); + } + + @Override + protected void updateWorkerMemoryProperties() + { + super.updateWorkerMemoryProperties(); + } + } + private NativeExecutionProcessFactory createNativeExecutionProcessFactory() { TaskId taskId = new TaskId("testid", 0, 0, 0, 0); From 493639063b0f5d7bcc570dbabb448b33fcc70093 Mon Sep 17 00:00:00 2001 From: Pratik Joseph Dabre Date: Wed, 10 Sep 2025 13:41:10 -0700 Subject: [PATCH 079/113] Split Provisio plugin packaging into plugins and native-plugins directory --- .../sphinx/plugin/native-sidecar-plugin.rst | 2 + .../conf/docker/common/compose-commons.sh | 10 --- presto-server/src/main/provisio/presto.xml | 85 +++++++++++++++++-- 3 files changed, 81 insertions(+), 16 deletions(-) diff --git a/presto-docs/src/main/sphinx/plugin/native-sidecar-plugin.rst b/presto-docs/src/main/sphinx/plugin/native-sidecar-plugin.rst index 49440a816288b..760176afbe635 100644 --- a/presto-docs/src/main/sphinx/plugin/native-sidecar-plugin.rst +++ b/presto-docs/src/main/sphinx/plugin/native-sidecar-plugin.rst @@ -26,6 +26,8 @@ Property Name Description ``coordinator-sidecar-enabled`` Enables sidecar in the coordinator true ``native-execution-enabled`` Enables native execution true ``presto.default-namespace`` Sets the default function namespace `native.default` +``plugin.dir`` Specifies which directory under installation root `{root-directory}/native-plugins/` + to scan for plugins at startup. ============================================ ===================================================================== ============================== .. _sidecar-worker-properties: diff --git a/presto-product-tests/conf/docker/common/compose-commons.sh b/presto-product-tests/conf/docker/common/compose-commons.sh index 5c20783716b60..eae9f18ce9583 100644 --- a/presto-product-tests/conf/docker/common/compose-commons.sh +++ b/presto-product-tests/conf/docker/common/compose-commons.sh @@ -39,16 +39,6 @@ if [[ -z "${PRESTO_SERVER_DIR:-}" ]]; then source "${PRODUCT_TESTS_ROOT}/target/classes/presto.env" PRESTO_SERVER_DIR="${PROJECT_ROOT}/presto-server/target/presto-server-${PRESTO_VERSION}/" fi - -# The following plugin results in a function signature conflict when loaded in Java/ sidecar disabled native clusters. -# This plugin is only meant for sidecar enabled native clusters, hence exclude it. -PLUGIN_TO_EXCLUDE="native-sql-invoked-functions-plugin" - -if [[ -d "${PRESTO_SERVER_DIR}/plugin/${PLUGIN_TO_EXCLUDE}" ]]; then - echo "Excluding plugin: $PLUGIN_TO_EXCLUDE" - rm -rf "${PRESTO_SERVER_DIR}/plugin/${PLUGIN_TO_EXCLUDE}" -fi - export_canonical_path PRESTO_SERVER_DIR if [[ -z "${PRESTO_CLI_JAR:-}" ]]; then diff --git a/presto-server/src/main/provisio/presto.xml b/presto-server/src/main/provisio/presto.xml index d15a041c7d1f5..59dd851e6436e 100644 --- a/presto-server/src/main/provisio/presto.xml +++ b/presto-server/src/main/provisio/presto.xml @@ -28,7 +28,7 @@ - + @@ -281,19 +281,92 @@ - - + + - - + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + - + From 6732f41193932b2bc3ef69cff26939f78b869c4a Mon Sep 17 00:00:00 2001 From: Ke Wang Date: Wed, 10 Sep 2025 17:52:30 -0700 Subject: [PATCH 080/113] [native] Add runtime stats taskCreationTime --- presto-native-execution/presto_cpp/main/PrestoTask.cpp | 1 + presto-native-execution/presto_cpp/main/PrestoTask.h | 3 +++ presto-native-execution/presto_cpp/main/TaskManager.cpp | 1 + 3 files changed, 5 insertions(+) diff --git a/presto-native-execution/presto_cpp/main/PrestoTask.cpp b/presto-native-execution/presto_cpp/main/PrestoTask.cpp index 3b0cb9f03006f..c48ace86bad8e 100644 --- a/presto-native-execution/presto_cpp/main/PrestoTask.cpp +++ b/presto-native-execution/presto_cpp/main/PrestoTask.cpp @@ -769,6 +769,7 @@ void PrestoTask::updateTimeInfoLocked( taskRuntimeStats["endTime"].addValue(veloxTaskStats.endTimeMs); } taskRuntimeStats.insert({"nativeProcessCpuTime", fromNanos(processCpuTime_)}); + taskRuntimeStats.insert({"taskCreationTime", fromNanos((createFinishTimeMs - createTimeMs) * 1'000'000)}); } void PrestoTask::updateMemoryInfoLocked( diff --git a/presto-native-execution/presto_cpp/main/PrestoTask.h b/presto-native-execution/presto_cpp/main/PrestoTask.h index 297f877f3793e..3732ccffedde2 100644 --- a/presto-native-execution/presto_cpp/main/PrestoTask.h +++ b/presto-native-execution/presto_cpp/main/PrestoTask.h @@ -121,7 +121,10 @@ struct PrestoTask { uint64_t lastTaskStatsUpdateMs{0}; uint64_t lastMemoryReservation{0}; + /// Time point (in ms) when the time we start task creating. uint64_t createTimeMs{0}; + /// Time point (in ms) when the time we finish task creating. + uint64_t createFinishTimeMs{0}; uint64_t startTimeMs{0}; uint64_t firstSplitStartTimeMs{0}; uint64_t lastEndTimeMs{0}; diff --git a/presto-native-execution/presto_cpp/main/TaskManager.cpp b/presto-native-execution/presto_cpp/main/TaskManager.cpp index 10ed292bac4f6..1d893c4a110de 100644 --- a/presto-native-execution/presto_cpp/main/TaskManager.cpp +++ b/presto-native-execution/presto_cpp/main/TaskManager.cpp @@ -576,6 +576,7 @@ std::unique_ptr TaskManager::createOrUpdateTaskImpl( } execTask = prestoTask->task; } + prestoTask->createFinishTimeMs = getCurrentTimeMs(); // Outside of prestoTask->mutex. VELOX_CHECK_NOT_NULL( execTask, From d5abf92ed89fb82a5ddf664aafaa9e36f158173a Mon Sep 17 00:00:00 2001 From: Sumi Mathew Date: Mon, 1 Sep 2025 18:05:26 +0530 Subject: [PATCH 081/113] Upgrade org.reflections version --- presto-pinot-toolkit/pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/presto-pinot-toolkit/pom.xml b/presto-pinot-toolkit/pom.xml index 88367dbd9fde3..5abe5d9005762 100644 --- a/presto-pinot-toolkit/pom.xml +++ b/presto-pinot-toolkit/pom.xml @@ -515,7 +515,7 @@ org.reflections reflections - 0.9.9 + 0.10.2 com.google.code.findbugs From 1fc6b0c2154751a1c7f2cb7c719326c81b3f5c34 Mon Sep 17 00:00:00 2001 From: Pratik Joseph Dabre Date: Wed, 10 Sep 2025 16:22:14 -0700 Subject: [PATCH 082/113] [native] Add parameterized varchar type in the list of supported types in NativeTypeManager --- .../typemanager/NativeTypeManager.java | 15 +++------ .../sidecar/TestNativeSidecarPlugin.java | 6 ++++ .../presto/nativetests/TestWindowQueries.java | 32 ------------------- 3 files changed, 10 insertions(+), 43 deletions(-) diff --git a/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/typemanager/NativeTypeManager.java b/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/typemanager/NativeTypeManager.java index e86ce277c0f31..dd4122a04cc6d 100644 --- a/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/typemanager/NativeTypeManager.java +++ b/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/typemanager/NativeTypeManager.java @@ -20,7 +20,6 @@ import com.facebook.presto.common.type.TypeManager; import com.facebook.presto.common.type.TypeSignature; import com.facebook.presto.common.type.TypeSignatureParameter; -import com.facebook.presto.common.type.VarcharType; import com.facebook.presto.spi.type.UnknownTypeException; import com.google.common.cache.CacheBuilder; import com.google.common.cache.CacheLoader; @@ -67,7 +66,6 @@ import static com.facebook.presto.common.type.StandardTypes.UUID; import static com.facebook.presto.common.type.StandardTypes.VARBINARY; import static com.facebook.presto.common.type.StandardTypes.VARCHAR; -import static com.facebook.presto.common.type.VarcharType.createUnboundedVarcharType; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Preconditions.checkState; import static com.google.common.base.Throwables.throwIfUnchecked; @@ -109,7 +107,10 @@ public class NativeTypeManager QDIGEST, ROW, TDIGEST, - FunctionType.NAME); + FunctionType.NAME, + // todo: fix this hack: parametrized varchar isn't supported in native execution yet + // Currently, native execution without sidecar enabled works with parameterized varchar, hence adding this to the list of supported types. + VARCHAR); private final TypeManager typeManager; private final LoadingCache parametricTypeCache; @@ -134,10 +135,6 @@ public NativeTypeManager(TypeManager typeManager) @Override public Type getType(TypeSignature typeSignature) { - // Todo: Fix this hack, native execution does not support parameterized varchar type signatures. - if (typeSignature.getBase().equals(VARCHAR)) { - typeSignature = createUnboundedVarcharType().getTypeSignature(); - } Type type = types.get(typeSignature); if (type != null) { return type; @@ -190,10 +187,6 @@ public boolean canCoerce(Type actualType, Type expectedType) private void addAllTypes(List typesList, List parametricTypesList) { typesList.forEach(this::addType); - // todo: Fix this hack - // Native engine does not support parameterized varchar, and varchar isn't in the lists of types returned from the engine - addType(VarcharType.VARCHAR); - parametricTypesList.forEach(this::addParametricType); } diff --git a/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java b/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java index 40ade2247847d..4cd909d6464f4 100644 --- a/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java +++ b/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java @@ -220,6 +220,7 @@ public void testGeneralQueries() "date_trunc('second', from_unixtime(orderkey, '+00:00')) FROM orders"); assertQuery("SELECT mod(orderkey, linenumber) FROM lineitem"); assertQueryFails("SELECT IF(true, 0/0, 1)", "[\\s\\S]*/ by zero native.default.fail[\\s\\S]*"); + assertQuery("select CASE WHEN true THEN 'Yes' ELSE 'No' END"); } @Test @@ -353,6 +354,11 @@ public void testInformationSchemaTables() { assertQuery("select lower(table_name) from information_schema.tables " + "where table_name = 'lineitem' or table_name = 'LINEITEM' "); + assertQuery("SELECT table_name, CASE WHEN abs(ordinal_position) > 3 THEN 'high' WHEN abs(ordinal_position) > 1 THEN 'medium' ELSE 'low' END as position_category, COUNT(*) \n" + + "FROM information_schema.columns " + + "WHERE table_catalog = 'hive' AND table_name IN ('nation', 'region', 'lineitem', 'orders') " + + "GROUP BY table_name, CASE WHEN abs(ordinal_position) > 3 THEN 'high' WHEN abs(ordinal_position) > 1 THEN 'medium' ELSE 'low' END " + + "ORDER BY table_name, position_category"); } @Test diff --git a/presto-native-tests/src/test/java/com/facebook/presto/nativetests/TestWindowQueries.java b/presto-native-tests/src/test/java/com/facebook/presto/nativetests/TestWindowQueries.java index 6db7af6e559ed..3b2a2a8c44d05 100644 --- a/presto-native-tests/src/test/java/com/facebook/presto/nativetests/TestWindowQueries.java +++ b/presto-native-tests/src/test/java/com/facebook/presto/nativetests/TestWindowQueries.java @@ -13,18 +13,11 @@ */ package com.facebook.presto.nativetests; -import com.facebook.presto.testing.MaterializedResult; import com.facebook.presto.testing.QueryRunner; import com.facebook.presto.tests.AbstractTestWindowQueries; -import org.intellij.lang.annotations.Language; import org.testng.annotations.BeforeClass; import org.testng.annotations.Test; -import static com.facebook.presto.common.type.BigintType.BIGINT; -import static com.facebook.presto.common.type.VarcharType.createUnboundedVarcharType; -import static com.facebook.presto.common.type.VarcharType.createVarcharType; -import static com.facebook.presto.testing.MaterializedResult.resultBuilder; -import static com.facebook.presto.testing.assertions.Assert.assertEquals; import static java.lang.Boolean.parseBoolean; public class TestWindowQueries @@ -234,29 +227,4 @@ public void testTypes() "(INTERVAL '2' month, ARRAY[INTERVAL '1' month, INTERVAL '2' month]), " + "(INTERVAL '5' year, ARRAY[INTERVAL '5' year])"); } - - // Todo: Refactor this test case when support for varchar(N) is added in native execution. - // The return types do not match on the native query runner : types=[varchar, bigint] and the java query runner: types=[varchar(3), bigint]. - @Override - @Test - public void testWindowFunctionWithGroupBy() - { - @Language("SQL") String sql = "SELECT *, rank() OVER (PARTITION BY x)\n" + - "FROM (SELECT 'foo' x)\n" + - "GROUP BY 1"; - - MaterializedResult actual = computeActual(sql); - MaterializedResult expected; - if (sidecarEnabled) { - expected = resultBuilder(getSession(), createUnboundedVarcharType(), BIGINT) - .row("foo", 1L) - .build(); - } - else { - expected = resultBuilder(getSession(), createVarcharType(3), BIGINT) - .row("foo", 1L) - .build(); - } - assertEquals(actual, expected); - } } From 6b5dbf80c57426039979c2ad41fae3674277ab89 Mon Sep 17 00:00:00 2001 From: Sumi Mathew Date: Tue, 9 Sep 2025 12:11:02 +0530 Subject: [PATCH 083/113] Upgrade org.antlr version --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index c2de59f18bcdf..368a195f21996 100644 --- a/pom.xml +++ b/pom.xml @@ -43,7 +43,7 @@ 17 3.3.9 - 4.7.1 + 4.13.2 0.221 ${dep.airlift.version} 0.38 From f31246e1ecce7cc1ea6b2fff34a0dc701fa9c5a9 Mon Sep 17 00:00:00 2001 From: adkharat Date: Thu, 28 Aug 2025 12:45:23 +0530 Subject: [PATCH 084/113] enhance csp security with form-action directive and stricter img-src --- .../main/java/com/facebook/presto/server/CoordinatorModule.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/presto-main/src/main/java/com/facebook/presto/server/CoordinatorModule.java b/presto-main/src/main/java/com/facebook/presto/server/CoordinatorModule.java index a165616c09f4c..1813c5eb06995 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/CoordinatorModule.java +++ b/presto-main/src/main/java/com/facebook/presto/server/CoordinatorModule.java @@ -144,7 +144,7 @@ public class CoordinatorModule { private static final String DEFAULT_WEBUI_CSP = "default-src 'self'; style-src 'self' 'unsafe-inline' https://fonts.googleapis.com; " + - "font-src 'self' https://fonts.gstatic.com; frame-ancestors 'self'; img-src http: https: data:"; + "font-src 'self' https://fonts.gstatic.com; frame-ancestors 'self'; img-src 'self' http: https: data:; form-action 'self'"; public static HttpResourceBinding webUIBinder(Binder binder, String path, String classPathResourceBase) { From 0bdb471a9893d10d68abdf2a2268013c8b51988f Mon Sep 17 00:00:00 2001 From: Amit Dutta Date: Wed, 10 Sep 2025 22:44:54 -0700 Subject: [PATCH 085/113] [native] Advance velox --- presto-native-execution/velox | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/presto-native-execution/velox b/presto-native-execution/velox index 7992424ce4d78..63f9644ea7839 160000 --- a/presto-native-execution/velox +++ b/presto-native-execution/velox @@ -1 +1 @@ -Subproject commit 7992424ce4d7886681270019b1729bb44a91e7ab +Subproject commit 63f9644ea78390f6631b193263edbc2b6b56b268 From 809ae7bcfb6581a4d36bc4fca833c5d9944caad1 Mon Sep 17 00:00:00 2001 From: Christian Zentgraf Date: Thu, 11 Sep 2025 19:13:35 -0700 Subject: [PATCH 086/113] [native] Disable Velox mono-library build --- presto-native-execution/CMakeLists.txt | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/presto-native-execution/CMakeLists.txt b/presto-native-execution/CMakeLists.txt index aea12fd9a8456..1e7fce6c7717b 100644 --- a/presto-native-execution/CMakeLists.txt +++ b/presto-native-execution/CMakeLists.txt @@ -227,6 +227,11 @@ include_directories(${CMAKE_BINARY_DIR}) # set this for backwards compatibility, will be overwritten in velox/ set(VELOX_GTEST_INCUDE_DIR "velox/third_party/googletest/googletest/include") +# Do not use the Mono library because it causes link errors. +set(VELOX_MONO_LIBRARY + OFF + CACHE BOOL "Build Velox mono library") + add_subdirectory(velox) if(PRESTO_ENABLE_TESTING) From d7193c7ee7c83687de11cf85e45d45e714e553d2 Mon Sep 17 00:00:00 2001 From: beinan Date: Tue, 9 Sep 2025 20:21:51 +0000 Subject: [PATCH 087/113] Bump up iceberg version to 1.8.1 Upgrade nessie to make it compatible with iceberg 1.8.x Fix HTTP client dependency versions for Iceberg 1.8.1 compatibility - Update httpcore5 from 5.2.4 to 5.3.1 - Resolves NoSuchMethodError in DefaultHttpRequestWriterFactory constructor --- presto-iceberg/pom.xml | 6 ++--- .../rest/IcebergRestCatalogServlet.java | 27 +++++++++++++------ 2 files changed, 22 insertions(+), 11 deletions(-) diff --git a/presto-iceberg/pom.xml b/presto-iceberg/pom.xml index 66445b698a25c..59a435bbff114 100644 --- a/presto-iceberg/pom.xml +++ b/presto-iceberg/pom.xml @@ -14,8 +14,8 @@ ${project.parent.basedir} 17 - 1.6.1 - 0.95.0 + 1.8.1 + 0.103.0 true @@ -613,7 +613,7 @@ org.apache.httpcomponents.core5 httpcore5 - 5.2.4 + 5.3.1 diff --git a/presto-iceberg/src/test/java/org/apache/iceberg/rest/IcebergRestCatalogServlet.java b/presto-iceberg/src/test/java/org/apache/iceberg/rest/IcebergRestCatalogServlet.java index 588f556065616..0fca1ffae07aa 100644 --- a/presto-iceberg/src/test/java/org/apache/iceberg/rest/IcebergRestCatalogServlet.java +++ b/presto-iceberg/src/test/java/org/apache/iceberg/rest/IcebergRestCatalogServlet.java @@ -25,7 +25,7 @@ import org.apache.iceberg.exceptions.RESTException; import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap; import org.apache.iceberg.relocated.com.google.common.io.CharStreams; -import org.apache.iceberg.rest.RESTCatalogAdapter.HTTPMethod; +import org.apache.iceberg.rest.HTTPRequest.HTTPMethod; import org.apache.iceberg.rest.RESTCatalogAdapter.Route; import org.apache.iceberg.rest.responses.ErrorResponse; import org.apache.iceberg.util.Pair; @@ -111,15 +111,18 @@ protected void execute(ServletRequestContext context, HttpServletResponse respon } try { + HTTPRequest request = restCatalogAdapter.buildRequest( + context.method(), + context.path(), + context.queryParams(), + context.headers(), + context.body()); Object responseBody = restCatalogAdapter.execute( - context.method(), - context.path(), - context.queryParams(), - context.body(), + request, context.route().responseClass(), - context.headers(), - handle(response)); + handleResponseError(response), + handleResponseHeader(response)); if (responseBody != null) { RESTObjectMapper.mapper().writeValue(response.getWriter(), responseBody); @@ -143,7 +146,15 @@ protected void execute(ServletRequestContext context, HttpServletResponse respon } } - protected Consumer handle(HttpServletResponse response) + private Consumer> handleResponseHeader(HttpServletResponse response) + { + return (responseHeaders) -> { + LOG.error("Unexpected response header: %s", responseHeaders); + throw new RuntimeException("Unexpected response header: " + responseHeaders); + }; + } + + protected Consumer handleResponseError(HttpServletResponse response) { return (errorResponse) -> { response.setStatus(errorResponse.code()); From 0a029e742f5d8dc1d551545dc63bc275c8722fef Mon Sep 17 00:00:00 2001 From: Shrinidhi Joshi Date: Thu, 11 Sep 2025 12:57:44 -0700 Subject: [PATCH 088/113] [pos][native] Unit testing fixes Makes 2 changes to pos native tests to run them on Java 17 1. Disable spark ui in tests. SparkUI needs to instantiate JavaX servlets and this causes runtime error due to incompatibility with glassfish servlets used by new airflit.jaxrs 2. Update NativeExecutionModule instantiation to have hive.properties in catalog properties due to a recent change - https://github.com/prestodb/presto/pull/25943 --- presto-native-execution/pom.xml | 13 +++++++++++++ .../spark/PrestoSparkNativeQueryRunnerUtils.java | 4 +++- .../nativeprocess/NativeExecutionModule.java | 6 ------ .../presto/spark/PrestoSparkQueryRunner.java | 3 ++- 4 files changed, 18 insertions(+), 8 deletions(-) diff --git a/presto-native-execution/pom.xml b/presto-native-execution/pom.xml index 4f48109d76937..f9150f63a4daf 100644 --- a/presto-native-execution/pom.xml +++ b/presto-native-execution/pom.xml @@ -30,6 +30,19 @@ + + + javax.ws.rs + javax.ws.rs-api + test + + + javax.servlet + javax.servlet-api + 3.1.0 + test + + com.facebook.presto presto-tpch diff --git a/presto-native-execution/src/test/java/com/facebook/presto/spark/PrestoSparkNativeQueryRunnerUtils.java b/presto-native-execution/src/test/java/com/facebook/presto/spark/PrestoSparkNativeQueryRunnerUtils.java index 21ef570381795..a0856f8c6a530 100644 --- a/presto-native-execution/src/test/java/com/facebook/presto/spark/PrestoSparkNativeQueryRunnerUtils.java +++ b/presto-native-execution/src/test/java/com/facebook/presto/spark/PrestoSparkNativeQueryRunnerUtils.java @@ -99,7 +99,8 @@ public static Map getNativeExecutionSparkConfigs() public static PrestoSparkQueryRunner createHiveRunner() { - PrestoSparkQueryRunner queryRunner = createRunner("hive", new NativeExecutionModule()); + PrestoSparkQueryRunner queryRunner = createRunner("hive", new NativeExecutionModule( + Optional.of(ImmutableMap.of("hive", ImmutableMap.of("connector.name", "hive"))))); PrestoNativeQueryRunnerUtils.setupJsonFunctionNamespaceManager(queryRunner, "external_functions.json", "json"); return queryRunner; @@ -186,6 +187,7 @@ private static Database createDatabaseMetastoreObject(String name) private static Map getNativeExecutionShuffleConfigs() { ImmutableMap.Builder sparkConfigs = ImmutableMap.builder(); + sparkConfigs.put("spark.ui.enabled", "false"); sparkConfigs.put(SPARK_SHUFFLE_MANAGER, "com.facebook.presto.spark.classloader_interface.PrestoSparkNativeExecutionShuffleManager"); sparkConfigs.put(FALLBACK_SPARK_SHUFFLE_MANAGER, "org.apache.spark.shuffle.sort.SortShuffleManager"); return sparkConfigs.build(); diff --git a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/nativeprocess/NativeExecutionModule.java b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/nativeprocess/NativeExecutionModule.java index 5636681a85da5..8df6a83395b35 100644 --- a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/nativeprocess/NativeExecutionModule.java +++ b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/nativeprocess/NativeExecutionModule.java @@ -41,12 +41,6 @@ public class NativeExecutionModule { private Optional>> catalogProperties; - // For use by production system where the configurations can only be tuned via configurations. - public NativeExecutionModule() - { - this.catalogProperties = Optional.empty(); - } - // In the future, we would make more bindings injected into NativeExecutionModule // to be able to test various configuration parameters @VisibleForTesting diff --git a/presto-spark-base/src/test/java/com/facebook/presto/spark/PrestoSparkQueryRunner.java b/presto-spark-base/src/test/java/com/facebook/presto/spark/PrestoSparkQueryRunner.java index bd8f8c006c9d7..0d5945fdb5a82 100644 --- a/presto-spark-base/src/test/java/com/facebook/presto/spark/PrestoSparkQueryRunner.java +++ b/presto-spark-base/src/test/java/com/facebook/presto/spark/PrestoSparkQueryRunner.java @@ -247,7 +247,8 @@ public static PrestoSparkQueryRunner createHivePrestoSparkQueryRunner(Iterable Date: Fri, 12 Sep 2025 08:14:12 -0700 Subject: [PATCH 089/113] [pos][native] Add presto-on-spark native tests step to Github CI on PRs --- .../prestocpp-linux-build-and-unit-test.yml | 78 +++++++++++++++++++ 1 file changed, 78 insertions(+) diff --git a/.github/workflows/prestocpp-linux-build-and-unit-test.yml b/.github/workflows/prestocpp-linux-build-and-unit-test.yml index f3fc4843c3321..7917102a6b06b 100644 --- a/.github/workflows/prestocpp-linux-build-and-unit-test.yml +++ b/.github/workflows/prestocpp-linux-build-and-unit-test.yml @@ -312,6 +312,84 @@ jobs: -Duser.timezone=America/Bahia_Banderas \ -T1C + prestocpp-linux-presto-on-spark-e2e-tests: + needs: prestocpp-linux-build-for-test + runs-on: ubuntu-22.04 + strategy: + fail-fast: false + matrix: + storage-format: [ "PARQUET", "DWRF" ] + enable-sidecar: [ "true", "false" ] + container: + image: prestodb/presto-native-dependency:0.293-20250522140509-484b00e + env: + MAVEN_OPTS: "-Xmx4G -XX:+ExitOnOutOfMemoryError" + MAVEN_FAST_INSTALL: "-B -V --quiet -T 1C -DskipTests -Dair.check.skip-all -Dmaven.javadoc.skip=true" + MAVEN_TEST: "-B -Dair.check.skip-all -Dmaven.javadoc.skip=true -DLogTestDurationListener.enabled=true --fail-at-end" + steps: + - uses: actions/checkout@v4 + + - name: Fix git permissions + # Usually actions/checkout does this but as we run in a container + # it doesn't work + run: git config --global --add safe.directory ${GITHUB_WORKSPACE} + + - name: Download artifacts + uses: actions/download-artifact@v4 + with: + name: presto-native-build + path: presto-native-execution/_build/release + + # Permissions are lost when uploading. Details here: https://github.com/actions/upload-artifact/issues/38 + - name: Restore execute permissions and library path + run: | + chmod +x ${GITHUB_WORKSPACE}/presto-native-execution/_build/release/presto_cpp/main/presto_server + chmod +x ${GITHUB_WORKSPACE}/presto-native-execution/_build/release/velox/velox/functions/remote/server/velox_functions_remote_server_main + # Ensure transitive dependency libboost-iostreams is found. + ldconfig /usr/local/lib + + - name: Install OpenJDK17 + uses: actions/setup-java@v4 + with: + distribution: 'temurin' + java-version: 17.0.15 + cache: 'maven' + - name: Download nodejs to maven cache + run: .github/bin/download_nodejs + + - name: Maven install + env: + # Use different Maven options to install. + MAVEN_OPTS: "-Xmx2G -XX:+ExitOnOutOfMemoryError" + run: | + for i in $(seq 1 3); do ./mvnw clean install $MAVEN_FAST_INSTALL -pl 'presto-native-tests' -am && s=0 && break || s=$? && sleep 10; done; (exit $s) + + - name: Run presto-on-spark native tests + run: | + export PRESTO_SERVER_PATH="${GITHUB_WORKSPACE}/presto-native-execution/_build/release/presto_cpp/main/presto_server" + export TESTFILES=`find ./presto-native-execution/src/test -type f -name 'TestPrestoSpark*.java'` + # Convert file paths to comma separated class names + export TESTCLASSES=TestPrestoSparkExpressionCompiler,TestPrestoSparkNativeBitwiseFunctionQueries,TestPrestoSparkNativeTpchConnectorQueries,TestPrestoSparkNativeSimpleQueries,TestPrestoSparkSqlFunctions,TestPrestoSparkNativeTpchQueries + for test_file in $TESTFILES + do + tmp=${test_file##*/} + test_class=${tmp%%\.*} + export TESTCLASSES="${TESTCLASSES},$test_class" + done + export TESTCLASSES=${TESTCLASSES#,} + echo "TESTCLASSES = $TESTCLASSES" + + mvn test \ + ${MAVEN_TEST} \ + -pl 'presto-native-execution' \ + -DstorageFormat=${{ matrix.storage-format }} \ + -DsidecarEnabled=${{ matrix.enable-sidecar }} \ + -Dtest="${TESTCLASSES}" \ + -DPRESTO_SERVER=${PRESTO_SERVER_PATH} \ + -DDATA_DIR=${RUNNER_TEMP} \ + -Duser.timezone=America/Bahia_Banderas \ + -T1C + prestocpp-linux-presto-sidecar-tests: needs: [changes, prestocpp-linux-build-for-test] runs-on: ubuntu-22.04 From 39ea2c08aa1da956a0f41a1874283aa501f394cd Mon Sep 17 00:00:00 2001 From: Shrinidhi Joshi Date: Wed, 23 Jul 2025 11:02:01 -0700 Subject: [PATCH 090/113] [pos][native] Disable failing presto-on-spark native tests Tracking disabled test fixed in https://github.com/prestodb/presto/issues/25511 --- .../AbstractTestExpressionCompiler.java | 68 +++++++++++++++++++ ...PrestoSparkNativeArrayFunctionQueries.java | 6 ++ .../TestPrestoSparkNativeGeneralQueries.java | 20 ++++++ ...PrestoSparkNativeTpchConnectorQueries.java | 5 ++ 4 files changed, 99 insertions(+) diff --git a/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/AbstractTestExpressionCompiler.java b/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/AbstractTestExpressionCompiler.java index 9ed6cbf0e9981..dec795a4b6fca 100644 --- a/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/AbstractTestExpressionCompiler.java +++ b/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/AbstractTestExpressionCompiler.java @@ -84,12 +84,80 @@ public void testNestedColumnFilter() { } + // Caused by: com.facebook.presto.spi.PrestoException: Catalog does not exist: hive @Override @Ignore + public void testBinaryOperatorsBigintDecimal() + { + } + + // Caused by: com.facebook.presto.spi.PrestoException: Catalog does not exist: hive + @Override + @Ignore + public void testBinaryOperatorsBoolean() + { + } + + // Caused by: com.facebook.presto.spi.PrestoException: Catalog does not exist: hive + @Override + @Ignore + public void testBinaryOperatorsDoubleDecimal() + { + } + + // Caused by: com.facebook.presto.spi.PrestoException: Catalog does not exist: hive + @Override + @Ignore + public void testBinaryOperatorsDoubleDouble() + { + } + public void testTernaryOperatorsLongLong() { } + // Caused by: com.facebook.presto.spi.PrestoException: Catalog does not exist: hive + @Override + @Ignore + public void testBinaryOperatorsDoubleIntegral() + { + } + + // Caused by: com.facebook.presto.spi.PrestoException: Catalog does not exist: hive + @Override + @Ignore + public void testBinaryOperatorsIntegerDecimal() + { + } + + // Caused by: com.facebook.presto.spi.PrestoException: Catalog does not exist: hive + @Override + @Ignore + public void testBinaryOperatorsIntegralDouble() + { + } + + // Caused by: com.facebook.presto.spi.PrestoException: Catalog does not exist: hive + @Override + @Ignore + public void testBinaryOperatorsIntegralIntegral() + { + } + + // Caused by: com.facebook.presto.spi.PrestoException: Catalog does not exist: hive + @Override + @Ignore + public void testBinaryOperatorsString() + { + } + + // Caused by: com.facebook.presto.spi.PrestoException: Catalog does not exist: hive + @Override + @Ignore + public void testNullif() + { + } + @Override @Ignore public void testTernaryOperatorsLongDouble() diff --git a/presto-native-execution/src/test/java/com/facebook/presto/spark/TestPrestoSparkNativeArrayFunctionQueries.java b/presto-native-execution/src/test/java/com/facebook/presto/spark/TestPrestoSparkNativeArrayFunctionQueries.java index 40db300489142..89871668882e5 100644 --- a/presto-native-execution/src/test/java/com/facebook/presto/spark/TestPrestoSparkNativeArrayFunctionQueries.java +++ b/presto-native-execution/src/test/java/com/facebook/presto/spark/TestPrestoSparkNativeArrayFunctionQueries.java @@ -16,6 +16,7 @@ import com.facebook.presto.nativeworker.AbstractTestNativeArrayFunctionQueries; import com.facebook.presto.testing.ExpectedQueryRunner; import com.facebook.presto.testing.QueryRunner; +import org.testng.annotations.Ignore; public class TestPrestoSparkNativeArrayFunctionQueries extends AbstractTestNativeArrayFunctionQueries @@ -32,4 +33,9 @@ protected ExpectedQueryRunner createExpectedQueryRunner() { return PrestoSparkNativeQueryRunnerUtils.createJavaQueryRunner(); } + + // Caused by: com.facebook.presto.sql.analyzer.SemanticException: line 1:32: Function array_sort_desc not registered + @Override + @Ignore + public void testArraySort() {} } diff --git a/presto-native-execution/src/test/java/com/facebook/presto/spark/TestPrestoSparkNativeGeneralQueries.java b/presto-native-execution/src/test/java/com/facebook/presto/spark/TestPrestoSparkNativeGeneralQueries.java index 902de47b2fd79..458323383d2dc 100644 --- a/presto-native-execution/src/test/java/com/facebook/presto/spark/TestPrestoSparkNativeGeneralQueries.java +++ b/presto-native-execution/src/test/java/com/facebook/presto/spark/TestPrestoSparkNativeGeneralQueries.java @@ -106,4 +106,24 @@ public void testRowWiseExchange() {} @Override @Ignore public void testAnalyzeStatsOnDecimals() {} + + //Caused by: com.facebook.presto.sql.analyzer.SemanticException: line 1:31: Function array_duplicates not registered + @Override + @Ignore + public void testArrayAndMapFunctions() {} + + // VeloxRuntimeError: it != connectors().end() Connector with ID 'hivecached' not registered + @Override + @Ignore + public void testCatalogWithCacheEnabled() {} + + // Caused by: com.facebook.presto.spi.PrestoException: Sampling function: key_sampling_percent not cannot be resolved + @Override + @Ignore + public void testKeyBasedSamplingInlined() {} + + // VeloxRuntimeError: !noMoreSplits_ + @Override + @Ignore + public void testUnionAllInsert() {} } diff --git a/presto-native-execution/src/test/java/com/facebook/presto/spark/TestPrestoSparkNativeTpchConnectorQueries.java b/presto-native-execution/src/test/java/com/facebook/presto/spark/TestPrestoSparkNativeTpchConnectorQueries.java index b2a0bb1ad5d51..05d194f6f6754 100644 --- a/presto-native-execution/src/test/java/com/facebook/presto/spark/TestPrestoSparkNativeTpchConnectorQueries.java +++ b/presto-native-execution/src/test/java/com/facebook/presto/spark/TestPrestoSparkNativeTpchConnectorQueries.java @@ -43,4 +43,9 @@ public void testMissingTpchConnector() @Override @Ignore public void testTpchTinyTables() {} + + // VeloxRuntimeError: it != connectors().end() Connector with ID 'tpchstandard' not registered + @Override + @Ignore + public void testTpchDateFilter() {} } From 4865fc95b6d29af18c7e009200e4462ac4d605c2 Mon Sep 17 00:00:00 2001 From: haneel-kumar Date: Tue, 12 Aug 2025 23:30:20 -0700 Subject: [PATCH 091/113] feat(connector): Add mTLS support for Arrow Flight This commit introduces mutual TLS authentication for the Arrow Flight connector, including necessary configuration options for both the client and server. It also includes fixes to the CI pipeline and C++ tests to ensure the new mTLS functionality is properly validated. co-authored-by: Ajas Mangal co-authored-by: Elbin Pallimalil co-authored-by: Thanzeel Hassan --- .github/workflows/arrow-flight-tests.yml | 4 +- .../plugin/arrow/ArrowFlightQueryRunner.java | 54 +++--- .../arrow/TestArrowFlightEchoQueries.java | 6 +- .../TestArrowFlightIntegrationSmokeTest.java | 4 +- .../plugin/arrow/TestArrowFlightMtls.java | 22 +-- .../arrow/TestArrowFlightNativeQueries.java | 132 +++++++++------ .../TestArrowFlightNativeQueriesMtls.java | 33 ++++ .../plugin/arrow/TestArrowFlightQueries.java | 4 +- .../src/test/resources/{mtls => certs}/ca.crt | 0 .../test/resources/{mtls => certs}/client.crt | 0 .../test/resources/{mtls => certs}/client.key | 0 .../{mtls => certs}/invalid_cert.crt | 0 .../test/resources/{mtls => certs}/server.crt | 0 .../test/resources/{mtls => certs}/server.key | 0 .../src/test/resources/server.crt | 32 ---- .../src/test/resources/server.key | 52 ------ .../sphinx/connector/base-arrow-flight.rst | 30 ++++ .../arrow_flight/ArrowFlightConfig.cpp | 10 ++ .../arrow_flight/ArrowFlightConfig.h | 10 ++ .../arrow_flight/ArrowFlightConnector.cpp | 20 +++ .../tests/ArrowFlightConfigTest.cpp | 8 +- .../tests/ArrowFlightConnectorMTlsTest.cpp | 158 ++++++++++++++++++ .../tests/ArrowFlightConnectorTlsTest.cpp | 6 +- .../arrow_flight/tests/CMakeLists.txt | 11 +- .../arrow_flight/tests/data/README.md | 10 +- .../arrow_flight/tests/data/certs/ca.crt | 22 +++ .../arrow_flight/tests/data/certs/client.crt | 22 +++ .../arrow_flight/tests/data/certs/client.key | 28 ++++ .../arrow_flight/tests/data/certs/server.crt | 23 +++ .../arrow_flight/tests/data/certs/server.key | 28 ++++ .../arrow_flight/tests/data/generate_certs.sh | 55 ++++++ .../tests/data/generate_tls_certs.sh | 40 ----- .../arrow_flight/tests/data/tls_certs/ca.crt | 22 --- .../tests/data/tls_certs/server.crt | 22 --- .../tests/data/tls_certs/server.key | 28 ---- 35 files changed, 591 insertions(+), 305 deletions(-) create mode 100644 presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightNativeQueriesMtls.java rename presto-base-arrow-flight/src/test/resources/{mtls => certs}/ca.crt (100%) rename presto-base-arrow-flight/src/test/resources/{mtls => certs}/client.crt (100%) rename presto-base-arrow-flight/src/test/resources/{mtls => certs}/client.key (100%) rename presto-base-arrow-flight/src/test/resources/{mtls => certs}/invalid_cert.crt (100%) rename presto-base-arrow-flight/src/test/resources/{mtls => certs}/server.crt (100%) rename presto-base-arrow-flight/src/test/resources/{mtls => certs}/server.key (100%) delete mode 100644 presto-base-arrow-flight/src/test/resources/server.crt delete mode 100644 presto-base-arrow-flight/src/test/resources/server.key create mode 100644 presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/ArrowFlightConnectorMTlsTest.cpp create mode 100644 presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/certs/ca.crt create mode 100644 presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/certs/client.crt create mode 100644 presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/certs/client.key create mode 100644 presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/certs/server.crt create mode 100644 presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/certs/server.key create mode 100755 presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/generate_certs.sh delete mode 100755 presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/generate_tls_certs.sh delete mode 100644 presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/tls_certs/ca.crt delete mode 100644 presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/tls_certs/server.crt delete mode 100644 presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/tls_certs/server.key diff --git a/.github/workflows/arrow-flight-tests.yml b/.github/workflows/arrow-flight-tests.yml index f6f39252bed78..22ebf9439fcb2 100644 --- a/.github/workflows/arrow-flight-tests.yml +++ b/.github/workflows/arrow-flight-tests.yml @@ -51,7 +51,7 @@ jobs: # Run Maven tests for the target module, excluding native tests - name: Maven Tests - run: ./mvnw test ${MAVEN_TEST} -pl ${{ matrix.modules }} -Dtest="*,!TestArrowFlightNativeQueries" + run: ./mvnw test ${MAVEN_TEST} -pl ${{ matrix.modules }} -Dtest="*,!TestArrowFlightNativeQueries*" prestocpp-linux-build-for-test: runs-on: ubuntu-22.04 @@ -207,7 +207,7 @@ jobs: mvn test \ ${MAVEN_TEST} \ -pl ${{ matrix.modules }} \ - -Dtest="TestArrowFlightNativeQueries" \ + -Dtest="TestArrowFlightNativeQueries*" \ -DPRESTO_SERVER=${PRESTO_SERVER_PATH} \ -DDATA_DIR=${RUNNER_TEMP} \ -Duser.timezone=America/Bahia_Banderas \ diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/ArrowFlightQueryRunner.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/ArrowFlightQueryRunner.java index eea7ff598da72..66db074b419b8 100644 --- a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/ArrowFlightQueryRunner.java +++ b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/ArrowFlightQueryRunner.java @@ -39,6 +39,8 @@ public class ArrowFlightQueryRunner { + private static final Logger log = Logger.get(ArrowFlightQueryRunner.class); + private ArrowFlightQueryRunner() { throw new UnsupportedOperationException("This is a utility class and cannot be instantiated"); @@ -54,24 +56,15 @@ public static int findUnusedPort() public static DistributedQueryRunner createQueryRunner(int flightServerPort) throws Exception { - return createQueryRunner(flightServerPort, ImmutableMap.of(), ImmutableMap.of(), Optional.empty()); + return createQueryRunner(flightServerPort, ImmutableMap.of(), ImmutableMap.of(), Optional.empty(), Optional.empty()); } public static DistributedQueryRunner createQueryRunner( int flightServerPort, Map extraProperties, Map coordinatorProperties, - Optional> externalWorkerLauncher) - throws Exception - { - return createQueryRunner(extraProperties, ImmutableMap.of("arrow-flight.server.port", String.valueOf(flightServerPort)), coordinatorProperties, externalWorkerLauncher); - } - - private static DistributedQueryRunner createQueryRunner( - Map extraProperties, - Map catalogProperties, - Map coordinatorProperties, - Optional> externalWorkerLauncher) + Optional> externalWorkerLauncher, + Optional mTLSEnabled) throws Exception { Session session = testSessionBuilder() @@ -86,17 +79,24 @@ private static DistributedQueryRunner createQueryRunner( DistributedQueryRunner queryRunner = queryRunnerBuilder .setExtraProperties(extraProperties) .setCoordinatorProperties(coordinatorProperties) - .setExternalWorkerLauncher(externalWorkerLauncher).build(); + .setExternalWorkerLauncher(externalWorkerLauncher) + .build(); try { boolean nativeExecution = externalWorkerLauncher.isPresent(); queryRunner.installPlugin(new TestingArrowFlightPlugin(nativeExecution)); + Map catalogProperties = ImmutableMap.of("arrow-flight.server.port", String.valueOf(flightServerPort)); ImmutableMap.Builder properties = ImmutableMap.builder() .putAll(catalogProperties) .put("arrow-flight.server", "localhost") .put("arrow-flight.server-ssl-enabled", "true") - .put("arrow-flight.server-ssl-certificate", "src/test/resources/server.crt"); + .put("arrow-flight.server-ssl-certificate", "src/test/resources/certs/ca.crt"); + + if (mTLSEnabled.orElse(false)) { + properties.put("arrow-flight.client-ssl-certificate", "src/test/resources/certs/client.crt"); + properties.put("arrow-flight.client-ssl-key", "src/test/resources/certs/client.key"); + } queryRunner.createCatalog(ARROW_FLIGHT_CATALOG, ARROW_FLIGHT_CONNECTOR, properties.build()); @@ -125,25 +125,33 @@ public static void main(String[] args) { Logging.initialize(); + boolean mTLSenabled = Boolean.parseBoolean(System.getProperty("flight.mtls.enabled", "false")); + RootAllocator allocator = new RootAllocator(Long.MAX_VALUE); Location serverLocation = Location.forGrpcTls("localhost", 9443); - File certChainFile = new File("src/test/resources/server.crt"); - File privateKeyFile = new File("src/test/resources/server.key"); - FlightServer server = FlightServer.builder(allocator, serverLocation, new TestingArrowProducer(allocator)) - .useTls(certChainFile, privateKeyFile) - .build(); + FlightServer.Builder serverBuilder = FlightServer.builder(allocator, serverLocation, new TestingArrowProducer(allocator)); - server.start(); + File serverCert = new File("src/test/resources/certs/server.crt"); + File serverKey = new File("src/test/resources/certs/server.key"); + serverBuilder.useTls(serverCert, serverKey); - Logger log = Logger.get(ArrowFlightQueryRunner.class); + if (mTLSenabled) { + File caCert = new File("src/test/resources/certs/ca.crt"); + serverBuilder.useMTlsClientVerification(caCert); + } + + FlightServer server = serverBuilder.build(); + server.start(); log.info("Server listening on port " + server.getPort()); DistributedQueryRunner queryRunner = createQueryRunner( + server.getPort(), ImmutableMap.of("http-server.http.port", "8080"), - ImmutableMap.of("arrow-flight.server.port", String.valueOf(9443)), ImmutableMap.of(), - Optional.empty()); + Optional.empty(), + Optional.of(mTLSenabled)); + Thread.sleep(10); log.info("======== SERVER STARTED ========"); log.info("\n====\n%s\n====", queryRunner.getCoordinator().getBaseUrl()); diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightEchoQueries.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightEchoQueries.java index 8715d413f687f..f658fb6f805c1 100644 --- a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightEchoQueries.java +++ b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightEchoQueries.java @@ -131,8 +131,8 @@ public void setup() throws Exception { arrowFlightQueryRunner = getDistributedQueryRunner(); - File certChainFile = new File("src/test/resources/server.crt"); - File privateKeyFile = new File("src/test/resources/server.key"); + File certChainFile = new File("src/test/resources/certs/server.crt"); + File privateKeyFile = new File("src/test/resources/certs/server.key"); allocator = new RootAllocator(Long.MAX_VALUE); @@ -407,7 +407,7 @@ private static MapType createMapType(Type keyType, Type valueType) private static FlightClient createFlightClient(BufferAllocator allocator, int serverPort) throws IOException { - InputStream trustedCertificate = new ByteArrayInputStream(Files.readAllBytes(Paths.get("src/test/resources/server.crt"))); + InputStream trustedCertificate = new ByteArrayInputStream(Files.readAllBytes(Paths.get("src/test/resources/certs/server.crt"))); Location location = Location.forGrpcTls("localhost", serverPort); return FlightClient.builder(allocator, location).useTls().trustedCertificates(trustedCertificate).build(); } diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightIntegrationSmokeTest.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightIntegrationSmokeTest.java index 1bca76c23a9f0..916dff5e964e2 100644 --- a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightIntegrationSmokeTest.java +++ b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightIntegrationSmokeTest.java @@ -48,8 +48,8 @@ public void setup() throws Exception { arrowFlightQueryRunner = getDistributedQueryRunner(); - File certChainFile = new File("src/test/resources/server.crt"); - File privateKeyFile = new File("src/test/resources/server.key"); + File certChainFile = new File("src/test/resources/certs/server.crt"); + File privateKeyFile = new File("src/test/resources/certs/server.key"); allocator = new RootAllocator(Long.MAX_VALUE); Location location = Location.forGrpcTls("127.0.0.1", serverPort); diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightMtls.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightMtls.java index 169f08db4ddb6..f0f3b8aa76ed0 100644 --- a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightMtls.java +++ b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightMtls.java @@ -60,9 +60,9 @@ private void setup() arrowFlightQueryRunner.createCatalog(ARROW_FLIGHT_CATALOG_WITH_NO_MTLS_CERTS, ARROW_FLIGHT_CONNECTOR, getNoMtlsCatalogProperties()); arrowFlightQueryRunner.createCatalog(ARROW_FLIGHT_CATALOG_WITH_MTLS_CERTS, ARROW_FLIGHT_CONNECTOR, getMtlsCatalogProperties()); - File certChainFile = new File("src/test/resources/mtls/server.crt"); - File privateKeyFile = new File("src/test/resources/mtls/server.key"); - File caCertFile = new File("src/test/resources/mtls/ca.crt"); + File certChainFile = new File("src/test/resources/certs/server.crt"); + File privateKeyFile = new File("src/test/resources/certs/server.key"); + File caCertFile = new File("src/test/resources/certs/ca.crt"); allocator = new RootAllocator(Long.MAX_VALUE); @@ -89,7 +89,7 @@ private void tearDown() protected QueryRunner createQueryRunner() throws Exception { - return ArrowFlightQueryRunner.createQueryRunner(serverPort, ImmutableMap.of(), ImmutableMap.of(), Optional.empty()); + return ArrowFlightQueryRunner.createQueryRunner(serverPort, ImmutableMap.of(), ImmutableMap.of(), Optional.empty(), Optional.empty()); } private Map getInvalidCertCatalogProperties() @@ -98,9 +98,9 @@ private Map getInvalidCertCatalogProperties() .put("arrow-flight.server.port", String.valueOf(serverPort)) .put("arrow-flight.server", "localhost") .put("arrow-flight.server-ssl-enabled", "true") - .put("arrow-flight.server-ssl-certificate", "src/test/resources/mtls/server.crt") - .put("arrow-flight.client-ssl-certificate", "src/test/resources/mtls/invalid_cert.crt") - .put("arrow-flight.client-ssl-key", "src/test/resources/mtls/client.key"); + .put("arrow-flight.server-ssl-certificate", "src/test/resources/certs/server.crt") + .put("arrow-flight.client-ssl-certificate", "src/test/resources/certs/invalid_cert.crt") + .put("arrow-flight.client-ssl-key", "src/test/resources/certs/client.key"); return catalogProperties.build(); } @@ -110,7 +110,7 @@ private Map getNoMtlsCatalogProperties() .put("arrow-flight.server.port", String.valueOf(serverPort)) .put("arrow-flight.server", "localhost") .put("arrow-flight.server-ssl-enabled", "true") - .put("arrow-flight.server-ssl-certificate", "src/test/resources/mtls/server.crt"); + .put("arrow-flight.server-ssl-certificate", "src/test/resources/certs/server.crt"); return catalogProperties.build(); } @@ -120,9 +120,9 @@ private Map getMtlsCatalogProperties() .put("arrow-flight.server.port", String.valueOf(serverPort)) .put("arrow-flight.server", "localhost") .put("arrow-flight.server-ssl-enabled", "true") - .put("arrow-flight.server-ssl-certificate", "src/test/resources/mtls/server.crt") - .put("arrow-flight.client-ssl-certificate", "src/test/resources/mtls/client.crt") - .put("arrow-flight.client-ssl-key", "src/test/resources/mtls/client.key"); + .put("arrow-flight.server-ssl-certificate", "src/test/resources/certs/server.crt") + .put("arrow-flight.client-ssl-certificate", "src/test/resources/certs/client.crt") + .put("arrow-flight.client-ssl-key", "src/test/resources/certs/client.key"); return catalogProperties.build(); } diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightNativeQueries.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightNativeQueries.java index 21ee7f53ed78a..d6f945c45853c 100644 --- a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightNativeQueries.java +++ b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightNativeQueries.java @@ -60,23 +60,32 @@ public TestArrowFlightNativeQueries() this.serverPort = ArrowFlightQueryRunner.findUnusedPort(); } + protected boolean ismTLSEnabled() + { + return false; + } + @BeforeClass public void setup() throws Exception { arrowFlightQueryRunner = getDistributedQueryRunner(); - allocator = new RootAllocator(Long.MAX_VALUE); Location location = Location.forGrpcTls("localhost", serverPort); - File certChainFile = new File("src/test/resources/server.crt"); - File privateKeyFile = new File("src/test/resources/server.key"); + FlightServer.Builder serverBuilder = FlightServer.builder(allocator, location, new TestingArrowProducer(allocator)); - server = FlightServer.builder(allocator, location, new TestingArrowProducer(allocator)) - .useTls(certChainFile, privateKeyFile) - .build(); + File serverCert = new File("src/test/resources/certs/server.crt"); + File serverKey = new File("src/test/resources/certs/server.key"); + serverBuilder.useTls(serverCert, serverKey); + + if (ismTLSEnabled()) { + File caCert = new File("src/test/resources/certs/ca.crt"); + serverBuilder.useMTlsClientVerification(caCert); + } + server = serverBuilder.build(); server.start(); - log.info("Server listening on port %s", server.getPort()); + log.info("Server listening on port %s (%s)", server.getPort(), ismTLSEnabled() ? "mTLS" : "TLS"); } @AfterClass(alwaysRun = true) @@ -99,9 +108,13 @@ protected QueryRunner createQueryRunner() log.info("Using PRESTO_SERVER binary at %s", prestoServerPath); ImmutableMap coordinatorProperties = ImmutableMap.of("native-execution-enabled", "true"); - String flightCertPath = Paths.get("src/test/resources/server.crt").toAbsolutePath().toString(); - return ArrowFlightQueryRunner.createQueryRunner(serverPort, getNativeWorkerSystemProperties(), coordinatorProperties, getExternalWorkerLauncher(prestoServerPath.toString(), serverPort, flightCertPath)); + return ArrowFlightQueryRunner.createQueryRunner( + serverPort, + getNativeWorkerSystemProperties(), + coordinatorProperties, + getExternalWorkerLauncher(prestoServerPath.toString(), serverPort, ismTLSEnabled()), + Optional.of(ismTLSEnabled())); } @Override @@ -338,50 +351,63 @@ public static Map getNativeWorkerSystemProperties() .build(); } - public static Optional> getExternalWorkerLauncher(String prestoServerPath, int flightServerPort, String flightCertPath) + public static Optional> getExternalWorkerLauncher(String prestoServerPath, int flightServerPort, boolean ismTLSEnabled) { - return - Optional.of((workerIndex, discoveryUri) -> { - try { - Path dir = Paths.get("/tmp", TestArrowFlightNativeQueries.class.getSimpleName()); - Files.createDirectories(dir); - Path tempDirectoryPath = Files.createTempDirectory(dir, "worker"); - log.info("Temp directory for Worker #%d: %s", workerIndex, tempDirectoryPath.toString()); - - // Write config file - use an ephemeral port for the worker. - String configProperties = format("discovery.uri=%s%n" + - "presto.version=testversion%n" + - "system-memory-gb=4%n" + - "http-server.http.port=0%n", discoveryUri); - - Files.write(tempDirectoryPath.resolve("config.properties"), configProperties.getBytes()); - Files.write(tempDirectoryPath.resolve("node.properties"), - format("node.id=%s%n" + - "node.internal-address=127.0.0.1%n" + - "node.environment=testing%n" + - "node.location=test-location", UUID.randomUUID()).getBytes()); - - Path catalogDirectoryPath = tempDirectoryPath.resolve("catalog"); - Files.createDirectory(catalogDirectoryPath); - - Files.write(catalogDirectoryPath.resolve(format("%s.properties", ARROW_FLIGHT_CATALOG)), - format("connector.name=%s\n" + - "arrow-flight.server=localhost\n" + - "arrow-flight.server.port=%d\n" + - "arrow-flight.server-ssl-enabled=true\n" + - "arrow-flight.server-ssl-certificate=%s", ARROW_FLIGHT_CONNECTOR, flightServerPort, flightCertPath).getBytes()); - - // Disable stack trace capturing as some queries (using TRY) generate a lot of exceptions. - return new ProcessBuilder(prestoServerPath, "--logtostderr=1", "--v=1") - .directory(tempDirectoryPath.toFile()) - .redirectErrorStream(true) - .redirectOutput(ProcessBuilder.Redirect.to(tempDirectoryPath.resolve("worker." + workerIndex + ".out").toFile())) - .redirectError(ProcessBuilder.Redirect.to(tempDirectoryPath.resolve("worker." + workerIndex + ".out").toFile())) - .start(); - } - catch (IOException e) { - throw new UncheckedIOException(e); - } - }); + return Optional.of((workerIndex, discoveryUri) -> { + try { + Path dir = Paths.get("/tmp", TestArrowFlightNativeQueries.class.getSimpleName()); + Files.createDirectories(dir); + Path tempDirectoryPath = Files.createTempDirectory(dir, "worker"); + log.info("Temp directory for Worker #%d: %s", workerIndex, tempDirectoryPath.toString()); + + // Write config file - use an ephemeral port for the worker. + String configProperties = format("discovery.uri=%s%n" + + "presto.version=testversion%n" + + "system-memory-gb=4%n" + + "http-server.http.port=0%n", discoveryUri); + + Files.write(tempDirectoryPath.resolve("config.properties"), configProperties.getBytes()); + Files.write(tempDirectoryPath.resolve("node.properties"), + format("node.id=%s%n" + + "node.internal-address=127.0.0.1%n" + + "node.environment=testing%n" + + "node.location=test-location", UUID.randomUUID()).getBytes()); + + Path catalogDirectoryPath = tempDirectoryPath.resolve("catalog"); + Files.createDirectory(catalogDirectoryPath); + + String caCertPath = Paths.get("src/test/resources/certs/ca.crt").toAbsolutePath().toString(); + + StringBuilder catalogBuilder = new StringBuilder(); + catalogBuilder.append(format( + "connector.name=%s\n" + + "arrow-flight.server=localhost\n" + + "arrow-flight.server.port=%d\n" + + "arrow-flight.server-ssl-enabled=true\n" + + "arrow-flight.server-ssl-certificate=%s\n", + ARROW_FLIGHT_CONNECTOR, flightServerPort, caCertPath)); + + if (ismTLSEnabled) { + String clientCertPath = Paths.get("src/test/resources/certs/client.crt").toAbsolutePath().toString(); + String clientKeyPath = Paths.get("src/test/resources/certs/client.key").toAbsolutePath().toString(); + catalogBuilder.append(format("arrow-flight.client-ssl-certificate=%s\n", clientCertPath)); + catalogBuilder.append(format("arrow-flight.client-ssl-key=%s\n", clientKeyPath)); + } + + Files.write( + catalogDirectoryPath.resolve(format("%s.properties", ARROW_FLIGHT_CATALOG)), + catalogBuilder.toString().getBytes()); + + return new ProcessBuilder(prestoServerPath, "--logtostderr=1", "--v=1") + .directory(tempDirectoryPath.toFile()) + .redirectErrorStream(true) + .redirectOutput(ProcessBuilder.Redirect.to(tempDirectoryPath.resolve("worker." + workerIndex + ".out").toFile())) + .redirectError(ProcessBuilder.Redirect.to(tempDirectoryPath.resolve("worker." + workerIndex + ".out").toFile())) + .start(); + } + catch (IOException e) { + throw new UncheckedIOException(e); + } + }); } } diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightNativeQueriesMtls.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightNativeQueriesMtls.java new file mode 100644 index 0000000000000..25079773e7c02 --- /dev/null +++ b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightNativeQueriesMtls.java @@ -0,0 +1,33 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.facebook.plugin.arrow; + +import java.io.IOException; + +public class TestArrowFlightNativeQueriesMtls + extends TestArrowFlightNativeQueries +{ + public TestArrowFlightNativeQueriesMtls() + throws IOException + { + super(); + } + + @Override + protected boolean ismTLSEnabled() + { + return true; + } +} diff --git a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightQueries.java b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightQueries.java index 9be967eab2c08..cc1edc0735ce0 100644 --- a/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightQueries.java +++ b/presto-base-arrow-flight/src/test/java/com/facebook/plugin/arrow/TestArrowFlightQueries.java @@ -70,8 +70,8 @@ public void setup() throws Exception { arrowFlightQueryRunner = getDistributedQueryRunner(); - File certChainFile = new File("src/test/resources/server.crt"); - File privateKeyFile = new File("src/test/resources/server.key"); + File certChainFile = new File("src/test/resources/certs/server.crt"); + File privateKeyFile = new File("src/test/resources/certs/server.key"); allocator = new RootAllocator(Long.MAX_VALUE); Location location = Location.forGrpcTls("localhost", serverPort); diff --git a/presto-base-arrow-flight/src/test/resources/mtls/ca.crt b/presto-base-arrow-flight/src/test/resources/certs/ca.crt similarity index 100% rename from presto-base-arrow-flight/src/test/resources/mtls/ca.crt rename to presto-base-arrow-flight/src/test/resources/certs/ca.crt diff --git a/presto-base-arrow-flight/src/test/resources/mtls/client.crt b/presto-base-arrow-flight/src/test/resources/certs/client.crt similarity index 100% rename from presto-base-arrow-flight/src/test/resources/mtls/client.crt rename to presto-base-arrow-flight/src/test/resources/certs/client.crt diff --git a/presto-base-arrow-flight/src/test/resources/mtls/client.key b/presto-base-arrow-flight/src/test/resources/certs/client.key similarity index 100% rename from presto-base-arrow-flight/src/test/resources/mtls/client.key rename to presto-base-arrow-flight/src/test/resources/certs/client.key diff --git a/presto-base-arrow-flight/src/test/resources/mtls/invalid_cert.crt b/presto-base-arrow-flight/src/test/resources/certs/invalid_cert.crt similarity index 100% rename from presto-base-arrow-flight/src/test/resources/mtls/invalid_cert.crt rename to presto-base-arrow-flight/src/test/resources/certs/invalid_cert.crt diff --git a/presto-base-arrow-flight/src/test/resources/mtls/server.crt b/presto-base-arrow-flight/src/test/resources/certs/server.crt similarity index 100% rename from presto-base-arrow-flight/src/test/resources/mtls/server.crt rename to presto-base-arrow-flight/src/test/resources/certs/server.crt diff --git a/presto-base-arrow-flight/src/test/resources/mtls/server.key b/presto-base-arrow-flight/src/test/resources/certs/server.key similarity index 100% rename from presto-base-arrow-flight/src/test/resources/mtls/server.key rename to presto-base-arrow-flight/src/test/resources/certs/server.key diff --git a/presto-base-arrow-flight/src/test/resources/server.crt b/presto-base-arrow-flight/src/test/resources/server.crt deleted file mode 100644 index f070253c9e119..0000000000000 --- a/presto-base-arrow-flight/src/test/resources/server.crt +++ /dev/null @@ -1,32 +0,0 @@ ------BEGIN CERTIFICATE----- -MIIFlTCCA32gAwIBAgIUf2p3qdxzpsOofYDpXNXDJA1fjbgwDQYJKoZIhvcNAQEL -BQAwWTELMAkGA1UEBhMCQVUxEzARBgNVBAgMClNvbWUtU3RhdGUxITAfBgNVBAoM -GEludGVybmV0IFdpZGdpdHMgUHR5IEx0ZDESMBAGA1UEAwwJbG9jYWxob3N0MCAX -DTI0MTIyMTA3MzUwMFoYDzIxMjQxMTI3MDczNTAwWjBZMQswCQYDVQQGEwJBVTET -MBEGA1UECAwKU29tZS1TdGF0ZTEhMB8GA1UECgwYSW50ZXJuZXQgV2lkZ2l0cyBQ -dHkgTHRkMRIwEAYDVQQDDAlsb2NhbGhvc3QwggIiMA0GCSqGSIb3DQEBAQUAA4IC -DwAwggIKAoICAQCmsVs+OHGQ/GlR9+KpgcKkBvXZQAUtCJzFG27GXBYoltJDAxsi -o47KbGkUWxf5E+4oFbbz7s8OHSuQnjP2XVHbKvmgicf3uP98fKkQtjvl+xn8vT18 -s6a8Kv8hj+f6MRWGwpHa/sQKU9uZmRYTRh+32vDZRtAvmpzkf6+2K8B1fFbwVmuC -j3ULb+0iYHnomC3aMWBFXkxjEmsamx4YK74NtQU98+EjQZwhWgWXhW5shS1kSs4r -3N6++tonBz+tDKAhCMueRRJAQXGjKqL7qDZn7wpk53L/fZT9mgRYyA+PN2ND9L0H -nMGMjJl71p42kkgIGOpllsmK6+g0Bj6aC/uCEnX2AtM0g57Th7U2aLwgOmRshC0s -uuBHxMWUgzJsB1dscXrFPPB+XVcUlgcwRbGsQG5VzK/rYRV4y1FmF5vSLY60mF43 -hFDAcnh4mnnBkobca5Dl6PSUpFmDkF56IUgYQCTrE6hPYnrIJKhPcfpmr/tm/Acr -ra1sPp/QPSFIxI9j7Nzm/QsOBF3Zy4AbbbOmhJOjNtwEi59r9za/FhxV5kSN/YM9 -HyyYYebxW/jXF/7hzQzWYfJBz1SgdD9prl4ml8VMVJdhZmBTzhVciPXizRUKkbD+ -LvQKw8q4a24/VruUvW15J39qalhdyWf3vqGuORWowpG7oYnGXik3kYVT5QIDAQAB -o1MwUTAdBgNVHQ4EFgQUkkKtG5568IDoFtn3AK0Yes0CqGgwHwYDVR0jBBgwFoAU -kkKtG5568IDoFtn3AK0Yes0CqGgwDwYDVR0TAQH/BAUwAwEB/zANBgkqhkiG9w0B -AQsFAAOCAgEAVjwxyO5N4EGQHf9euKiTjksUMeJrJkPMFYS/1UmQB5/yc0TkOxv3 -M8yvTzkzxWAQV2X+CPbOLgqO2eCoSev2GHVUyHLLUEFATzAgTpF79Ovzq6JkSgFo -GFxTL0TuA+d++430wYyts7iYY9uT5Xt8xp4+SwXrCvLWxHosP3xGUmNvY4sP0hdS -beIvdGE/J8izgy5DRt2fWZ03mmwfKqiz/qhKGj9DDsHkA/1jyKIivP/nufr9dzDr -41lhk1N7qFWkOjbMd06NYySIe0MaapIkenjT1IOgqGw2f98RfSEomoaXuxN0VoSM -6dZ4rN97cER25X7/zE0zCZurjCLHzPTyuTspYGEK+9U4plOeWh0keQEcdpcHTAr+ -NqU3VlhXVxz91nVREpRJmKk+r4c+xdrfY3YkDcSJ1dazbd3eS05Ggx51KOqes8Zb -hFQfhIDqvaqXDNlkBezLpr4v/MU69+cp7SOn5uPnOccS6sd7fzl/PUZhEjd5ZIws -8SX79OwhQjbYZYRHJSPPasb8B1amULtoo0pJ5izSkXxileiGuXhRO5stiBux7SL+ -oJAztjuRf0IvP6LWOMHgqquzc2JiEDCz0DPnTCqoXGZlT2HPGzXOoDSTOmRdFx+L -qi/DeY+MpIMVov/rplqjydXw6AuQDxcV1GvyjMvaHxJG5MEBC/mVeqQ= ------END CERTIFICATE----- \ No newline at end of file diff --git a/presto-base-arrow-flight/src/test/resources/server.key b/presto-base-arrow-flight/src/test/resources/server.key deleted file mode 100644 index 4524c26943f2d..0000000000000 --- a/presto-base-arrow-flight/src/test/resources/server.key +++ /dev/null @@ -1,52 +0,0 @@ ------BEGIN PRIVATE KEY----- -MIIJQwIBADANBgkqhkiG9w0BAQEFAASCCS0wggkpAgEAAoICAQCmsVs+OHGQ/GlR -9+KpgcKkBvXZQAUtCJzFG27GXBYoltJDAxsio47KbGkUWxf5E+4oFbbz7s8OHSuQ -njP2XVHbKvmgicf3uP98fKkQtjvl+xn8vT18s6a8Kv8hj+f6MRWGwpHa/sQKU9uZ -mRYTRh+32vDZRtAvmpzkf6+2K8B1fFbwVmuCj3ULb+0iYHnomC3aMWBFXkxjEmsa -mx4YK74NtQU98+EjQZwhWgWXhW5shS1kSs4r3N6++tonBz+tDKAhCMueRRJAQXGj -KqL7qDZn7wpk53L/fZT9mgRYyA+PN2ND9L0HnMGMjJl71p42kkgIGOpllsmK6+g0 -Bj6aC/uCEnX2AtM0g57Th7U2aLwgOmRshC0suuBHxMWUgzJsB1dscXrFPPB+XVcU -lgcwRbGsQG5VzK/rYRV4y1FmF5vSLY60mF43hFDAcnh4mnnBkobca5Dl6PSUpFmD -kF56IUgYQCTrE6hPYnrIJKhPcfpmr/tm/Acrra1sPp/QPSFIxI9j7Nzm/QsOBF3Z -y4AbbbOmhJOjNtwEi59r9za/FhxV5kSN/YM9HyyYYebxW/jXF/7hzQzWYfJBz1Sg -dD9prl4ml8VMVJdhZmBTzhVciPXizRUKkbD+LvQKw8q4a24/VruUvW15J39qalhd -yWf3vqGuORWowpG7oYnGXik3kYVT5QIDAQABAoICAE0MeIzTgSbPjQz+w82u9V1k -/DlNfrb4nqH7EqJsSS+8uvaPlnjV2fgV0SJAEt4mCLSNiPHKpfkzoYHopkMPknj4 -Lcc3OG94GtubMXhQi3I7tSDeBfBAh+a9Bw2n20WJb5ZJFCsCDHJrnXsrSAljpeCR -OjdsJGmEkVWK8ZiGM6D6dqMDhxEjpynAs/7qUh8hTDxpC0M1GaDHkCsNnQT2HxRt -4jznH97wgi7mUeReIBLYIgmUDCU5I9ppz/EvSA8AYXmze46uBYge19xgJlKlR3SW -CJtoYf7XOMlZ6f1xh8OeiesM0l0U51/EU2Nq6dl2lwXrIlkPsBve/AckBcalmDwh -tJgAOHM8VEMp0imJ8+SPGZupiIt+sTVSt4aq2uB5dtTBDcXGx4gVUYJoekI4yQtv -OWKAQAobq0Cutogx8hyXr2tSizp5pGpJDKx8/G/3YcNmRqCq+HQgwshlHQy/hPqs -QC/jcuOkr5GOZASF3wLCuLmAwFCELL96iVEJvPLe5Yb1u4HbCrsy9/N8rHV1s1qa -xMcGiIS4m/fwLbD+TjgM1foESeQciHH7GWVZ3Osn6Inpp6vLJKDhLhGBgpaYF72S -0fft6CkIy/CqbwXPe/3Wm0PL0SLTsSk0lKhw6Zeuvg0Z0jUrmDmlTrGo8j4zU7cA -Gc8uiZuFAU2ZoS9R7SZXAoIBAQDSfF8H4C8jn9IvTw/S3APkP7VyLKFSTitEfagt -qcAZr7sYWYGdYDL1hX+jBopAWdzbJ+N3MZBccH67f3XDZWVH4GDQWF6tyOpyseBe -42aU/yCr6ZSoxi91W9V+oyiTFh+t2cEkNyGey/pbXVBw45jncHpzIAaERm5gD8DD -MxOadIvDxEGquGs9MUgZg6ABFmqvQIfiZ7I3PvfoHJGrHsPyM9i6AqAkJe1rKiN3 -tt2wiDk6zDor4vFKn+5LaRQ5INbpgBuCJrvDsRh3MeKLb0ddCYcAStHeAjHJZbH4 -OiC5y9hpaIcSQN3Z4gEhdMd+Lhi1+z36uN5k7KCvbGmO3sG/AoIBAQDKvMiVqTgm -a36FFswwZJd5o+Hx6G8Nt5y7ZYLJgoq4wyTPYA5ZIHBp2Jtb//IGJMKpeLcRigxx -nKu4zrLZPHesgYzHbCH9i5fOBFl8yF9pZnbCU/zGOYhiVdn7MHZgJw8nlstWr1Jr -cMnxBBXEjgL65Lrse1noiPATNrvBsirFaJky6HhxRPbkijKKd0Xqc0FeHMxt6IHU -y6yZRzguI0M1A8RGR4CrqsOGrdv0vkMSiumkUdz8JW4w9R69n6ax/qNAZSMql0ss -PIPOqGtYRPhibOhXKPl0p30X32YXx/SKm8+L9Sr1ny76dab/bSnxStOdihGJLCs4 -l7vFkuJgTMtbAoIBAQCYA3CSfIsu3EbtGdlgvLsmxggh7C+aBJBlB6dFSzpMksi5 -njLo2MgU35Q9xgRk00GZGWbC94295RTyDuya8IjD7z2cWqYONnNz4BkeDndQli0f -WzOc7Hzr8iXvLqCoEatRYFmH8TUbvU8TWwI0dXtBcs9Mg82RDFi8kcPyddnri85A -1WVjiYsRh5z9qD0PbAQii6VXkvJ3ycc64B8oCbEUI/Oa6ziCws2DvswcsnnK+6bx -WvuMJHuFHJn55mrPk3MC8h1r0tN6UlVMCEAH2ZcdjzrrsB1/i/Au9n4gusJVzO1/ -uxkJysUujXWplvBYpav9CfVKNOeQ1gB6kP5vS1t7AoIBAB1DquCPkJ9bHOQxKkBC -BOt2EINOvdkJDAKw4HQd99A7uvCEOQ38dL2Smrpo85KXc9HqruJFPw6XQuJmU8Kv -y8aG3L9ciHuEzuDaF+C/O6aHN9VNMkuaumkXY2Oy1yOB/9oDFk7o98iyezPjFxFM -Pnng0mqYU54RRjY/zFJlWW8tbg+/JsOS5OCQYkNCfEEfaewf1BJ5YWRKEhv9/8oJ -JQZeCNLsN1KQT7D9H6bwX9YpXxhtCK0M6h7/AvT0OqeuzfnZn33iYON9yLjn7rbL -Hd93QQJz065XDuOHR8FfB5mKbCcTuKPD2pAks3pjU46U8n7nEyjtyz9cB6q5TRwB -eckCggEBAMd/3riFnoc2VF6LsA0VQzZiJleN0G9g3AF/5acqHHjMv6aJs/AjV7ZF -hFdiqTcQFEg87HSOzrtV2DVNWJc7DMcMIuwGhVec+D4wedJ2wrQv+MEgySltSBZq -wcPVn5IQiml38NnG/tPIrHETb0PIoa8+iu/Jg80o7j8M3+DKVKfCyfh334PjFK6W -B/mkgC9PfcfeA/Doby9pJsRnqmAJTeWjxbefksckI4PcRCLMEwggB2ReiIWef8Q+ -IooNIxypWtBWtpNEl5lhO7Y2f65Whp34TjooXmGBl1lj4szO8PNnf9QA865J86OS -kOoFda4Sn7LajkbSX0wTGMuXDpmx34M= ------END PRIVATE KEY----- \ No newline at end of file diff --git a/presto-docs/src/main/sphinx/connector/base-arrow-flight.rst b/presto-docs/src/main/sphinx/connector/base-arrow-flight.rst index e1676d3e05249..a68b4308a9d73 100644 --- a/presto-docs/src/main/sphinx/connector/base-arrow-flight.rst +++ b/presto-docs/src/main/sphinx/connector/base-arrow-flight.rst @@ -59,6 +59,36 @@ Property Name Description ``arrow-flight.server-ssl-enabled`` Port is ssl enabled ========================================== ============================================================== +Mutual TLS (mTLS) Support +------------------------- + + +To connect the Presto client to an Arrow Flight server with mutual TLS (mTLS) enabled, you must configure the client to present a valid certificate and key that the server can validate. This enhances security by ensuring both the client and server authenticate each other. + +To enable mTLS, the following properties must be configured: + +- ``arrow-flight.server-ssl-enabled=true``: Explicitly enables TLS for the connection. +- ``arrow-flight.server-ssl-certificate``: Path to the server's SSL certificate. +- ``arrow-flight.client-ssl-certificate``: Path to the client's SSL certificate. +- ``arrow-flight.client-ssl-key``: Path to the client's SSL private key. + +These properties must be used alongside the existing SSL configurations for the server, such as ``arrow-flight.server-ssl-certificate`` and ``arrow-flight.server-ssl-enabled=true``. Make sure the server is configured to trust the client certificates (typically via a shared CA). + +Below is an example code snippet to configure the Arrow Flight server with mTLS: + +.. code-block:: java + + File certChainFile = new File("src/test/resources/certs/server.crt"); + File privateKeyFile = new File("src/test/resources/certs/server.key"); + File caCertFile = new File("src/test/resources/certs/ca.crt"); + + server = FlightServer.builder(allocator, location, new TestingArrowProducer(allocator)) + .useTls(certChainFile, privateKeyFile) + .useMTlsClientVerification(caCertFile) + .build(); + + server.start(); + Querying Arrow-Flight --------------------- diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/ArrowFlightConfig.cpp b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/ArrowFlightConfig.cpp index c0900c6c41de4..18d637df02db8 100644 --- a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/ArrowFlightConfig.cpp +++ b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/ArrowFlightConfig.cpp @@ -43,4 +43,14 @@ std::optional ArrowFlightConfig::serverSslCertificate() const { config_->get(kServerSslCertificate)); } +std::optional ArrowFlightConfig::clientSslCertificate() const { + return static_cast>( + config_->get(kClientSslCertificate)); +} + +std::optional ArrowFlightConfig::clientSslKey() const { + return static_cast>( + config_->get(kClientSslKey)); +} + } // namespace facebook::presto diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/ArrowFlightConfig.h b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/ArrowFlightConfig.h index 77ad8e9379cf3..4839733baaab0 100644 --- a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/ArrowFlightConfig.h +++ b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/ArrowFlightConfig.h @@ -38,6 +38,12 @@ class ArrowFlightConfig { static constexpr const char* kServerSslCertificate = "arrow-flight.server-ssl-certificate"; + static constexpr const char* kClientSslCertificate = + "arrow-flight.client-ssl-certificate"; + + static constexpr const char* kClientSslKey = + "arrow-flight.client-ssl-key"; + std::string authenticatorName() const; std::optional defaultServerHostname() const; @@ -50,6 +56,10 @@ class ArrowFlightConfig { std::optional serverSslCertificate() const; + std::optional clientSslCertificate() const; + + std::optional clientSslKey() const; + private: const std::shared_ptr config_; }; diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/ArrowFlightConnector.cpp b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/ArrowFlightConnector.cpp index 28e3e5b5f074c..5c02ac051e076 100644 --- a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/ArrowFlightConnector.cpp +++ b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/ArrowFlightConnector.cpp @@ -71,6 +71,26 @@ ArrowFlightConnector::initClientOpts( clientOpts->tls_root_certs = cert; } + auto clientCertPath = config->clientSslCertificate(); + if (clientCertPath.has_value()) { + std::ifstream certFile(clientCertPath.value()); + VELOX_CHECK( + certFile.is_open(), "Could not open client certificate at {}", clientCertPath.value()); + clientOpts->cert_chain.assign( + (std::istreambuf_iterator(certFile)), + (std::istreambuf_iterator())); + } + + auto clientKeyPath = config->clientSslKey(); + if (clientKeyPath.has_value()) { + std::ifstream keyFile(clientKeyPath.value()); + VELOX_CHECK( + keyFile.is_open(), "Could not open client key at {}", clientKeyPath.value()); + clientOpts->private_key.assign( + (std::istreambuf_iterator(keyFile)), + (std::istreambuf_iterator())); + } + return clientOpts; } diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/ArrowFlightConfigTest.cpp b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/ArrowFlightConfigTest.cpp index eb946f1fcae76..8ec912b417650 100644 --- a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/ArrowFlightConfigTest.cpp +++ b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/ArrowFlightConfigTest.cpp @@ -27,6 +27,8 @@ TEST(ArrowFlightConfigTest, defaultConfig) { ASSERT_EQ(config.defaultServerSslEnabled(), false); ASSERT_EQ(config.serverVerify(), true); ASSERT_EQ(config.serverSslCertificate(), std::nullopt); + ASSERT_EQ(config.clientSslCertificate(), std::nullopt); + ASSERT_EQ(config.clientSslKey(), std::nullopt); } TEST(ArrowFlightConfigTest, overrideConfig) { @@ -36,7 +38,9 @@ TEST(ArrowFlightConfigTest, overrideConfig) { {ArrowFlightConfig::kDefaultServerPort, "9000"}, {ArrowFlightConfig::kDefaultServerSslEnabled, "true"}, {ArrowFlightConfig::kServerVerify, "false"}, - {ArrowFlightConfig::kServerSslCertificate, "my-cert.crt"}}; + {ArrowFlightConfig::kServerSslCertificate, "my-cert.crt"}, + {ArrowFlightConfig::kClientSslCertificate, "/path/to/client.crt"}, + {ArrowFlightConfig::kClientSslKey, "/path/to/client.key"}}; auto config = ArrowFlightConfig( std::make_shared(std::move(configMap))); ASSERT_EQ(config.authenticatorName(), "my-authenticator"); @@ -45,4 +49,6 @@ TEST(ArrowFlightConfigTest, overrideConfig) { ASSERT_EQ(config.defaultServerSslEnabled(), true); ASSERT_EQ(config.serverVerify(), false); ASSERT_EQ(config.serverSslCertificate(), "my-cert.crt"); + ASSERT_EQ(config.clientSslCertificate(), "/path/to/client.crt"); + ASSERT_EQ(config.clientSslKey(), "/path/to/client.key"); } diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/ArrowFlightConnectorMTlsTest.cpp b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/ArrowFlightConnectorMTlsTest.cpp new file mode 100644 index 0000000000000..8225172c8b526 --- /dev/null +++ b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/ArrowFlightConnectorMTlsTest.cpp @@ -0,0 +1,158 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include +#include +#include "presto_cpp/main/connectors/arrow_flight/ArrowFlightConnector.h" +#include "presto_cpp/main/connectors/arrow_flight/tests/utils/ArrowFlightConnectorTestBase.h" +#include "presto_cpp/main/connectors/arrow_flight/tests/utils/ArrowFlightPlanBuilder.h" +#include "presto_cpp/main/connectors/arrow_flight/tests/utils/Utils.h" +#include "velox/common/base/tests/GTestUtils.h" +#include "velox/common/config/Config.h" +#include "velox/exec/tests/utils/AssertQueryBuilder.h" + +using namespace arrow; +using namespace facebook::velox; +using namespace facebook::velox::exec::test; + +namespace facebook::presto::test { + +class ArrowFlightConnectorMtlsTestBase : public ArrowFlightConnectorTestBase { + protected: + explicit ArrowFlightConnectorMtlsTestBase( + std::shared_ptr config) + : ArrowFlightConnectorTestBase(std::move(config)) {} + + void setFlightServerOptions( + flight::FlightServerOptions* serverOptions) override { + flight::CertKeyPair tlsCertificate{ + .pem_cert = readFile("./data/certs/server.crt"), + .pem_key = readFile("./data/certs/server.key")}; + serverOptions->tls_certificates.push_back(tlsCertificate); + serverOptions->verify_client = true; + serverOptions->root_certificates = readFile("./data/certs/ca.crt"); + } + + void executeSuccessfulQuery() { + std::vector idData = { + 1, 12, 2, std::numeric_limits::max()}; + + updateTable( + "sample-data", + makeArrowTable({"id"}, {makeNumericArray(idData)})); + + auto idVec = makeFlatVector(idData); + + auto plan = ArrowFlightPlanBuilder() + .flightTableScan(ROW({"id"}, {BIGINT()})) + .planNode(); + + AssertQueryBuilder(plan) + .splits(makeSplits({"sample-data"})) + .assertResults(makeRowVector({idVec})); + } + + std::function createQueryFunction() { + std::vector idData = { + 1, 12, 2, std::numeric_limits::max()}; + + updateTable( + "sample-data", + makeArrowTable({"id"}, {makeNumericArray(idData)})); + + auto idVec = makeFlatVector(idData); + + auto plan = ArrowFlightPlanBuilder() + .flightTableScan(ROW({"id"}, {BIGINT()})) + .planNode(); + + return [this, plan, idVec]() { + AssertQueryBuilder(plan) + .splits(makeSplits({"sample-data"})) + .assertResults(makeRowVector({idVec})); + }; + } +}; + +class ArrowFlightConnectorMtlsTest : public ArrowFlightConnectorMtlsTestBase { + protected: + explicit ArrowFlightConnectorMtlsTest() + : ArrowFlightConnectorMtlsTestBase( + std::make_shared( + std::unordered_map{ + {ArrowFlightConfig::kDefaultServerSslEnabled, "true"}, + {ArrowFlightConfig::kServerVerify, "true"}, + {ArrowFlightConfig::kServerSslCertificate, + "./data/certs/ca.crt"}, + {ArrowFlightConfig::kClientSslCertificate, "./data/certs/client.crt"}, + {ArrowFlightConfig::kClientSslKey, "./data/certs/client.key"}})) {} +}; + +TEST_F(ArrowFlightConnectorMtlsTest, successfulMtlsConnection) { + executeSuccessfulQuery(); +} + +class ArrowFlightMtlsNoClientCertTest : public ArrowFlightConnectorMtlsTestBase { + protected: + ArrowFlightMtlsNoClientCertTest() + : ArrowFlightConnectorMtlsTestBase( + std::make_shared( + std::unordered_map{ + {ArrowFlightConfig::kDefaultServerSslEnabled, "true"}, + {ArrowFlightConfig::kServerVerify, "true"}, + {ArrowFlightConfig::kServerSslCertificate, "./data/certs/ca.crt"}})) {} +}; + +TEST_F(ArrowFlightMtlsNoClientCertTest, mtlsFailsWithoutClientCert) { + auto queryFunction = createQueryFunction(); + VELOX_ASSERT_THROW(queryFunction(), "failed to connect"); +} + +class ArrowFlightConnectorImplicitSslTest : public ArrowFlightConnectorMtlsTestBase { + protected: + ArrowFlightConnectorImplicitSslTest() + : ArrowFlightConnectorMtlsTestBase( + std::make_shared( + std::unordered_map{ + {ArrowFlightConfig::kServerVerify, "true"}, + {ArrowFlightConfig::kServerSslCertificate, + "./data/certs/ca.crt"}, + {ArrowFlightConfig::kClientSslCertificate, "./data/certs/client.crt"}, + {ArrowFlightConfig::kClientSslKey, "./data/certs/client.key"}})) {} +}; + +TEST_F(ArrowFlightConnectorImplicitSslTest, successfulImplicitSslConnection) { + executeSuccessfulQuery(); +} + +class ArrowFlightImplicitSslNoClientCertTest : public ArrowFlightConnectorMtlsTestBase { + protected: + ArrowFlightImplicitSslNoClientCertTest() + : ArrowFlightConnectorMtlsTestBase( + std::make_shared( + std::unordered_map{ + {ArrowFlightConfig::kDefaultServerSslEnabled, "true"}, + {ArrowFlightConfig::kServerVerify, "true"}, + {ArrowFlightConfig::kServerSslCertificate, "./data/certs/ca.crt"} + })) {} +}; + +TEST_F(ArrowFlightImplicitSslNoClientCertTest, mtlsFailsWithoutClientCertOnImplicitSsl) { + auto queryFunction = createQueryFunction(); + VELOX_ASSERT_THROW(queryFunction(), "failed to connect"); +} + +} // namespace facebook::presto::test diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/ArrowFlightConnectorTlsTest.cpp b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/ArrowFlightConnectorTlsTest.cpp index e0c611225bd65..0be7df437e5e4 100644 --- a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/ArrowFlightConnectorTlsTest.cpp +++ b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/ArrowFlightConnectorTlsTest.cpp @@ -39,8 +39,8 @@ class ArrowFlightConnectorTlsTestBase : public ArrowFlightConnectorTestBase { void setFlightServerOptions( flight::FlightServerOptions* serverOptions) override { flight::CertKeyPair tlsCertificate{ - .pem_cert = readFile("./data/tls_certs/server.crt"), - .pem_key = readFile("./data/tls_certs/server.key")}; + .pem_cert = readFile("./data/certs/server.crt"), + .pem_key = readFile("./data/certs/server.key")}; serverOptions->tls_certificates.push_back(tlsCertificate); } @@ -83,7 +83,7 @@ class ArrowFlightConnectorTlsTest : public ArrowFlightConnectorTlsTestBase { {ArrowFlightConfig::kDefaultServerSslEnabled, "true"}, {ArrowFlightConfig::kServerVerify, "true"}, {ArrowFlightConfig::kServerSslCertificate, - "./data/tls_certs/ca.crt"}})) {} + "./data/certs/ca.crt"}})) {} }; TEST_F(ArrowFlightConnectorTlsTest, tlsEnabled) { diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/CMakeLists.txt b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/CMakeLists.txt index 9af596a913973..f346d1212a1ad 100644 --- a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/CMakeLists.txt +++ b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/CMakeLists.txt @@ -23,16 +23,19 @@ target_link_libraries( add_executable( presto_flight_connector_test ArrowFlightConnectorTest.cpp ArrowFlightConnectorAuthTest.cpp + ArrowFlightConnectorMTlsTest.cpp ArrowFlightConnectorTlsTest.cpp ArrowFlightConnectorDataTypeTest.cpp ArrowFlightConfigTest.cpp) -set(DATA_DIR "${CMAKE_CURRENT_SOURCE_DIR}/data/tls_certs") - add_custom_target( copy_flight_test_data ALL - COMMAND ${CMAKE_COMMAND} -E copy_directory ${DATA_DIR} - $/data/tls_certs) + COMMAND ${CMAKE_COMMAND} -E copy_directory_if_different + "${CMAKE_CURRENT_SOURCE_DIR}/data" + "${CMAKE_CURRENT_BINARY_DIR}/data" + COMMENT "Copying test data files..." +) +add_dependencies(presto_flight_connector_test copy_flight_test_data) add_test(presto_flight_connector_test presto_flight_connector_test) target_link_libraries( diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/README.md b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/README.md index 3a5f2e5786c67..17e778a84750c 100644 --- a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/README.md +++ b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/README.md @@ -1,7 +1,7 @@ -### Placeholder TLS Certificates for Arrow Flight Connector Unit Testing -The `tls_certs` directory contains placeholder TLS certificates generated for unit testing the Arrow Flight Connector with TLS enabled. These certificates are not intended for production use and should only be used in the context of unit tests. +### Placeholder TLS & mTLS Certificates for Arrow Flight Connector Unit Testing +The `certs/` directory contains placeholder certificates used for unit and integration testing of the Arrow Flight Connector with **TLS** and **mutual TLS (mTLS)** enabled. These certificates are **not intended for production use** and should only be used for local development or automated testing scenarios. -### Generating TLS Certificates -To create the TLS certificates and keys inside the `tls_certs` folder, run the following command: +### Generating Certificates for TLS and mTLS Testing +To generate the necessary certificates and keys for both TLS and mTLS testing, run the following script: -`./generate_tls_certs.sh` +`./generate_certs.sh` diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/certs/ca.crt b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/certs/ca.crt new file mode 100644 index 0000000000000..bc6036c52aa12 --- /dev/null +++ b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/certs/ca.crt @@ -0,0 +1,22 @@ +-----BEGIN CERTIFICATE----- +MIIDmzCCAoOgAwIBAgIUVB6MDVAXGLccJY5XrssVNLKbH/gwDQYJKoZIhvcNAQEL +BQAwXDELMAkGA1UEBhMCVVMxDjAMBgNVBAgMBVN0YXRlMQ0wCwYDVQQHDARDaXR5 +MQ4wDAYDVQQKDAVNeU9yZzEPMA0GA1UECwwGTXlVbml0MQ0wCwYDVQQDDARNeUNB +MCAXDTI1MDczMDE2NTU1M1oYDzMwMjQxMTMwMTY1NTUzWjBcMQswCQYDVQQGEwJV +UzEOMAwGA1UECAwFU3RhdGUxDTALBgNVBAcMBENpdHkxDjAMBgNVBAoMBU15T3Jn +MQ8wDQYDVQQLDAZNeVVuaXQxDTALBgNVBAMMBE15Q0EwggEiMA0GCSqGSIb3DQEB +AQUAA4IBDwAwggEKAoIBAQCQ5D8+R5/HesO9BrgWF92wGOesBcOPQL7Nemzh1qYT +fhxNV7HGN1W5fByWsVelZ9326CTa7Yr4TlRXC8GHt9YbRvU68LBU1kAGUqGmPdHY +LgPZJkQqFhYdPdD++Gw2+ecfYp+Ls9pZm/pNOeCWBV0F3RaoPZ19m/C/lDdZz8OX +V6t2to+Yh0GXJyyJfO+w7qG6B/j8UiRYtnnMq3ywTcWqsUmYp4+uwmkEEN9/eSo9 +JanjpEiv6o9Yb9J5StXRPmbAoXOl45o87A6qo0vzgYdP6uKkPQxo6wSdltb5qctM +CKHm4bYFT6IHoVzUGVHiq2iRna87OiDPqnFaX3ktWJB9AgMBAAGjUzBRMB0GA1Ud +DgQWBBRGccIvJxNjWYY8hGZ3HIlut0p17jAfBgNVHSMEGDAWgBRGccIvJxNjWYY8 +hGZ3HIlut0p17jAPBgNVHRMBAf8EBTADAQH/MA0GCSqGSIb3DQEBCwUAA4IBAQAI +sUZOfL8ZJEu0Oq7AHgOmxMYhdRQmFr45C54fiNb6QLs5bPJZ9T1/nbPQZvF1hkFh +wB3pDn8rNntiYXkCkrQ6PAiQfn0WHl4jLYCoGYxFbSP1QViZNid7dPmpaxccjMhL +Zk7htfCS1HtHWWBZPMDDA8hsUvBf4qusVonO71XGL22Z2ZKtgvDJYAyoxm7xwIo2 +mqSH9TfOnHYE0hUpo3u4PdmVAfCXzSDRccLALnVlzt50ColmAQgzj3MnwWfXJmdv +kjBhIZ9Obt7Hf6FcBhX9/qQN1t3u6mLIjf2akokRFblcW5Xgv4X1c7pEtdtM9V+o +MVqiDfm49mpMcOdHDeb1 +-----END CERTIFICATE----- diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/certs/client.crt b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/certs/client.crt new file mode 100644 index 0000000000000..7313fbf6babc6 --- /dev/null +++ b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/certs/client.crt @@ -0,0 +1,22 @@ +-----BEGIN CERTIFICATE----- +MIIDmTCCAoGgAwIBAgIUHZmKZs2+ejbJGEezgoSUTPKF42YwDQYJKoZIhvcNAQEL +BQAwXDELMAkGA1UEBhMCVVMxDjAMBgNVBAgMBVN0YXRlMQ0wCwYDVQQHDARDaXR5 +MQ4wDAYDVQQKDAVNeU9yZzEPMA0GA1UECwwGTXlVbml0MQ0wCwYDVQQDDARNeUNB +MCAXDTI1MDczMDE2NTU1M1oYDzMwMjQxMTMwMTY1NTUzWjBrMQswCQYDVQQGEwJV +UzEOMAwGA1UECAwFU3RhdGUxDTALBgNVBAcMBENpdHkxDjAMBgNVBAoMBU15T3Jn +MQ8wDQYDVQQLDAZNeVVuaXQxHDAaBgNVBAMME2NsaWVudC5teWRvbWFpbi5jb20w +ggEiMA0GCSqGSIb3DQEBAQUAA4IBDwAwggEKAoIBAQDESuInkBMIwJ02rNuO1DPN +3BfuNzF2SJq5/UgH9soWxVIYqLOwinBmWRsrRVOps9mJjN/GYbj27GT1LFXZHIwh +MYf5JLE/Y8CL/JwnMyxe3wkiBb1s2d+urOrIPMLQTVLGtClyPVT1PnmOd/h7vJ+k +I5ciK9G+8krAq5NoWM/OvhdTK2pOns1jl6Lq2c4IGDgQHt65uKRqPTPtc1xe/Vxd +4sTT8LJQEHu6zdl5nGFdGq4eIJpr+9LuM8ZCYkTAYpzHJ2XKP7jEj6ZRmsIgwwTf +VMON+yAF06Tmk9VCRkLUGjNMA4Xh24MejieE0JSq7mIAhSYWE+nalG8HsxtsQWDt +AgMBAAGjQjBAMB0GA1UdDgQWBBRPc2AWGMwuJEbnnMXvzgNAizik1jAfBgNVHSME +GDAWgBRGccIvJxNjWYY8hGZ3HIlut0p17jANBgkqhkiG9w0BAQsFAAOCAQEACJde +pXZvE04uw6tv+iGplmYNfasMQ9JXbvi0JMlnt9Y7ajf0F3g6yw9xWQfp8mCWVuof +aS/X8qw3loundeprxVq/2V6pFXStLFXJCXX+YL0Wl8AMv8VOxdZ8+hYlkfMoiKnx +ZKaWgrVtI/idFRUJLg8aHLRk5qVOwACBg9DAxMBC4V4MCQBfvDDdY6Y5qAM8o7PN +YGsPv5JIQI/3jsG2ZNQ/A8Ar+BNKWqnwRg2jjXysjPJPaU8TExvFXUypQaUn+a8X +/Y+CXhGabfbodEbEvBny6tiFQMb3YWhd0kHkjrYylY8GOpQ/ziEC8s7JYcaLSiu/ +ko1ErmRY/vhQIvrylw== +-----END CERTIFICATE----- diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/certs/client.key b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/certs/client.key new file mode 100644 index 0000000000000..2c285f776ba2d --- /dev/null +++ b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/certs/client.key @@ -0,0 +1,28 @@ +-----BEGIN PRIVATE KEY----- +MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQDESuInkBMIwJ02 +rNuO1DPN3BfuNzF2SJq5/UgH9soWxVIYqLOwinBmWRsrRVOps9mJjN/GYbj27GT1 +LFXZHIwhMYf5JLE/Y8CL/JwnMyxe3wkiBb1s2d+urOrIPMLQTVLGtClyPVT1PnmO +d/h7vJ+kI5ciK9G+8krAq5NoWM/OvhdTK2pOns1jl6Lq2c4IGDgQHt65uKRqPTPt +c1xe/Vxd4sTT8LJQEHu6zdl5nGFdGq4eIJpr+9LuM8ZCYkTAYpzHJ2XKP7jEj6ZR +msIgwwTfVMON+yAF06Tmk9VCRkLUGjNMA4Xh24MejieE0JSq7mIAhSYWE+nalG8H +sxtsQWDtAgMBAAECggEAKjnqszSS56nZ2Bpw4+Wh3EnZywjMBuRBBrwmC/KK0EGz +8rKN7y8k1VubVOBlykayiBzKQciha9r4L+bQ8/LocTaQx+edCqQolmSp6ePgCmuj +8RH3iSxIanDv09IAXaOYqD63AMiRV22QZDXKOkIePIbcewEerps8Ofze6c5bK9/H +jWWScpphiJcTVSNKqTzVMJ27DcpTn/bQ+nPYQNdyMj51d1Vaebesmv1ieQfvdMTA +FFTCYUfNrWWjB9CEhZAg3sgr3dOVVTl8LC8VF6ug5EvNeuW2l8FX+fgVaZWKzGpf +veItiHIeMAEw6ZQPTumSvIO2fhg5Iue9fu4hR08twQKBgQDuqsRTJNG8DhR2CZrH +KWeW8ac7pOiHSHxqpERfwkxtpUEsLXib8foPRJYdEzvDDgL3SxaRmTnwtyZvhZgl +RE7ebRvAHkHy+nvim6wBgFYe2YwTTm9d3gCRxOHknEXfXSxzbwJznDlZrHlJ9XDQ +OFBgLbhheF694y1l605E/8s/QQKBgQDSjEztXmrzlQvXKhOp+eL6o+jRrfrv5hmC +ORy063ecVD/bA3OsgOKd5Q0v7dEyCC5IesQA0uY4GZ7zjsxcA4tUp20fZoTEeMpZ +DznTz5Lqkpa8DCeMLj6fPWKnb4zZeduvBE4BRzS4Mdpzp/OrLAGgE2hwWNIB56uq +pyloAkEirQKBgQDPZjg7JFjaOcYQGSKWheWOJyszSogC37u2lE8Sg+8UrTGoaU9Q +/QNXdzuXwpoBU9DCA092cRgHlbDh4s8nO2fqJBikZ+bZdlBnyO29VEACiPwP3u4q +PPxzsAq5NhAGHZq+KS6RNqYjxhyUZ6SEXRuDqNd8ZDS4gI137vZSQZLmwQKBgDcs +mQQjF/fY+Q9bcWe7miWASoSYCQhQziJ4APPQOLn4wfsMvoVYCQrDeV8z/PwVdLt9 +oFtu6PGOlT7SDu+V5i866Levz98EoFISUV8WKDPcUi/ZJ4vumm50UaP68XgUHOOS +RzbCiCg0uEBSpOIYWBywuU+nlvD02uGPiKQ+4v7JAoGAdyREhzfyz+WO5wdHLRIB +k9cz448Lp/Uh2pcgAkFvtBwCAGZn86YJZ7JFD/HkQbw7amVEn+nxoCeWHkbSnzt0 +8Gu0hdqo6SaxOdf1wKBel79r/GH4ZwYRzWJoO/Iy//JJJO7g9GylAo0rBdQR9n60 +ZPMIZJLfKmGQWahl2EqyMG8= +-----END PRIVATE KEY----- diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/certs/server.crt b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/certs/server.crt new file mode 100644 index 0000000000000..bdae53b6bd114 --- /dev/null +++ b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/certs/server.crt @@ -0,0 +1,23 @@ +-----BEGIN CERTIFICATE----- +MIIDxDCCAqygAwIBAgIUHZmKZs2+ejbJGEezgoSUTPKF42UwDQYJKoZIhvcNAQEL +BQAwXDELMAkGA1UEBhMCVVMxDjAMBgNVBAgMBVN0YXRlMQ0wCwYDVQQHDARDaXR5 +MQ4wDAYDVQQKDAVNeU9yZzEPMA0GA1UECwwGTXlVbml0MQ0wCwYDVQQDDARNeUNB +MCAXDTI1MDczMDE2NTU1M1oYDzMwMjQxMTMwMTY1NTUzWjBrMQswCQYDVQQGEwJV +UzEOMAwGA1UECAwFU3RhdGUxDTALBgNVBAcMBENpdHkxDjAMBgNVBAoMBU15T3Jn +MQ8wDQYDVQQLDAZNeVVuaXQxHDAaBgNVBAMME3NlcnZlci5teWRvbWFpbi5jb20w +ggEiMA0GCSqGSIb3DQEBAQUAA4IBDwAwggEKAoIBAQDS9+yttKuydwcc4te8CtJS +yldr163zuSemazYUDY07ZMc/VFfSISYBX3yrW8uveR4zt8H8XhdkPHV6t3K3GxcY +28CpUXr9TkCH2xMd4xn2DuewaTuf+yCcM6TLFh00nyqcYPhZGn3Bd5jGXKxI8PsH +butswFKM/t9VTRtTpBgMsw8SEy0vyvsJyPTHm8aWY7tCDSI+vI8bHcG9sO8cuin8 +0JcZ4rRTgZpmDDlcY0OniRwowB5ph3eU0uaHIX+EWuht3+1trjyrSFuQ6y82f/0N +MchcA3vTG0qF4ulcZn0yng1wuC1YorInYiDUahxDXHvmsoYz+0yOkR5fVReHLfdp +AgMBAAGjbTBrMCkGA1UdEQQiMCCCE3NlcnZlci5teWRvbWFpbi5jb22CCWxvY2Fs +aG9zdDAdBgNVHQ4EFgQUuxBCAwXgk1cEC3myPm4Q/b8+oWMwHwYDVR0jBBgwFoAU +RnHCLycTY1mGPIRmdxyJbrdKde4wDQYJKoZIhvcNAQELBQADggEBABlwupm4YIBd +dPRoX6S/Ta8oAvkz+gv/s9vIekG7fTzfGDmh9hQEtVU4OVT0ifVAX7x5bf0v7KKp +jXDctZqyttXAtu6e1nQBE+MC9MXi68YU2hSApJF+z7WTPXEbcuIOQPpvXfiD7s2j +455tMF90iKWjcFGdgB1usiQeNeNDjBwcvkGJhLhKO6UqsHLh2BBUTPhDnpyBkt1y +zsZk3YDKr2ipaIj3OvBwueJNL8I3oU04eZMAFeSYbnJ49GdSDLTxglACoOkZrbju +4qwV1bbIB97AtSdhcg3OxBZTAgVep8kOl3N/spNfgq4N/jTdsl+hoEKlChH/7kI3 +oYZfv+nc6s0= +-----END CERTIFICATE----- diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/certs/server.key b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/certs/server.key new file mode 100644 index 0000000000000..1906728090bb7 --- /dev/null +++ b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/certs/server.key @@ -0,0 +1,28 @@ +-----BEGIN PRIVATE KEY----- +MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQDS9+yttKuydwcc +4te8CtJSyldr163zuSemazYUDY07ZMc/VFfSISYBX3yrW8uveR4zt8H8XhdkPHV6 +t3K3GxcY28CpUXr9TkCH2xMd4xn2DuewaTuf+yCcM6TLFh00nyqcYPhZGn3Bd5jG +XKxI8PsHbutswFKM/t9VTRtTpBgMsw8SEy0vyvsJyPTHm8aWY7tCDSI+vI8bHcG9 +sO8cuin80JcZ4rRTgZpmDDlcY0OniRwowB5ph3eU0uaHIX+EWuht3+1trjyrSFuQ +6y82f/0NMchcA3vTG0qF4ulcZn0yng1wuC1YorInYiDUahxDXHvmsoYz+0yOkR5f +VReHLfdpAgMBAAECggEAGOjA/zmH1EiNhHGcO02jy7asX8VVeqNv9QxPlEqNVGfv +xqB0xhC35g2aMLlj8VIBqOWXd+68IE+rJ1QlrUz7iynXM6a1ONdWczQAq9S2qgDU +hlXGfnsuPIM0f+4agK4SX+hrKkogcwll9nXWub4KRbRpA6wpkxA82luCUHvdgxIi +PW37HHYWUKoAu828PpG1wGUf6wDFnSEhEuYfgHGkcWyhpRw5QZP563QTMCKWPjcW +Nm+XJ10gpAB4Q+zcqxji7r+5uIbF9Zobkd9VuaWHavQfrgHpjfLnZjHr3r3BABi0 +U1Zz5x5R19r54aY93J0Fq4DOlJ1Gf80eBLJEs1b8CQKBgQD53IfUp4ScfIs1I3Vw +JOclRY3UY/uZDQ+edXLqT91UxRpmvq8pmL82Y+idWGrx1WtCk2uOFGAl37Y+8goW +bMHEZ8Wv8NtV29sqduuCE210miZN3q7EmTEj0AOOKd/Skwyte3rxvFw0didQvAhE ++uZEIZ+XaUNF102hT2BaTIm8fwKBgQDYJsgiqGRzGtP3sOeLKxYmGDE3fGB4JHYU +U9kZ0pTQWiYsl+F7lXdwqkUApgU1rFuA3oR7dV4a8zS7BbyLK2WYfWE4yAoRm52W +VnIGsdG4Z4wFGuNR7d+m+MouP4HYSJFtUoJFJxYXU4Kc3H88Ob5UmerVNFN7VSaU +W/jtek34FwKBgH+c4sL5zAEgmvjI43IjZuriW03ewuGoihGkaszBfYmOIa3YNh5I +pWBiJqw2PGjHV8DpCkXGolS1rZ74f650XYKyfYUevudbItTNZ/tHcN/c2zNqSFig +5TglRauWN3qVICR6rJBKY81niyzw3Ehe3Lxvb9MlL/a7wCpjIBL+hFqBAoGBAIax +jA+EzauosSPtWiwv+kpc0vaXi+nyFp7OLUBZKCC5vIYXUwxW9KoBgKRJ0H9E23Rv +tTDVz4GNwnM0vOwga9vdbaMbjKKyTT4sujuPvXdjFy7rNXKNf8wlxp+RNZGYjv8H +5mO/WpXIlWC4SpU2CnPfwiV/yPHW+waCVZlumH2bAoGAJscnMS4YQXiSoLQweiyQ ++mev9+i9h57/M0bk6WQHWzhAjxmZDtFUPO8P4vi3akAhCnt0yswh8e3G6h3F9TEz +PQ4AIFMajb0PWUEcGsZhBh7TCJwxPXo34SbXGkIB5RYogT/mDVMjEfduAMen/0sp +Te0BRozzkKHtIcGAqozt7ZU= +-----END PRIVATE KEY----- diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/generate_certs.sh b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/generate_certs.sh new file mode 100755 index 0000000000000..11875d6fe462a --- /dev/null +++ b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/generate_certs.sh @@ -0,0 +1,55 @@ +#!/bin/bash + +# Set directory for certificates and keys. +CERT_DIR="./certs" +mkdir -p $CERT_DIR + +# Dummy values for the certificates. +COUNTRY="US" +STATE="State" +LOCALITY="City" +ORGANIZATION="MyOrg" +ORG_UNIT="MyUnit" +COMMON_NAME="MyCA" +SERVER_CN="server.mydomain.com" +CLIENT_CN="client.mydomain.com" + +# Step 1: Generate CA private key and self-signed certificate. +openssl genpkey -algorithm RSA -out $CERT_DIR/ca.key +openssl req -key $CERT_DIR/ca.key -new -x509 -out $CERT_DIR/ca.crt -days 365000 \ + -subj "/C=$COUNTRY/ST=$STATE/L=$LOCALITY/O=$ORGANIZATION/OU=$ORG_UNIT/CN=$COMMON_NAME" + +# Step 2: Generate server private key. +openssl genpkey -algorithm RSA -out $CERT_DIR/server.key + +# Step 3: Generate server certificate signing request (CSR). +openssl req -new -key $CERT_DIR/server.key -out $CERT_DIR/server.csr \ + -subj "/C=$COUNTRY/ST=$STATE/L=$LOCALITY/O=$ORGANIZATION/OU=$ORG_UNIT/CN=$SERVER_CN" \ + -addext "subjectAltName=DNS:$SERVER_CN,DNS:localhost" + +# Step 4: Sign server CSR with the CA certificate to generate the server certificate. +openssl x509 -req -in $CERT_DIR/server.csr -CA $CERT_DIR/ca.crt -CAkey $CERT_DIR/ca.key \ + -CAcreateserial -out $CERT_DIR/server.crt -days 365000 \ + -extfile <(printf "subjectAltName=DNS:$SERVER_CN,DNS:localhost") + +# Step 5: Generate client private key. +openssl genpkey -algorithm RSA -out $CERT_DIR/client.key + +# Step 6: Generate client certificate signing request (CSR). +openssl req -new -key $CERT_DIR/client.key -out $CERT_DIR/client.csr \ + -subj "/C=$COUNTRY/ST=$STATE/L=$LOCALITY/O=$ORGANIZATION/OU=$ORG_UNIT/CN=$CLIENT_CN" + +# Step 7: Sign client CSR with the CA certificate to generate the client certificate. +openssl x509 -req -in $CERT_DIR/client.csr -CA $CERT_DIR/ca.crt -CAkey $CERT_DIR/ca.key \ + -CAcreateserial -out $CERT_DIR/client.crt -days 365000 + +# Step 8: Output summary. +echo "Certificate Authority (CA) certificate : $CERT_DIR/ca.crt" +echo "Server certificate : $CERT_DIR/server.crt" +echo "Server private key : $CERT_DIR/server.key" +echo "Client certificate : $CERT_DIR/client.crt" +echo "Client private key : $CERT_DIR/client.key" + +# Step 9: Remove unused files. +rm -f $CERT_DIR/*.csr $CERT_DIR/ca.srl $CERT_DIR/ca.key + diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/generate_tls_certs.sh b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/generate_tls_certs.sh deleted file mode 100755 index 718f313c70a75..0000000000000 --- a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/generate_tls_certs.sh +++ /dev/null @@ -1,40 +0,0 @@ -#!/bin/bash - -# Set directory for certificates and keys. -CERT_DIR="./tls_certs" -mkdir -p $CERT_DIR - -# Dummy values for the certificates. -COUNTRY="US" -STATE="State" -LOCALITY="City" -ORGANIZATION="MyOrg" -ORG_UNIT="MyUnit" -COMMON_NAME="MyCA" -SERVER_CN="server.mydomain.com" - -# Step 1: Generate CA private key and self-signed certificate. -openssl genpkey -algorithm RSA -out $CERT_DIR/ca.key -openssl req -key $CERT_DIR/ca.key -new -x509 -out $CERT_DIR/ca.crt -days 365000 \ - -subj "/C=$COUNTRY/ST=$STATE/L=$LOCALITY/O=$ORGANIZATION/OU=$ORG_UNIT/CN=$COMMON_NAME" - -# Step 2: Generate server private key. -openssl genpkey -algorithm RSA -out $CERT_DIR/server.key - -# Step 3: Generate server certificate signing request (CSR). -openssl req -new -key $CERT_DIR/server.key -out $CERT_DIR/server.csr \ - -subj "/C=$COUNTRY/ST=$STATE/L=$LOCALITY/O=$ORGANIZATION/OU=$ORG_UNIT/CN=$SERVER_CN" \ - -addext "subjectAltName=DNS:$COMMON_NAME,DNS:localhost" \ - -# Step 4: Sign server CSR with the CA certificate to generate the server certificate. -openssl x509 -req -in $CERT_DIR/server.csr -CA $CERT_DIR/ca.crt -CAkey $CERT_DIR/ca.key \ - -CAcreateserial -out $CERT_DIR/server.crt -days 365000 \ - -extfile <(printf "subjectAltName=DNS:$COMMON_NAME,DNS:localhost") - -# Step 5: Output the generated files. -echo "Certificate Authority (CA) certificate: $CERT_DIR/ca.crt" -echo "Server certificate: $CERT_DIR/server.crt" -echo "Server private key: $CERT_DIR/server.key" - -# Step 6: Remove unused files. -rm -rf $CERT_DIR/server.csr $CERT_DIR/ca.srl $CERT_DIR/ca.key diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/tls_certs/ca.crt b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/tls_certs/ca.crt deleted file mode 100644 index 6740e89c54e17..0000000000000 --- a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/tls_certs/ca.crt +++ /dev/null @@ -1,22 +0,0 @@ ------BEGIN CERTIFICATE----- -MIIDmzCCAoOgAwIBAgIUf+rP48iL39yGlAfFQTIp5bmM4uQwDQYJKoZIhvcNAQEL -BQAwXDELMAkGA1UEBhMCVVMxDjAMBgNVBAgMBVN0YXRlMQ0wCwYDVQQHDARDaXR5 -MQ4wDAYDVQQKDAVNeU9yZzEPMA0GA1UECwwGTXlVbml0MQ0wCwYDVQQDDARNeUNB -MCAXDTI0MTIwMzExMDQxMVoYDzMwMjQwNDA1MTEwNDExWjBcMQswCQYDVQQGEwJV -UzEOMAwGA1UECAwFU3RhdGUxDTALBgNVBAcMBENpdHkxDjAMBgNVBAoMBU15T3Jn -MQ8wDQYDVQQLDAZNeVVuaXQxDTALBgNVBAMMBE15Q0EwggEiMA0GCSqGSIb3DQEB -AQUAA4IBDwAwggEKAoIBAQCliiXIcSmxXAAq2k/XjcZniDgEDCxWKZGiV8JBiJwY -MMBJtqcVzWfiDpO2u6d1dfGb6utlRW+1dnwupzURCMmZff4bqlPx4ZejRXDrWzKz -08WSpDVZwC2H5XOllwK36Cn4gvPRe3YWVcdDGHy7GL+zsJENvawJj0BH952MU4bk -sV52zEkN291bfN9sSYfT1NCJuLPM0Qsf97DeQ+wHXEw+t4XVMF3FQbciQp0y6CnA -wfFFN14WDiWxukP1I3kuDYYA6h/WJCQMp5rU2NCB9nIQrulYRxFaepMYENLxgAyj -gFaoRh2Kt2k7XKv6WOa6CmYm2dZERPlbA+oNAHkaHw6lAgMBAAGjUzBRMB0GA1Ud -DgQWBBSN+3vRlXGjs6c+rN94qgEnkPLl3DAfBgNVHSMEGDAWgBSN+3vRlXGjs6c+ -rN94qgEnkPLl3DAPBgNVHRMBAf8EBTADAQH/MA0GCSqGSIb3DQEBCwUAA4IBAQAb -L40Oe2b/2xdUSyjqXJceVxaeA291fCpnu1C1JliP0hNI3fu9jjZhXHQoYub/4mod -8lriEDIcOCCiUfmi404akpqQHuBmOHaKEOtaaQkezjPsYnUra+O2ssqUo2zto5bK -gR0LGsb+4AO0bDvq+QVI6kEQqAAIf6qC+kpg/jV4iKJ1J6Qw4R3QppYBm6SQcfvI -hfUfDSO6SNfy0f/ZVCavbJIP9zG/BfAD9DEERocw03PiN5bm4IXJ3HH8rxyuBfJ5 -Eg/fPP5TlZ2H7Kqb3VgVBGWJtNXWmJphHyraBJTEuxgXWvl6AaW0P/3dsJi3rfdD -zDIT7AmENLCom8Gl0bgM ------END CERTIFICATE----- diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/tls_certs/server.crt b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/tls_certs/server.crt deleted file mode 100644 index 92c91f2d613b0..0000000000000 --- a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/tls_certs/server.crt +++ /dev/null @@ -1,22 +0,0 @@ ------BEGIN CERTIFICATE----- -MIIDtTCCAp2gAwIBAgIUUhmhZP94nIowrg2EarzfEBp6W1EwDQYJKoZIhvcNAQEL -BQAwXDELMAkGA1UEBhMCVVMxDjAMBgNVBAgMBVN0YXRlMQ0wCwYDVQQHDARDaXR5 -MQ4wDAYDVQQKDAVNeU9yZzEPMA0GA1UECwwGTXlVbml0MQ0wCwYDVQQDDARNeUNB -MCAXDTI0MTIwMzExMDQxMVoYDzMwMjQwNDA1MTEwNDExWjBrMQswCQYDVQQGEwJV -UzEOMAwGA1UECAwFU3RhdGUxDTALBgNVBAcMBENpdHkxDjAMBgNVBAoMBU15T3Jn -MQ8wDQYDVQQLDAZNeVVuaXQxHDAaBgNVBAMME3NlcnZlci5teWRvbWFpbi5jb20w -ggEiMA0GCSqGSIb3DQEBAQUAA4IBDwAwggEKAoIBAQDSxC4zCC4GFZbX+fdFgWbL -sj4PortyOM7mzRjNaQ3M0FTSEy5xET9C2qFlBCJ7AL7DlbSLmSckYY/FkdfMqNN2 -+NZ0Dy2d6bZN+ly5N/QBVnyS/5HVC3MXa6Y2BmFXiBnczWfGBwj+uVHlKOUWUyNi -EyUkhuPwtYXkFmJoqBxJSPC6cxX6NzMujnwCF18dUf0Vra44osu4moaovmg3c9jM -cBtmafFs9F54FoAEuLotjISVEa7VY6th5RxXJHpgas+0R5EBddGYKbTRiUYjht7r -pS+An0ey02oOjEWdqLnQSg/SUGKuRXULyE5l1A0HfNQtvepUQotb9ull1F7OrbfB -AgMBAAGjXjBcMBoGA1UdEQQTMBGCBE15Q0GCCWxvY2FsaG9zdDAdBgNVHQ4EFgQU -vnCLWjre4jqkKzC24psCPh1oIQwwHwYDVR0jBBgwFoAUjft70ZVxo7OnPqzfeKoB -J5Dy5dwwDQYJKoZIhvcNAQELBQADggEBAJCiJgtTw/7b/g3QnJfM4teQkFS440Ii -weqQJMoP6als8Fc3opPKv9eC5w0wqaLlIdwJjzGM5PmCAtGVafo22TbqhZyQdzQu -TUKv1DaVF0JBVAGVxTSDIK9r5Ww4mDAQnQENLC6soS3AvYDEi+8667YLoNNdhRCX -q2D5v76UN45idiShppxOw53whsvpHv+wyqcdse7DhgM9boCbx51Uvv3l/AEToyaj -S1xeIkBwNpSYU0ax2Lr1j2yoKbzAa3MHy8Php+T5CGji02+HwwlvlPDLtw8q5gHw -BLSwlAHgclPxUTWNNoCqjfX8Bi083+QDCLm0rgQ45xljNDbFAF1Y5hA= ------END CERTIFICATE----- diff --git a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/tls_certs/server.key b/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/tls_certs/server.key deleted file mode 100644 index 2cdf5750a4753..0000000000000 --- a/presto-native-execution/presto_cpp/main/connectors/arrow_flight/tests/data/tls_certs/server.key +++ /dev/null @@ -1,28 +0,0 @@ ------BEGIN PRIVATE KEY----- -MIIEvAIBADANBgkqhkiG9w0BAQEFAASCBKYwggSiAgEAAoIBAQDSxC4zCC4GFZbX -+fdFgWbLsj4PortyOM7mzRjNaQ3M0FTSEy5xET9C2qFlBCJ7AL7DlbSLmSckYY/F -kdfMqNN2+NZ0Dy2d6bZN+ly5N/QBVnyS/5HVC3MXa6Y2BmFXiBnczWfGBwj+uVHl -KOUWUyNiEyUkhuPwtYXkFmJoqBxJSPC6cxX6NzMujnwCF18dUf0Vra44osu4moao -vmg3c9jMcBtmafFs9F54FoAEuLotjISVEa7VY6th5RxXJHpgas+0R5EBddGYKbTR -iUYjht7rpS+An0ey02oOjEWdqLnQSg/SUGKuRXULyE5l1A0HfNQtvepUQotb9ull -1F7OrbfBAgMBAAECggEAAxbZuuESGGAMMm9HLGhKHgbHU8gnv2Phdbrka+SYBYg5 -UYzTHLh3FwEsjd4VnaweJ7CN1WDb1NvWmTum/DCebJ1HKqtjKLAZfk8q2TLGmXdL -pzWOdQ8MX1fKP2sIlcl0kFbNCE8vprjneDyBLtqOK36eiAh/fl6BQ12QAMLjyv/L -OwXSY4ESs/RzxRzFgdT98cDZFL7y0FVIjJo/Q5lfW9UwwSfw8tOLNXKTYwPHqIfJ -NjfWD7IqztQlnanyRXv5dScp80i8p9qgH0i8YfVBHZDeOmHGLcltilLRZ0dQ/X0g -Lrr0aIO3iLhmTIkJRzUnGeyvDjxcPINvRSBBwXy04QKBgQDpFJa/EwSsWj8586oh -xgm0Z3q+FiEeCe7aLLPcXAS2EDvix5ibJDT2y1Aa/kXq25S53npa/7Ov6TJs5H4g -eyshDtR1wVhz+rIggREiX/sagkhwnNsssUZFv5t9PdnaFXpVnH49m5Qc8HO3owtN -t8EGSRcAQ4o/fLWLs51qd38cIQKBgQDnfd8YPyDQ03xDC/3+Qrypyc/xhGnCuj7w -ZeA5iEyTnnNxL0a0B6PWcSk2BZReMNQKgYtipnsOQKtwHMttxtXYs/VQpeB4KoWE -zEwW0fV3MMsXN+nVJlEZnVaTbmYXknjeZrh/rNjsY96yxw8NtvAuYSpnqtr3N2nd -iMQ3G/QnoQKBgGMi+bdNvIgeXpQkmrGAzTHpbaCaQv3G1cwAhYPts6dIomAj6znZ -nZl3ApxomI57VPf1s+8uoVvqASOl0Cu6l66Y4y8uzJOQBuGiZApN7rzouy0C2opY -4H3cMKOFgjqrNfxh8qP7n3TrpRxvgehNhxFIVzsqfwvf3EwOWp8lMnBhAoGAZ25E -Ge9K2ENGCCb5i3uCFFLJiF3ja1AQAxVhxBL0NBjd97pp2tJ3D79r7Gk9y4ABndgX -0TIVVV7ruqIC+r+WmMZ/W1NiIg7NrXIipSeWh3TTqUIgRk5iehFkt2biUrHtM2Gu -Gc2+9pAA1tw+C6CrW+2qJrueLksiEAulsAHba0ECgYBIgIiY+Gx+XecEgCwAhWcn -GzNDAAlA4IgBjHpUtIByflzQDqlECKXjPbVBKfyq6eLt40upFmQCLsn+AkiQau8A -3cFAK9wJOAHv9KuWDrbHyhRE9CrJ6BqsY2goC3LiFCTgJy1TrRl6CDaFzHivONwF -LNPflYk5s376UWqxC+HtIA== ------END PRIVATE KEY----- From 95215807d26610abf277abb1850297153313300b Mon Sep 17 00:00:00 2001 From: Xiaoxuan Meng Date: Sat, 13 Sep 2025 23:31:00 -0700 Subject: [PATCH 092/113] [sapphire-velox]fix: Fix to handle multiple TaskSource have the same source node (#26031) Summary: Sapphire-Velox might send multiple task sources with the same source node. Task manager doesn't expect this and directly send splits of each task source to velox task. Since Sapphire-Velox send all splits once for each velox task, then all such task sources have no more splits set. This hit the check failure in the recent added no-more split check in Velox task split add API. This PR fixes the issue by merge the splits from multiple task sources if they share the same source node id. Reviewed By: zacw7, tanjialiang Differential Revision: D82367224 --- .../presto_cpp/main/TaskManager.cpp | 28 +++++++++++++++- .../presto_cpp/main/tests/TaskManagerTest.cpp | 33 +++++++++++++++++++ .../TestPrestoSparkNativeGeneralQueries.java | 5 --- 3 files changed, 60 insertions(+), 6 deletions(-) diff --git a/presto-native-execution/presto_cpp/main/TaskManager.cpp b/presto-native-execution/presto_cpp/main/TaskManager.cpp index 1d893c4a110de..cafac1c0dcf6d 100644 --- a/presto-native-execution/presto_cpp/main/TaskManager.cpp +++ b/presto-native-execution/presto_cpp/main/TaskManager.cpp @@ -542,7 +542,7 @@ std::unique_ptr TaskManager::createOrUpdateTaskImpl( { std::lock_guard l(prestoTask->mutex); prestoTask->updateCoordinatorHeartbeatLocked(); - if (not prestoTask->task && planFragment.planNode) { + if ((prestoTask->task == nullptr) && (planFragment.planNode != nullptr)) { // If the task is aborted, no need to do anything else. // This takes care of DELETE task message coming before CREATE task. if (prestoTask->info.taskStatus.state == protocol::TaskState::ABORTED) { @@ -611,7 +611,33 @@ std::unique_ptr TaskManager::createOrUpdateTaskImpl( VLOG(1) << "Failed to update output buffers for task: " << taskId; } + folly::F14FastMap sourcesMap; for (const auto& source : sources) { + auto it = sourcesMap.find(source.planNodeId); + if (it == sourcesMap.end()) { + // No existing source with same planNodeId, add as new + sourcesMap.emplace(source.planNodeId, source); + continue; + } + + // Merge with existing source that has the same planNodeId + auto& merged = it->second; + + // Merge splits + merged.splits.insert( + merged.splits.end(), source.splits.begin(), source.splits.end()); + + // Merge noMoreSplitsForLifespan + merged.noMoreSplitsForLifespan.insert( + merged.noMoreSplitsForLifespan.end(), + source.noMoreSplitsForLifespan.begin(), + source.noMoreSplitsForLifespan.end()); + + // Use OR logic for noMoreSplits flag + merged.noMoreSplits = merged.noMoreSplits || source.noMoreSplits; + } + + for (const auto& [_, source] : sourcesMap) { // Add all splits from the source to the task. VLOG(1) << "Adding " << source.splits.size() << " splits to " << taskId << " for node " << source.planNodeId; diff --git a/presto-native-execution/presto_cpp/main/tests/TaskManagerTest.cpp b/presto-native-execution/presto_cpp/main/tests/TaskManagerTest.cpp index c5eba098be4e5..02f84d83a77ed 100644 --- a/presto-native-execution/presto_cpp/main/tests/TaskManagerTest.cpp +++ b/presto-native-execution/presto_cpp/main/tests/TaskManagerTest.cpp @@ -715,6 +715,39 @@ TEST_P(TaskManagerTest, tableScanAllSplitsAtOnce) { assertResults(taskId, rowType_, "SELECT * FROM tmp WHERE c0 % 5 = 0"); } +TEST_P(TaskManagerTest, addSplitsWithSameSourceNode) { + const auto tableDir = exec::test::TempDirectoryPath::create(); + auto filePaths = makeFilePaths(tableDir, 5); + auto vectors = makeVectors(filePaths.size(), 1'000); + for (int i = 0; i < filePaths.size(); i++) { + writeToFile(filePaths[i], vectors[i]); + } + duckDbQueryRunner_.createTable("tmp", vectors); + + const auto planFragment = exec::test::PlanBuilder() + .tableScan(rowType_) + .filter("c0 % 5 = 0") + .partitionedOutput({}, 1, {"c0", "c1"}, GetParam()) + .planFragment(); + + protocol::TaskUpdateRequest updateRequest; + // Create multiple task sources with the same source node id. + std::vector taskSources; + taskSources.reserve(filePaths.size()); + long splitSequenceId{0}; + for (const auto& filePath : filePaths) { + taskSources.push_back(makeSource("0", {filePath}, /*noMoreSplits=*/true, splitSequenceId)); + } + taskSources.reserve(filePaths.size()); + updateRequest.sources = std::move(taskSources); + + protocol::TaskId taskId = "scan.0.0.1.0"; + auto taskInfo = createOrUpdateTask(taskId, updateRequest, planFragment); + + ASSERT_GE(taskInfo->stats.queuedTimeInNanos, 0); + assertResults(taskId, rowType_, "SELECT * FROM tmp WHERE c0 % 5 = 0"); +} + TEST_P(TaskManagerTest, fecthFromFinishedTask) { const auto tableDir = exec::test::TempDirectoryPath::create(); auto filePaths = makeFilePaths(tableDir, 5); diff --git a/presto-native-execution/src/test/java/com/facebook/presto/spark/TestPrestoSparkNativeGeneralQueries.java b/presto-native-execution/src/test/java/com/facebook/presto/spark/TestPrestoSparkNativeGeneralQueries.java index 458323383d2dc..448ddde08bd1f 100644 --- a/presto-native-execution/src/test/java/com/facebook/presto/spark/TestPrestoSparkNativeGeneralQueries.java +++ b/presto-native-execution/src/test/java/com/facebook/presto/spark/TestPrestoSparkNativeGeneralQueries.java @@ -121,9 +121,4 @@ public void testCatalogWithCacheEnabled() {} @Override @Ignore public void testKeyBasedSamplingInlined() {} - - // VeloxRuntimeError: !noMoreSplits_ - @Override - @Ignore - public void testUnionAllInsert() {} } From 003cda687413b5d9be451165a36af7e7aea6970c Mon Sep 17 00:00:00 2001 From: Ke Wang Date: Sat, 13 Sep 2025 11:40:12 -0700 Subject: [PATCH 093/113] [native] Assign createFinishTimeMs at point of task completion in prestoTask Previously, prestoTask->createFinishTimeMs was set after the lock scope, potentially not reflecting the actual task creation finish time. Now, the assignment is moved inside the lock, right after the task is created and assigned, to more accurately capture when the task creation completes. --- presto-native-execution/presto_cpp/main/TaskManager.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/presto-native-execution/presto_cpp/main/TaskManager.cpp b/presto-native-execution/presto_cpp/main/TaskManager.cpp index cafac1c0dcf6d..da39aefd316f1 100644 --- a/presto-native-execution/presto_cpp/main/TaskManager.cpp +++ b/presto-native-execution/presto_cpp/main/TaskManager.cpp @@ -573,10 +573,10 @@ std::unique_ptr TaskManager::createOrUpdateTaskImpl( prestoTask->task = std::move(newExecTask); prestoTask->info.needsPlan = false; startTask = true; + prestoTask->createFinishTimeMs = getCurrentTimeMs(); } execTask = prestoTask->task; } - prestoTask->createFinishTimeMs = getCurrentTimeMs(); // Outside of prestoTask->mutex. VELOX_CHECK_NOT_NULL( execTask, From 8f96415095e4c9b90802306e88df1975bd8ac628 Mon Sep 17 00:00:00 2001 From: Nishitha-Bhaskaran Date: Tue, 9 Sep 2025 15:03:24 +0530 Subject: [PATCH 094/113] Upgrade org.fusesource.jansi:jansi version to 2.4.2 --- pom.xml | 2 +- .../src/main/java/com/facebook/presto/cli/KeyReader.java | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/pom.xml b/pom.xml index 368a195f21996..31b39909c3c5c 100644 --- a/pom.xml +++ b/pom.xml @@ -1537,7 +1537,7 @@ org.fusesource.jansi jansi - 1.18 + 2.4.2 diff --git a/presto-cli/src/main/java/com/facebook/presto/cli/KeyReader.java b/presto-cli/src/main/java/com/facebook/presto/cli/KeyReader.java index 7a75fa24603f3..a242a98511c92 100644 --- a/presto-cli/src/main/java/com/facebook/presto/cli/KeyReader.java +++ b/presto-cli/src/main/java/com/facebook/presto/cli/KeyReader.java @@ -18,11 +18,12 @@ import java.io.IOException; import java.io.InputStream; -import static org.fusesource.jansi.internal.CLibrary.STDIN_FILENO; import static org.fusesource.jansi.internal.CLibrary.isatty; public final class KeyReader { + private static final int STDIN_FILENO = 0; + private KeyReader() {} @SuppressWarnings("resource") From 4261dcf3e33b2685419690a692e45c686a3fed74 Mon Sep 17 00:00:00 2001 From: Li Zhou Date: Mon, 15 Sep 2025 17:50:21 +0100 Subject: [PATCH 095/113] Enable more features for prestissimo image (#25712) ## Description 1. Enables features for prestissimo image by default, added flags below when building the image: ``` -DPRESTO_ENABLE_REMOTE_FUNCTIONS=ON -DPRESTO_ENABLE_JWT=ON -DPRESTO_STATS_REPORTER_TYPE=PROMETHEUS -DPRESTO_MEMORY_CHECKER_TYPE=LINUX_MEMORY_CHECKER -DPRESTO_ENABLE_SPATIAL=ON ``` 2. Use cache mount on ccache directory to accelerate local build 3. Added ARM_BUILD_TARGET for arm build 4. Fixed error in centos dependency image when building arrow ## Motivation and Context By default the image is built with ``` -DPRESTO_ENABLE_TESTING=OFF -DPRESTO_ENABLE_PARQUET=ON -DPRESTO_ENABLE_S3=ON ``` Add more features so that user can try without rebuild the image ## Impact Release ## Test Plan Build and test ## Contributor checklist - [ ] Please make sure your submission complies with our [contributing guide](https://github.com/prestodb/presto/blob/master/CONTRIBUTING.md), in particular [code style](https://github.com/prestodb/presto/blob/master/CONTRIBUTING.md#code-style) and [commit standards](https://github.com/prestodb/presto/blob/master/CONTRIBUTING.md#commit-standards). - [ ] PR description addresses the issue accurately and concisely. If the change is non-trivial, a GitHub Issue is referenced. - [ ] Documented new properties (with its default value), SQL syntax, functions, or other functionality. - [ ] If release notes are required, they follow the [release notes guidelines](https://github.com/prestodb/presto/wiki/Release-Notes-Guidelines). - [ ] Adequate tests were added if applicable. - [ ] CI passed. ## Release Notes Please follow [release notes guidelines](https://github.com/prestodb/presto/wiki/Release-Notes-Guidelines) and fill in the release notes below. ``` == NO RELEASE NOTE == ``` --- .github/workflows/presto-release-publish.yml | 12 +++++++++++- .../scripts/dockerfiles/centos-dependency.dockerfile | 7 ++++++- .../dockerfiles/prestissimo-runtime.dockerfile | 6 ++++-- .../dockerfiles/ubuntu-22.04-dependency.dockerfile | 8 +++++++- presto-native-execution/scripts/setup-adapters.sh | 2 +- 5 files changed, 29 insertions(+), 6 deletions(-) diff --git a/.github/workflows/presto-release-publish.yml b/.github/workflows/presto-release-publish.yml index bcd408405feb5..78e156eeea97a 100644 --- a/.github/workflows/presto-release-publish.yml +++ b/.github/workflows/presto-release-publish.yml @@ -398,7 +398,17 @@ jobs: working-directory: presto-native-execution run: | df -h - docker compose build centos-native-runtime + docker compose build --build-arg EXTRA_CMAKE_FLAGS=" + -DPRESTO_ENABLE_PARQUET=ON \ + -DPRESTO_ENABLE_REMOTE_FUNCTIONS=ON \ + -DPRESTO_ENABLE_JWT=ON \ + -DPRESTO_STATS_REPORTER_TYPE=PROMETHEUS \ + -DPRESTO_MEMORY_CHECKER_TYPE=LINUX_MEMORY_CHECKER \ + -DPRESTO_ENABLE_SPATIAL=ON \ + -DPRESTO_ENABLE_TESTING=OFF \ + -DPRESTO_ENABLE_S3=ON" \ + --build-arg NUM_THREADS=2 \ + centos-native-runtime - name: Add release tag working-directory: presto-native-execution diff --git a/presto-native-execution/scripts/dockerfiles/centos-dependency.dockerfile b/presto-native-execution/scripts/dockerfiles/centos-dependency.dockerfile index e67770e78a0b8..fc9c0c810f265 100644 --- a/presto-native-execution/scripts/dockerfiles/centos-dependency.dockerfile +++ b/presto-native-execution/scripts/dockerfiles/centos-dependency.dockerfile @@ -12,9 +12,14 @@ FROM quay.io/centos/centos:stream9 -ENV PROMPT_ALWAYS_RESPOND=n +# Set this when build arm with common flags +# from https://github.com/facebookincubator/velox/pull/14366 +ARG ARM_BUILD_TARGET + +ENV PROMPT_ALWAYS_RESPOND=y ENV CC=/opt/rh/gcc-toolset-12/root/bin/gcc ENV CXX=/opt/rh/gcc-toolset-12/root/bin/g++ +ENV ARM_BUILD_TARGET=${ARM_BUILD_TARGET} RUN mkdir -p /scripts /velox/scripts COPY scripts /scripts diff --git a/presto-native-execution/scripts/dockerfiles/prestissimo-runtime.dockerfile b/presto-native-execution/scripts/dockerfiles/prestissimo-runtime.dockerfile index 30f299172d7fe..826a818b0d7f5 100644 --- a/presto-native-execution/scripts/dockerfiles/prestissimo-runtime.dockerfile +++ b/presto-native-execution/scripts/dockerfiles/prestissimo-runtime.dockerfile @@ -26,8 +26,10 @@ ENV BUILD_DIR="" RUN mkdir -p /prestissimo /runtime-libraries COPY . /prestissimo/ -RUN EXTRA_CMAKE_FLAGS=${EXTRA_CMAKE_FLAGS} \ - NUM_THREADS=${NUM_THREADS} make --directory="/prestissimo/" cmake-and-build BUILD_TYPE=${BUILD_TYPE} BUILD_DIR=${BUILD_DIR} BUILD_BASE_DIR=${BUILD_BASE_DIR} +RUN --mount=type=cache,target=/root/.ccache,sharing=locked \ + EXTRA_CMAKE_FLAGS=${EXTRA_CMAKE_FLAGS} \ + NUM_THREADS=${NUM_THREADS} make --directory="/prestissimo/" cmake-and-build BUILD_TYPE=${BUILD_TYPE} BUILD_DIR=${BUILD_DIR} BUILD_BASE_DIR=${BUILD_BASE_DIR} && \ + ccache -sz -v RUN !(LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/local/lib:/usr/local/lib64 ldd /prestissimo/${BUILD_BASE_DIR}/${BUILD_DIR}/presto_cpp/main/presto_server | grep "not found") && \ LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/local/lib:/usr/local/lib64 ldd /prestissimo/${BUILD_BASE_DIR}/${BUILD_DIR}/presto_cpp/main/presto_server | awk 'NF == 4 { system("cp " $3 " /runtime-libraries") }' diff --git a/presto-native-execution/scripts/dockerfiles/ubuntu-22.04-dependency.dockerfile b/presto-native-execution/scripts/dockerfiles/ubuntu-22.04-dependency.dockerfile index 2738552101ba3..8c548c2679adb 100644 --- a/presto-native-execution/scripts/dockerfiles/ubuntu-22.04-dependency.dockerfile +++ b/presto-native-execution/scripts/dockerfiles/ubuntu-22.04-dependency.dockerfile @@ -16,12 +16,18 @@ FROM ${base} # Set a default timezone, can be overriden via ARG ARG tz="America/New_York" ARG DEBIAN_FRONTEND="noninteractive" -ENV PROMPT_ALWAYS_RESPOND=n + +# Set this when build arm with common flags +# from https://github.com/facebookincubator/velox/pull/14366 +ARG ARM_BUILD_TARGET + +ENV PROMPT_ALWAYS_RESPOND=y ENV SUDO=" " # TZ and DEBIAN_FRONTEND="noninteractive" # are required to avoid tzdata installation # to prompt for region selection. ENV TZ=${tz} +ENV ARM_BUILD_TARGET=${ARM_BUILD_TARGET} RUN mkdir -p /scripts /velox/scripts COPY scripts /scripts diff --git a/presto-native-execution/scripts/setup-adapters.sh b/presto-native-execution/scripts/setup-adapters.sh index 516b28c2302b7..42596d1c059b1 100755 --- a/presto-native-execution/scripts/setup-adapters.sh +++ b/presto-native-execution/scripts/setup-adapters.sh @@ -47,7 +47,7 @@ function install_arrow_flight { # Arrow Flight enabled. The Velox version of Arrow is used. # NOTE: benchmarks are on due to a compilation error with v15.0.0, once updated that can be removed # see https://github.com/apache/arrow/issues/41617 - EXTRA_ARROW_OPTIONS=" -DARROW_FLIGHT=ON -DARROW_BUILD_BENCHMARKS=ON " + EXTRA_ARROW_OPTIONS=" -DARROW_FLIGHT=ON -DARROW_BUILD_BENCHMARKS=ON -DgRPC_SOURCE=BUNDLED -DProtobuf_SOURCE=BUNDLED " install_arrow } From 0494998b3888fd858f30962803aeca2e6e0bc093 Mon Sep 17 00:00:00 2001 From: Pratik Joseph Dabre Date: Mon, 15 Sep 2025 11:35:57 -0700 Subject: [PATCH 096/113] Package memory connector plugin under native-plugin/ using Provisio --- presto-server/src/main/provisio/presto.xml | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/presto-server/src/main/provisio/presto.xml b/presto-server/src/main/provisio/presto.xml index 59dd851e6436e..b7176785fedd4 100644 --- a/presto-server/src/main/provisio/presto.xml +++ b/presto-server/src/main/provisio/presto.xml @@ -360,6 +360,12 @@ + + + + + + From d23167db6f8fa53ffaf427c3d7c594b412f5c9d2 Mon Sep 17 00:00:00 2001 From: Dilli-Babu-Godari Date: Thu, 21 Aug 2025 15:26:45 +0530 Subject: [PATCH 097/113] Add array_sort() and array_sort_desc() with lambda support for key extraction This commit introduces a new overloaded functions 1. array_sort() that accepts an array and a lambda expression to extract sort keys, then sorts the array in ascending order based on those keys. 2. array_sort_desc() that accepts an array and a lambda expression to extract sort keys, then sorts the array in descending order based on those keys. Such as, array_sort(ARRAY['hello', 'hi', 'world'], x -> length(x)) -- Returns: ['hi', 'hello', 'world'] array_sort(ARRAY[row('apples', 23), row('bananas', 12)], x -> x[2]) -- Returns: [row('bananas', 12), row('apples', 23)] array_sort_desc(ARRAY['hello', 'hi', 'world'], x -> length(x)) -- Returns: ['hello', 'world', 'hi'] array_sort_desc(ARRAY[row('apples', 23), row('bananas', 12)], x -> x[2]) -- Returns: [row('apples', 23), row('bananas', 12)] The implementation leverages the same code generation approach to optimize key extraction based on element and key types. --- .../src/main/sphinx/functions/array.rst | 24 + ...uiltInTypeAndFunctionNamespaceManager.java | 3 + .../AbstractArraySortByKeyFunction.java | 420 ++++++++++++++++++ .../scalar/TestArraySortByKeyFunctions.java | 309 +++++++++++++ .../sidecar/TestNativeSidecarPlugin.java | 64 +++ ...arPluginWithoutLoadingFunctionalities.java | 4 - 6 files changed, 820 insertions(+), 4 deletions(-) create mode 100644 presto-main-base/src/main/java/com/facebook/presto/operator/scalar/AbstractArraySortByKeyFunction.java create mode 100644 presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestArraySortByKeyFunctions.java diff --git a/presto-docs/src/main/sphinx/functions/array.rst b/presto-docs/src/main/sphinx/functions/array.rst index 7cd3c60589dde..fde1e78f07581 100644 --- a/presto-docs/src/main/sphinx/functions/array.rst +++ b/presto-docs/src/main/sphinx/functions/array.rst @@ -198,6 +198,18 @@ Array Functions -1, IF(cardinality(x) = cardinality(y), 0, 1))); -- [[1, 2], [2, 3, 1], [4, 2, 1, 4]] +.. function:: array_sort(array(T), function(T,U)) -> array(T) + + Sorts and returns the ``array`` using a lambda function to extract sorting keys. The function is applied + to each element of the array to produce a key, and the array is sorted based on these keys in ascending order. + Null array elements and null keys are placed at the end. :: + + SELECT array_sort(ARRAY['pear', 'apple', 'banana', 'kiwi'], x -> length(x)); -- ['pear', 'kiwi', 'apple', 'banana'] + SELECT array_sort(ARRAY[5, 20, 3, 9, 100], x -> x); -- [3, 5, 9, 20, 100] + SELECT array_sort(ARRAY['apple', NULL, 'banana', NULL], x -> length(x)); -- ['apple', 'banana', NULL, NULL] + SELECT array_sort(ARRAY[CAST(0.0 AS DOUBLE), CAST('NaN' AS DOUBLE), CAST('Infinity' AS DOUBLE), CAST('-Infinity' AS DOUBLE)], x -> x); -- [-Infinity, 0.0, Infinity, NaN] + SELECT array_sort(ARRAY[ROW('a', 3), ROW('b', 1), ROW('c', 2)], x -> x[2]); -- [ROW('b', 1), ROW('c', 2), ROW('a', 3)] + .. function:: array_sort_desc(x) -> array Returns the ``array`` sorted in the descending order. Elements of the ``array`` must be orderable. @@ -207,6 +219,18 @@ Array Functions SELECT array_sort_desc(ARRAY [null, 100, null, 1, 10, 50]); -- [100, 50, 10, 1, null, null] SELECT array_sort_desc(ARRAY [ARRAY ["a", null], null, ARRAY ["a"]); -- [["a", null], ["a"], null] +.. function:: array_sort_desc(array(T), function(T,U)) -> array(T) + + Sorts and returns the ``array`` in descending order using a lambda function to extract sorting keys. + The function is applied to each element of the array to produce a key, and the array is sorted based + on these keys in descending order. Null array elements and null keys are placed at the end. :: + + SELECT array_sort_desc(ARRAY['pear', 'apple', 'banana', 'kiwi'], x -> length(x)); -- ['banana', 'apple', 'pear', 'kiwi'] + SELECT array_sort_desc(ARRAY[5, 20, 3, 9, 100], x -> x); -- [100, 20, 9, 5, 3] + SELECT array_sort_desc(ARRAY['apple', NULL, 'banana', NULL], x -> length(x)); -- ['banana', 'apple', NULL, NULL] + SELECT array_sort_desc(ARRAY[CAST(0.0 AS DOUBLE), CAST('NaN' AS DOUBLE), CAST('Infinity' AS DOUBLE), CAST('-Infinity' AS DOUBLE)], x -> x); -- [NaN, Infinity, 0.0, -Infinity] + SELECT array_sort_desc(ARRAY[ROW('a', 3), ROW('b', 1), ROW('c', 2)], x -> x[2]); -- [ROW('a', 3), ROW('c', 2), ROW('b', 1)] + .. function:: array_split_into_chunks(array(T), int) -> array(array(T)) Returns an ``array`` of arrays splitting the input ``array`` into chunks of given length. diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInTypeAndFunctionNamespaceManager.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInTypeAndFunctionNamespaceManager.java index d37d658e69204..4984cd8fe065d 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInTypeAndFunctionNamespaceManager.java +++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/BuiltInTypeAndFunctionNamespaceManager.java @@ -110,6 +110,7 @@ import com.facebook.presto.operator.aggregation.sketch.kll.KllSketchAggregationFunction; import com.facebook.presto.operator.aggregation.sketch.kll.KllSketchWithKAggregationFunction; import com.facebook.presto.operator.aggregation.sketch.theta.ThetaSketchAggregationFunction; +import com.facebook.presto.operator.scalar.AbstractArraySortByKeyFunction; import com.facebook.presto.operator.scalar.ArrayAllMatchFunction; import com.facebook.presto.operator.scalar.ArrayAnyMatchFunction; import com.facebook.presto.operator.scalar.ArrayCardinalityFunction; @@ -882,6 +883,8 @@ private List getBuiltInFunctions(FunctionsConfig function .scalar(ArrayGreaterThanOrEqualOperator.class) .scalar(ArrayElementAtFunction.class) .scalar(ArraySortFunction.class) + .function(AbstractArraySortByKeyFunction.ArraySortByKeyFunction.ARRAY_SORT_BY_KEY_FUNCTION) + .function(AbstractArraySortByKeyFunction.ArraySortDescByKeyFunction.ARRAY_SORT_DESC_BY_KEY_FUNCTION) .scalar(MapSubsetFunction.class) .scalar(ArraySortComparatorFunction.class) .scalar(ArrayShuffleFunction.class) diff --git a/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/AbstractArraySortByKeyFunction.java b/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/AbstractArraySortByKeyFunction.java new file mode 100644 index 0000000000000..30313bca7d438 --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/operator/scalar/AbstractArraySortByKeyFunction.java @@ -0,0 +1,420 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.operator.scalar; + +import com.facebook.presto.bytecode.BytecodeBlock; +import com.facebook.presto.bytecode.CallSiteBinder; +import com.facebook.presto.bytecode.ClassDefinition; +import com.facebook.presto.bytecode.MethodDefinition; +import com.facebook.presto.bytecode.Parameter; +import com.facebook.presto.bytecode.Scope; +import com.facebook.presto.bytecode.Variable; +import com.facebook.presto.bytecode.control.IfStatement; +import com.facebook.presto.common.NotSupportedException; +import com.facebook.presto.common.QualifiedObjectName; +import com.facebook.presto.common.block.Block; +import com.facebook.presto.common.block.BlockBuilder; +import com.facebook.presto.common.function.SqlFunctionProperties; +import com.facebook.presto.common.type.Type; +import com.facebook.presto.metadata.BoundVariables; +import com.facebook.presto.metadata.FunctionAndTypeManager; +import com.facebook.presto.metadata.SqlScalarFunction; +import com.facebook.presto.spi.PrestoException; +import com.facebook.presto.spi.function.ComplexTypeFunctionDescriptor; +import com.facebook.presto.spi.function.FunctionKind; +import com.facebook.presto.spi.function.LambdaArgumentDescriptor; +import com.facebook.presto.spi.function.LambdaDescriptor; +import com.facebook.presto.spi.function.Signature; +import com.facebook.presto.spi.function.SqlFunctionVisibility; +import com.facebook.presto.sql.gen.lambda.UnaryFunctionInterface; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; +import com.google.common.primitives.Primitives; +import it.unimi.dsi.fastutil.ints.IntComparator; + +import java.lang.invoke.MethodHandle; +import java.lang.invoke.MethodHandles; +import java.util.Optional; + +import static com.facebook.presto.bytecode.Access.FINAL; +import static com.facebook.presto.bytecode.Access.PUBLIC; +import static com.facebook.presto.bytecode.Access.a; +import static com.facebook.presto.bytecode.Parameter.arg; +import static com.facebook.presto.bytecode.ParameterizedType.type; +import static com.facebook.presto.bytecode.expression.BytecodeExpressions.constantNull; +import static com.facebook.presto.bytecode.expression.BytecodeExpressions.equal; +import static com.facebook.presto.common.type.TypeSignature.parseTypeSignature; +import static com.facebook.presto.common.type.UnknownType.UNKNOWN; +import static com.facebook.presto.metadata.BuiltInTypeAndFunctionNamespaceManager.JAVA_BUILTIN_NAMESPACE; +import static com.facebook.presto.operator.scalar.ScalarFunctionImplementationChoice.ArgumentProperty.functionTypeArgumentProperty; +import static com.facebook.presto.operator.scalar.ScalarFunctionImplementationChoice.ArgumentProperty.valueTypeArgumentProperty; +import static com.facebook.presto.operator.scalar.ScalarFunctionImplementationChoice.NullConvention.RETURN_NULL_ON_NULL; +import static com.facebook.presto.spi.StandardErrorCode.INVALID_FUNCTION_ARGUMENT; +import static com.facebook.presto.spi.StandardErrorCode.NOT_SUPPORTED; +import static com.facebook.presto.spi.function.Signature.typeVariable; +import static com.facebook.presto.sql.gen.SqlTypeBytecodeExpression.constantType; +import static com.facebook.presto.util.CompilerUtils.defineClass; +import static com.facebook.presto.util.CompilerUtils.makeClassName; +import static com.facebook.presto.util.Reflection.methodHandle; +import static it.unimi.dsi.fastutil.ints.IntArrays.quickSort; + +public abstract class AbstractArraySortByKeyFunction + extends SqlScalarFunction +{ + private final ComplexTypeFunctionDescriptor descriptor; + + protected AbstractArraySortByKeyFunction(String functionName) + { + super(new Signature( + QualifiedObjectName.valueOf(JAVA_BUILTIN_NAMESPACE, functionName), + FunctionKind.SCALAR, + ImmutableList.of(typeVariable("T"), typeVariable("K")), + ImmutableList.of(), + parseTypeSignature("array(T)"), + ImmutableList.of(parseTypeSignature("array(T)"), parseTypeSignature("function(T,K)")), + false)); + descriptor = new ComplexTypeFunctionDescriptor( + true, + ImmutableList.of(new LambdaDescriptor(1, ImmutableMap.of(0, new LambdaArgumentDescriptor(0, ComplexTypeFunctionDescriptor::prependAllSubscripts)))), + Optional.of(ImmutableSet.of(0)), + Optional.of(ComplexTypeFunctionDescriptor::clearRequiredSubfields), + getSignature()); + } + + @Override + public SqlFunctionVisibility getVisibility() + { + return SqlFunctionVisibility.PUBLIC; + } + + @Override + public boolean isDeterministic() + { + return true; + } + + @Override + public BuiltInScalarFunctionImplementation specialize(BoundVariables boundVariables, int arity, FunctionAndTypeManager functionAndTypeManager) + { + Type elementType = boundVariables.getTypeVariable("T"); + Type keyType = boundVariables.getTypeVariable("K"); + + // Generate the specialized key extractor instance once + KeyExtractor keyExtractor = generateKeyExtractor(elementType, keyType); + + MethodHandle raw = methodHandle( + AbstractArraySortByKeyFunction.class, + "sortByKey", + AbstractArraySortByKeyFunction.class, + Type.class, + Type.class, + KeyExtractor.class, + SqlFunctionProperties.class, + Block.class, + UnaryFunctionInterface.class); + + MethodHandle bound = MethodHandles.insertArguments(raw, 0, this, elementType, keyType, keyExtractor); + + return new BuiltInScalarFunctionImplementation( + false, + ImmutableList.of( + valueTypeArgumentProperty(RETURN_NULL_ON_NULL), // array parameter + functionTypeArgumentProperty(UnaryFunctionInterface.class)), // keyFunction parameter + bound); + } + + @Override + public ComplexTypeFunctionDescriptor getComplexTypeFunctionDescriptor() + { + return descriptor; + } + + public static Block sortByKey( + AbstractArraySortByKeyFunction function, + Type elementType, + Type keyType, + KeyExtractor keyExtractor, + SqlFunctionProperties properties, + Block array, + UnaryFunctionInterface keyFunction) + { + int arrayLength = array.getPositionCount(); + if (arrayLength < 2) { + return array; + } + + // Create array of indices and extracted keys + int[] indices = new int[arrayLength]; + BlockBuilder keyBlockBuilder = keyType.createBlockBuilder(null, arrayLength); + + // Extract keys for all elements + for (int i = 0; i < arrayLength; i++) { + indices[i] = i; + if (array.isNull(i)) { + keyBlockBuilder.appendNull(); + } + else { + try { + // Use the generated KeyExtractor implementation (direct virtual call) + keyExtractor.extract(properties, array, i, keyFunction, keyBlockBuilder); + } + catch (Throwable t) { + throw new PrestoException(INVALID_FUNCTION_ARGUMENT, String.format("Error applying key function to element at position %d", i), t); + } + } + } + + Block keysBlock = keyBlockBuilder.build(); + + // Sort indices based on extracted keys using Type's compareTo + try { + if (array.mayHaveNull() || keysBlock.mayHaveNull()) { + quickSort(indices, new NullableComparator(array, keysBlock, keyType, function)); + } + else { + quickSort(indices, new NonNullableComparator(keysBlock, keyType, function)); + } + } + catch (NotSupportedException | UnsupportedOperationException e) { + throw new PrestoException(INVALID_FUNCTION_ARGUMENT, "Key type does not support comparison", e); + } + catch (PrestoException e) { + if (e.getErrorCode() == NOT_SUPPORTED.toErrorCode()) { + throw new PrestoException(INVALID_FUNCTION_ARGUMENT, "Key type does not support comparison", e); + } + throw e; + } + + // Build result block with sorted elements + BlockBuilder resultBuilder = elementType.createBlockBuilder(null, arrayLength); + for (int i = 0; i < arrayLength; i++) { + elementType.appendTo(array, indices[i], resultBuilder); + } + + return resultBuilder.build(); + } + + /** + * KeyExtractor is a simple interface implemented by generated classes. + * Implementations must write the extracted key into the provided BlockBuilder + * (or appendNull) for the given position. + */ + public interface KeyExtractor + { + void extract(SqlFunctionProperties properties, Block array, int position, UnaryFunctionInterface keyFunction, BlockBuilder keyBlockBuilder) throws Throwable; + } + + // Generate just the key extraction logic + public static KeyExtractor generateKeyExtractor(Type elementType, Type keyType) + { + CallSiteBinder binder = new CallSiteBinder(); + Class elementJavaType = Primitives.wrap(elementType.getJavaType()); + Class keyJavaType = Primitives.wrap(keyType.getJavaType()); + + String className = "ArraySortKeyExtractorImpl_" + elementType.getTypeSignature() + "_" + keyType.getTypeSignature(); + ClassDefinition definition = new ClassDefinition( + a(PUBLIC, FINAL), + makeClassName(className), + type(Object.class), + type(KeyExtractor.class)); + definition.declareDefaultConstructor(a(PUBLIC)); + + Parameter properties = arg("properties", SqlFunctionProperties.class); + Parameter array = arg("array", Block.class); + Parameter position = arg("position", int.class); + Parameter keyFunction = arg("keyFunction", UnaryFunctionInterface.class); + Parameter keyBlockBuilder = arg("keyBlockBuilder", BlockBuilder.class); + + MethodDefinition method = definition.declareMethod( + a(PUBLIC), + "extract", + type(void.class), + ImmutableList.of(properties, array, position, keyFunction, keyBlockBuilder)); + + BytecodeBlock body = method.getBody(); + Scope scope = method.getScope(); + Variable element = scope.declareVariable(elementJavaType, "element"); + Variable key = scope.declareVariable(keyJavaType, "key"); + + // Load element with correct primitive handling + if (!elementType.equals(UNKNOWN)) { + // generates the correct getLong/getDouble/getBoolean/getSlice/getObject call + body.append(element.set(constantType(binder, elementType).getValue(array, position).cast(elementJavaType))); + } + else { + body.append(element.set(constantNull(elementJavaType))); + } + + body.append(key.set(keyFunction.invoke("apply", Object.class, element.cast(Object.class)).cast(keyJavaType))); + + // Write the key to the block builder + if (!keyType.equals(UNKNOWN)) { + body.append(new IfStatement() + .condition(equal(key, constantNull(keyJavaType))) + .ifTrue(keyBlockBuilder.invoke("appendNull", BlockBuilder.class).pop()) + .ifFalse(constantType(binder, keyType).writeValue(keyBlockBuilder, key.cast(keyType.getJavaType())))); + } + else { + body.append(keyBlockBuilder.invoke("appendNull", BlockBuilder.class).pop()); + } + + body.ret(); + + Class generatedClass = defineClass(definition, Object.class, binder.getBindings(), AbstractArraySortByKeyFunction.class.getClassLoader()); + + try { + // instantiate generated class and cast to KeyExtractor for direct virtual call + return (KeyExtractor) generatedClass.getConstructor().newInstance(); + } + catch (ReflectiveOperationException e) { + throw new RuntimeException("Failed to instantiate generated key extractor", e); + } + } + + // Abstract method to be implemented by subclasses to define comparison direction + protected abstract int compareKeys(Type keyType, Block keysBlock, int leftIndex, int rightIndex); + + private static class NullableComparator + implements IntComparator + { + private final Block array; + private final Block keysBlock; + private final Type keyType; + private final AbstractArraySortByKeyFunction function; + + public NullableComparator(Block array, Block keysBlock, Type keyType, AbstractArraySortByKeyFunction function) + { + this.array = array; + this.keysBlock = keysBlock; + this.keyType = keyType; + this.function = function; + } + + @Override + public int compare(int leftIndex, int rightIndex) + { + boolean leftArrayNull = array.isNull(leftIndex); + boolean rightArrayNull = array.isNull(rightIndex); + + if (leftArrayNull && rightArrayNull) { + return 0; + } + if (leftArrayNull) { + return 1; + } + if (rightArrayNull) { + return -1; + } + + boolean leftKeyNull = keysBlock.isNull(leftIndex); + boolean rightKeyNull = keysBlock.isNull(rightIndex); + + if (leftKeyNull && rightKeyNull) { + return 0; + } + if (leftKeyNull) { + return 1; + } + if (rightKeyNull) { + return -1; + } + + int result = function.compareKeys(keyType, keysBlock, leftIndex, rightIndex); + + // If keys are equal, maintain original order + if (result == 0) { + return Integer.compare(leftIndex, rightIndex); + } + + return result; + } + } + + private static class NonNullableComparator + implements IntComparator + { + private final Block keysBlock; + private final Type keyType; + private final AbstractArraySortByKeyFunction function; + + public NonNullableComparator(Block keysBlock, Type keyType, AbstractArraySortByKeyFunction function) + { + this.keysBlock = keysBlock; + this.keyType = keyType; + this.function = function; + } + + @Override + public int compare(int leftIndex, int rightIndex) + { + int result = function.compareKeys(keyType, keysBlock, leftIndex, rightIndex); + + // If keys are equal, maintain original order + if (result == 0) { + return Integer.compare(leftIndex, rightIndex); + } + + return result; + } + } + + public static class ArraySortByKeyFunction + extends AbstractArraySortByKeyFunction + { + public static final ArraySortByKeyFunction ARRAY_SORT_BY_KEY_FUNCTION = new ArraySortByKeyFunction(); + + private ArraySortByKeyFunction() + { + super("array_sort"); + } + + @Override + public String getDescription() + { + return "Sorts the given array using a lambda function to extract sorting keys. " + + "Null array elements and null keys are placed at the end. " + + "Example: array_sort(ARRAY['apple', 'banana', 'cherry'], x -> length(x))"; + } + + @Override + protected int compareKeys(Type keyType, Block keysBlock, int leftIndex, int rightIndex) + { + return keyType.compareTo(keysBlock, leftIndex, keysBlock, rightIndex); + } + } + + public static class ArraySortDescByKeyFunction + extends AbstractArraySortByKeyFunction + { + public static final ArraySortDescByKeyFunction ARRAY_SORT_DESC_BY_KEY_FUNCTION = new ArraySortDescByKeyFunction(); + + private ArraySortDescByKeyFunction() + { + super("array_sort_desc"); + } + + @Override + public String getDescription() + { + return "Sorts the given array in descending order using a lambda function to extract sorting keys"; + } + + @Override + protected int compareKeys(Type keyType, Block keysBlock, int leftIndex, int rightIndex) + { + return keyType.compareTo(keysBlock, rightIndex, keysBlock, leftIndex); + } + } +} diff --git a/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestArraySortByKeyFunctions.java b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestArraySortByKeyFunctions.java new file mode 100644 index 0000000000000..dee4f22e42a96 --- /dev/null +++ b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/TestArraySortByKeyFunctions.java @@ -0,0 +1,309 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.operator.scalar; + +import com.facebook.presto.common.type.ArrayType; +import com.facebook.presto.common.type.RowType; +import com.google.common.collect.ImmutableList; +import org.testng.annotations.Test; + +import static com.facebook.presto.common.type.BigintType.BIGINT; +import static com.facebook.presto.common.type.BooleanType.BOOLEAN; +import static com.facebook.presto.common.type.DoubleType.DOUBLE; +import static com.facebook.presto.common.type.IntegerType.INTEGER; +import static com.facebook.presto.common.type.UnknownType.UNKNOWN; +import static com.facebook.presto.common.type.VarcharType.createVarcharType; +import static java.util.Arrays.asList; + +public class TestArraySortByKeyFunctions + extends AbstractTestFunctions +{ + @Test + public void testBasic() + { + // Test array_sort + assertFunction( + "array_sort(ARRAY['pear', 'apple', 'banana', 'kiwi'], x -> length(x))", + new ArrayType(createVarcharType(6)), + asList("pear", "kiwi", "apple", "banana")); + + assertFunction( + "array_sort(ARRAY['pear', 'apple', 'banana', 'kiwi'], x -> substr(x, length(x), 1))", + new ArrayType(createVarcharType(6)), + asList("banana", "apple", "kiwi", "pear")); + + // Test array_sort_desc + assertFunction( + "array_sort_desc(ARRAY['pear', 'apple', 'banana', 'kiwi'], x -> length(x))", + new ArrayType(createVarcharType(6)), + asList("banana", "apple", "pear", "kiwi")); + + assertFunction( + "array_sort_desc(ARRAY['pear', 'apple', 'banana', 'kiwi'], x -> substr(x, length(x), 1))", + new ArrayType(createVarcharType(6)), + asList("pear", "kiwi", "apple", "banana")); + } + + @Test + public void testNulls() + { + // Test array_sort + assertFunction( + "array_sort(ARRAY['apple', NULL, 'banana', NULL], x -> length(x))", + new ArrayType(createVarcharType(6)), + asList("apple", "banana", null, null)); + + assertFunction( + "array_sort(ARRAY['apple', 'banana', 'pear'], x -> IF(x = 'banana', NULL, length(x)))", + new ArrayType(createVarcharType(6)), + asList("pear", "apple", "banana")); + + assertFunction( + "array_sort(ARRAY['apple', NULL, 'banana', 'pear', NULL], x -> length(x))", + new ArrayType(createVarcharType(6)), + asList("pear", "apple", "banana", null, null)); + + // Test array_sort_desc + assertFunction( + "array_sort_desc(ARRAY['apple', NULL, 'banana', NULL], x -> length(x))", + new ArrayType(createVarcharType(6)), + asList("banana", "apple", null, null)); + + assertFunction( + "array_sort_desc(ARRAY['apple', 'banana', 'pear'], x -> IF(x = 'banana', NULL, length(x)))", + new ArrayType(createVarcharType(6)), + asList("apple", "pear", "banana")); + + assertFunction( + "array_sort_desc(ARRAY['apple', NULL, 'banana', 'pear', NULL], x -> length(x))", + new ArrayType(createVarcharType(6)), + asList("banana", "apple", "pear", null, null)); + } + + @Test + public void testSpecialDoubleValues() + { + // Test array_sort + assertFunction( + "array_sort(ARRAY[CAST(0.0 AS DOUBLE), CAST('NaN' AS DOUBLE), CAST('Infinity' AS DOUBLE), CAST('-Infinity' AS DOUBLE)], x -> x)", + new ArrayType(DOUBLE), + asList(Double.NEGATIVE_INFINITY, 0.0, Double.POSITIVE_INFINITY, Double.NaN)); + + // Test array_sort_desc + assertFunction( + "array_sort_desc(ARRAY[CAST(0.0 AS DOUBLE), CAST('NaN' AS DOUBLE), CAST('Infinity' AS DOUBLE), CAST('-Infinity' AS DOUBLE)], x -> x)", + new ArrayType(DOUBLE), + asList(Double.NaN, Double.POSITIVE_INFINITY, 0.0, Double.NEGATIVE_INFINITY)); + } + + @Test + public void testNumericKeys() + { + // Test array_sort + assertFunction( + "array_sort(ARRAY[5, 20, 3, 9, 100], x -> x)", + new ArrayType(INTEGER), + asList(3, 5, 9, 20, 100)); + + assertFunction( + "array_sort(ARRAY[CAST(5000000000 AS BIGINT), CAST(20000000000 AS BIGINT), CAST(3000000000 AS BIGINT), CAST(9000000000 AS BIGINT), CAST(100000000000 AS BIGINT)], x -> x)", + new ArrayType(BIGINT), + asList(3000000000L, 5000000000L, 9000000000L, 20000000000L, 100000000000L)); + + assertFunction( + "array_sort(ARRAY[CAST(5.5 AS DOUBLE), CAST(20.1 AS DOUBLE), CAST(3.9 AS DOUBLE), CAST(9.0 AS DOUBLE), CAST(100.0 AS DOUBLE)], x -> x)", + new ArrayType(DOUBLE), + asList(3.9, 5.5, 9.0, 20.1, 100.0)); + + assertFunction( + "array_sort(ARRAY[5, 20, 3, 9, 100], x -> x % 10)", + new ArrayType(INTEGER), + asList(20, 100, 3, 5, 9)); + + assertFunction( + "array_sort(ARRAY[CAST(5000000000 AS BIGINT), CAST(20000000000 AS BIGINT), CAST(3000000000 AS BIGINT), CAST(9000000000 AS BIGINT), CAST(100000000000 AS BIGINT)], x -> x % CAST(10000000000 AS BIGINT))", + new ArrayType(BIGINT), + asList(20000000000L, 100000000000L, 3000000000L, 5000000000L, 9000000000L)); + + // Test array_sort_desc + assertFunction( + "array_sort_desc(ARRAY[5, 20, 3, 9, 100], x -> x)", + new ArrayType(INTEGER), + asList(100, 20, 9, 5, 3)); + + assertFunction( + "array_sort_desc(ARRAY[CAST(5000000000 AS BIGINT), CAST(20000000000 AS BIGINT), CAST(3000000000 AS BIGINT), CAST(9000000000 AS BIGINT), CAST(100000000000 AS BIGINT)], x -> x)", + new ArrayType(BIGINT), + asList(100000000000L, 20000000000L, 9000000000L, 5000000000L, 3000000000L)); + + assertFunction( + "array_sort_desc(ARRAY[CAST(5.5 AS DOUBLE), CAST(20.1 AS DOUBLE), CAST(3.9 AS DOUBLE), CAST(9.0 AS DOUBLE), CAST(100.0 AS DOUBLE)], x -> x)", + new ArrayType(DOUBLE), + asList(100.0, 20.1, 9.0, 5.5, 3.9)); + + assertFunction( + "array_sort_desc(ARRAY[5, 20, 3, 9, 100], x -> x % 10)", + new ArrayType(INTEGER), + asList(9, 5, 3, 20, 100)); + + assertFunction( + "array_sort_desc(ARRAY[CAST(5000000000 AS BIGINT), CAST(20000000000 AS BIGINT), CAST(3000000000 AS BIGINT), CAST(9000000000 AS BIGINT), CAST(100000000000 AS BIGINT)], x -> x % CAST(10000000000 AS BIGINT))", + new ArrayType(BIGINT), + asList(9000000000L, 5000000000L, 3000000000L, 20000000000L, 100000000000L)); + } + + @Test + public void testBooleanKeys() + { + // Test array_sort + assertFunction( + "array_sort(ARRAY[true, false, true, false], x -> x)", + new ArrayType(BOOLEAN), + asList(false, false, true, true)); + + assertFunction( + "array_sort(ARRAY[true, false, true, false], x -> NOT x)", + new ArrayType(BOOLEAN), + asList(true, true, false, false)); + + // Test array_sort_desc + assertFunction( + "array_sort_desc(ARRAY[true, false, true, false], x -> x)", + new ArrayType(BOOLEAN), + asList(true, true, false, false)); + + assertFunction( + "array_sort_desc(ARRAY[true, false, true, false], x -> NOT x)", + new ArrayType(BOOLEAN), + asList(false, false, true, true)); + } + + @Test + public void testComplexTypes() + { + // Test array_sort + assertFunction( + "array_sort(ARRAY[ARRAY[1, 2, 3], ARRAY[4, 5], ARRAY[6, 7, 8, 9]], x -> cardinality(x))", + new ArrayType(new ArrayType(INTEGER)), + asList(asList(4, 5), asList(1, 2, 3), asList(6, 7, 8, 9))); + + assertFunction( + "array_sort(ARRAY[ROW('a', 3), ROW('b', 1), ROW('c', 2)], x -> x[2])", + new ArrayType(RowType.anonymous(ImmutableList.of(createVarcharType(1), INTEGER))), + asList(asList("b", 1), asList("c", 2), asList("a", 3))); + + assertFunction( + "array_sort(ARRAY[ROW('a', CAST(3000000000 AS BIGINT)), ROW('b', CAST(1000000000 AS BIGINT)), ROW('c', CAST(2000000000 AS BIGINT))], x -> x[2])", + new ArrayType(RowType.anonymous(ImmutableList.of(createVarcharType(1), BIGINT))), + asList(asList("b", 1000000000L), asList("c", 2000000000L), asList("a", 3000000000L))); + + // Test array_sort_desc + assertFunction( + "array_sort_desc(ARRAY[ARRAY[1, 2, 3], ARRAY[4, 5], ARRAY[6, 7, 8, 9]], x -> cardinality(x))", + new ArrayType(new ArrayType(INTEGER)), + asList(asList(6, 7, 8, 9), asList(1, 2, 3), asList(4, 5))); + + assertFunction( + "array_sort_desc(ARRAY[ROW('a', 3), ROW('b', 1), ROW('c', 2)], x -> x[2])", + new ArrayType(RowType.anonymous(ImmutableList.of(createVarcharType(1), INTEGER))), + asList(asList("a", 3), asList("c", 2), asList("b", 1))); + + assertFunction( + "array_sort_desc(ARRAY[ROW('a', CAST(3000000000 AS BIGINT)), ROW('b', CAST(1000000000 AS BIGINT)), ROW('c', CAST(2000000000 AS BIGINT))], x -> x[2])", + new ArrayType(RowType.anonymous(ImmutableList.of(createVarcharType(1), BIGINT))), + asList(asList("a", 3000000000L), asList("c", 2000000000L), asList("b", 1000000000L))); + } + + @Test + public void testEdgeCases() + { + // Test array_sort + assertFunction( + "array_sort(ARRAY[], x -> x)", + new ArrayType(UNKNOWN), + ImmutableList.of()); + + assertFunction( + "array_sort(ARRAY[5], x -> x)", + new ArrayType(INTEGER), + asList(5)); + + assertFunction( + "array_sort(ARRAY[NULL, NULL, NULL], x -> x)", + new ArrayType(UNKNOWN), + asList(null, null, null)); + + // Test array_sort_desc + assertFunction( + "array_sort_desc(ARRAY[], x -> x)", + new ArrayType(UNKNOWN), + ImmutableList.of()); + + assertFunction( + "array_sort_desc(ARRAY[5], x -> x)", + new ArrayType(INTEGER), + asList(5)); + + assertFunction( + "array_sort_desc(ARRAY[NULL, NULL, NULL], x -> x)", + new ArrayType(UNKNOWN), + asList(null, null, null)); + } + + @Test + public void testTypeCoercion() + { + // Test array_sort + assertFunction( + "array_sort(ARRAY[5, 20, 3, 9, 100], x -> x + CAST(0.5 AS DOUBLE))", + new ArrayType(INTEGER), + asList(3, 5, 9, 20, 100)); + + assertFunction( + "array_sort(ARRAY[5, 20, 3, 9, 100], x -> x * CAST(1000000000 AS BIGINT))", + new ArrayType(INTEGER), + asList(3, 5, 9, 20, 100)); + + assertFunction( + "array_sort(ARRAY['5', '20', '3', '9', '100'], x -> cast(x as integer))", + new ArrayType(createVarcharType(3)), + asList("3", "5", "9", "20", "100")); + + assertFunction( + "array_sort(ARRAY['5000000000', '20000000000', '3000000000', '9000000000', '100000000000'], x -> cast(x as bigint))", + new ArrayType(createVarcharType(12)), + asList("3000000000", "5000000000", "9000000000", "20000000000", "100000000000")); + + // Test array_sort_desc + assertFunction( + "array_sort_desc(ARRAY[5, 20, 3, 9, 100], x -> x + CAST(0.5 AS DOUBLE))", + new ArrayType(INTEGER), + asList(100, 20, 9, 5, 3)); + + assertFunction( + "array_sort_desc(ARRAY[5, 20, 3, 9, 100], x -> x * CAST(1000000000 AS BIGINT))", + new ArrayType(INTEGER), + asList(100, 20, 9, 5, 3)); + + assertFunction( + "array_sort_desc(ARRAY['5', '20', '3', '9', '100'], x -> cast(x as integer))", + new ArrayType(createVarcharType(3)), + asList("100", "20", "9", "5", "3")); + + assertFunction( + "array_sort_desc(ARRAY['5000000000', '20000000000', '3000000000', '9000000000', '100000000000'], x -> cast(x as bigint))", + new ArrayType(createVarcharType(12)), + asList("100000000000", "20000000000", "9000000000", "5000000000", "3000000000")); + } +} diff --git a/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java b/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java index 4cd909d6464f4..d3f584222c3b9 100644 --- a/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java +++ b/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPlugin.java @@ -289,6 +289,70 @@ public void testLambdaFunctions() " Exception 2: line 1:31: Expected a lambda that takes ([12])" + Pattern.quote(" argument(s) but got 3\n")); } + @Test + public void testArraySortByKeyFunction() + { + // Basic string sorting by length + assertQuerySucceeds("SELECT array_sort(ARRAY['pear', 'apple', 'banana', 'kiwi'], x -> length(x))"); + assertQuerySucceeds("SELECT array_sort(ARRAY['pear', 'apple', 'banana', 'kiwi'], x -> substr(x, length(x), 1))"); + // Sorting with nulls + assertQuerySucceeds("SELECT array_sort(ARRAY['apple', NULL, 'banana', NULL], x -> length(x))"); + assertQuerySucceeds("SELECT array_sort(ARRAY['apple', 'banana', 'pear'], x -> IF(x = 'banana', NULL, length(x)))"); + assertQuerySucceeds("SELECT array_sort(ARRAY['apple', NULL, 'banana', 'pear', NULL], x -> length(x))"); + // Special double values + assertQuerySucceeds("SELECT array_sort(ARRAY[CAST(0.0 AS DOUBLE), CAST('NaN' AS DOUBLE), CAST('Infinity' AS DOUBLE), CAST('-Infinity' AS DOUBLE)], x -> x)"); + // Numeric keys + assertQuerySucceeds("SELECT array_sort(ARRAY[5, 20, 3, 9, 100], x -> x)"); + assertQuerySucceeds("SELECT array_sort(ARRAY[CAST(5000000000 AS BIGINT), CAST(20000000000 AS BIGINT), CAST(3000000000 AS BIGINT), CAST(9000000000 AS BIGINT), CAST(100000000000 AS BIGINT)], x -> x)"); + assertQuerySucceeds("SELECT array_sort(ARRAY[CAST(5.5 AS DOUBLE), CAST(20.1 AS DOUBLE), CAST(3.9 AS DOUBLE), CAST(9.0 AS DOUBLE), CAST(100.0 AS DOUBLE)], x -> x)"); + assertQuerySucceeds("SELECT array_sort(ARRAY[5, 20, 3, 9, 100], x -> x % 10)"); + // Boolean keys + assertQuerySucceeds("SELECT array_sort(ARRAY[true, false, true, false], x -> x)"); + assertQuerySucceeds("SELECT array_sort(ARRAY[true, false, true, false], x -> NOT x)"); + // Complex types + assertQuerySucceeds("SELECT array_sort(ARRAY[ARRAY[1, 2, 3], ARRAY[4, 5], ARRAY[6, 7, 8, 9]], x -> cardinality(x))"); + assertQuerySucceeds("SELECT array_sort(ARRAY[ROW('a', 3), ROW('b', 1), ROW('c', 2)], x -> x[2])"); + // Edge cases + assertQuerySucceeds("SELECT array_sort(ARRAY[], x -> x)"); + assertQuerySucceeds("SELECT array_sort(ARRAY[5], x -> x)"); + assertQuerySucceeds("SELECT array_sort(ARRAY[NULL, NULL, NULL], x -> x)"); + // Type coercion + assertQuerySucceeds("SELECT array_sort(ARRAY[5, 20, 3, 9, 100], x -> x + CAST(0.5 AS DOUBLE))"); + assertQuerySucceeds("SELECT array_sort(ARRAY[5, 20, 3, 9, 100], x -> x * CAST(1000000000 AS BIGINT))"); + } + + @Test + public void testArraySortDescByKeyFunction() + { + // Basic string sorting by length in descending order + assertQuerySucceeds("SELECT array_sort_desc(ARRAY['pear', 'apple', 'banana', 'kiwi'], x -> length(x))"); + assertQuerySucceeds("SELECT array_sort_desc(ARRAY['pear', 'apple', 'banana', 'kiwi'], x -> substr(x, length(x), 1))"); + // Sorting with nulls + assertQuerySucceeds("SELECT array_sort_desc(ARRAY['apple', NULL, 'banana', NULL], x -> length(x))"); + assertQuerySucceeds("SELECT array_sort_desc(ARRAY['apple', 'banana', 'pear'], x -> IF(x = 'banana', NULL, length(x)))"); + assertQuerySucceeds("SELECT array_sort_desc(ARRAY['apple', NULL, 'banana', 'pear', NULL], x -> length(x))"); + // Special double values + assertQuerySucceeds("SELECT array_sort_desc(ARRAY[CAST(0.0 AS DOUBLE), CAST('NaN' AS DOUBLE), CAST('Infinity' AS DOUBLE), CAST('-Infinity' AS DOUBLE)], x -> x)"); + // Numeric keys + assertQuerySucceeds("SELECT array_sort_desc(ARRAY[5, 20, 3, 9, 100], x -> x)"); + assertQuerySucceeds("SELECT array_sort_desc(ARRAY[CAST(5000000000 AS BIGINT), CAST(20000000000 AS BIGINT), CAST(3000000000 AS BIGINT), CAST(9000000000 AS BIGINT), CAST(100000000000 AS BIGINT)], x -> x)"); + assertQuerySucceeds("SELECT array_sort_desc(ARRAY[CAST(5.5 AS DOUBLE), CAST(20.1 AS DOUBLE), CAST(3.9 AS DOUBLE), CAST(9.0 AS DOUBLE), CAST(100.0 AS DOUBLE)], x -> x)"); + assertQuerySucceeds("SELECT array_sort_desc(ARRAY[5, 20, 3, 9, 100], x -> x % 10)"); + // Boolean keys + assertQuerySucceeds("SELECT array_sort_desc(ARRAY[true, false, true, false], x -> x)"); + assertQuerySucceeds("SELECT array_sort_desc(ARRAY[true, false, true, false], x -> NOT x)"); + // Complex types + assertQuerySucceeds("SELECT array_sort_desc(ARRAY[ARRAY[1, 2, 3], ARRAY[4, 5], ARRAY[6, 7, 8, 9]], x -> cardinality(x))"); + assertQuerySucceeds("SELECT array_sort_desc(ARRAY[ROW('a', 3), ROW('b', 1), ROW('c', 2)], x -> x[2])"); + // Edge cases + assertQuerySucceeds("SELECT array_sort_desc(ARRAY[], x -> x)"); + assertQuerySucceeds("SELECT array_sort_desc(ARRAY[5], x -> x)"); + assertQuerySucceeds("SELECT array_sort_desc(ARRAY[NULL, NULL, NULL], x -> x)"); + // Type coercion + assertQuerySucceeds("SELECT array_sort_desc(ARRAY[5, 20, 3, 9, 100], x -> x + CAST(0.5 AS DOUBLE))"); + assertQuerySucceeds("SELECT array_sort_desc(ARRAY[5, 20, 3, 9, 100], x -> x * CAST(1000000000 AS BIGINT))"); + } + @Test public void testApproxPercentile() { diff --git a/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPluginWithoutLoadingFunctionalities.java b/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPluginWithoutLoadingFunctionalities.java index 01e786a4a0f66..a0997d9ab365c 100644 --- a/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPluginWithoutLoadingFunctionalities.java +++ b/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestNativeSidecarPluginWithoutLoadingFunctionalities.java @@ -75,9 +75,5 @@ public void testBasicQueries() assertQuery("select count(*) from nation"); assertQuery("select count(abs(orderkey) between 1 and 60000) from orders group by orderkey"); assertQuery("SELECT count(orderkey) FROM orders WHERE orderkey < 0 GROUP BY GROUPING SETS (())"); - // This query will work when sidecar is enabled, should fail without it. - assertQueryFails( - "select array_sort(array[row('apples', 23), row('bananas', 12), row('grapes', 44)], x -> x[2])", - "line 1:84: Expected a lambda that takes 2 argument\\(s\\) but got 1"); } } From abad3add7684bfd65bd2648776cf13fdddd648ff Mon Sep 17 00:00:00 2001 From: Dilli-Babu-Godari Date: Thu, 11 Sep 2025 14:27:59 +0530 Subject: [PATCH 098/113] Bump jdbi-core and jdbi-sqlclient to latest version Upgrade org.jdbi:jdbi3-core:3.4.0 to org.jdbi:jdbi3-core:3.49.5 org.jdbi:jdbi3-sqlobject:3.4.0 to org.jdbi:jdbi3-sqlobject:3.49.5 This upgrade will fix below vulnerabilities CVE-2024-1597, CVE-2023-32697 CVE-2023-2976, CVE-2022-41946 CVE-2022-41853,CVE-2022-31197 CVE-2022-26520,CVE-2022-23221 CVE-2022-21724,CVE-2021-42392 CVE-2020-8908, CVE-2020-13692 CVE-2018-10237. --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 31b39909c3c5c..221226ef50479 100644 --- a/pom.xml +++ b/pom.xml @@ -50,7 +50,7 @@ 0.6 1.12.782 4.12.0 - 3.4.0 + 3.49.0 19.3.0.0 ${dep.airlift.version} + 1.0.2 diff --git a/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiRecordCursors.java b/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiRecordCursors.java index dd73e684fc370..fbab341ade696 100644 --- a/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiRecordCursors.java +++ b/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiRecordCursors.java @@ -55,6 +55,7 @@ import static org.apache.hadoop.hive.serde2.ColumnProjectionUtils.READ_ALL_COLUMNS; import static org.apache.hadoop.hive.serde2.ColumnProjectionUtils.READ_COLUMN_IDS_CONF_STR; import static org.apache.hadoop.hive.serde2.ColumnProjectionUtils.READ_COLUMN_NAMES_CONF_STR; +import static org.apache.hudi.common.config.HoodieReaderConfig.FILE_GROUP_READER_ENABLED; class HudiRecordCursors { @@ -105,6 +106,7 @@ public static RecordCursor createRealtimeRecordCursor( jobConf.setBoolean(READ_ALL_COLUMNS, false); jobConf.set(READ_COLUMN_IDS_CONF_STR, join(dataColumns, HudiColumnHandle::getId)); jobConf.set(READ_COLUMN_NAMES_CONF_STR, join(dataColumns, HudiColumnHandle::getName)); + jobConf.setBoolean(FILE_GROUP_READER_ENABLED.key(), false); schema.stringPropertyNames() .forEach(name -> jobConf.set(name, schema.getProperty(name))); refineCompressionCodecs(jobConf); diff --git a/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiSplitManager.java b/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiSplitManager.java index e1f38c291aa88..c6e3cae63915f 100644 --- a/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiSplitManager.java +++ b/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiSplitManager.java @@ -45,6 +45,7 @@ import org.apache.hudi.common.table.timeline.HoodieTimeline; import org.apache.hudi.common.table.view.HoodieTableFileSystemView; import org.apache.hudi.common.util.HoodieTimer; +import org.apache.hudi.storage.StorageConfiguration; import java.io.IOException; import java.util.List; @@ -62,6 +63,7 @@ import static java.lang.String.format; import static java.util.Objects.requireNonNull; import static org.apache.hudi.common.table.view.FileSystemViewManager.createInMemoryFileSystemViewWithTimeline; +import static org.apache.hudi.hadoop.fs.HadoopFSUtils.getStorageConfWithCopy; public class HudiSplitManager implements ConnectorSplitManager @@ -104,7 +106,7 @@ public ConnectorSplitSource getSplits( HudiTableHandle table = layout.getTable(); // Retrieve and prune partitions - HoodieTimer timer = new HoodieTimer().startTimer(); + HoodieTimer timer = HoodieTimer.start(); List partitions = hudiPartitionManager.getEffectivePartitions(session, metastore, table.getSchemaTableName(), layout.getTupleDomain()); log.debug("Took %d ms to get %d partitions", timer.endTimer(), partitions.size()); if (partitions.isEmpty()) { @@ -114,10 +116,10 @@ public ConnectorSplitSource getSplits( // Load Hudi metadata ExtendedFileSystem fs = getFileSystem(session, table); HoodieMetadataConfig metadataConfig = HoodieMetadataConfig.newBuilder().enable(isHudiMetadataTableEnabled(session)).build(); - Configuration conf = fs.getConf(); + StorageConfiguration conf = getStorageConfWithCopy(fs.getConf()); HoodieTableMetaClient metaClient = HoodieTableMetaClient.builder().setConf(conf).setBasePath(table.getPath()).build(); HoodieTimeline timeline = metaClient.getActiveTimeline().getCommitsTimeline().filterCompletedInstants(); - String timestamp = timeline.lastInstant().map(HoodieInstant::getTimestamp).orElse(null); + String timestamp = timeline.lastInstant().map(HoodieInstant::requestedTime).orElse(null); if (timestamp == null) { // no completed instant for current table return new FixedSplitSource(ImmutableList.of()); diff --git a/presto-hudi/src/main/java/com/facebook/presto/hudi/split/HudiPartitionSplitGenerator.java b/presto-hudi/src/main/java/com/facebook/presto/hudi/split/HudiPartitionSplitGenerator.java index 54fd0f4bda8a0..21ca4f7c36f16 100644 --- a/presto-hudi/src/main/java/com/facebook/presto/hudi/split/HudiPartitionSplitGenerator.java +++ b/presto-hudi/src/main/java/com/facebook/presto/hudi/split/HudiPartitionSplitGenerator.java @@ -30,7 +30,6 @@ import com.facebook.presto.spi.schedule.NodeSelectionStrategy; import com.google.common.collect.ImmutableList; import org.apache.hadoop.fs.Path; -import org.apache.hudi.common.fs.FSUtils; import org.apache.hudi.common.model.FileSlice; import org.apache.hudi.common.table.view.HoodieTableFileSystemView; import org.apache.hudi.common.util.HoodieTimer; @@ -47,6 +46,8 @@ import static com.facebook.presto.hudi.HudiSplitManager.getHudiPartition; import static com.google.common.collect.ImmutableList.toImmutableList; import static java.util.Objects.requireNonNull; +import static org.apache.hudi.common.fs.FSUtils.getRelativePartitionPath; +import static org.apache.hudi.hadoop.fs.HadoopFSUtils.convertToStoragePath; /** * A runnable to take partition names from a queue of partitions to process, @@ -93,7 +94,7 @@ public HudiPartitionSplitGenerator( @Override public void run() { - HoodieTimer timer = new HoodieTimer().startTimer(); + HoodieTimer timer = HoodieTimer.start(); while (!concurrentPartitionQueue.isEmpty()) { String partitionName = concurrentPartitionQueue.poll(); if (partitionName != null) { @@ -107,7 +108,7 @@ private void generateSplitsFromPartition(String partitionName) { HudiPartition hudiPartition = getHudiPartition(metastore, metastoreContext, layout, partitionName); Path partitionPath = new Path(hudiPartition.getStorage().getLocation()); - String relativePartitionPath = FSUtils.getRelativePartitionPath(tablePath, partitionPath); + String relativePartitionPath = getRelativePartitionPath(convertToStoragePath(tablePath), convertToStoragePath(partitionPath)); Stream fileSlices = HudiTableType.MOR.equals(table.getTableType()) ? fsView.getLatestMergedFileSlicesBeforeOrOn(relativePartitionPath, latestInstant) : fsView.getLatestFileSlicesBeforeOrOn(relativePartitionPath, latestInstant, false); From de552585313abee4e4f5a055d57628b1bd5ec2ed Mon Sep 17 00:00:00 2001 From: Zac Wen Date: Mon, 15 Sep 2025 22:11:24 -0700 Subject: [PATCH 102/113] [native] Add protocol support for non-index-lookup-join-condition --- .../main/types/PrestoToVeloxQueryPlan.cpp | 166 ++++++++++++------ .../main/types/PrestoToVeloxQueryPlan.h | 27 ++- .../tests/PrestoToVeloxQueryPlanTest.cpp | 61 +++++++ presto-native-execution/velox | 2 +- 4 files changed, 189 insertions(+), 67 deletions(-) diff --git a/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.cpp b/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.cpp index 6c6efc582ac04..04dd0b48f3cc6 100644 --- a/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.cpp +++ b/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.cpp @@ -419,10 +419,9 @@ core::PlanNodePtr VeloxQueryPlanConverterBase::toVeloxQueryPlan( const auto desiredSourceOutput = toRowType(node->inputs[i], typeParser_); - for (auto j = 0; j < outputType->size(); j++) { - projections.emplace_back( - std::make_shared( - outputType->childAt(j), desiredSourceOutput->nameOf(j))); + for (auto j = 0; j < outputType->size(); ++j) { + projections.emplace_back(std::make_shared( + outputType->childAt(j), desiredSourceOutput->nameOf(j))); } sourceNodes[i] = std::make_shared( @@ -574,6 +573,13 @@ std::shared_ptr isAnd( return isSpecialForm(expression, protocol::Form::AND); } +// Check if input RowExpression is an 'or' expression and returns it as +// SpecialFormExpression. Returns nullptr if input expression is something else. +std::shared_ptr isOr( + const std::shared_ptr& expression) { + return isSpecialForm(expression, protocol::Form::OR); +} + // Checks if input PlanNode represents a local exchange with single source and // returns it as ExchangeNode. Returns nullptr if input node is something else. std::shared_ptr isLocalSingleSourceExchange( @@ -662,9 +668,8 @@ core::PlanNodePtr VeloxQueryPlanConverterBase::toVeloxQueryPlan( std::vector projections; projections.reserve(leftNames.size() + 1); for (auto i = 0; i < leftNames.size(); i++) { - projections.emplace_back( - std::make_shared( - leftTypes[i], leftNames[i])); + projections.emplace_back(std::make_shared( + leftTypes[i], leftNames[i])); } const bool constantValue = joinType.value() == core::JoinType::kLeftSemiFilter; @@ -1130,9 +1135,8 @@ VeloxQueryPlanConverterBase::toVeloxQueryPlan( std::vector groupingKeys; groupingKeys.reserve(node->groupingColumns.size()); for (const auto& [output, input] : node->groupingColumns) { - groupingKeys.emplace_back( - core::GroupIdNode::GroupingKeyInfo{ - output.name, exprConverter_.toVeloxExpr(input)}); + groupingKeys.emplace_back(core::GroupIdNode::GroupingKeyInfo{ + output.name, exprConverter_.toVeloxExpr(input)}); } return std::make_shared( @@ -1287,7 +1291,8 @@ core::PlanNodePtr VeloxQueryPlanConverterBase::toVeloxQueryPlan( exprConverter_.toVeloxExpr(node->filter), toVeloxQueryPlan(node->left, tableWriteInfo, taskId), toVeloxQueryPlan(node->right, tableWriteInfo, taskId), - toRowType(node->outputVariables, typeParser_));} + toRowType(node->outputVariables, typeParser_)); +} std::shared_ptr VeloxQueryPlanConverterBase::toVeloxQueryPlan( @@ -1314,12 +1319,25 @@ VeloxQueryPlanConverterBase::toVeloxQueryPlan( toVeloxQueryPlan(node->indexSource, tableWriteInfo, taskId); std::vector joinConditionPtrs{}; + std::vector unsupportedConditions{}; if (node->filter) { parseIndexLookupCondition( *node->filter, exprConverter_, /*acceptConstant=*/false, - joinConditionPtrs); + joinConditionPtrs, + unsupportedConditions); + } + + // Combine unsupported conditions into a single filter using AND + core::TypedExprPtr joinFilter; + if (!unsupportedConditions.empty()) { + if (unsupportedConditions.size() == 1) { + joinFilter = unsupportedConditions[0]; + } else { + joinFilter = std::make_shared( + BOOLEAN(), unsupportedConditions, "and"); + } } return std::make_shared( @@ -1328,6 +1346,7 @@ VeloxQueryPlanConverterBase::toVeloxQueryPlan( leftKeys, rightKeys, joinConditionPtrs, + joinFilter, false, left, std::dynamic_pointer_cast(right), @@ -2343,16 +2362,49 @@ void parseSqlFunctionHandle( } } +void handleUnsupportedIndexLookupCondition( + const std::shared_ptr& filter, + const VeloxExprConverter& exprConverter, + std::vector& unsupportedConditions) { + // OR conditions cannot be converted to index lookup conditions + if (const auto orForm = isOr(filter)) { + VELOX_UNSUPPORTED( + "Unsupported index lookup condition: {}", toJsonString(filter)); + } + unsupportedConditions.push_back(exprConverter.toVeloxExpr(filter)); +} + +#ifdef VELOX_ENABLE_BACKWARD_COMPATIBILITY void parseIndexLookupCondition( const std::shared_ptr& filter, const VeloxExprConverter& exprConverter, bool acceptConstant, std::vector& joinConditionPtrs) { + std::vector unsupportedConditions{}; + parseIndexLookupCondition( + filter, + exprConverter, + acceptConstant, + joinConditionPtrs, + unsupportedConditions); +} +#endif + +void parseIndexLookupCondition( + const std::shared_ptr& filter, + const VeloxExprConverter& exprConverter, + bool acceptConstant, + std::vector& joinConditionPtrs, + std::vector& unsupportedConditions) { if (const auto andForm = isAnd(filter)) { VELOX_CHECK_EQ(andForm->arguments.size(), 2); for (const auto& child : andForm->arguments) { parseIndexLookupCondition( - child, exprConverter, acceptConstant, joinConditionPtrs); + child, + exprConverter, + acceptConstant, + joinConditionPtrs, + unsupportedConditions); } return; } @@ -2368,17 +2420,14 @@ void parseIndexLookupCondition( const auto lowerExpr = exprConverter.toVeloxExpr(between->arguments[1]); const auto upperExpr = exprConverter.toVeloxExpr(between->arguments[2]); - VELOX_CHECK( - acceptConstant || - !(core::TypedExprs::isConstant(lowerExpr) && - core::TypedExprs::isConstant(upperExpr)), - "At least one of the between condition bounds needs to be not constant: {}", - toJsonString(filter)); - - joinConditionPtrs.push_back( - std::make_shared( - keyColumnExpr, lowerExpr, upperExpr)); - return; + if (acceptConstant || + !(core::TypedExprs::isConstant(lowerExpr) && + core::TypedExprs::isConstant(upperExpr))) { + joinConditionPtrs.push_back( + std::make_shared( + keyColumnExpr, lowerExpr, upperExpr)); + return; + } } if (const auto contains = isContains(filter)) { @@ -2391,15 +2440,13 @@ void parseIndexLookupCondition( const auto conditionColumnExpr = exprConverter.toVeloxExpr(contains->arguments[0]); - VELOX_CHECK( - acceptConstant || !core::TypedExprs::isConstant(conditionColumnExpr), - "The condition column needs to be not constant: {}", - toJsonString(filter)); - joinConditionPtrs.push_back( - std::make_shared( - keyColumnExpr, conditionColumnExpr)); - return; + if (acceptConstant || !core::TypedExprs::isConstant(conditionColumnExpr)) { + joinConditionPtrs.push_back( + std::make_shared( + keyColumnExpr, conditionColumnExpr)); + return; + } } if (const auto equals = isEqual(filter)) { @@ -2410,34 +2457,39 @@ void parseIndexLookupCondition( const bool leftIsConstant = core::TypedExprs::isConstant(leftExpr); const bool rightIsConstant = core::TypedExprs::isConstant(rightExpr); - VELOX_CHECK_NE( - leftIsConstant, - rightIsConstant, - "The equal condition must have one key and one constant: {}", - toJsonString(filter)); - - // Determine which argument is the key (non-constant) and which is the value - // (constant) - const auto& keyArgument = - leftIsConstant ? equals->arguments[1] : equals->arguments[0]; - const auto& constantExpr = leftIsConstant ? leftExpr : rightExpr; - - const auto keyColumnExpr = exprConverter.toVeloxExpr( - std::dynamic_pointer_cast( - keyArgument)); + VELOX_CHECK( + !(leftIsConstant && rightIsConstant), + "The equal condition must have at least one side to be non-constant: {}", + toJsonString(filter)); - VELOX_CHECK_NOT_NULL( - keyColumnExpr, - "Key argument must be a variable reference: {}", - toJsonString(keyArgument)); + // Check if the equal condition is a constant express and a field access. + std::shared_ptr keyArgument; + core::TypedExprPtr constantExpr; + if (core::TypedExprs::isFieldAccess(leftExpr) && rightIsConstant) { + keyArgument = equals->arguments[0]; + constantExpr = rightExpr; + } else if (core::TypedExprs::isFieldAccess(rightExpr) && leftIsConstant) { + keyArgument = equals->arguments[1]; + constantExpr = leftExpr; + } - joinConditionPtrs.push_back( - std::make_shared( - keyColumnExpr, constantExpr)); - return; + if (keyArgument != nullptr && constantExpr != nullptr) { + const auto keyColumnExpr = exprConverter.toVeloxExpr( + checked_pointer_cast( + keyArgument)); + VELOX_CHECK_NOT_NULL( + keyColumnExpr, + "Key argument must be a variable reference: {}", + toJsonString(keyArgument)); + joinConditionPtrs.push_back( + std::make_shared( + keyColumnExpr, constantExpr)); + return; + } } - VELOX_UNSUPPORTED( - "Unsupported index lookup condition: {}", toJsonString(filter)); + // For unsupported conditions, add to the vector or throw error. + handleUnsupportedIndexLookupCondition( + filter, exprConverter, unsupportedConditions); } } // namespace facebook::presto diff --git a/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.h b/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.h index 15c87c94e0984..1e1880cc9e345 100644 --- a/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.h +++ b/presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.h @@ -111,9 +111,9 @@ class VeloxQueryPlanConverterBase { const protocol::TaskId& taskId); velox::core::PlanNodePtr toVeloxQueryPlan( - const std::shared_ptr& node, - const std::shared_ptr& tableWriteInfo, - const protocol::TaskId& taskId); + const std::shared_ptr& node, + const std::shared_ptr& tableWriteInfo, + const protocol::TaskId& taskId); std::shared_ptr toVeloxQueryPlan( const std::shared_ptr& node, @@ -121,9 +121,9 @@ class VeloxQueryPlanConverterBase { const protocol::TaskId& taskId); std::shared_ptr toVeloxQueryPlan( - const std::shared_ptr& node, - const std::shared_ptr& tableWriteInfo, - const protocol::TaskId& taskId); + const std::shared_ptr& node, + const std::shared_ptr& tableWriteInfo, + const protocol::TaskId& taskId); velox::core::PlanNodePtr toVeloxQueryPlan( const std::shared_ptr& node, @@ -156,9 +156,9 @@ class VeloxQueryPlanConverterBase { const protocol::TaskId& taskId); std::shared_ptr toVeloxQueryPlan( - const std::shared_ptr& node, - const std::shared_ptr& tableWriteInfo, - const protocol::TaskId& taskId); + const std::shared_ptr& node, + const std::shared_ptr& tableWriteInfo, + const protocol::TaskId& taskId); std::shared_ptr toVeloxQueryPlan( const std::shared_ptr& node, @@ -297,9 +297,18 @@ void parseSqlFunctionHandle( std::vector& rawInputTypes, TypeParser& typeParser); +#ifdef VELOX_ENABLE_BACKWARD_COMPATIBILITY void parseIndexLookupCondition( const std::shared_ptr& filter, const VeloxExprConverter& exprConverter, bool acceptConstant, std::vector& joinConditionPtrs); +#endif + +void parseIndexLookupCondition( + const std::shared_ptr& filter, + const VeloxExprConverter& exprConverter, + bool acceptConstant, + std::vector& joinConditionPtrs, + std::vector& unsupportedConditions); } // namespace facebook::presto diff --git a/presto-native-execution/presto_cpp/main/types/tests/PrestoToVeloxQueryPlanTest.cpp b/presto-native-execution/presto_cpp/main/types/tests/PrestoToVeloxQueryPlanTest.cpp index 240e9cbd6d75a..3c4045a89b9e9 100644 --- a/presto-native-execution/presto_cpp/main/types/tests/PrestoToVeloxQueryPlanTest.cpp +++ b/presto-native-execution/presto_cpp/main/types/tests/PrestoToVeloxQueryPlanTest.cpp @@ -231,6 +231,67 @@ TEST_F(PrestoToVeloxQueryPlanTest, parseIndexJoinNode) { "builtInFunctionKind": "ENGINE" }, "returnType": "boolean" + }, + { + "@type": "call", + "arguments": [ + { + "@type": "constant", + "type": "bigint", + "valueBlock": "CgAAAExPTkdfQVJSQVkBAAAAAAAAAAAAAAAA" + }, + { + "@type": "call", + "arguments": [ + { + "@type": "variable", + "name": "c1", + "type": "bigint" + }, + { + "@type": "constant", + "type": "bigint", + "valueBlock": "CgAAAExPTkdfQVJSQVkBAAAAAGQAAAAAAAAA" + } + ], + "displayName": "MODULUS", + "functionHandle": { + "@type": "$static", + "builtInFunctionKind": "ENGINE", + "signature": { + "argumentTypes": [ + "bigint", + "bigint" + ], + "kind": "SCALAR", + "longVariableConstraints": [], + "name": "presto.default.$operator$modulus", + "returnType": "bigint", + "typeVariableConstraints": [], + "variableArity": false + } + }, + "returnType": "bigint" + } + ], + "displayName": "$operator$equal", + "functionHandle": { + "@type": "$static", + "builtInFunctionKind": "ENGINE", + "signature": { + "argumentTypes": [ + "bigint", + "bigint" + ], + "kind": "SCALAR", + "longVariableConstraints": [], + "name": "presto.default.$operator$equal", + "returnType": "boolean", + "typeVariableConstraints": [], + "variableArity": false + } + }, + "returnType": "boolean" } ], "form": "AND", diff --git a/presto-native-execution/velox b/presto-native-execution/velox index 63f9644ea7839..71a310307a734 160000 --- a/presto-native-execution/velox +++ b/presto-native-execution/velox @@ -1 +1 @@ -Subproject commit 63f9644ea78390f6631b193263edbc2b6b56b268 +Subproject commit 71a310307a734e170206163efb504afb7e6d834c From ed6e4a990dea4d627c034e1c378af3ff872f4107 Mon Sep 17 00:00:00 2001 From: mohsaka <135669458+mohsaka@users.noreply.github.com> Date: Fri, 25 Jul 2025 10:36:56 +0100 Subject: [PATCH 103/113] Plan table function invocation and add tests Changes adapted from trino/PR#11336, 12951, 14175 Original commit: d4c73389bbdb6b48c24a0969b259286b05a99ade 565700985baff0c4b29fdb1e3e26139a29318b9e ec8b9fd2b2cc9c8bc78c0ca1317dc34fcf2c48c7 98fc1ee8b29fca86f2a1b3abe4989524940333a6 1aea489884346822c812b1a242acc286e3e1248e 8bd17171a8469b9351e2fd7d9f2f49f4af9ea209 Author: kasiafi Modifications were made to adapt to Presto including: Change CatalogName to ConnectorId Change Symbol to VariableReferenceExpression TableFunctionNode extends InternalPlanNode instead of PlanNode. Add applyTableFunction to all implementations of Metadata Add empty ConnectorTableLayoutHandle to TableHandle in MetadataManger::applyTableFunction Removal of PlannerContext and replaced with Metadata Co-authored-by: kasiafi <30203062+kasiafi@users.noreply.github.com> Co-authored-by: Pratik Joseph Dabre Co-authored-by: Xin Zhang --- .../jdbc/TestJdbcDistributedQueries.java | 14 + .../metadata/DelegatingMetadataManager.java | 7 + .../facebook/presto/metadata/Metadata.java | 3 + .../presto/metadata/MetadataManager.java | 13 + .../presto/metadata/TableFunctionHandle.java | 58 +++ .../sql/planner/LocalExecutionPlanner.java | 7 + .../presto/sql/planner/PlanOptimizers.java | 9 + .../presto/sql/planner/RelationPlanner.java | 41 ++ .../rule/RewriteTableFunctionToTableScan.java | 97 ++++ .../planner/optimizations/AddExchanges.java | 7 + .../UnaliasSymbolReferences.java | 16 + .../sql/planner/plan/InternalPlanVisitor.java | 5 + .../presto/sql/planner/plan/Patterns.java | 5 + .../sql/planner/plan/TableFunctionNode.java | 178 +++++++ .../sql/planner/planPrinter/PlanPrinter.java | 15 + .../sanity/ValidateDependenciesChecker.java | 7 + ...java => TestTVFConnectorColumnHandle.java} | 6 +- .../tvf/TestTVFConnectorFactory.java | 452 ++++++++++++++++++ .../tvf/TestTVFConnectorPageSource.java | 82 ++++ .../connector/tvf/TestTVFConnectorPlugin.java | 37 ++ ....java => TestTVFConnectorTableHandle.java} | 29 +- .../TestTVFConnectorTransactionHandle.java | 22 + .../connector/tvf/TestTVFHandleResolver.java | 62 +++ .../tvf/TestTVFPartitioningHandle.java | 77 +++ .../tvf/TestTVFTableLayoutHandle.java | 66 +++ .../connector/tvf/TestingTableFunctions.java | 21 +- .../presto/metadata/AbstractMockMetadata.java | 7 + .../tests/TestTableFunctionInvocation.java | 94 ++++ 28 files changed, 1410 insertions(+), 27 deletions(-) create mode 100644 presto-main-base/src/main/java/com/facebook/presto/metadata/TableFunctionHandle.java create mode 100644 presto-main-base/src/main/java/com/facebook/presto/sql/planner/iterative/rule/RewriteTableFunctionToTableScan.java create mode 100644 presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/TableFunctionNode.java rename presto-main-base/src/test/java/com/facebook/presto/connector/tvf/{MockConnectorColumnHandle.java => TestTVFConnectorColumnHandle.java} (92%) create mode 100644 presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorFactory.java create mode 100644 presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorPageSource.java create mode 100644 presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorPlugin.java rename presto-main-base/src/test/java/com/facebook/presto/connector/tvf/{MockConnectorTableHandle.java => TestTVFConnectorTableHandle.java} (84%) create mode 100644 presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorTransactionHandle.java create mode 100644 presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFHandleResolver.java create mode 100644 presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFPartitioningHandle.java create mode 100644 presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFTableLayoutHandle.java create mode 100644 presto-tests/src/main/java/com/facebook/presto/tests/TestTableFunctionInvocation.java diff --git a/presto-base-jdbc/src/test/java/com/facebook/presto/plugin/jdbc/TestJdbcDistributedQueries.java b/presto-base-jdbc/src/test/java/com/facebook/presto/plugin/jdbc/TestJdbcDistributedQueries.java index 2d2962742a6e3..8b73ea66dc137 100644 --- a/presto-base-jdbc/src/test/java/com/facebook/presto/plugin/jdbc/TestJdbcDistributedQueries.java +++ b/presto-base-jdbc/src/test/java/com/facebook/presto/plugin/jdbc/TestJdbcDistributedQueries.java @@ -13,11 +13,14 @@ */ package com.facebook.presto.plugin.jdbc; +import com.facebook.presto.Session; import com.facebook.presto.testing.QueryRunner; import com.facebook.presto.tests.AbstractTestQueries; import io.airlift.tpch.TpchTable; +import org.testng.annotations.Test; import static com.facebook.presto.plugin.jdbc.JdbcQueryRunner.createJdbcQueryRunner; +import static com.facebook.presto.testing.TestingSession.testSessionBuilder; public class TestJdbcDistributedQueries extends AbstractTestQueries @@ -33,4 +36,15 @@ protected QueryRunner createQueryRunner() public void testLargeIn() { } + + @Test + public void testNativeQueryParameters() + { + Session session = testSessionBuilder() + .addPreparedStatement("my_query_simple", "SELECT * FROM TABLE(system.query(query => ?))") + .addPreparedStatement("my_query", "SELECT * FROM TABLE(system.query(query => format('SELECT %s FROM %s', ?, ?)))") + .build(); + assertQueryFails(session, "EXECUTE my_query_simple USING 'SELECT 1 a'", "line 1:21: Table function system.query not registered"); + assertQueryFails(session, "EXECUTE my_query USING 'a', '(SELECT 2 a) t'", "line 1:21: Table function system.query not registered"); + } } diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/DelegatingMetadataManager.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/DelegatingMetadataManager.java index 7302236671319..818e342cd8887 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/metadata/DelegatingMetadataManager.java +++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/DelegatingMetadataManager.java @@ -34,6 +34,7 @@ import com.facebook.presto.spi.connector.ConnectorCapabilities; import com.facebook.presto.spi.connector.ConnectorOutputMetadata; import com.facebook.presto.spi.connector.ConnectorTableVersion; +import com.facebook.presto.spi.connector.TableFunctionApplicationResult; import com.facebook.presto.spi.constraints.TableConstraint; import com.facebook.presto.spi.function.SqlFunction; import com.facebook.presto.spi.plan.PartitioningHandle; @@ -657,4 +658,10 @@ public String normalizeIdentifier(Session session, String catalogName, String id { return delegate.normalizeIdentifier(session, catalogName, identifier); } + + @Override + public Optional> applyTableFunction(Session session, TableFunctionHandle handle) + { + return delegate.applyTableFunction(session, handle); + } } diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/Metadata.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/Metadata.java index 2a0c62d7d4b0e..f53d1b7022746 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/metadata/Metadata.java +++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/Metadata.java @@ -41,6 +41,7 @@ import com.facebook.presto.spi.connector.ConnectorOutputMetadata; import com.facebook.presto.spi.connector.ConnectorPartitioningHandle; import com.facebook.presto.spi.connector.ConnectorTableVersion; +import com.facebook.presto.spi.connector.TableFunctionApplicationResult; import com.facebook.presto.spi.constraints.TableConstraint; import com.facebook.presto.spi.function.SqlFunction; import com.facebook.presto.spi.plan.PartitioningHandle; @@ -547,4 +548,6 @@ default boolean isPushdownSupportedForFilter(Session session, TableHandle tableH } String normalizeIdentifier(Session session, String catalogName, String identifier); + + Optional> applyTableFunction(Session session, TableFunctionHandle handle); } diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/MetadataManager.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/MetadataManager.java index bfcafe2e98339..6e714c5013f3e 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/metadata/MetadataManager.java +++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/MetadataManager.java @@ -61,6 +61,7 @@ import com.facebook.presto.spi.connector.ConnectorPartitioningMetadata; import com.facebook.presto.spi.connector.ConnectorTableVersion; import com.facebook.presto.spi.connector.ConnectorTransactionHandle; +import com.facebook.presto.spi.connector.TableFunctionApplicationResult; import com.facebook.presto.spi.constraints.TableConstraint; import com.facebook.presto.spi.function.SqlFunction; import com.facebook.presto.spi.plan.PartitioningHandle; @@ -1539,6 +1540,18 @@ private ColumnMetadata normalizedColumnMetadata(Session session, String catalogN .build(); } + @Override + public Optional> applyTableFunction(Session session, TableFunctionHandle handle) + { + ConnectorId connectorId = handle.getConnectorId(); + ConnectorMetadata metadata = getMetadata(session, connectorId); + + return metadata.applyTableFunction(session.toConnectorSession(connectorId), handle.getFunctionHandle()) + .map(result -> new TableFunctionApplicationResult<>( + new TableHandle(connectorId, result.getTableHandle(), handle.getTransactionHandle(), Optional.empty()), + result.getColumnHandles())); + } + private ViewDefinition deserializeView(String data) { try { diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/TableFunctionHandle.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/TableFunctionHandle.java new file mode 100644 index 0000000000000..20e65066d4320 --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/TableFunctionHandle.java @@ -0,0 +1,58 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.metadata; + +import com.facebook.presto.spi.ConnectorId; +import com.facebook.presto.spi.connector.ConnectorTransactionHandle; +import com.facebook.presto.spi.function.table.ConnectorTableFunctionHandle; +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; + +import static java.util.Objects.requireNonNull; + +public class TableFunctionHandle +{ + private final ConnectorId connectorId; + private final ConnectorTableFunctionHandle functionHandle; + private final ConnectorTransactionHandle transactionHandle; + + @JsonCreator + public TableFunctionHandle( + @JsonProperty("connectorId") ConnectorId connectorId, + @JsonProperty("functionHandle") ConnectorTableFunctionHandle functionHandle, + @JsonProperty("transactionHandle") ConnectorTransactionHandle transactionHandle) + { + this.connectorId = requireNonNull(connectorId, "connectorId is null"); + this.functionHandle = requireNonNull(functionHandle, "functionHandle is null"); + this.transactionHandle = requireNonNull(transactionHandle, "transactionHandle is null"); + } + + @JsonProperty + public ConnectorId getConnectorId() + { + return connectorId; + } + + @JsonProperty + public ConnectorTableFunctionHandle getFunctionHandle() + { + return functionHandle; + } + + @JsonProperty + public ConnectorTransactionHandle getTransactionHandle() + { + return transactionHandle; + } +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/LocalExecutionPlanner.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/LocalExecutionPlanner.java index c07fd31994d6b..85a52a298bcb7 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/LocalExecutionPlanner.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/LocalExecutionPlanner.java @@ -210,6 +210,7 @@ import com.facebook.presto.sql.planner.plan.RowNumberNode; import com.facebook.presto.sql.planner.plan.SampleNode; import com.facebook.presto.sql.planner.plan.StatisticsWriterNode; +import com.facebook.presto.sql.planner.plan.TableFunctionNode; import com.facebook.presto.sql.planner.plan.TableWriterMergeNode; import com.facebook.presto.sql.planner.plan.TopNRowNumberNode; import com.facebook.presto.sql.planner.plan.UpdateNode; @@ -1204,6 +1205,12 @@ public PhysicalOperation visitWindow(WindowNode node, LocalExecutionPlanContext return new PhysicalOperation(operatorFactory, outputMappings.build(), context, source); } + @Override + public PhysicalOperation visitTableFunction(TableFunctionNode node, LocalExecutionPlanContext context) + { + throw new UnsupportedOperationException("execution by operator is not yet implemented for table function " + node.getName()); + } + @Override public PhysicalOperation visitTopN(TopNNode node, LocalExecutionPlanContext context) { diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/PlanOptimizers.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/PlanOptimizers.java index 4d08e779911c9..927fda540e07a 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/PlanOptimizers.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/PlanOptimizers.java @@ -133,6 +133,7 @@ import com.facebook.presto.sql.planner.iterative.rule.RewriteConstantArrayContainsToInExpression; import com.facebook.presto.sql.planner.iterative.rule.RewriteFilterWithExternalFunctionToProject; import com.facebook.presto.sql.planner.iterative.rule.RewriteSpatialPartitioningAggregation; +import com.facebook.presto.sql.planner.iterative.rule.RewriteTableFunctionToTableScan; import com.facebook.presto.sql.planner.iterative.rule.RuntimeReorderJoinSides; import com.facebook.presto.sql.planner.iterative.rule.ScaledWriterRule; import com.facebook.presto.sql.planner.iterative.rule.SimplifyCardinalityMap; @@ -855,6 +856,14 @@ public PlanOptimizers( costCalculator, ImmutableSet.of(new ScaledWriterRule()))); + builder.add( + new IterativeOptimizer( + metadata, + ruleStats, + statsCalculator, + costCalculator, + ImmutableSet.of(new RewriteTableFunctionToTableScan(metadata)))); + if (!noExchange) { builder.add(new ReplicateSemiJoinInDelete()); // Must run before AddExchanges diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/RelationPlanner.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/RelationPlanner.java index 81e6e0bd1598d..11241f2f3dce1 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/RelationPlanner.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/RelationPlanner.java @@ -21,6 +21,7 @@ import com.facebook.presto.common.type.RowType; import com.facebook.presto.common.type.Type; import com.facebook.presto.metadata.Metadata; +import com.facebook.presto.metadata.TableFunctionHandle; import com.facebook.presto.spi.ColumnHandle; import com.facebook.presto.spi.PrestoException; import com.facebook.presto.spi.TableHandle; @@ -55,6 +56,7 @@ import com.facebook.presto.sql.planner.optimizations.SampleNodeUtil; import com.facebook.presto.sql.planner.plan.LateralJoinNode; import com.facebook.presto.sql.planner.plan.SampleNode; +import com.facebook.presto.sql.planner.plan.TableFunctionNode; import com.facebook.presto.sql.tree.AliasedRelation; import com.facebook.presto.sql.tree.Cast; import com.facebook.presto.sql.tree.CoalesceExpression; @@ -84,6 +86,7 @@ import com.facebook.presto.sql.tree.SetOperation; import com.facebook.presto.sql.tree.SymbolReference; import com.facebook.presto.sql.tree.Table; +import com.facebook.presto.sql.tree.TableFunctionInvocation; import com.facebook.presto.sql.tree.TableSubquery; import com.facebook.presto.sql.tree.Union; import com.facebook.presto.sql.tree.Unnest; @@ -297,6 +300,44 @@ private RelationPlan addColumnMasks(Table table, RelationPlan plan, SqlPlannerCo return new RelationPlan(planBuilder.getRoot(), plan.getScope(), newMappings.build()); } + @Override + protected RelationPlan visitTableFunctionInvocation(TableFunctionInvocation node, SqlPlannerContext context) + { + Analysis.TableFunctionInvocationAnalysis functionAnalysis = analysis.getTableFunctionAnalysis(node); + + // TODO handle input relations: + // 1. extract the input relations from node.getArguments() and plan them. Apply relation coercions if requested. + // 2. for each input relation, prepare the TableArgumentProperties record, consisting of: + // - row or set semantics (from the actualArgument) + // - prune when empty property (from the actualArgument) + // - pass through columns property (from the actualArgument) + // - optional Specification: ordering scheme and partitioning (from the node's argument) <- planned upon the source's RelationPlan (or combined RelationPlan from all sources) + // TODO add - argument name + // TODO add - mapping column name => Symbol // TODO mind the fields without names and duplicate field names in RelationType + List sources = ImmutableList.of(); + List inputRelationsProperties = ImmutableList.of(); + + Scope scope = analysis.getScope(node); + + ImmutableList.Builder outputVariablesBuilder = ImmutableList.builder(); + for (Field field : scope.getRelationType().getAllFields()) { + VariableReferenceExpression variable = variableAllocator.newVariable(getSourceLocation(node), field.getName().get(), field.getType()); + outputVariablesBuilder.add(variable); + } + + List outputVariables = outputVariablesBuilder.build(); + PlanNode root = new TableFunctionNode( + idAllocator.getNextId(), + functionAnalysis.getFunctionName(), + functionAnalysis.getArguments(), + outputVariablesBuilder.build(), + sources.stream().map(RelationPlan::getRoot).collect(toImmutableList()), + inputRelationsProperties, + new TableFunctionHandle(functionAnalysis.getConnectorId(), functionAnalysis.getConnectorTableFunctionHandle(), functionAnalysis.getTransactionHandle())); + + return new RelationPlan(root, scope, outputVariables); + } + @Override protected RelationPlan visitAliasedRelation(AliasedRelation node, SqlPlannerContext context) { diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/iterative/rule/RewriteTableFunctionToTableScan.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/iterative/rule/RewriteTableFunctionToTableScan.java new file mode 100644 index 0000000000000..2418377c7ac53 --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/iterative/rule/RewriteTableFunctionToTableScan.java @@ -0,0 +1,97 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.sql.planner.iterative.rule; + +import com.facebook.presto.common.predicate.TupleDomain; +import com.facebook.presto.matching.Captures; +import com.facebook.presto.matching.Pattern; +import com.facebook.presto.metadata.Metadata; +import com.facebook.presto.spi.ColumnHandle; +import com.facebook.presto.spi.TableHandle; +import com.facebook.presto.spi.connector.TableFunctionApplicationResult; +import com.facebook.presto.spi.plan.TableScanNode; +import com.facebook.presto.spi.relation.VariableReferenceExpression; +import com.facebook.presto.sql.planner.iterative.Rule; +import com.facebook.presto.sql.planner.plan.TableFunctionNode; +import com.google.common.collect.ImmutableMap; + +import java.util.List; +import java.util.Optional; + +import static com.facebook.presto.matching.Pattern.empty; +import static com.facebook.presto.sql.planner.plan.Patterns.sources; +import static com.facebook.presto.sql.planner.plan.Patterns.tableFunction; +import static com.google.common.base.Preconditions.checkState; +import static java.util.Objects.requireNonNull; + +/* + * This process converts connector-resolvable TableFunctionNodes into equivalent + * TableScanNodes by invoking the connector’s applyTableFunction() during planning. + * It allows table-valued functions whose results can be expressed as a ConnectorTableHandle + * to be treated like regular scans and benefit from normal scan optimizations. + * + * Example: + * Before Transformation: + * TableFunction(my_function(arg1, arg2)) + * + * After Transformation: + * TableScan(my_function(arg1, arg2)).applyTableFunction_tableHandle) + * assignments: {outputVar1 -> my_function(arg1, arg2)).applyTableFunction_colHandle1, + * outputVar2 -> my_function(arg1, arg2)).applyTableFunction_colHandle2} + */ +public class RewriteTableFunctionToTableScan + implements Rule +{ + private static final Pattern PATTERN = tableFunction() + .with(empty(sources())); + + private final Metadata metadata; + + public RewriteTableFunctionToTableScan(Metadata metadata) + { + this.metadata = requireNonNull(metadata, "metadata is null"); + } + + @Override + public Pattern getPattern() + { + return PATTERN; + } + + @Override + public Result apply(TableFunctionNode tableFunctionNode, Captures captures, Context context) + { + Optional> result = metadata.applyTableFunction(context.getSession(), tableFunctionNode.getHandle()); + + if (!result.isPresent()) { + return Result.empty(); + } + + List columnHandles = result.get().getColumnHandles(); + checkState(tableFunctionNode.getOutputVariables().size() == columnHandles.size(), "returned table does not match the node's output"); + ImmutableMap.Builder assignments = ImmutableMap.builder(); + for (int i = 0; i < columnHandles.size(); i++) { + assignments.put(tableFunctionNode.getOutputVariables().get(i), columnHandles.get(i)); + } + + return Result.ofPlanNode(new TableScanNode( + tableFunctionNode.getSourceLocation(), + tableFunctionNode.getId(), + result.get().getTableHandle(), + tableFunctionNode.getOutputVariables(), + assignments.buildOrThrow(), + TupleDomain.all(), + TupleDomain.all(), Optional.empty())); + } +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/AddExchanges.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/AddExchanges.java index 9749dd023c58e..84c151ac6e7a0 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/AddExchanges.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/AddExchanges.java @@ -77,6 +77,7 @@ import com.facebook.presto.sql.planner.plan.RowNumberNode; import com.facebook.presto.sql.planner.plan.SequenceNode; import com.facebook.presto.sql.planner.plan.StatisticsWriterNode; +import com.facebook.presto.sql.planner.plan.TableFunctionNode; import com.facebook.presto.sql.planner.plan.TopNRowNumberNode; import com.google.common.annotations.VisibleForTesting; import com.google.common.cache.CacheBuilder; @@ -410,6 +411,12 @@ public PlanWithProperties visitWindow(WindowNode node, PreferredProperties prefe return rebaseAndDeriveProperties(node, child); } + @Override + public PlanWithProperties visitTableFunction(TableFunctionNode node, PreferredProperties preferredProperties) + { + throw new UnsupportedOperationException("execution by operator is not yet implemented for table function " + node.getName()); + } + @Override public PlanWithProperties visitRowNumber(RowNumberNode node, PreferredProperties preferredProperties) { diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/UnaliasSymbolReferences.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/UnaliasSymbolReferences.java index eb61dd1bf42cb..c279a0675ec7f 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/UnaliasSymbolReferences.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/UnaliasSymbolReferences.java @@ -74,6 +74,7 @@ import com.facebook.presto.sql.planner.plan.SequenceNode; import com.facebook.presto.sql.planner.plan.SimplePlanRewriter; import com.facebook.presto.sql.planner.plan.StatisticsWriterNode; +import com.facebook.presto.sql.planner.plan.TableFunctionNode; import com.facebook.presto.sql.planner.plan.TableWriterMergeNode; import com.facebook.presto.sql.planner.plan.TopNRowNumberNode; import com.facebook.presto.sql.planner.plan.UpdateNode; @@ -477,6 +478,21 @@ public PlanNode visitTableFinish(TableFinishNode node, RewriteContext cont return mapper.map(node, source); } + @Override + public PlanNode visitTableFunction(TableFunctionNode node, RewriteContext context) + { + return new TableFunctionNode( + node.getSourceLocation(), + node.getId(), + Optional.empty(), + node.getName(), + node.getArguments(), + node.getOutputVariables(), + node.getSources(), + node.getTableArgumentProperties(), + node.getHandle()); + } + @Override public PlanNode visitRowNumber(RowNumberNode node, RewriteContext context) { diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/InternalPlanVisitor.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/InternalPlanVisitor.java index 15b53e87ac787..58ef334a0f01e 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/InternalPlanVisitor.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/InternalPlanVisitor.java @@ -126,4 +126,9 @@ public R visitSequence(SequenceNode node, C context) { return visitPlan(node, context); } + + public R visitTableFunction(TableFunctionNode node, C context) + { + return visitPlan(node, context); + } } diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/Patterns.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/Patterns.java index a2d18830e9d90..f1a00c4b1a128 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/Patterns.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/Patterns.java @@ -229,6 +229,11 @@ public static Pattern window() return typeOf(WindowNode.class); } + public static Pattern tableFunction() + { + return typeOf(TableFunctionNode.class); + } + public static Pattern rowNumber() { return typeOf(RowNumberNode.class); diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/TableFunctionNode.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/TableFunctionNode.java new file mode 100644 index 0000000000000..97892523498c0 --- /dev/null +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/plan/TableFunctionNode.java @@ -0,0 +1,178 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.sql.planner.plan; + +import com.facebook.presto.metadata.TableFunctionHandle; +import com.facebook.presto.spi.SourceLocation; +import com.facebook.presto.spi.function.table.Argument; +import com.facebook.presto.spi.plan.DataOrganizationSpecification; +import com.facebook.presto.spi.plan.PlanNode; +import com.facebook.presto.spi.plan.PlanNodeId; +import com.facebook.presto.spi.relation.VariableReferenceExpression; +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.errorprone.annotations.Immutable; + +import java.util.List; +import java.util.Map; +import java.util.Optional; + +import static com.google.common.base.Preconditions.checkArgument; +import static java.util.Objects.requireNonNull; + +@Immutable +public class TableFunctionNode + extends InternalPlanNode +{ + private final String name; + private final Map arguments; + private final List outputVariables; + private final List sources; + private final List tableArgumentProperties; + private final TableFunctionHandle handle; + + @JsonCreator + public TableFunctionNode( + @JsonProperty("id") PlanNodeId id, + @JsonProperty("name") String name, + @JsonProperty("arguments") Map arguments, + @JsonProperty("outputVariables") List outputVariables, + @JsonProperty("sources") List sources, + @JsonProperty("tableArgumentProperties") List tableArgumentProperties, + @JsonProperty("handle") TableFunctionHandle handle) + { + this(Optional.empty(), id, Optional.empty(), name, arguments, outputVariables, sources, tableArgumentProperties, handle); + } + + public TableFunctionNode( + Optional sourceLocation, + PlanNodeId id, + Optional statsEquivalentPlanNode, + String name, + Map arguments, + List outputVariables, + List sources, + List tableArgumentProperties, + TableFunctionHandle handle) + { + super(sourceLocation, id, statsEquivalentPlanNode); + this.name = requireNonNull(name, "name is null"); + this.arguments = requireNonNull(arguments, "arguments is null"); + this.outputVariables = requireNonNull(outputVariables, "outputVariables is null"); + this.sources = requireNonNull(sources, "sources is null"); + this.tableArgumentProperties = requireNonNull(tableArgumentProperties, "tableArgumentProperties is null"); + this.handle = requireNonNull(handle, "handle is null"); + } + + @JsonProperty + public String getName() + { + return name; + } + + @JsonProperty + public Map getArguments() + { + return arguments; + } + + @JsonProperty + public List getOutputVariables() + { + return outputVariables; + } + + @JsonProperty + public List getTableArgumentProperties() + { + return tableArgumentProperties; + } + + @JsonProperty + public TableFunctionHandle getHandle() + { + return handle; + } + + @JsonProperty + @Override + public List getSources() + { + return sources; + } + + @Override + public R accept(InternalPlanVisitor visitor, C context) + { + return visitor.visitTableFunction(this, context); + } + + @Override + public PlanNode replaceChildren(List newSources) + { + checkArgument(sources.size() == newSources.size(), "wrong number of new children"); + return new TableFunctionNode(getId(), name, arguments, outputVariables, newSources, tableArgumentProperties, handle); + } + + @Override + public PlanNode assignStatsEquivalentPlanNode(Optional statsEquivalentPlanNode) + { + return new TableFunctionNode(getSourceLocation(), getId(), statsEquivalentPlanNode, name, arguments, outputVariables, sources, tableArgumentProperties, handle); + } + + public static class TableArgumentProperties + { + private final boolean rowSemantics; + private final boolean pruneWhenEmpty; + private final boolean passThroughColumns; + private final Optional specification; + + @JsonCreator + public TableArgumentProperties( + @JsonProperty("rowSemantics") boolean rowSemantics, + @JsonProperty("pruneWhenEmpty") boolean pruneWhenEmpty, + @JsonProperty("passThroughColumns") boolean passThroughColumns, + @JsonProperty("specification") Optional specification) + { + this.rowSemantics = rowSemantics; + this.pruneWhenEmpty = pruneWhenEmpty; + this.passThroughColumns = passThroughColumns; + this.specification = requireNonNull(specification, "specification is null"); + } + + @JsonProperty + public boolean isRowSemantics() + { + return rowSemantics; + } + + @JsonProperty + public boolean isPruneWhenEmpty() + { + return pruneWhenEmpty; + } + + @JsonProperty + public boolean isPassThroughColumns() + { + return passThroughColumns; + } + + @JsonProperty + public Optional specification() + { + return specification; + } + } +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/planPrinter/PlanPrinter.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/planPrinter/PlanPrinter.java index 8eaacb43ff006..e536f99c5f67f 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/planPrinter/PlanPrinter.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/planPrinter/PlanPrinter.java @@ -98,6 +98,7 @@ import com.facebook.presto.sql.planner.plan.SampleNode; import com.facebook.presto.sql.planner.plan.SequenceNode; import com.facebook.presto.sql.planner.plan.StatisticsWriterNode; +import com.facebook.presto.sql.planner.plan.TableFunctionNode; import com.facebook.presto.sql.planner.plan.TableWriterMergeNode; import com.facebook.presto.sql.planner.plan.TopNRowNumberNode; import com.facebook.presto.sql.planner.plan.UpdateNode; @@ -1317,6 +1318,20 @@ public Void visitLateralJoin(LateralJoinNode node, Void context) public Void visitOffset(OffsetNode node, Void context) { addNode(node, "Offset", format("[%s]", node.getCount())); + return processChildren(node, context); + } + + @Override + public Void visitTableFunction(TableFunctionNode node, Void context) + { + NodeRepresentation nodeOutput = addNode( + node, + "TableFunction", + node.getName()); + + checkArgument( + node.getSources().isEmpty() && node.getTableArgumentProperties().isEmpty(), + "Table or descriptor arguments are not yet supported in PlanPrinter"); return processChildren(node, context); } diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/sanity/ValidateDependenciesChecker.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/sanity/ValidateDependenciesChecker.java index 60f7140716018..333b4c9c4fd76 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/sanity/ValidateDependenciesChecker.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/sanity/ValidateDependenciesChecker.java @@ -69,6 +69,7 @@ import com.facebook.presto.sql.planner.plan.SampleNode; import com.facebook.presto.sql.planner.plan.SequenceNode; import com.facebook.presto.sql.planner.plan.StatisticsWriterNode; +import com.facebook.presto.sql.planner.plan.TableFunctionNode; import com.facebook.presto.sql.planner.plan.TableWriterMergeNode; import com.facebook.presto.sql.planner.plan.TopNRowNumberNode; import com.facebook.presto.sql.planner.plan.UpdateNode; @@ -116,6 +117,12 @@ public Void visitPlan(PlanNode node, Set boundVaria throw new UnsupportedOperationException("not yet implemented: " + node.getClass().getName()); } + @Override + public Void visitTableFunction(TableFunctionNode node, Set boundSymbols) + { + return null; + } + @Override public Void visitExplainAnalyze(ExplainAnalyzeNode node, Set boundVariables) { diff --git a/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/MockConnectorColumnHandle.java b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorColumnHandle.java similarity index 92% rename from presto-main-base/src/test/java/com/facebook/presto/connector/tvf/MockConnectorColumnHandle.java rename to presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorColumnHandle.java index 84fbd45bfa9b9..063d4e37a839c 100644 --- a/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/MockConnectorColumnHandle.java +++ b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorColumnHandle.java @@ -23,14 +23,14 @@ import static com.google.common.base.MoreObjects.toStringHelper; import static java.util.Objects.requireNonNull; -public class MockConnectorColumnHandle +public class TestTVFConnectorColumnHandle implements ColumnHandle { private final String name; private final Type type; @JsonCreator - public MockConnectorColumnHandle( + public TestTVFConnectorColumnHandle( @JsonProperty("name") String name, @JsonProperty("type") Type type) { @@ -68,7 +68,7 @@ public boolean equals(Object o) if ((o == null) || (getClass() != o.getClass())) { return false; } - MockConnectorColumnHandle other = (MockConnectorColumnHandle) o; + TestTVFConnectorColumnHandle other = (TestTVFConnectorColumnHandle) o; return Objects.equals(name, other.name) && Objects.equals(type, other.type); } diff --git a/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorFactory.java b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorFactory.java new file mode 100644 index 0000000000000..2674c87d28cc6 --- /dev/null +++ b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorFactory.java @@ -0,0 +1,452 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.connector.tvf; + +import com.facebook.presto.common.RuntimeStats; +import com.facebook.presto.common.predicate.TupleDomain; +import com.facebook.presto.common.type.Type; +import com.facebook.presto.spi.ColumnHandle; +import com.facebook.presto.spi.ColumnMetadata; +import com.facebook.presto.spi.ConnectorHandleResolver; +import com.facebook.presto.spi.ConnectorPageSource; +import com.facebook.presto.spi.ConnectorSession; +import com.facebook.presto.spi.ConnectorSplit; +import com.facebook.presto.spi.ConnectorSplitSource; +import com.facebook.presto.spi.ConnectorTableHandle; +import com.facebook.presto.spi.ConnectorTableLayout; +import com.facebook.presto.spi.ConnectorTableLayoutHandle; +import com.facebook.presto.spi.ConnectorTableLayoutResult; +import com.facebook.presto.spi.ConnectorTableMetadata; +import com.facebook.presto.spi.ConnectorViewDefinition; +import com.facebook.presto.spi.Constraint; +import com.facebook.presto.spi.FixedSplitSource; +import com.facebook.presto.spi.HostAddress; +import com.facebook.presto.spi.InMemoryRecordSet; +import com.facebook.presto.spi.NodeProvider; +import com.facebook.presto.spi.RecordPageSource; +import com.facebook.presto.spi.SchemaTableName; +import com.facebook.presto.spi.SchemaTablePrefix; +import com.facebook.presto.spi.SplitContext; +import com.facebook.presto.spi.connector.Connector; +import com.facebook.presto.spi.connector.ConnectorContext; +import com.facebook.presto.spi.connector.ConnectorFactory; +import com.facebook.presto.spi.connector.ConnectorMetadata; +import com.facebook.presto.spi.connector.ConnectorPageSourceProvider; +import com.facebook.presto.spi.connector.ConnectorRecordSetProvider; +import com.facebook.presto.spi.connector.ConnectorSplitManager; +import com.facebook.presto.spi.connector.ConnectorTransactionHandle; +import com.facebook.presto.spi.connector.TableFunctionApplicationResult; +import com.facebook.presto.spi.function.table.ConnectorTableFunction; +import com.facebook.presto.spi.function.table.ConnectorTableFunctionHandle; +import com.facebook.presto.spi.schedule.NodeSelectionStrategy; +import com.facebook.presto.spi.statistics.TableStatistics; +import com.facebook.presto.spi.transaction.IsolationLevel; +import com.facebook.presto.tpch.TpchColumnHandle; +import com.facebook.presto.tpch.TpchRecordSetProvider; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; + +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.function.BiFunction; +import java.util.function.Function; +import java.util.function.Supplier; +import java.util.stream.IntStream; + +import static com.facebook.presto.common.type.VarcharType.createUnboundedVarcharType; +import static com.facebook.presto.spi.schedule.NodeSelectionStrategy.NO_PREFERENCE; +import static com.google.common.collect.ImmutableList.toImmutableList; +import static com.google.common.collect.ImmutableMap.toImmutableMap; +import static java.util.Objects.requireNonNull; + +public class TestTVFConnectorFactory + implements ConnectorFactory +{ + private final Function> listSchemaNames; + private final BiFunction> listTables; + private final BiFunction> getViews; + private final BiFunction> getColumnHandles; + private final Supplier getTableStatistics; + private final ApplyTableFunction applyTableFunction; + private final Set tableFunctions; + + private TestTVFConnectorFactory( + Function> listSchemaNames, + BiFunction> listTables, + BiFunction> getViews, + BiFunction> getColumnHandles, + Supplier getTableStatistics, + ApplyTableFunction applyTableFunction, + Set tableFunctions) + { + this.listSchemaNames = requireNonNull(listSchemaNames, "listSchemaNames is null"); + this.listTables = requireNonNull(listTables, "listTables is null"); + this.getViews = requireNonNull(getViews, "getViews is null"); + this.getColumnHandles = requireNonNull(getColumnHandles, "getColumnHandles is null"); + this.getTableStatistics = requireNonNull(getTableStatistics, "getTableStatistics is null"); + this.applyTableFunction = requireNonNull(applyTableFunction, "applyTableFunction is null"); + this.tableFunctions = requireNonNull(tableFunctions, "tableFunctions is null"); + } + + @Override + public String getName() + { + return "testTVF"; + } + + @Override + public ConnectorHandleResolver getHandleResolver() + { + return new TestTVFHandleResolver(); + } + + @Override + public Connector create(String catalogName, Map config, ConnectorContext context) + { + return new TestTVFConnector(context, listSchemaNames, listTables, getViews, getColumnHandles, getTableStatistics, applyTableFunction, tableFunctions); + } + + public static Builder builder() + { + return new Builder(); + } + + public static Function> defaultGetColumns() + { + return table -> IntStream.range(0, 100) + .boxed() + .map(i -> ColumnMetadata.builder().setName("column_" + i).setType(createUnboundedVarcharType()).build()) + .collect(toImmutableList()); + } + + @FunctionalInterface + public interface ApplyTableFunction + { + Optional> apply(ConnectorSession session, ConnectorTableFunctionHandle handle); + } + + public static class TestTVFConnector + implements Connector + { + private static final String DELETE_ROW_ID = "delete_row_id"; + private static final String UPDATE_ROW_ID = "update_row_id"; + private static final String MERGE_ROW_ID = "merge_row_id"; + + private final ConnectorContext context; + private final Function> listSchemaNames; + private final BiFunction> listTables; + private final BiFunction> getViews; + private final BiFunction> getColumnHandles; + private final Supplier getTableStatistics; + private final ApplyTableFunction applyTableFunction; + private final Set tableFunctions; + + public TestTVFConnector( + ConnectorContext context, + Function> listSchemaNames, + BiFunction> listTables, + BiFunction> getViews, + BiFunction> getColumnHandles, + Supplier getTableStatistics, + ApplyTableFunction applyTableFunction, + Set tableFunctions) + { + this.context = requireNonNull(context, "context is null"); + this.listSchemaNames = requireNonNull(listSchemaNames, "listSchemaNames is null"); + this.listTables = requireNonNull(listTables, "listTables is null"); + this.getViews = requireNonNull(getViews, "getViews is null"); + this.getColumnHandles = requireNonNull(getColumnHandles, "getColumnHandles is null"); + this.getTableStatistics = requireNonNull(getTableStatistics, "getTableStatistics is null"); + this.applyTableFunction = requireNonNull(applyTableFunction, "applyTableFunction is null"); + this.tableFunctions = requireNonNull(tableFunctions, "tableFunctions is null"); + } + + @Override + public ConnectorTransactionHandle beginTransaction(IsolationLevel isolationLevel, boolean readOnly) + { + return TestTVFConnectorTransactionHandle.INSTANCE; + } + + @Override + public ConnectorMetadata getMetadata(ConnectorTransactionHandle transaction) + { + return new TestTVFConnectorMetadata(); + } + + public enum TestTVFConnectorSplit + implements ConnectorSplit + { + TEST_TVF_CONNECTOR_SPLIT; + + @Override + public NodeSelectionStrategy getNodeSelectionStrategy() + { + return NO_PREFERENCE; + } + + @Override + public List getPreferredNodes(NodeProvider nodeProvider) + { + return Collections.emptyList(); + } + + @Override + public Object getInfo() + { + return null; + } + } + + @Override + public ConnectorSplitManager getSplitManager() + { + return new ConnectorSplitManager() + { + @Override + public ConnectorSplitSource getSplits(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorTableLayoutHandle layout, SplitSchedulingContext splitSchedulingContext) + { + return new FixedSplitSource(Collections.singleton(TestTVFConnectorSplit.TEST_TVF_CONNECTOR_SPLIT)); + } + }; + } + + @Override + public ConnectorRecordSetProvider getRecordSetProvider() + { + return new TpchRecordSetProvider(); + } + + @Override + public ConnectorPageSourceProvider getPageSourceProvider() + { + return new TestTVFConnectorPageSourceProvider(); + } + + @Override + public Set getTableFunctions() + { + return tableFunctions; + } + + private class TestTVFConnectorMetadata + implements ConnectorMetadata + { + @Override + public List listSchemaNames(ConnectorSession session) + { + return listSchemaNames.apply(session); + } + + @Override + public ConnectorTableHandle getTableHandle(ConnectorSession session, SchemaTableName tableName) + { + return new ConnectorTableHandle() {}; + } + + @Override + public ConnectorTableMetadata getTableMetadata(ConnectorSession session, ConnectorTableHandle tableHandle) + { + TestTVFConnectorTableHandle table = (TestTVFConnectorTableHandle) tableHandle; + return new ConnectorTableMetadata( + table.getTableName(), + defaultGetColumns().apply(table.getTableName()), + ImmutableMap.of()); + } + + @Override + public List listTables(ConnectorSession session, String schemaNameOrNull) + { + return listTables.apply(session, schemaNameOrNull); + } + + public void setTableProperties(ConnectorSession session, ConnectorTableHandle tableHandle, Map properties) + { + } + + @Override + public Map getColumnHandles(ConnectorSession session, ConnectorTableHandle tableHandle) + { + return (Map) (Map) getColumnHandles.apply(session, tableHandle); + } + + @Override + public ColumnMetadata getColumnMetadata(ConnectorSession session, ConnectorTableHandle tableHandle, ColumnHandle columnHandle) + { + if (columnHandle instanceof TestTVFConnectorColumnHandle) { + TestTVFConnectorColumnHandle testTVFColumnHandle = (TestTVFConnectorColumnHandle) columnHandle; + return ColumnMetadata.builder().setName(testTVFColumnHandle.getName()).setType(testTVFColumnHandle.getType()).build(); + } + else { + TpchColumnHandle tpchColumnHandle = (TpchColumnHandle) columnHandle; + return ColumnMetadata.builder().setName(tpchColumnHandle.getColumnName()).setType(tpchColumnHandle.getType()).build(); + } + } + + @Override + public Map> listTableColumns(ConnectorSession session, SchemaTablePrefix prefix) + { + return listTables(session, prefix.getSchemaName()).stream() + .collect(toImmutableMap(table -> table, table -> IntStream.range(0, 100) + .boxed() + .map(i -> ColumnMetadata.builder().setName("column_" + i).setType(createUnboundedVarcharType()).build()) + .collect(toImmutableList()))); + } + + @Override + public ConnectorTableLayoutResult getTableLayoutForConstraint( + ConnectorSession session, + ConnectorTableHandle table, + Constraint constraint, + Optional> desiredColumns) + { + // TODO: Currently not supporting constraints + TestTVFTableLayoutHandle tvfLayout = new TestTVFTableLayoutHandle((TestTVFConnectorTableHandle) table, TupleDomain.none()); + return new ConnectorTableLayoutResult(new ConnectorTableLayout(tvfLayout, + Optional.empty(), + tvfLayout.getPredicate(), + Optional.empty(), + Optional.empty(), + Optional.empty(), + Collections.emptyList(), + Optional.empty()), TupleDomain.none()); + } + + @Override + public ConnectorTableLayout getTableLayout(ConnectorSession session, ConnectorTableLayoutHandle handle) + { + TestTVFTableLayoutHandle tvfTableLayout = (TestTVFTableLayoutHandle) handle; + return new ConnectorTableLayout(tvfTableLayout); + } + + @Override + public Map getViews(ConnectorSession session, SchemaTablePrefix prefix) + { + return getViews.apply(session, prefix); + } + + @Override + public TableStatistics getTableStatistics(ConnectorSession session, ConnectorTableHandle tableHandle, Optional tableLayoutHandle, List columnHandles, Constraint constraint) + { + return getTableStatistics.get(); + } + + @Override + public Optional> applyTableFunction(ConnectorSession session, ConnectorTableFunctionHandle handle) + { + return applyTableFunction.apply(session, handle); + } + } + + private class TestTVFConnectorPageSourceProvider + implements ConnectorPageSourceProvider + { + @Override + public ConnectorPageSource createPageSource(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorSplit split, ConnectorTableLayoutHandle layout, List columns, SplitContext splitContext, RuntimeStats runtimeStats) + { + TestTVFConnectorTableHandle handle = ((TestTVFTableLayoutHandle) layout).getTable(); + SchemaTableName tableName = handle.getTableName(); + List projection = columns.stream() + .map(TestTVFConnectorColumnHandle.class::cast) + .collect(toImmutableList()); + List types = columns.stream() + .map(TestTVFConnectorColumnHandle.class::cast) + .map(TestTVFConnectorColumnHandle::getType) + .collect(toImmutableList()); + return new TestTVFConnectorPageSource(new RecordPageSource(new InMemoryRecordSet(types, ImmutableList.of()))); + } + + private Map getColumnIndexes(SchemaTableName tableName) + { + ImmutableMap.Builder columnIndexes = ImmutableMap.builder(); + List columnMetadata = defaultGetColumns().apply(tableName); + for (int index = 0; index < columnMetadata.size(); index++) { + columnIndexes.put(columnMetadata.get(index).getName(), index); + } + return columnIndexes.buildOrThrow(); + } + } + } + + public static final class Builder + { + private Function> listSchemaNames = (session) -> ImmutableList.of(); + private BiFunction> listTables = (session, schemaName) -> ImmutableList.of(); + private BiFunction> getViews = (session, schemaTablePrefix) -> ImmutableMap.of(); + private BiFunction> getColumnHandles = (session, tableHandle) -> { + TestTVFConnectorTableHandle table = (TestTVFConnectorTableHandle) tableHandle; + return defaultGetColumns().apply(table.getTableName()).stream() + .collect(toImmutableMap(ColumnMetadata::getName, column -> + new TestTVFConnectorColumnHandle(column.getName(), column.getType()))); + }; + private Supplier getTableStatistics = TableStatistics::empty; + private ApplyTableFunction applyTableFunction = (session, handle) -> Optional.empty(); + private Set tableFunctions = ImmutableSet.of(); + + public Builder withListSchemaNames(Function> listSchemaNames) + { + this.listSchemaNames = requireNonNull(listSchemaNames, "listSchemaNames is null"); + return this; + } + + public Builder withListTables(BiFunction> listTables) + { + this.listTables = requireNonNull(listTables, "listTables is null"); + return this; + } + + public Builder withGetViews(BiFunction> getViews) + { + this.getViews = requireNonNull(getViews, "getViews is null"); + return this; + } + + public Builder withGetColumnHandles(BiFunction> getColumnHandles) + { + this.getColumnHandles = requireNonNull(getColumnHandles, "getColumnHandles is null"); + return this; + } + + public Builder withGetTableStatistics(Supplier getTableStatistics) + { + this.getTableStatistics = requireNonNull(getTableStatistics, "getTableStatistics is null"); + return this; + } + + public Builder withApplyTableFunction(ApplyTableFunction applyTableFunction) + { + this.applyTableFunction = applyTableFunction; + return this; + } + + public Builder withTableFunctions(Iterable tableFunctions) + { + this.tableFunctions = ImmutableSet.copyOf(tableFunctions); + return this; + } + + public TestTVFConnectorFactory build() + { + return new TestTVFConnectorFactory(listSchemaNames, listTables, getViews, getColumnHandles, getTableStatistics, applyTableFunction, tableFunctions); + } + + private static T notSupported() + { + throw new UnsupportedOperationException(); + } + } +} diff --git a/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorPageSource.java b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorPageSource.java new file mode 100644 index 0000000000000..126b23a48e796 --- /dev/null +++ b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorPageSource.java @@ -0,0 +1,82 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.connector.tvf; + +import com.facebook.presto.common.Page; +import com.facebook.presto.spi.ConnectorPageSource; + +import java.io.IOException; +import java.util.concurrent.CompletableFuture; + +import static java.util.Objects.requireNonNull; + +public class TestTVFConnectorPageSource + implements ConnectorPageSource +{ + private final ConnectorPageSource delegate; + + public TestTVFConnectorPageSource(ConnectorPageSource delegate) + { + this.delegate = requireNonNull(delegate, "delegate is null"); + } + + @Override + public long getCompletedBytes() + { + return delegate.getCompletedBytes(); + } + + @Override + public long getCompletedPositions() + { + return 0; + } + + @Override + public long getReadTimeNanos() + { + return delegate.getReadTimeNanos(); + } + + @Override + public boolean isFinished() + { + return delegate.isFinished(); + } + + @Override + public Page getNextPage() + { + return delegate.getNextPage(); + } + + @Override + public long getSystemMemoryUsage() + { + return 0; + } + + @Override + public void close() + throws IOException + { + delegate.close(); + } + + @Override + public CompletableFuture isBlocked() + { + return delegate.isBlocked(); + } +} diff --git a/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorPlugin.java b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorPlugin.java new file mode 100644 index 0000000000000..a26ef9c34466a --- /dev/null +++ b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorPlugin.java @@ -0,0 +1,37 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.connector.tvf; + +import com.facebook.presto.spi.Plugin; +import com.facebook.presto.spi.connector.ConnectorFactory; +import com.google.common.collect.ImmutableList; + +import static java.util.Objects.requireNonNull; + +public class TestTVFConnectorPlugin + implements Plugin +{ + private final ConnectorFactory connectorFactory; + + public TestTVFConnectorPlugin(ConnectorFactory connectorFactory) + { + this.connectorFactory = requireNonNull(connectorFactory, "connectorFactory is null"); + } + + @Override + public Iterable getConnectorFactories() + { + return ImmutableList.of(connectorFactory); + } +} diff --git a/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/MockConnectorTableHandle.java b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorTableHandle.java similarity index 84% rename from presto-main-base/src/test/java/com/facebook/presto/connector/tvf/MockConnectorTableHandle.java rename to presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorTableHandle.java index b29826b684dd7..9c90975cbc82b 100644 --- a/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/MockConnectorTableHandle.java +++ b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorTableHandle.java @@ -27,28 +27,31 @@ import static java.util.Objects.requireNonNull; -public class MockConnectorTableHandle +public class TestTVFConnectorTableHandle implements ConnectorTableHandle { + // These are example fields for a Connector's Table Handle. + // For other examples, see TpchTableHandle or any other implementations + // of ConnectorTableHandle. private final SchemaTableName tableName; - private final TupleDomain constraint; private final Optional> columns; + private final TupleDomain constraint; - public MockConnectorTableHandle(SchemaTableName tableName) + public TestTVFConnectorTableHandle(SchemaTableName tableName) { - this(tableName, TupleDomain.all(), Optional.empty()); + this(tableName, Optional.empty(), TupleDomain.all()); } @JsonCreator - public MockConnectorTableHandle( + public TestTVFConnectorTableHandle( @JsonProperty SchemaTableName tableName, - @JsonProperty("constraint") TupleDomain constraint, - @JsonProperty("columns") Optional> columns) + @JsonProperty("columns") Optional> columns, + @JsonProperty("constraint") TupleDomain constraint) { this.tableName = requireNonNull(tableName, "tableName is null"); - this.constraint = requireNonNull(constraint, "constraint is null"); requireNonNull(columns, "columns is null"); this.columns = columns.map(ImmutableList::copyOf); + this.constraint = requireNonNull(constraint, "constraint is null"); } @JsonProperty @@ -58,15 +61,15 @@ public SchemaTableName getTableName() } @JsonProperty - public TupleDomain getConstraint() + public Optional> getColumns() { - return constraint; + return columns; } @JsonProperty - public Optional> getColumns() + public TupleDomain getConstraint() { - return columns; + return constraint; } @Override @@ -78,7 +81,7 @@ public boolean equals(Object o) if (o == null || getClass() != o.getClass()) { return false; } - MockConnectorTableHandle other = (MockConnectorTableHandle) o; + TestTVFConnectorTableHandle other = (TestTVFConnectorTableHandle) o; return Objects.equals(tableName, other.tableName) && Objects.equals(constraint, other.constraint) && Objects.equals(columns, other.columns); diff --git a/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorTransactionHandle.java b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorTransactionHandle.java new file mode 100644 index 0000000000000..668d3fb6575b2 --- /dev/null +++ b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFConnectorTransactionHandle.java @@ -0,0 +1,22 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.connector.tvf; + +import com.facebook.presto.spi.connector.ConnectorTransactionHandle; + +public enum TestTVFConnectorTransactionHandle + implements ConnectorTransactionHandle +{ + INSTANCE +} diff --git a/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFHandleResolver.java b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFHandleResolver.java new file mode 100644 index 0000000000000..1d6771ef0bfb7 --- /dev/null +++ b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFHandleResolver.java @@ -0,0 +1,62 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.connector.tvf; + +import com.facebook.presto.spi.ColumnHandle; +import com.facebook.presto.spi.ConnectorHandleResolver; +import com.facebook.presto.spi.ConnectorSplit; +import com.facebook.presto.spi.ConnectorTableHandle; +import com.facebook.presto.spi.ConnectorTableLayoutHandle; +import com.facebook.presto.spi.connector.ConnectorPartitioningHandle; +import com.facebook.presto.spi.connector.ConnectorTransactionHandle; + +public class TestTVFHandleResolver + implements ConnectorHandleResolver +{ + @Override + public Class getTableHandleClass() + { + return TestTVFConnectorTableHandle.class; + } + + @Override + public Class getColumnHandleClass() + { + return TestTVFConnectorColumnHandle.class; + } + + @Override + public Class getSplitClass() + { + return TestTVFConnectorFactory.TestTVFConnector.TestTVFConnectorSplit.class; + } + + @Override + public Class getTableLayoutHandleClass() + { + return TestTVFTableLayoutHandle.class; + } + + @Override + public Class getTransactionHandleClass() + { + return TestTVFConnectorTransactionHandle.class; + } + + @Override + public Class getPartitioningHandleClass() + { + return TestTVFPartitioningHandle.class; + } +} diff --git a/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFPartitioningHandle.java b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFPartitioningHandle.java new file mode 100644 index 0000000000000..cca6fc0a48bb4 --- /dev/null +++ b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFPartitioningHandle.java @@ -0,0 +1,77 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.connector.tvf; + +import com.facebook.presto.spi.connector.ConnectorPartitioningHandle; +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; + +import java.util.Objects; + +import static com.google.common.base.Preconditions.checkArgument; +import static java.util.Objects.requireNonNull; + +public class TestTVFPartitioningHandle + implements ConnectorPartitioningHandle +{ + private final String table; + private final long totalRows; + + @JsonCreator + public TestTVFPartitioningHandle(@JsonProperty("table") String table, @JsonProperty("totalRows") long totalRows) + { + this.table = requireNonNull(table, "table is null"); + + checkArgument(totalRows > 0, "totalRows must be at least 1"); + this.totalRows = totalRows; + } + + @JsonProperty + public String getTable() + { + return table; + } + + @JsonProperty + public long getTotalRows() + { + return totalRows; + } + + @Override + public boolean equals(Object o) + { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + TestTVFPartitioningHandle that = (TestTVFPartitioningHandle) o; + return Objects.equals(table, that.table) && + totalRows == that.totalRows; + } + + @Override + public int hashCode() + { + return Objects.hash(table, totalRows); + } + + @Override + public String toString() + { + return table + ":" + totalRows; + } +} diff --git a/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFTableLayoutHandle.java b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFTableLayoutHandle.java new file mode 100644 index 0000000000000..c9f61203cf1de --- /dev/null +++ b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestTVFTableLayoutHandle.java @@ -0,0 +1,66 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.connector.tvf; + +import com.facebook.presto.common.plan.PlanCanonicalizationStrategy; +import com.facebook.presto.common.predicate.TupleDomain; +import com.facebook.presto.spi.ColumnHandle; +import com.facebook.presto.spi.ConnectorSplit; +import com.facebook.presto.spi.ConnectorTableLayoutHandle; +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.common.collect.ImmutableMap; + +import java.util.Optional; + +public class TestTVFTableLayoutHandle + implements ConnectorTableLayoutHandle +{ + private final TestTVFConnectorTableHandle table; + private final TupleDomain predicate; + + @JsonCreator + public TestTVFTableLayoutHandle(@JsonProperty("table") TestTVFConnectorTableHandle table, @JsonProperty("predicate") TupleDomain predicate) + { + this.table = table; + this.predicate = predicate; + } + + @JsonProperty + public TestTVFConnectorTableHandle getTable() + { + return table; + } + + @JsonProperty + public TupleDomain getPredicate() + { + return predicate; + } + + @Override + public String toString() + { + return table.toString(); + } + + @Override + public Object getIdentifier(Optional split, PlanCanonicalizationStrategy strategy) + { + return ImmutableMap.builder() + .put("table", table) + .put("predicate", predicate.canonicalize(ignored -> false)) + .build(); + } +} diff --git a/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestingTableFunctions.java b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestingTableFunctions.java index 1986401eac75a..f2831f4bf10e1 100644 --- a/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestingTableFunctions.java +++ b/presto-main-base/src/test/java/com/facebook/presto/connector/tvf/TestingTableFunctions.java @@ -163,7 +163,6 @@ public TableFunctionAnalysis analyze(ConnectorSession session, ConnectorTransact public static class SimpleTableFunction extends AbstractConnectorTableFunction { - private static final String SCHEMA_NAME = "system"; private static final String FUNCTION_NAME = "simple_table_function"; private static final String TABLE_NAME = "simple_table"; @@ -201,17 +200,17 @@ public TableFunctionAnalysis analyze(ConnectorSession session, ConnectorTransact public static class SimpleTableFunctionHandle implements ConnectorTableFunctionHandle { - private final MockConnectorTableHandle tableHandle; + private final TestTVFConnectorTableHandle tableHandle; public SimpleTableFunctionHandle(String schema, String table, String column) { - this.tableHandle = new MockConnectorTableHandle( + this.tableHandle = new TestTVFConnectorTableHandle( new SchemaTableName(schema, table), - TupleDomain.all(), - Optional.of(ImmutableList.of(new MockConnectorColumnHandle(column, BOOLEAN)))); + Optional.of(ImmutableList.of(new TestTVFConnectorColumnHandle(column, BOOLEAN))), + TupleDomain.all()); } - public MockConnectorTableHandle getTableHandle() + public TestTVFConnectorTableHandle getTableHandle() { return tableHandle; } @@ -301,17 +300,17 @@ public TableFunctionAnalysis analyze(ConnectorSession session, ConnectorTransact public static class TestingTableFunctionPushdownHandle implements ConnectorTableFunctionHandle { - private final MockConnectorTableHandle tableHandle; + private final TestTVFConnectorTableHandle tableHandle; public TestingTableFunctionPushdownHandle() { - this.tableHandle = new MockConnectorTableHandle( + this.tableHandle = new TestTVFConnectorTableHandle( new SchemaTableName(SCHEMA_NAME, TABLE_NAME), - TupleDomain.all(), - Optional.of(ImmutableList.of(new MockConnectorColumnHandle(COLUMN_NAME, BOOLEAN)))); + Optional.of(ImmutableList.of(new TestTVFConnectorColumnHandle(COLUMN_NAME, BOOLEAN))), + TupleDomain.all()); } - public MockConnectorTableHandle getTableHandle() + public TestTVFConnectorTableHandle getTableHandle() { return tableHandle; } diff --git a/presto-main-base/src/test/java/com/facebook/presto/metadata/AbstractMockMetadata.java b/presto-main-base/src/test/java/com/facebook/presto/metadata/AbstractMockMetadata.java index e13cbb5a80c0c..64d33822ed2dc 100644 --- a/presto-main-base/src/test/java/com/facebook/presto/metadata/AbstractMockMetadata.java +++ b/presto-main-base/src/test/java/com/facebook/presto/metadata/AbstractMockMetadata.java @@ -35,6 +35,7 @@ import com.facebook.presto.spi.connector.ConnectorCapabilities; import com.facebook.presto.spi.connector.ConnectorOutputMetadata; import com.facebook.presto.spi.connector.ConnectorTableVersion; +import com.facebook.presto.spi.connector.TableFunctionApplicationResult; import com.facebook.presto.spi.constraints.TableConstraint; import com.facebook.presto.spi.function.SqlFunction; import com.facebook.presto.spi.plan.PartitioningHandle; @@ -685,4 +686,10 @@ public String normalizeIdentifier(Session session, String catalogName, String id { return identifier.toLowerCase(ENGLISH); } + + @Override + public Optional> applyTableFunction(Session session, TableFunctionHandle handle) + { + return Optional.empty(); + } } diff --git a/presto-tests/src/main/java/com/facebook/presto/tests/TestTableFunctionInvocation.java b/presto-tests/src/main/java/com/facebook/presto/tests/TestTableFunctionInvocation.java new file mode 100644 index 0000000000000..067c750f09497 --- /dev/null +++ b/presto-tests/src/main/java/com/facebook/presto/tests/TestTableFunctionInvocation.java @@ -0,0 +1,94 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.tests; + +import com.facebook.presto.connector.tvf.TestTVFConnectorColumnHandle; +import com.facebook.presto.connector.tvf.TestTVFConnectorFactory; +import com.facebook.presto.connector.tvf.TestTVFConnectorPlugin; +import com.facebook.presto.connector.tvf.TestingTableFunctions.SimpleTableFunction; +import com.facebook.presto.connector.tvf.TestingTableFunctions.SimpleTableFunction.SimpleTableFunctionHandle; +import com.facebook.presto.spi.ConnectorSession; +import com.facebook.presto.spi.ConnectorTableHandle; +import com.facebook.presto.spi.connector.TableFunctionApplicationResult; +import com.facebook.presto.testing.QueryRunner; +import com.google.common.collect.ImmutableSet; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + +import java.util.Map; +import java.util.Optional; +import java.util.function.BiFunction; +import java.util.stream.IntStream; + +import static com.facebook.presto.common.type.VarcharType.createUnboundedVarcharType; +import static com.facebook.presto.testing.TestingSession.testSessionBuilder; +import static com.google.common.collect.ImmutableMap.toImmutableMap; + +public class TestTableFunctionInvocation + extends AbstractTestQueryFramework +{ + private static final String TESTING_CATALOG = "testing_catalog1"; + private static final String TABLE_FUNCTION_SCHEMA = "table_function_schema"; + + @Override + protected QueryRunner createQueryRunner() + throws Exception + { + return DistributedQueryRunner.builder(testSessionBuilder() + .setCatalog(TESTING_CATALOG) + .setSchema(TABLE_FUNCTION_SCHEMA) + .build()) + .build(); + } + + @BeforeClass + public void setUp() + { + DistributedQueryRunner queryRunner = getDistributedQueryRunner(); + + BiFunction> getColumnHandles = (session, tableHandle) -> IntStream.range(0, 100) + .boxed() + .map(i -> "column_" + i) + .collect(toImmutableMap(column -> column, column -> new TestTVFConnectorColumnHandle(column, createUnboundedVarcharType()) {})); + + queryRunner.installPlugin(new TestTVFConnectorPlugin(TestTVFConnectorFactory.builder() + .withTableFunctions(ImmutableSet.of(new SimpleTableFunction())) + .withApplyTableFunction((session, handle) -> { + if (handle instanceof SimpleTableFunctionHandle) { + SimpleTableFunctionHandle functionHandle = (SimpleTableFunctionHandle) handle; + return Optional.of(new TableFunctionApplicationResult<>(functionHandle.getTableHandle(), functionHandle.getTableHandle().getColumns().orElseThrow(() -> new IllegalStateException("Columns are missing")))); + } + throw new IllegalStateException("Unsupported table function handle: " + handle.getClass().getSimpleName()); + }).withGetColumnHandles(getColumnHandles) + .build())); + queryRunner.createCatalog(TESTING_CATALOG, "testTVF"); + } + + @Test + public void testPrimitiveDefaultArgument() + { + assertQuery("SELECT boolean_column FROM TABLE(system.simple_table_function(column => 'boolean_column', ignored => 1))", "SELECT true WHERE false"); + + // skip the `ignored` argument. + assertQuery("SELECT boolean_column FROM TABLE(system.simple_table_function(column => 'boolean_column'))", + "SELECT true WHERE false"); + } + + @Test + public void testNoArgumentsPassed() + { + assertQuery("SELECT col FROM TABLE(system.simple_table_function())", + "SELECT true WHERE false"); + } +} From 877721081036a4562017008f7844170faf7f040f Mon Sep 17 00:00:00 2001 From: Anant Aneja <1797669+aaneja@users.noreply.github.com> Date: Wed, 30 Apr 2025 10:42:53 +0530 Subject: [PATCH 104/113] External auth support for presto-jdbc and presto-client Partial cherry pick of - https://github.com/trinodb/trino/commit/693cfb624d7a72c12a349ea4157b66bc8a43b93c https://github.com/trinodb/trino/commit/0031497cee8502b26fa1b44b9bd455d6abd59f98 https://github.com/trinodb/trino/commit/ceb2d1b8238c4b6949abf2e7f506aab83bf1ea02 https://github.com/trinodb/trino/commit/fe5335b0baf1b32fa189d7201d4df4ca03da0504 https://github.com/trinodb/trino/commit/050d089c518b45d8c9937e7b9a80711ea5f20219 Co-authored-by: Stephen Yugel Co-authored-by: David Phillips Co-authored-by: s2lomon Co-authored-by: Mateusz Gajewski --- presto-client/pom.xml | 22 + .../external/CompositeRedirectHandler.java | 53 ++ .../DesktopBrowserRedirectHandler.java | 41 ++ .../auth/external/ExternalAuthentication.java | 77 +++ .../auth/external/ExternalAuthenticator.java | 121 ++++ .../external/ExternalRedirectStrategy.java | 40 ++ .../client/auth/external/HttpTokenPoller.java | 193 +++++++ .../client/auth/external/KnownToken.java | 34 ++ .../client/auth/external/LocalKnownToken.java | 46 ++ .../auth/external/MemoryCachedKnownToken.java | 83 +++ .../auth/external/RedirectException.java | 28 + .../client/auth/external/RedirectHandler.java | 22 + .../external/SystemOpenRedirectHandler.java | 102 ++++ .../SystemOutPrintRedirectHandler.java | 27 + .../presto/client/auth/external/Token.java | 31 ++ .../client/auth/external/TokenPollResult.java | 80 +++ .../client/auth/external/TokenPoller.java | 24 + .../auth/external/MockRedirectHandler.java | 59 ++ .../client/auth/external/MockTokenPoller.java | 81 +++ .../external/TestExternalAuthentication.java | 130 +++++ .../external/TestExternalAuthenticator.java | 370 ++++++++++++ .../auth/external/TestHttpTokenPoller.java | 224 ++++++++ presto-jdbc/pom.xml | 41 +- .../presto/jdbc/ConnectionProperties.java | 64 +++ .../facebook/presto/jdbc/KnownTokenCache.java | 29 + .../facebook/presto/jdbc/PrestoDriverUri.java | 44 +- .../jdbc/TestJdbcExternalAuthentication.java | 525 ++++++++++++++++++ .../server/security/SecurityConfig.java | 5 +- 28 files changed, 2587 insertions(+), 9 deletions(-) create mode 100644 presto-client/src/main/java/com/facebook/presto/client/auth/external/CompositeRedirectHandler.java create mode 100644 presto-client/src/main/java/com/facebook/presto/client/auth/external/DesktopBrowserRedirectHandler.java create mode 100644 presto-client/src/main/java/com/facebook/presto/client/auth/external/ExternalAuthentication.java create mode 100644 presto-client/src/main/java/com/facebook/presto/client/auth/external/ExternalAuthenticator.java create mode 100644 presto-client/src/main/java/com/facebook/presto/client/auth/external/ExternalRedirectStrategy.java create mode 100644 presto-client/src/main/java/com/facebook/presto/client/auth/external/HttpTokenPoller.java create mode 100644 presto-client/src/main/java/com/facebook/presto/client/auth/external/KnownToken.java create mode 100644 presto-client/src/main/java/com/facebook/presto/client/auth/external/LocalKnownToken.java create mode 100644 presto-client/src/main/java/com/facebook/presto/client/auth/external/MemoryCachedKnownToken.java create mode 100644 presto-client/src/main/java/com/facebook/presto/client/auth/external/RedirectException.java create mode 100644 presto-client/src/main/java/com/facebook/presto/client/auth/external/RedirectHandler.java create mode 100644 presto-client/src/main/java/com/facebook/presto/client/auth/external/SystemOpenRedirectHandler.java create mode 100644 presto-client/src/main/java/com/facebook/presto/client/auth/external/SystemOutPrintRedirectHandler.java create mode 100644 presto-client/src/main/java/com/facebook/presto/client/auth/external/Token.java create mode 100644 presto-client/src/main/java/com/facebook/presto/client/auth/external/TokenPollResult.java create mode 100644 presto-client/src/main/java/com/facebook/presto/client/auth/external/TokenPoller.java create mode 100644 presto-client/src/test/java/com/facebook/presto/client/auth/external/MockRedirectHandler.java create mode 100644 presto-client/src/test/java/com/facebook/presto/client/auth/external/MockTokenPoller.java create mode 100644 presto-client/src/test/java/com/facebook/presto/client/auth/external/TestExternalAuthentication.java create mode 100644 presto-client/src/test/java/com/facebook/presto/client/auth/external/TestExternalAuthenticator.java create mode 100644 presto-client/src/test/java/com/facebook/presto/client/auth/external/TestHttpTokenPoller.java create mode 100644 presto-jdbc/src/main/java/com/facebook/presto/jdbc/KnownTokenCache.java create mode 100644 presto-jdbc/src/test/java/com/facebook/presto/jdbc/TestJdbcExternalAuthentication.java diff --git a/presto-client/pom.xml b/presto-client/pom.xml index e61853a18f6c4..3547e6c17c001 100644 --- a/presto-client/pom.xml +++ b/presto-client/pom.xml @@ -56,6 +56,11 @@ jackson-databind + + com.facebook.airlift + concurrent + + com.facebook.airlift json @@ -107,6 +112,11 @@ kotlin-stdlib-jdk8 + + net.jodah + failsafe + + com.facebook.presto @@ -114,6 +124,12 @@ test + + org.assertj + assertj-core + test + + org.testng testng @@ -142,6 +158,12 @@ drift-codec-utils test + + + com.squareup.okhttp3 + mockwebserver + test + diff --git a/presto-client/src/main/java/com/facebook/presto/client/auth/external/CompositeRedirectHandler.java b/presto-client/src/main/java/com/facebook/presto/client/auth/external/CompositeRedirectHandler.java new file mode 100644 index 0000000000000..bc1ebb0355a3a --- /dev/null +++ b/presto-client/src/main/java/com/facebook/presto/client/auth/external/CompositeRedirectHandler.java @@ -0,0 +1,53 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.client.auth.external; + +import java.net.URI; +import java.util.List; + +import static com.google.common.base.Preconditions.checkState; +import static com.google.common.collect.ImmutableList.toImmutableList; +import static java.util.Objects.requireNonNull; + +public class CompositeRedirectHandler + implements RedirectHandler +{ + private final List handlers; + + public CompositeRedirectHandler(List strategies) + { + this.handlers = requireNonNull(strategies, "strategies is null") + .stream() + .map(ExternalRedirectStrategy::getHandler) + .collect(toImmutableList()); + checkState(!handlers.isEmpty(), "Expected at least one external redirect handler"); + } + + @Override + public void redirectTo(URI uri) throws RedirectException + { + RedirectException redirectException = new RedirectException("Could not redirect to " + uri); + for (RedirectHandler handler : handlers) { + try { + handler.redirectTo(uri); + return; + } + catch (RedirectException e) { + redirectException.addSuppressed(e); + } + } + + throw redirectException; + } +} diff --git a/presto-client/src/main/java/com/facebook/presto/client/auth/external/DesktopBrowserRedirectHandler.java b/presto-client/src/main/java/com/facebook/presto/client/auth/external/DesktopBrowserRedirectHandler.java new file mode 100644 index 0000000000000..6918cba955c4b --- /dev/null +++ b/presto-client/src/main/java/com/facebook/presto/client/auth/external/DesktopBrowserRedirectHandler.java @@ -0,0 +1,41 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.client.auth.external; + +import java.io.IOException; +import java.net.URI; + +import static java.awt.Desktop.Action.BROWSE; +import static java.awt.Desktop.getDesktop; +import static java.awt.Desktop.isDesktopSupported; + +public final class DesktopBrowserRedirectHandler + implements RedirectHandler +{ + @Override + public void redirectTo(URI uri) + throws RedirectException + { + if (!isDesktopSupported() || !getDesktop().isSupported(BROWSE)) { + throw new RedirectException("Desktop Browser is not available. Make sure your Java process is not in headless mode (-Djava.awt.headless=false)"); + } + + try { + getDesktop().browse(uri); + } + catch (IOException e) { + throw new RedirectException("Failed to redirect", e); + } + } +} diff --git a/presto-client/src/main/java/com/facebook/presto/client/auth/external/ExternalAuthentication.java b/presto-client/src/main/java/com/facebook/presto/client/auth/external/ExternalAuthentication.java new file mode 100644 index 0000000000000..afee9bd899b42 --- /dev/null +++ b/presto-client/src/main/java/com/facebook/presto/client/auth/external/ExternalAuthentication.java @@ -0,0 +1,77 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.client.auth.external; + +import com.facebook.presto.client.ClientException; +import com.google.common.annotations.VisibleForTesting; + +import java.net.URI; +import java.time.Duration; +import java.util.Optional; + +import static java.util.Objects.requireNonNull; + +class ExternalAuthentication +{ + private final URI tokenUri; + private final Optional redirectUri; + + public ExternalAuthentication(URI tokenUri, Optional redirectUri) + { + this.tokenUri = requireNonNull(tokenUri, "tokenUri is null"); + this.redirectUri = requireNonNull(redirectUri, "redirectUri is null"); + } + + public Optional obtainToken(Duration timeout, RedirectHandler handler, TokenPoller poller) + { + redirectUri.ifPresent(handler::redirectTo); + + URI currentUri = tokenUri; + + long start = System.nanoTime(); + long timeoutNanos = timeout.toNanos(); + + while (true) { + long remaining = timeoutNanos - (System.nanoTime() - start); + if (remaining < 0) { + return Optional.empty(); + } + + TokenPollResult result = poller.pollForToken(currentUri, Duration.ofNanos(remaining)); + + if (result.isFailed()) { + throw new ClientException(result.getError()); + } + + if (result.isPending()) { + currentUri = result.getNextTokenUri(); + continue; + } + poller.tokenReceived(currentUri); + return Optional.of(result.getToken()); + } + } + + @VisibleForTesting + Optional getRedirectUri() + { + return redirectUri; + } + + @VisibleForTesting + URI getTokenUri() + { + return tokenUri; + } +} diff --git a/presto-client/src/main/java/com/facebook/presto/client/auth/external/ExternalAuthenticator.java b/presto-client/src/main/java/com/facebook/presto/client/auth/external/ExternalAuthenticator.java new file mode 100644 index 0000000000000..69bd0faf72e0f --- /dev/null +++ b/presto-client/src/main/java/com/facebook/presto/client/auth/external/ExternalAuthenticator.java @@ -0,0 +1,121 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.client.auth.external; + +import com.facebook.presto.client.ClientException; +import com.google.common.annotations.VisibleForTesting; +import okhttp3.Authenticator; +import okhttp3.Challenge; +import okhttp3.Interceptor; +import okhttp3.Request; +import okhttp3.Response; +import okhttp3.Route; + +import javax.annotation.Nullable; + +import java.io.IOException; +import java.net.URI; +import java.net.URISyntaxException; +import java.time.Duration; +import java.util.Map; +import java.util.Optional; + +import static com.google.common.net.HttpHeaders.AUTHORIZATION; +import static java.lang.String.format; +import static java.util.Objects.requireNonNull; + +public class ExternalAuthenticator + implements Authenticator, Interceptor +{ + public static final String TOKEN_URI_FIELD = "x_token_server"; + public static final String REDIRECT_URI_FIELD = "x_redirect_server"; + + private final TokenPoller tokenPoller; + private final RedirectHandler redirectHandler; + private final Duration timeout; + private final KnownToken knownToken; + + public ExternalAuthenticator(RedirectHandler redirect, TokenPoller tokenPoller, KnownToken knownToken, Duration timeout) + { + this.tokenPoller = requireNonNull(tokenPoller, "tokenPoller is null"); + this.redirectHandler = requireNonNull(redirect, "redirect is null"); + this.knownToken = requireNonNull(knownToken, "knownToken is null"); + this.timeout = requireNonNull(timeout, "timeout is null"); + } + + @Nullable + @Override + public Request authenticate(Route route, Response response) + { + knownToken.setupToken(() -> { + Optional authentication = toAuthentication(response); + if (!authentication.isPresent()) { + return Optional.empty(); + } + + return authentication.get().obtainToken(timeout, redirectHandler, tokenPoller); + }); + + return knownToken.getToken() + .map(token -> withBearerToken(response.request(), token)) + .orElse(null); + } + + @Override + public Response intercept(Chain chain) + throws IOException + { + Optional token = knownToken.getToken(); + if (token.isPresent()) { + return chain.proceed(withBearerToken(chain.request(), token.get())); + } + + return chain.proceed(chain.request()); + } + + private static Request withBearerToken(Request request, Token token) + { + return request.newBuilder() + .header(AUTHORIZATION, "Bearer " + token.token()) + .build(); + } + + @VisibleForTesting + static Optional toAuthentication(Response response) + { + for (Challenge challenge : response.challenges()) { + if (challenge.scheme().equalsIgnoreCase("Bearer")) { + Optional tokenUri = parseField(challenge.authParams(), TOKEN_URI_FIELD); + Optional redirectUri = parseField(challenge.authParams(), REDIRECT_URI_FIELD); + if (tokenUri.isPresent()) { + return Optional.of(new ExternalAuthentication(tokenUri.get(), redirectUri)); + } + } + } + + return Optional.empty(); + } + + private static Optional parseField(Map fields, String key) + { + return Optional.ofNullable(fields.get(key)).map(value -> { + try { + return new URI(value); + } + catch (URISyntaxException e) { + throw new ClientException(format("Failed to parse URI for field '%s'", key), e); + } + }); + } +} diff --git a/presto-client/src/main/java/com/facebook/presto/client/auth/external/ExternalRedirectStrategy.java b/presto-client/src/main/java/com/facebook/presto/client/auth/external/ExternalRedirectStrategy.java new file mode 100644 index 0000000000000..07b004f0c15fe --- /dev/null +++ b/presto-client/src/main/java/com/facebook/presto/client/auth/external/ExternalRedirectStrategy.java @@ -0,0 +1,40 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.client.auth.external; + +import com.google.common.collect.ImmutableList; + +import static java.util.Objects.requireNonNull; + +public enum ExternalRedirectStrategy +{ + DESKTOP_OPEN(new DesktopBrowserRedirectHandler()), + SYSTEM_OPEN(new SystemOpenRedirectHandler()), + PRINT(new SystemOutPrintRedirectHandler()), + OPEN(new CompositeRedirectHandler(ImmutableList.of(SYSTEM_OPEN, DESKTOP_OPEN))), + ALL(new CompositeRedirectHandler(ImmutableList.of(OPEN, PRINT))) + /**/; + + private final RedirectHandler handler; + + ExternalRedirectStrategy(RedirectHandler handler) + { + this.handler = requireNonNull(handler, "handler is null"); + } + + public RedirectHandler getHandler() + { + return handler; + } +} diff --git a/presto-client/src/main/java/com/facebook/presto/client/auth/external/HttpTokenPoller.java b/presto-client/src/main/java/com/facebook/presto/client/auth/external/HttpTokenPoller.java new file mode 100644 index 0000000000000..b06b224ee80ac --- /dev/null +++ b/presto-client/src/main/java/com/facebook/presto/client/auth/external/HttpTokenPoller.java @@ -0,0 +1,193 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.client.auth.external; + +import com.facebook.airlift.json.JsonCodec; +import com.facebook.presto.client.JsonResponse; +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import net.jodah.failsafe.Failsafe; +import net.jodah.failsafe.FailsafeException; +import net.jodah.failsafe.RetryPolicy; +import okhttp3.HttpUrl; +import okhttp3.OkHttpClient; +import okhttp3.Request; +import okhttp3.Response; + +import java.io.IOException; +import java.io.UncheckedIOException; +import java.net.URI; +import java.time.Duration; +import java.util.Optional; +import java.util.function.Consumer; +import java.util.function.Supplier; + +import static com.facebook.airlift.json.JsonCodec.jsonCodec; +import static com.facebook.presto.client.JsonResponse.execute; +import static com.google.common.base.MoreObjects.firstNonNull; +import static com.google.common.net.HttpHeaders.USER_AGENT; +import static java.lang.String.format; +import static java.net.HttpURLConnection.HTTP_INTERNAL_ERROR; +import static java.net.HttpURLConnection.HTTP_OK; +import static java.net.HttpURLConnection.HTTP_UNAVAILABLE; +import static java.time.temporal.ChronoUnit.MILLIS; +import static java.util.Objects.requireNonNull; + +public class HttpTokenPoller + implements TokenPoller +{ + private static final JsonCodec TOKEN_POLL_CODEC = jsonCodec(TokenPollRepresentation.class); + private static final String USER_AGENT_VALUE = "PrestoTokenPoller/" + + firstNonNull(HttpTokenPoller.class.getPackage().getImplementationVersion(), "unknown"); + + private final Supplier client; + + public HttpTokenPoller(OkHttpClient client) + { + requireNonNull(client, "client is null"); + this.client = () -> client; + } + + public HttpTokenPoller(OkHttpClient client, Consumer refreshableClientConfig) + { + requireNonNull(client, "client is null"); + requireNonNull(refreshableClientConfig, "refreshableClientConfig is null"); + + this.client = () -> { + OkHttpClient.Builder builder = client.newBuilder(); + refreshableClientConfig.accept(builder); + return builder.build(); + }; + } + + @Override + public TokenPollResult pollForToken(URI tokenUri, Duration timeout) + { + try { + return Failsafe.with(new RetryPolicy() + .withMaxAttempts(-1) + .withMaxDuration(timeout) + .withBackoff(100, 500, MILLIS) + .handle(IOException.class)) + .get(() -> executePoll(prepareRequestBuilder(tokenUri).build())); + } + catch (FailsafeException e) { + if (e.getCause() instanceof IOException) { + throw new UncheckedIOException((IOException) e.getCause()); + } + throw e; + } + } + + @Override + public void tokenReceived(URI tokenUri) + { + try { + Failsafe.with(new RetryPolicy() + .withMaxAttempts(-1) + .withMaxDuration(Duration.ofSeconds(4)) + .withBackoff(100, 500, MILLIS) + .handleResultIf(code -> code >= HTTP_INTERNAL_ERROR)) + .get(() -> { + Request request = prepareRequestBuilder(tokenUri) + .delete() + .build(); + try (Response response = client.get().newCall(request) + .execute()) { + return response.code(); + } + }); + } + catch (FailsafeException e) { + if (e.getCause() instanceof IOException) { + throw new UncheckedIOException((IOException) e.getCause()); + } + throw e; + } + } + + private static Request.Builder prepareRequestBuilder(URI tokenUri) + { + HttpUrl url = HttpUrl.get(tokenUri); + if (url == null) { + throw new IllegalArgumentException("Invalid token URI: " + tokenUri); + } + + return new Request.Builder() + .url(url) + .addHeader(USER_AGENT, USER_AGENT_VALUE); + } + + private TokenPollResult executePoll(Request request) + throws IOException + { + JsonResponse response = executeRequest(request); + + if ((response.getStatusCode() == HTTP_OK) && response.hasValue()) { + return response.getValue().toResult(); + } + + Optional responseBody = Optional.ofNullable(response.getResponseBody()); + String message = format("Request to %s failed: %s [Error: %s]", request.url(), response, responseBody.orElse("")); + + if (response.getStatusCode() == HTTP_UNAVAILABLE) { + throw new IOException(message); + } + + return TokenPollResult.failed(message); + } + + private JsonResponse executeRequest(Request request) + throws IOException + { + try { + return execute(TOKEN_POLL_CODEC, client.get(), request); + } + catch (UncheckedIOException e) { + throw e.getCause(); + } + } + + public static class TokenPollRepresentation + { + private final String token; + private final URI nextUri; + private final String error; + + @JsonCreator + public TokenPollRepresentation( + @JsonProperty("token") String token, + @JsonProperty("nextUri") URI nextUri, + @JsonProperty("error") String error) + { + this.token = token; + this.nextUri = nextUri; + this.error = error; + } + + TokenPollResult toResult() + { + if (token != null) { + return TokenPollResult.successful(new Token(token)); + } + if (error != null) { + return TokenPollResult.failed(error); + } + if (nextUri != null) { + return TokenPollResult.pending(nextUri); + } + return TokenPollResult.failed("Failed to poll for token. No fields set in response."); + } + } +} diff --git a/presto-client/src/main/java/com/facebook/presto/client/auth/external/KnownToken.java b/presto-client/src/main/java/com/facebook/presto/client/auth/external/KnownToken.java new file mode 100644 index 0000000000000..ae57de609c02c --- /dev/null +++ b/presto-client/src/main/java/com/facebook/presto/client/auth/external/KnownToken.java @@ -0,0 +1,34 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.client.auth.external; + +import java.util.Optional; +import java.util.function.Supplier; + +public interface KnownToken +{ + Optional getToken(); + + void setupToken(Supplier> tokenSource); + + static KnownToken local() + { + return new LocalKnownToken(); + } + + static KnownToken memoryCached() + { + return MemoryCachedKnownToken.INSTANCE; + } +} diff --git a/presto-client/src/main/java/com/facebook/presto/client/auth/external/LocalKnownToken.java b/presto-client/src/main/java/com/facebook/presto/client/auth/external/LocalKnownToken.java new file mode 100644 index 0000000000000..31ff4bd89d525 --- /dev/null +++ b/presto-client/src/main/java/com/facebook/presto/client/auth/external/LocalKnownToken.java @@ -0,0 +1,46 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.client.auth.external; + +import javax.annotation.concurrent.NotThreadSafe; + +import java.util.Optional; +import java.util.function.Supplier; + +import static java.util.Objects.requireNonNull; + +/** + * LocalKnownToken class keeps the token on its field + * and it's designed to use it in fully serialized manner. + */ +@NotThreadSafe +class LocalKnownToken + implements KnownToken +{ + private Optional knownToken = Optional.empty(); + + @Override + public Optional getToken() + { + return knownToken; + } + + @Override + public void setupToken(Supplier> tokenSource) + { + requireNonNull(tokenSource, "tokenSource is null"); + + knownToken = tokenSource.get(); + } +} diff --git a/presto-client/src/main/java/com/facebook/presto/client/auth/external/MemoryCachedKnownToken.java b/presto-client/src/main/java/com/facebook/presto/client/auth/external/MemoryCachedKnownToken.java new file mode 100644 index 0000000000000..104841ac2dd14 --- /dev/null +++ b/presto-client/src/main/java/com/facebook/presto/client/auth/external/MemoryCachedKnownToken.java @@ -0,0 +1,83 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.client.auth.external; + +import javax.annotation.concurrent.ThreadSafe; + +import java.util.Optional; +import java.util.concurrent.locks.Lock; +import java.util.concurrent.locks.ReadWriteLock; +import java.util.concurrent.locks.ReentrantReadWriteLock; +import java.util.function.Supplier; + +/** + * This KnownToken instance forces all Connections to reuse same token. + * Every time an existing token is considered to be invalid each Connection + * will try to obtain a new token, but only the first one will actually do the job, + * where every other connection will be waiting on readLock + * until obtaining new token finishes. + *

+ * In general the game is to reuse same token and obtain it only once, no matter how + * many Connections will be actively using it. It's very important as obtaining the new token + * will take minutes, as it mostly requires user thinking time. + */ +@ThreadSafe +public class MemoryCachedKnownToken + implements KnownToken +{ + public static final MemoryCachedKnownToken INSTANCE = new MemoryCachedKnownToken(); + + private final ReadWriteLock lock = new ReentrantReadWriteLock(); + private final Lock readLock = lock.readLock(); + private final Lock writeLock = lock.writeLock(); + private Optional knownToken = Optional.empty(); + + private MemoryCachedKnownToken() + { + } + + @Override + public Optional getToken() + { + try { + readLock.lockInterruptibly(); + return knownToken; + } + catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + finally { + readLock.unlock(); + } + } + + @Override + public void setupToken(Supplier> tokenSource) + { + // Try to lock and generate new token. If some other thread (Connection) has + // already obtained writeLock and is generating new token, then skipp this + // to block on getToken() + if (writeLock.tryLock()) { + try { + // Clear knownToken before obtaining new token, as it might fail leaving old invalid token. + knownToken = Optional.empty(); + knownToken = tokenSource.get(); + } + finally { + writeLock.unlock(); + } + } + } +} diff --git a/presto-client/src/main/java/com/facebook/presto/client/auth/external/RedirectException.java b/presto-client/src/main/java/com/facebook/presto/client/auth/external/RedirectException.java new file mode 100644 index 0000000000000..0c5a1f5e149fa --- /dev/null +++ b/presto-client/src/main/java/com/facebook/presto/client/auth/external/RedirectException.java @@ -0,0 +1,28 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.client.auth.external; + +public class RedirectException + extends RuntimeException +{ + public RedirectException(String message, Throwable throwable) + { + super(message, throwable); + } + + public RedirectException(String message) + { + super(message); + } +} diff --git a/presto-client/src/main/java/com/facebook/presto/client/auth/external/RedirectHandler.java b/presto-client/src/main/java/com/facebook/presto/client/auth/external/RedirectHandler.java new file mode 100644 index 0000000000000..710be10238884 --- /dev/null +++ b/presto-client/src/main/java/com/facebook/presto/client/auth/external/RedirectHandler.java @@ -0,0 +1,22 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.client.auth.external; + +import java.net.URI; + +public interface RedirectHandler +{ + void redirectTo(URI uri) + throws RedirectException; +} diff --git a/presto-client/src/main/java/com/facebook/presto/client/auth/external/SystemOpenRedirectHandler.java b/presto-client/src/main/java/com/facebook/presto/client/auth/external/SystemOpenRedirectHandler.java new file mode 100644 index 0000000000000..d63841b26d830 --- /dev/null +++ b/presto-client/src/main/java/com/facebook/presto/client/auth/external/SystemOpenRedirectHandler.java @@ -0,0 +1,102 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.client.auth.external; + +import com.google.common.base.Splitter; +import com.google.common.collect.ImmutableList; + +import java.io.File; +import java.io.FileNotFoundException; +import java.io.IOException; +import java.net.URI; +import java.nio.file.Paths; +import java.util.List; +import java.util.Optional; + +import static java.lang.String.format; +import static java.util.Locale.ENGLISH; + +public class SystemOpenRedirectHandler + implements RedirectHandler +{ + private static final List LINUX_BROWSERS = ImmutableList.of( + "xdg-open", + "gnome-open", + "kde-open", + "chromium", + "google", + "google-chrome", + "firefox", + "mozilla", + "opera", + "epiphany", + "konqueror"); + + private static final String MACOS_OPEN_COMMAND = "open"; + private static final String WINDOWS_OPEN_COMMAND = "rundll32 url.dll,FileProtocolHandler"; + + private static final Splitter SPLITTER = Splitter.on(":") + .omitEmptyStrings() + .trimResults(); + + @Override + public void redirectTo(URI uri) + throws RedirectException + { + String operatingSystem = System.getProperty("os.name").toLowerCase(ENGLISH); + + try { + if (operatingSystem.contains("mac")) { + exec(uri, MACOS_OPEN_COMMAND); + } + else if (operatingSystem.contains("windows")) { + exec(uri, WINDOWS_OPEN_COMMAND); + } + else { + String executablePath = findLinuxBrowser() + .orElseThrow(() -> new FileNotFoundException("Could not find any known linux browser in $PATH")); + exec(uri, executablePath); + } + } + catch (IOException e) { + throw new RedirectException(format("Could not open uri %s", uri), e); + } + } + + private static Optional findLinuxBrowser() + { + List paths = SPLITTER.splitToList(System.getenv("PATH")); + for (String path : paths) { + File[] found = Paths.get(path) + .toFile() + .listFiles((dir, name) -> LINUX_BROWSERS.contains(name)); + + if (found == null) { + continue; + } + + if (found.length > 0) { + return Optional.of(found[0].getPath()); + } + } + + return Optional.empty(); + } + + private static void exec(URI uri, String openCommand) + throws IOException + { + Runtime.getRuntime().exec(openCommand + " " + uri.toString()); + } +} diff --git a/presto-client/src/main/java/com/facebook/presto/client/auth/external/SystemOutPrintRedirectHandler.java b/presto-client/src/main/java/com/facebook/presto/client/auth/external/SystemOutPrintRedirectHandler.java new file mode 100644 index 0000000000000..d0f3c7954dbd7 --- /dev/null +++ b/presto-client/src/main/java/com/facebook/presto/client/auth/external/SystemOutPrintRedirectHandler.java @@ -0,0 +1,27 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.client.auth.external; + +import java.net.URI; + +public class SystemOutPrintRedirectHandler + implements RedirectHandler +{ + @Override + public void redirectTo(URI uri) throws RedirectException + { + System.out.println("External authentication required. Please go to:"); + System.out.println(uri.toString()); + } +} diff --git a/presto-client/src/main/java/com/facebook/presto/client/auth/external/Token.java b/presto-client/src/main/java/com/facebook/presto/client/auth/external/Token.java new file mode 100644 index 0000000000000..82890f3f6b2df --- /dev/null +++ b/presto-client/src/main/java/com/facebook/presto/client/auth/external/Token.java @@ -0,0 +1,31 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.client.auth.external; + +import static java.util.Objects.requireNonNull; + +class Token +{ + private final String token; + + public Token(String token) + { + this.token = requireNonNull(token, "token is null"); + } + + public String token() + { + return token; + } +} diff --git a/presto-client/src/main/java/com/facebook/presto/client/auth/external/TokenPollResult.java b/presto-client/src/main/java/com/facebook/presto/client/auth/external/TokenPollResult.java new file mode 100644 index 0000000000000..3cd13b5dd3710 --- /dev/null +++ b/presto-client/src/main/java/com/facebook/presto/client/auth/external/TokenPollResult.java @@ -0,0 +1,80 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.client.auth.external; + +import java.net.URI; +import java.util.Optional; + +import static java.util.Objects.requireNonNull; + +public class TokenPollResult +{ + private enum State + { + PENDING, SUCCESSFUL, FAILED + } + + private final State state; + private final Optional errorMessage; + private final Optional nextTokenUri; + private final Optional token; + + public static TokenPollResult failed(String error) + { + return new TokenPollResult(State.FAILED, null, error, null); + } + + public static TokenPollResult pending(URI uri) + { + return new TokenPollResult(State.PENDING, null, null, uri); + } + + public static TokenPollResult successful(Token token) + { + return new TokenPollResult(State.SUCCESSFUL, token, null, null); + } + + private TokenPollResult(State state, Token token, String error, URI nextTokenUri) + { + this.state = requireNonNull(state, "state is null"); + this.token = Optional.ofNullable(token); + this.errorMessage = Optional.ofNullable(error); + this.nextTokenUri = Optional.ofNullable(nextTokenUri); + } + + public boolean isFailed() + { + return state == State.FAILED; + } + + public boolean isPending() + { + return state == State.PENDING; + } + + public Token getToken() + { + return token.orElseThrow(() -> new IllegalStateException("state is " + state)); + } + + public String getError() + { + return errorMessage.orElseThrow(() -> new IllegalStateException("state is " + state)); + } + + public URI getNextTokenUri() + { + return nextTokenUri.orElseThrow(() -> new IllegalStateException("state is " + state)); + } +} diff --git a/presto-client/src/main/java/com/facebook/presto/client/auth/external/TokenPoller.java b/presto-client/src/main/java/com/facebook/presto/client/auth/external/TokenPoller.java new file mode 100644 index 0000000000000..ea228137ab688 --- /dev/null +++ b/presto-client/src/main/java/com/facebook/presto/client/auth/external/TokenPoller.java @@ -0,0 +1,24 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.client.auth.external; + +import java.net.URI; +import java.time.Duration; + +public interface TokenPoller +{ + TokenPollResult pollForToken(URI tokenUri, Duration timeout); + + void tokenReceived(URI tokenUri); +} diff --git a/presto-client/src/test/java/com/facebook/presto/client/auth/external/MockRedirectHandler.java b/presto-client/src/test/java/com/facebook/presto/client/auth/external/MockRedirectHandler.java new file mode 100644 index 0000000000000..f4b8c6cd07633 --- /dev/null +++ b/presto-client/src/test/java/com/facebook/presto/client/auth/external/MockRedirectHandler.java @@ -0,0 +1,59 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.client.auth.external; + +import java.net.URI; +import java.time.Duration; +import java.util.concurrent.atomic.AtomicInteger; + +public class MockRedirectHandler + implements RedirectHandler +{ + private URI redirectedTo; + private AtomicInteger redirectionCount = new AtomicInteger(0); + private Duration redirectTime; + + @Override + public void redirectTo(URI uri) + throws RedirectException + { + redirectedTo = uri; + redirectionCount.incrementAndGet(); + try { + if (redirectTime != null) { + Thread.sleep(redirectTime.toMillis()); + } + } + catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + } + + public URI redirectedTo() + { + return redirectedTo; + } + + public int getRedirectionCount() + { + return redirectionCount.get(); + } + + public MockRedirectHandler sleepOnRedirect(Duration redirectTime) + { + this.redirectTime = redirectTime; + return this; + } +} diff --git a/presto-client/src/test/java/com/facebook/presto/client/auth/external/MockTokenPoller.java b/presto-client/src/test/java/com/facebook/presto/client/auth/external/MockTokenPoller.java new file mode 100644 index 0000000000000..6733a76ebb64d --- /dev/null +++ b/presto-client/src/test/java/com/facebook/presto/client/auth/external/MockTokenPoller.java @@ -0,0 +1,81 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.client.auth.external; + +import com.google.common.collect.ImmutableList; + +import java.net.URI; +import java.time.Duration; +import java.util.Map; +import java.util.concurrent.BlockingDeque; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.LinkedBlockingDeque; +import java.util.function.Function; + +public final class MockTokenPoller + implements TokenPoller +{ + private final Map> results = new ConcurrentHashMap<>(); + private URI tokenReceivedUri; + + public static TokenPoller onPoll(Function pollingStrategy) + { + return new TokenPoller() + { + @Override + public TokenPollResult pollForToken(URI tokenUri, Duration timeout) + { + return pollingStrategy.apply(tokenUri); + } + + @Override + public void tokenReceived(URI tokenUri) + { + } + }; + } + + public MockTokenPoller withResult(URI tokenUri, TokenPollResult result) + { + results.compute(tokenUri, (uri, queue) -> { + if (queue == null) { + return new LinkedBlockingDeque<>(ImmutableList.of(result)); + } + queue.add(result); + return queue; + }); + return this; + } + + @Override + public TokenPollResult pollForToken(URI tokenUri, Duration ignored) + { + BlockingDeque queue = results.get(tokenUri); + if (queue == null) { + throw new IllegalArgumentException("Unknown token URI: " + tokenUri); + } + return queue.remove(); + } + + @Override + public void tokenReceived(URI tokenUri) + { + this.tokenReceivedUri = tokenUri; + } + + public URI tokenReceivedUri() + { + return tokenReceivedUri; + } +} diff --git a/presto-client/src/test/java/com/facebook/presto/client/auth/external/TestExternalAuthentication.java b/presto-client/src/test/java/com/facebook/presto/client/auth/external/TestExternalAuthentication.java new file mode 100644 index 0000000000000..e7434c5ac242e --- /dev/null +++ b/presto-client/src/test/java/com/facebook/presto/client/auth/external/TestExternalAuthentication.java @@ -0,0 +1,130 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.client.auth.external; + +import com.facebook.presto.client.ClientException; +import org.testng.annotations.Test; + +import java.io.IOException; +import java.io.UncheckedIOException; +import java.net.URI; +import java.time.Duration; +import java.util.Optional; +import java.util.concurrent.TimeUnit; + +import static com.google.common.util.concurrent.Uninterruptibles.sleepUninterruptibly; +import static java.net.URI.create; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +public class TestExternalAuthentication +{ + private static final String AUTH_TOKEN = "authToken"; + private static final URI REDIRECT_URI = create("https://redirect.uri"); + private static final URI TOKEN_URI = create("https://token.uri"); + private static final Duration TIMEOUT = Duration.ofSeconds(1); + + @Test + public void testObtainTokenWhenTokenAlreadyExists() + { + MockRedirectHandler redirectHandler = new MockRedirectHandler(); + + MockTokenPoller poller = new MockTokenPoller() + .withResult(TOKEN_URI, TokenPollResult.successful(new Token(AUTH_TOKEN))); + + Optional token = new ExternalAuthentication(TOKEN_URI, Optional.of(REDIRECT_URI)) + .obtainToken(TIMEOUT, redirectHandler, poller); + + assertThat(redirectHandler.redirectedTo()).isEqualTo(REDIRECT_URI); + assertThat(token).map(Token::token).hasValue(AUTH_TOKEN); + assertThat(poller.tokenReceivedUri()).isEqualTo(TOKEN_URI); + } + + @Test + public void testObtainTokenWhenTokenIsReadyAtSecondAttempt() + { + RedirectHandler redirectHandler = new MockRedirectHandler(); + + URI nextTokenUri = TOKEN_URI.resolve("/next"); + MockTokenPoller poller = new MockTokenPoller() + .withResult(TOKEN_URI, TokenPollResult.pending(nextTokenUri)) + .withResult(nextTokenUri, TokenPollResult.successful(new Token(AUTH_TOKEN))); + + Optional token = new ExternalAuthentication(TOKEN_URI, Optional.of(REDIRECT_URI)) + .obtainToken(TIMEOUT, redirectHandler, poller); + + assertThat(token).map(Token::token).hasValue(AUTH_TOKEN); + assertThat(poller.tokenReceivedUri()).isEqualTo(nextTokenUri); + } + + @Test + public void testObtainTokenWhenTokenIsNeverAvailable() + { + RedirectHandler redirectHandler = new MockRedirectHandler(); + + TokenPoller poller = MockTokenPoller.onPoll(tokenUri -> { + sleepUninterruptibly(20, TimeUnit.MILLISECONDS); + return TokenPollResult.pending(TOKEN_URI); + }); + + Optional token = new ExternalAuthentication(TOKEN_URI, Optional.of(REDIRECT_URI)) + .obtainToken(TIMEOUT, redirectHandler, poller); + + assertThat(token).isEmpty(); + } + + @Test + public void testObtainTokenWhenPollingFails() + { + RedirectHandler redirectHandler = new MockRedirectHandler(); + + TokenPoller poller = new MockTokenPoller() + .withResult(TOKEN_URI, TokenPollResult.failed("error")); + + assertThatThrownBy(() -> new ExternalAuthentication(TOKEN_URI, Optional.of(REDIRECT_URI)) + .obtainToken(TIMEOUT, redirectHandler, poller)) + .isInstanceOf(ClientException.class) + .hasMessage("error"); + } + + @Test + public void testObtainTokenWhenPollingFailsWithException() + { + RedirectHandler redirectHandler = new MockRedirectHandler(); + + TokenPoller poller = MockTokenPoller.onPoll(uri -> { + throw new UncheckedIOException(new IOException("polling error")); + }); + + assertThatThrownBy(() -> new ExternalAuthentication(TOKEN_URI, Optional.of(REDIRECT_URI)) + .obtainToken(TIMEOUT, redirectHandler, poller)) + .isInstanceOf(UncheckedIOException.class) + .hasRootCauseInstanceOf(IOException.class); + } + + @Test + public void testObtainTokenWhenNoRedirectUriHasBeenProvided() + { + MockRedirectHandler redirectHandler = new MockRedirectHandler(); + + TokenPoller poller = new MockTokenPoller() + .withResult(TOKEN_URI, TokenPollResult.successful(new Token(AUTH_TOKEN))); + + Optional token = new ExternalAuthentication(TOKEN_URI, Optional.empty()) + .obtainToken(TIMEOUT, redirectHandler, poller); + + assertThat(redirectHandler.redirectedTo()).isNull(); + assertThat(token).map(Token::token).hasValue(AUTH_TOKEN); + } +} diff --git a/presto-client/src/test/java/com/facebook/presto/client/auth/external/TestExternalAuthenticator.java b/presto-client/src/test/java/com/facebook/presto/client/auth/external/TestExternalAuthenticator.java new file mode 100644 index 0000000000000..86cd89fc6dfa9 --- /dev/null +++ b/presto-client/src/test/java/com/facebook/presto/client/auth/external/TestExternalAuthenticator.java @@ -0,0 +1,370 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.client.auth.external; + +import com.facebook.presto.client.ClientException; +import com.google.common.collect.ImmutableList; +import okhttp3.HttpUrl; +import okhttp3.Protocol; +import okhttp3.Request; +import okhttp3.Response; +import org.assertj.core.api.ListAssert; +import org.assertj.core.api.ThrowableAssert; +import org.testng.annotations.AfterClass; +import org.testng.annotations.Test; + +import java.net.URI; +import java.net.URISyntaxException; +import java.time.Duration; +import java.util.ArrayList; +import java.util.List; +import java.util.Optional; +import java.util.concurrent.Callable; +import java.util.concurrent.CancellationException; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Future; +import java.util.stream.Stream; + +import static com.facebook.airlift.concurrent.Threads.daemonThreadsNamed; +import static com.facebook.presto.client.auth.external.ExternalAuthenticator.TOKEN_URI_FIELD; +import static com.facebook.presto.client.auth.external.ExternalAuthenticator.toAuthentication; +import static com.facebook.presto.client.auth.external.MockTokenPoller.onPoll; +import static com.facebook.presto.client.auth.external.TokenPollResult.successful; +import static com.google.common.base.Preconditions.checkState; +import static com.google.common.collect.ImmutableList.toImmutableList; +import static com.google.common.net.HttpHeaders.AUTHORIZATION; +import static com.google.common.net.HttpHeaders.WWW_AUTHENTICATE; +import static java.lang.String.format; +import static java.net.HttpURLConnection.HTTP_UNAUTHORIZED; +import static java.net.URI.create; +import static java.util.concurrent.Executors.newCachedThreadPool; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +@Test(singleThreaded = true) +public class TestExternalAuthenticator +{ + private static final ExecutorService executor = newCachedThreadPool(daemonThreadsNamed(TestExternalAuthenticator.class.getName() + "-%d")); + + @AfterClass(alwaysRun = true) + public void shutDownThreadPool() + { + executor.shutdownNow(); + } + + @Test + public void testChallengeWithOnlyTokenServerUri() + { + assertThat(buildAuthentication("Bearer x_token_server=\"http://token.uri\"")) + .hasValueSatisfying(authentication -> { + assertThat(authentication.getRedirectUri()).isEmpty(); + assertThat(authentication.getTokenUri()).isEqualTo(create("http://token.uri")); + }); + } + + @Test + public void testChallengeWithBothUri() + { + assertThat(buildAuthentication("Bearer x_redirect_server=\"http://redirect.uri\", x_token_server=\"http://token.uri\"")) + .hasValueSatisfying(authentication -> { + assertThat(authentication.getRedirectUri()).hasValue(create("http://redirect.uri")); + assertThat(authentication.getTokenUri()).isEqualTo(create("http://token.uri")); + }); + } + + @Test + public void testChallengeWithValuesWithoutQuotes() + { + // this is legal according to RFC 7235 + assertThat(buildAuthentication("Bearer x_redirect_server=http://redirect.uri, x_token_server=http://token.uri")) + .hasValueSatisfying(authentication -> { + assertThat(authentication.getRedirectUri()).hasValue(create("http://redirect.uri")); + assertThat(authentication.getTokenUri()).isEqualTo(create("http://token.uri")); + }); + } + + @Test + public void testChallengeWithAdditionalFields() + { + assertThat(buildAuthentication("Bearer type=\"token\", x_redirect_server=\"http://redirect.uri\", x_token_server=\"http://token.uri\", description=\"oauth challenge\"")) + .hasValueSatisfying(authentication -> { + assertThat(authentication.getRedirectUri()).hasValue(create("http://redirect.uri")); + assertThat(authentication.getTokenUri()).isEqualTo(create("http://token.uri")); + }); + } + + @Test + public void testInvalidChallenges() + { + // no authentication parameters + assertThat(buildAuthentication("Bearer")).isEmpty(); + + // no Bearer scheme prefix + assertThat(buildAuthentication("x_redirect_server=\"http://redirect.uri\", x_token_server=\"http://token.uri\"")).isEmpty(); + + // space instead of comma + assertThat(buildAuthentication("Bearer x_redirect_server=\"http://redirect.uri\" x_token_server=\"http://token.uri\"")).isEmpty(); + + // equals sign instead of comma + assertThat(buildAuthentication("Bearer x_redirect_server=\"http://redirect.uri\"=x_token_server=\"http://token.uri\"")).isEmpty(); + } + + @Test + public void testChallengeWithMalformedUri() + { + assertThatThrownBy(() -> buildAuthentication("Bearer x_token_server=\"http://[1.1.1.1]\"")) + .isInstanceOf(ClientException.class) + .hasMessageContaining(format("Failed to parse URI for field '%s'", TOKEN_URI_FIELD)) + .hasRootCauseInstanceOf(URISyntaxException.class); + } + + @Test + public void testAuthentication() + { + MockTokenPoller tokenPoller = new MockTokenPoller() + .withResult(URI.create("http://token.uri"), successful(new Token("valid-token"))); + ExternalAuthenticator authenticator = new ExternalAuthenticator(uri -> {}, tokenPoller, KnownToken.local(), Duration.ofSeconds(1)); + + Request authenticated = authenticator.authenticate(null, getUnauthorizedResponse("Bearer x_token_server=\"http://token.uri\"")); + + assertThat(authenticated.headers(AUTHORIZATION)) + .containsExactly("Bearer valid-token"); + } + + @Test + public void testReAuthenticationAfterRejectingToken() + { + MockTokenPoller tokenPoller = new MockTokenPoller() + .withResult(URI.create("http://token.uri"), successful(new Token("first-token"))) + .withResult(URI.create("http://token.uri"), successful(new Token("second-token"))); + ExternalAuthenticator authenticator = new ExternalAuthenticator(uri -> {}, tokenPoller, KnownToken.local(), Duration.ofSeconds(1)); + + Request request = authenticator.authenticate(null, getUnauthorizedResponse("Bearer x_token_server=\"http://token.uri\"")); + Request reAuthenticated = authenticator.authenticate(null, getUnauthorizedResponse("Bearer x_token_server=\"http://token.uri\"", request)); + + assertThat(reAuthenticated.headers(AUTHORIZATION)) + .containsExactly("Bearer second-token"); + } + + @Test(timeOut = 2000) + public void testAuthenticationFromMultipleThreadsWithLocallyStoredToken() + { + MockTokenPoller tokenPoller = new MockTokenPoller() + .withResult(URI.create("http://token.uri"), successful(new Token("valid-token-1"))) + .withResult(URI.create("http://token.uri"), successful(new Token("valid-token-2"))) + .withResult(URI.create("http://token.uri"), successful(new Token("valid-token-3"))) + .withResult(URI.create("http://token.uri"), successful(new Token("valid-token-4"))); + MockRedirectHandler redirectHandler = new MockRedirectHandler(); + + List> requests = times( + 4, + () -> new ExternalAuthenticator(redirectHandler, tokenPoller, KnownToken.local(), Duration.ofSeconds(1)) + .authenticate(null, getUnauthorizedResponse("Bearer x_token_server=\"http://token.uri\", x_redirect_server=\"http://redirect.uri\""))) + .map(executor::submit) + .collect(toImmutableList()); + + ConcurrentRequestAssertion assertion = new ConcurrentRequestAssertion(requests); + assertion.requests() + .extracting(Request::headers) + .extracting(headers -> headers.get(AUTHORIZATION)) + .contains("Bearer valid-token-1", "Bearer valid-token-2", "Bearer valid-token-3", "Bearer valid-token-4"); + assertion.assertThatNoExceptionsHasBeenThrown(); + assertThat(redirectHandler.getRedirectionCount()).isEqualTo(4); + } + + @Test(timeOut = 2000) + public void testAuthenticationFromMultipleThreadsWithCachedToken() + { + MockTokenPoller tokenPoller = new MockTokenPoller() + .withResult(URI.create("http://token.uri"), successful(new Token("valid-token"))); + MockRedirectHandler redirectHandler = new MockRedirectHandler() + .sleepOnRedirect(Duration.ofSeconds(1)); + + List> requests = times( + 2, + () -> new ExternalAuthenticator(redirectHandler, tokenPoller, KnownToken.memoryCached(), Duration.ofSeconds(1)) + .authenticate(null, getUnauthorizedResponse("Bearer x_token_server=\"http://token.uri\", x_redirect_server=\"http://redirect.uri\""))) + .map(executor::submit) + .collect(toImmutableList()); + + ConcurrentRequestAssertion assertion = new ConcurrentRequestAssertion(requests); + assertion.requests() + .extracting(Request::headers) + .extracting(headers -> headers.get(AUTHORIZATION)) + .containsOnly("Bearer valid-token"); + assertion.assertThatNoExceptionsHasBeenThrown(); + assertThat(redirectHandler.getRedirectionCount()).isEqualTo(1); + } + + @Test(timeOut = 2000) + public void testAuthenticationFromMultipleThreadsWithCachedTokenAfterAuthenticateFails() + { + MockTokenPoller tokenPoller = new MockTokenPoller() + .withResult(URI.create("http://token.uri"), TokenPollResult.successful(new Token("first-token"))) + .withResult(URI.create("http://token.uri"), TokenPollResult.failed("external authentication error")); + MockRedirectHandler redirectHandler = new MockRedirectHandler() + .sleepOnRedirect(Duration.ofMillis(500)); + + ExternalAuthenticator authenticator = new ExternalAuthenticator(redirectHandler, tokenPoller, KnownToken.memoryCached(), Duration.ofSeconds(1)); + Request firstRequest = authenticator.authenticate(null, getUnauthorizedResponse("Bearer x_token_server=\"http://token.uri\", x_redirect_server=\"http://redirect.uri\"")); + + List> requests = times( + 4, + () -> new ExternalAuthenticator(redirectHandler, tokenPoller, KnownToken.memoryCached(), Duration.ofSeconds(1)) + .authenticate(null, getUnauthorizedResponse("Bearer x_token_server=\"http://token.uri\", x_redirect_server=\"http://redirect.uri\"", firstRequest))) + .map(executor::submit) + .collect(toImmutableList()); + + ConcurrentRequestAssertion assertion = new ConcurrentRequestAssertion(requests); + assertion.firstException().hasMessage("external authentication error") + .isInstanceOf(ClientException.class); + + assertThat(redirectHandler.getRedirectionCount()).isEqualTo(2); + } + + @Test(timeOut = 2000) + public void testAuthenticationFromMultipleThreadsWithCachedTokenAfterAuthenticateTimesOut() + { + MockRedirectHandler redirectHandler = new MockRedirectHandler() + .sleepOnRedirect(Duration.ofSeconds(1)); + + List> requests = times( + 2, + () -> new ExternalAuthenticator(redirectHandler, onPoll(TokenPollResult::pending), KnownToken.memoryCached(), Duration.ofMillis(1)) + .authenticate(null, getUnauthorizedResponse("Bearer x_token_server=\"http://token.uri\", x_redirect_server=\"http://redirect.uri\""))) + .map(executor::submit) + .collect(toImmutableList()); + + ConcurrentRequestAssertion assertion = new ConcurrentRequestAssertion(requests); + assertion.requests() + .containsExactly(null, null); + assertion.assertThatNoExceptionsHasBeenThrown(); + assertThat(redirectHandler.getRedirectionCount()).isEqualTo(1); + } + + @Test(timeOut = 2000) + public void testAuthenticationFromMultipleThreadsWithCachedTokenAfterAuthenticateIsInterrupted() + throws Exception + { + ExecutorService interruptableThreadPool = newCachedThreadPool(daemonThreadsNamed(this.getClass().getName() + "-interruptable-%d")); + MockRedirectHandler redirectHandler = new MockRedirectHandler() + .sleepOnRedirect(Duration.ofMinutes(1)); + + ExternalAuthenticator authenticator = new ExternalAuthenticator(redirectHandler, onPoll(TokenPollResult::pending), KnownToken.memoryCached(), Duration.ofMillis(1)); + Future interruptedAuthentication = interruptableThreadPool.submit( + () -> authenticator.authenticate(null, getUnauthorizedResponse("Bearer x_token_server=\"http://token.uri\", x_redirect_server=\"http://redirect.uri\""))); + Thread.sleep(100); //It's here to make sure that authentication will start before the other threads. + List> requests = times( + 2, + () -> new ExternalAuthenticator(redirectHandler, onPoll(TokenPollResult::pending), KnownToken.memoryCached(), Duration.ofMillis(1)) + .authenticate(null, getUnauthorizedResponse("Bearer x_token_server=\"http://token.uri\", x_redirect_server=\"http://redirect.uri\""))) + .map(executor::submit) + .collect(toImmutableList()); + + Thread.sleep(100); + interruptableThreadPool.shutdownNow(); + + ConcurrentRequestAssertion assertion = new ConcurrentRequestAssertion(ImmutableList.>builder() + .addAll(requests) + .add(interruptedAuthentication) + .build()); + assertion.requests().containsExactly(null, null); + assertion.firstException().hasRootCauseInstanceOf(InterruptedException.class); + + assertThat(redirectHandler.getRedirectionCount()).isEqualTo(1); + } + + private static Stream> times(int times, Callable request) + { + return Stream.generate(() -> request) + .limit(times); + } + + private static Optional buildAuthentication(String challengeHeader) + { + return toAuthentication(getUnauthorizedResponse(challengeHeader)); + } + + private static Response getUnauthorizedResponse(String challengeHeader) + { + return getUnauthorizedResponse(challengeHeader, + new Request.Builder() + .url(HttpUrl.get("http://example.com")) + .build()); + } + + private static Response getUnauthorizedResponse(String challengeHeader, Request request) + { + return new Response.Builder() + .request(request) + .protocol(Protocol.HTTP_1_1) + .code(HTTP_UNAUTHORIZED) + .message("Unauthorized") + .header(WWW_AUTHENTICATE, challengeHeader) + .build(); + } + + static class ConcurrentRequestAssertion + { + private final List exceptions = new ArrayList<>(); + private final List requests = new ArrayList<>(); + + public ConcurrentRequestAssertion(List> requests) + { + for (Future request : requests) { + try { + this.requests.add(request.get()); + } + catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + catch (CancellationException ex) { + exceptions.add(ex); + } + catch (ExecutionException ex) { + checkState(ex.getCause() != null, "Missing cause on ExecutionException " + ex.getMessage()); + + exceptions.add(ex.getCause()); + } + } + } + + ThrowableAssert firstException() + { + return exceptions.stream() + .findFirst() + .map(ThrowableAssert::new) + .orElseGet(() -> new ThrowableAssert(() -> null)); + } + + void assertThatNoExceptionsHasBeenThrown() + { + if (!exceptions.isEmpty()) { + Throwable firstException = exceptions.get(0); + AssertionError assertionError = new AssertionError("Expected no exceptions, but some exceptions has been thrown", firstException); + for (int i = 1; i < exceptions.size(); i++) { + assertionError.addSuppressed(exceptions.get(i)); + } + throw assertionError; + } + } + + ListAssert requests() + { + return assertThat(requests); + } + } +} diff --git a/presto-client/src/test/java/com/facebook/presto/client/auth/external/TestHttpTokenPoller.java b/presto-client/src/test/java/com/facebook/presto/client/auth/external/TestHttpTokenPoller.java new file mode 100644 index 0000000000000..69f110d4ec917 --- /dev/null +++ b/presto-client/src/test/java/com/facebook/presto/client/auth/external/TestHttpTokenPoller.java @@ -0,0 +1,224 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.client.auth.external; + +import okhttp3.HttpUrl; +import okhttp3.OkHttpClient; +import okhttp3.mockwebserver.MockResponse; +import okhttp3.mockwebserver.MockWebServer; +import okhttp3.mockwebserver.RecordedRequest; +import org.testng.annotations.AfterMethod; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + +import java.io.IOException; +import java.io.UncheckedIOException; +import java.net.URI; +import java.time.Duration; + +import static com.google.common.net.HttpHeaders.CONTENT_TYPE; +import static com.google.common.net.MediaType.JSON_UTF_8; +import static java.lang.String.format; +import static java.net.HttpURLConnection.HTTP_GONE; +import static java.net.HttpURLConnection.HTTP_OK; +import static java.net.HttpURLConnection.HTTP_UNAVAILABLE; +import static java.net.URI.create; +import static java.util.concurrent.TimeUnit.MILLISECONDS; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +@Test(singleThreaded = true) +public class TestHttpTokenPoller +{ + private static final String TOKEN_PATH = "/v1/authentications/sso/test/token"; + private static final Duration ONE_SECOND = Duration.ofSeconds(1); + + private TokenPoller tokenPoller; + private MockWebServer server; + + @BeforeMethod(alwaysRun = true) + public void setup() + throws Exception + { + server = new MockWebServer(); + server.start(); + + tokenPoller = new HttpTokenPoller(new OkHttpClient.Builder() + .callTimeout(Duration.ofMillis(500)) + .build()); + } + + @AfterMethod(alwaysRun = true) + public void teardown() + throws IOException + { + server.close(); + server = null; + } + + @Test + public void testTokenReady() + { + server.enqueue(statusAndBody(HTTP_OK, jsonPair("token", "token"))); + + TokenPollResult result = tokenPoller.pollForToken(tokenUri(), ONE_SECOND); + + assertThat(result.getToken().token()).isEqualTo("token"); + assertThat(server.getRequestCount()).isEqualTo(1); + } + + @Test + public void testTokenNotReady() + { + server.enqueue(statusAndBody(HTTP_OK, jsonPair("nextUri", tokenUri()))); + + TokenPollResult result = tokenPoller.pollForToken(tokenUri(), ONE_SECOND); + + assertThat(result.isPending()).isTrue(); + assertThat(server.getRequestCount()).isEqualTo(1); + } + + @Test + public void testErrorResponse() + { + server.enqueue(statusAndBody(HTTP_OK, jsonPair("error", "test failure"))); + + TokenPollResult result = tokenPoller.pollForToken(tokenUri(), ONE_SECOND); + + assertThat(result.isFailed()).isTrue(); + assertThat(result.getError()).contains("test failure"); + } + + @Test + public void testBadHttpStatus() + { + server.enqueue(new MockResponse().setResponseCode(HTTP_GONE)); + + TokenPollResult result = tokenPoller.pollForToken(tokenUri(), ONE_SECOND); + + assertThat(result.isFailed()).isTrue(); + assertThat(result.getError()) + .matches("Request to http://.* failed: JsonResponse\\{statusCode=410, .*"); + } + + @Test + public void testInvalidJsonBody() + { + server.enqueue(statusAndBody(HTTP_OK, jsonPair("foo", "bar"))); + + TokenPollResult result = tokenPoller.pollForToken(tokenUri(), ONE_SECOND); + + assertThat(result.isFailed()).isTrue(); + assertThat(result.getError()) + .isEqualTo("Failed to poll for token. No fields set in response."); + } + + @Test + public void testInvalidNextUri() + { + server.enqueue(statusAndBody(HTTP_OK, jsonPair("nextUri", ":::"))); + + TokenPollResult result = tokenPoller.pollForToken(tokenUri(), ONE_SECOND); + + assertThat(result.isFailed()).isTrue(); + assertThat(result.getError()) + .matches("Request to http://.* failed: JsonResponse\\{statusCode=200, .*, hasValue=false} .*"); + } + + @Test + public void testHttpStatus503() + { + for (int i = 1; i <= 100; i++) { + server.enqueue(statusAndBody(HTTP_UNAVAILABLE, "Server failure #" + i)); + } + + assertThatThrownBy(() -> tokenPoller.pollForToken(tokenUri(), ONE_SECOND)) + .isInstanceOf(UncheckedIOException.class) + .hasRootCauseExactlyInstanceOf(IOException.class); + + assertThat(server.getRequestCount()).isGreaterThan(1); + } + + @Test + public void testHttpTimeout() + { + // force request to timeout by not enqueuing response + + assertThatThrownBy(() -> tokenPoller.pollForToken(tokenUri(), ONE_SECOND)) + .isInstanceOf(UncheckedIOException.class) + .hasMessageEndingWith(": timeout"); + } + + @Test + public void testTokenReceived() + throws InterruptedException + { + server.enqueue(status(HTTP_OK)); + + tokenPoller.tokenReceived(tokenUri()); + + RecordedRequest request = server.takeRequest(1, MILLISECONDS); + assertThat(request.getMethod()).isEqualTo("DELETE"); + assertThat(request.getRequestUrl()).isEqualTo(HttpUrl.get(tokenUri())); + } + + @Test + public void testTokenReceivedRetriesUntilNotErrorReturned() + { + server.enqueue(status(HTTP_UNAVAILABLE)); + server.enqueue(status(HTTP_UNAVAILABLE)); + server.enqueue(status(HTTP_UNAVAILABLE)); + server.enqueue(status(202)); + + tokenPoller.tokenReceived(tokenUri()); + + assertThat(server.getRequestCount()).isEqualTo(4); + } + + @Test + public void testTokenReceivedDoesNotRetriesIndefinitely() + { + for (int i = 1; i <= 100; i++) { + server.enqueue(status(HTTP_UNAVAILABLE)); + } + + tokenPoller.tokenReceived(tokenUri()); + + assertThat(server.getRequestCount()).isLessThan(100); + } + + private URI tokenUri() + { + return create("http://" + server.getHostName() + ":" + server.getPort() + TOKEN_PATH); + } + + private static String jsonPair(String key, Object value) + { + return format("{\"%s\": \"%s\"}", key, value); + } + + private static MockResponse statusAndBody(int status, String body) + { + return new MockResponse() + .setResponseCode(status) + .addHeader(CONTENT_TYPE, JSON_UTF_8) + .setBody(body); + } + + private static MockResponse status(int status) + { + return new MockResponse() + .setResponseCode(status); + } +} diff --git a/presto-jdbc/pom.xml b/presto-jdbc/pom.xml index c4fc6237d602c..071330ea73ed4 100644 --- a/presto-jdbc/pom.xml +++ b/presto-jdbc/pom.xml @@ -57,6 +57,11 @@ okhttp + + com.facebook.airlift + concurrent + + com.facebook.airlift security @@ -117,6 +122,24 @@ + + com.facebook.airlift + configuration + test + + + + com.facebook.airlift + http-server + test + + + + com.facebook.airlift + jaxrs + test + + com.facebook.presto presto-testng-services @@ -177,6 +200,18 @@ test + + javax.servlet + javax.servlet-api + test + + + + javax.ws.rs + javax.ws.rs-api + test + + org.testng testng @@ -201,12 +236,6 @@ test - - com.facebook.airlift - concurrent - test - - com.google.inject guice diff --git a/presto-jdbc/src/main/java/com/facebook/presto/jdbc/ConnectionProperties.java b/presto-jdbc/src/main/java/com/facebook/presto/jdbc/ConnectionProperties.java index c6c8e180820d7..823eccdae7c01 100644 --- a/presto-jdbc/src/main/java/com/facebook/presto/jdbc/ConnectionProperties.java +++ b/presto-jdbc/src/main/java/com/facebook/presto/jdbc/ConnectionProperties.java @@ -13,9 +13,12 @@ */ package com.facebook.presto.jdbc; +import com.facebook.presto.client.auth.external.ExternalRedirectStrategy; +import com.google.common.base.Splitter; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import com.google.common.net.HostAndPort; +import io.airlift.units.Duration; import okhttp3.Protocol; import java.io.File; @@ -26,12 +29,14 @@ import java.util.Properties; import java.util.Set; import java.util.function.Predicate; +import java.util.stream.StreamSupport; import static com.facebook.presto.jdbc.AbstractConnectionProperty.ClassListConverter.CLASS_LIST_CONVERTER; import static com.facebook.presto.jdbc.AbstractConnectionProperty.HttpProtocolConverter.HTTP_PROTOCOL_CONVERTER; import static com.facebook.presto.jdbc.AbstractConnectionProperty.ListValidateConvertor.LIST_VALIDATE_CONVERTOR; import static com.facebook.presto.jdbc.AbstractConnectionProperty.StringMapConverter.STRING_MAP_CONVERTER; import static com.facebook.presto.jdbc.AbstractConnectionProperty.checkedPredicate; +import static com.google.common.collect.ImmutableList.toImmutableList; import static java.util.Collections.unmodifiableMap; import static java.util.function.Function.identity; import static java.util.stream.Collectors.toMap; @@ -68,6 +73,11 @@ final class ConnectionProperties public static final ConnectionProperty FOLLOW_REDIRECTS = new FollowRedirects(); public static final ConnectionProperty SSL_KEY_STORE_TYPE = new SSLKeyStoreType(); public static final ConnectionProperty SSL_TRUST_STORE_TYPE = new SSLTrustStoreType(); + public static final ConnectionProperty EXTERNAL_AUTHENTICATION = new ExternalAuthentication(); + public static final ConnectionProperty EXTERNAL_AUTHENTICATION_TIMEOUT = new ExternalAuthenticationTimeout(); + public static final ConnectionProperty EXTERNAL_AUTHENTICATION_TOKEN_CACHE = new ExternalAuthenticationTokenCache(); + public static final ConnectionProperty> EXTERNAL_AUTHENTICATION_REDIRECT_HANDLERS = new ExternalAuthenticationRedirectHandlers(); + private static final Set> ALL_PROPERTIES = ImmutableSet.>builder() .add(USER) .add(PASSWORD) @@ -98,6 +108,10 @@ final class ConnectionProperties .add(QUERY_INTERCEPTORS) .add(VALIDATE_NEXTURI_SOURCE) .add(FOLLOW_REDIRECTS) + .add(EXTERNAL_AUTHENTICATION) + .add(EXTERNAL_AUTHENTICATION_TIMEOUT) + .add(EXTERNAL_AUTHENTICATION_TOKEN_CACHE) + .add(EXTERNAL_AUTHENTICATION_REDIRECT_HANDLERS) .build(); private static final Map> KEY_LOOKUP = unmodifiableMap(ALL_PROPERTIES.stream() @@ -411,4 +425,54 @@ public SSLKeyStoreType() super("SSLKeyStoreType", Optional.of(KeyStore.getDefaultType()), NOT_REQUIRED, ALLOWED, STRING_CONVERTER); } } + + private static Predicate isExternalAuthEnabled() + { + return checkedPredicate(properties -> EXTERNAL_AUTHENTICATION.getValue(properties).isPresent()); + } + + private static class ExternalAuthentication + extends AbstractConnectionProperty + { + public ExternalAuthentication() + { + super("externalAuthentication", Optional.of("false"), NOT_REQUIRED, ALLOWED, BOOLEAN_CONVERTER); + } + } + + private static class ExternalAuthenticationTimeout + extends AbstractConnectionProperty + { + public ExternalAuthenticationTimeout() + { + super("externalAuthenticationTimeout", NOT_REQUIRED, isExternalAuthEnabled(), Duration::valueOf); + } + } + + private static class ExternalAuthenticationTokenCache + extends AbstractConnectionProperty + { + public ExternalAuthenticationTokenCache() + { + super("externalAuthenticationTokenCache", Optional.of(KnownTokenCache.NONE.name()), NOT_REQUIRED, ALLOWED, KnownTokenCache::valueOf); + } + } + + private static class ExternalAuthenticationRedirectHandlers + extends AbstractConnectionProperty> + { + private static final Splitter ENUM_SPLITTER = Splitter.on(',').trimResults().omitEmptyStrings(); + + public ExternalAuthenticationRedirectHandlers() + { + super("externalAuthenticationRedirectHandlers", Optional.of("OPEN"), NOT_REQUIRED, ALLOWED, ExternalAuthenticationRedirectHandlers::parse); + } + + public static List parse(String value) + { + return StreamSupport.stream(ENUM_SPLITTER.split(value).spliterator(), false) + .map(ExternalRedirectStrategy::valueOf) + .collect(toImmutableList()); + } + } } diff --git a/presto-jdbc/src/main/java/com/facebook/presto/jdbc/KnownTokenCache.java b/presto-jdbc/src/main/java/com/facebook/presto/jdbc/KnownTokenCache.java new file mode 100644 index 0000000000000..8792075be56d1 --- /dev/null +++ b/presto-jdbc/src/main/java/com/facebook/presto/jdbc/KnownTokenCache.java @@ -0,0 +1,29 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.jdbc; + +import com.facebook.presto.client.auth.external.KnownToken; + +public enum KnownTokenCache +{ + NONE { + @Override + KnownToken create() + { + return KnownToken.local(); + } + }; + + abstract KnownToken create(); +} diff --git a/presto-jdbc/src/main/java/com/facebook/presto/jdbc/PrestoDriverUri.java b/presto-jdbc/src/main/java/com/facebook/presto/jdbc/PrestoDriverUri.java index e355b58f32331..1dcfc47b152d8 100644 --- a/presto-jdbc/src/main/java/com/facebook/presto/jdbc/PrestoDriverUri.java +++ b/presto-jdbc/src/main/java/com/facebook/presto/jdbc/PrestoDriverUri.java @@ -15,6 +15,12 @@ import com.facebook.presto.client.ClientException; import com.facebook.presto.client.OkHttpUtil; +import com.facebook.presto.client.auth.external.CompositeRedirectHandler; +import com.facebook.presto.client.auth.external.ExternalAuthenticator; +import com.facebook.presto.client.auth.external.HttpTokenPoller; +import com.facebook.presto.client.auth.external.RedirectHandler; +import com.facebook.presto.client.auth.external.TokenPoller; +import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Splitter; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; @@ -27,6 +33,7 @@ import java.net.URI; import java.net.URISyntaxException; import java.sql.SQLException; +import java.time.Duration; import java.util.Arrays; import java.util.HashMap; import java.util.List; @@ -35,6 +42,7 @@ import java.util.Optional; import java.util.Properties; import java.util.TimeZone; +import java.util.concurrent.atomic.AtomicReference; import static com.facebook.presto.client.GCSOAuthInterceptor.GCS_CREDENTIALS_PATH_KEY; import static com.facebook.presto.client.GCSOAuthInterceptor.GCS_OAUTH_SCOPES_KEY; @@ -51,6 +59,10 @@ import static com.facebook.presto.jdbc.ConnectionProperties.CLIENT_TAGS; import static com.facebook.presto.jdbc.ConnectionProperties.CUSTOM_HEADERS; import static com.facebook.presto.jdbc.ConnectionProperties.DISABLE_COMPRESSION; +import static com.facebook.presto.jdbc.ConnectionProperties.EXTERNAL_AUTHENTICATION; +import static com.facebook.presto.jdbc.ConnectionProperties.EXTERNAL_AUTHENTICATION_REDIRECT_HANDLERS; +import static com.facebook.presto.jdbc.ConnectionProperties.EXTERNAL_AUTHENTICATION_TIMEOUT; +import static com.facebook.presto.jdbc.ConnectionProperties.EXTERNAL_AUTHENTICATION_TOKEN_CACHE; import static com.facebook.presto.jdbc.ConnectionProperties.EXTRA_CREDENTIALS; import static com.facebook.presto.jdbc.ConnectionProperties.FOLLOW_REDIRECTS; import static com.facebook.presto.jdbc.ConnectionProperties.HTTP_PROTOCOLS; @@ -88,7 +100,7 @@ final class PrestoDriverUri private static final Splitter QUERY_SPLITTER = Splitter.on('&').omitEmptyStrings(); private static final Splitter ARG_SPLITTER = Splitter.on('=').limit(2); - + private static final AtomicReference REDIRECT_HANDLER = new AtomicReference<>(null); private final HostAndPort address; private final URI uri; @@ -282,6 +294,30 @@ public void setupClient(OkHttpClient.Builder builder) } builder.addInterceptor(tokenAuth(ACCESS_TOKEN.getValue(properties).get())); } + + if (EXTERNAL_AUTHENTICATION.getValue(properties).orElse(false)) { + if (!useSecureConnection) { + throw new SQLException("Authentication using external authorization requires SSL to be enabled"); + } + + // create HTTP client that shares the same settings, but without the external authenticator + TokenPoller poller = new HttpTokenPoller(builder.build()); + + Duration timeout = EXTERNAL_AUTHENTICATION_TIMEOUT.getValue(properties) + .map(value -> Duration.ofMillis(value.toMillis())) + .orElse(Duration.ofMinutes(2)); + + KnownTokenCache knownTokenCache = EXTERNAL_AUTHENTICATION_TOKEN_CACHE.getValue(properties).get(); + Optional configuredHandler = EXTERNAL_AUTHENTICATION_REDIRECT_HANDLERS.getValue(properties) + .map(CompositeRedirectHandler::new) + .map(RedirectHandler.class::cast); + RedirectHandler redirectHandler = Optional.ofNullable(REDIRECT_HANDLER.get()) + .orElseGet(() -> configuredHandler.orElseThrow(() -> new RuntimeException("External authentication redirect handler is not configured"))); + ExternalAuthenticator authenticator = new ExternalAuthenticator(redirectHandler, poller, knownTokenCache.create(), timeout); + + builder.authenticator(authenticator); + builder.addInterceptor(authenticator); + } } catch (ClientException e) { throw new SQLException(e.getMessage(), e); @@ -419,4 +455,10 @@ private static void validateConnectionProperties(Properties connectionProperties property.validate(connectionProperties); } } + + @VisibleForTesting + static void setRedirectHandler(RedirectHandler handler) + { + REDIRECT_HANDLER.set(requireNonNull(handler, "handler is null")); + } } diff --git a/presto-jdbc/src/test/java/com/facebook/presto/jdbc/TestJdbcExternalAuthentication.java b/presto-jdbc/src/test/java/com/facebook/presto/jdbc/TestJdbcExternalAuthentication.java new file mode 100644 index 0000000000000..408ef01386164 --- /dev/null +++ b/presto-jdbc/src/test/java/com/facebook/presto/jdbc/TestJdbcExternalAuthentication.java @@ -0,0 +1,525 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.jdbc; + +import com.facebook.airlift.configuration.AbstractConfigurationAwareModule; +import com.facebook.airlift.http.server.AuthenticationException; +import com.facebook.airlift.http.server.Authenticator; +import com.facebook.airlift.http.server.BasicPrincipal; +import com.facebook.airlift.log.Logging; +import com.facebook.presto.client.ClientException; +import com.facebook.presto.client.auth.external.DesktopBrowserRedirectHandler; +import com.facebook.presto.client.auth.external.RedirectException; +import com.facebook.presto.client.auth.external.RedirectHandler; +import com.facebook.presto.server.security.SecurityConfig; +import com.facebook.presto.server.testing.TestingPrestoServer; +import com.facebook.presto.sql.parser.SqlParserOptions; +import com.facebook.presto.tpch.TpchPlugin; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.inject.Binder; +import com.google.inject.Inject; +import com.google.inject.Key; +import com.google.inject.Module; +import com.google.inject.Scopes; +import com.google.inject.multibindings.Multibinder; +import okhttp3.HttpUrl; +import okhttp3.OkHttpClient; +import okhttp3.Request; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + +import javax.servlet.http.HttpServletRequest; +import javax.ws.rs.GET; +import javax.ws.rs.Path; +import javax.ws.rs.PathParam; +import javax.ws.rs.Produces; +import javax.ws.rs.core.Context; +import javax.ws.rs.core.Response; + +import java.io.File; +import java.io.IOException; +import java.net.URI; +import java.security.Principal; +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.ConcurrentModificationException; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Properties; +import java.util.Set; +import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.IntSupplier; + +import static com.facebook.airlift.configuration.ConditionalModule.installModuleIf; +import static com.facebook.airlift.jaxrs.JaxrsBinder.jaxrsBinder; +import static com.facebook.airlift.testing.Closeables.closeAll; +import static com.facebook.presto.jdbc.PrestoDriverUri.setRedirectHandler; +import static com.facebook.presto.jdbc.TestPrestoDriver.waitForNodeRefresh; +import static com.google.common.io.Resources.getResource; +import static com.google.inject.Scopes.SINGLETON; +import static com.google.inject.multibindings.Multibinder.newSetBinder; +import static com.google.inject.util.Modules.combine; +import static java.lang.String.format; +import static java.net.HttpURLConnection.HTTP_OK; +import static java.util.Objects.requireNonNull; +import static javax.ws.rs.core.HttpHeaders.AUTHORIZATION; +import static javax.ws.rs.core.MediaType.APPLICATION_JSON_TYPE; +import static javax.ws.rs.core.MediaType.TEXT_PLAIN; +import static javax.ws.rs.core.Response.Status.NOT_FOUND; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +@Test(singleThreaded = true) +public class TestJdbcExternalAuthentication +{ + private static final String TEST_CATALOG = "test_catalog"; + private TestingPrestoServer server; + + @BeforeClass + public void setup() + throws Exception + { + Logging.initialize(); + + Map properties = ImmutableMap.builder() + .put("http-server.authentication.type", "TEST_EXTERNAL") + .put("http-server.https.enabled", "true") + .put("http-server.https.keystore.path", new File(getResource("localhost.keystore").toURI()).getPath()) + .put("http-server.https.keystore.key", "changeit") + .build(); + List additionalModules = ImmutableList.builder() + .add(new DummyExternalAuthModule(() -> server.getAddress().getPort())) + .build(); + + server = new TestingPrestoServer(true, properties, null, null, new SqlParserOptions(), additionalModules); + server.installPlugin(new TpchPlugin()); + server.createCatalog(TEST_CATALOG, "tpch"); + waitForNodeRefresh(server); + } + + @AfterClass(alwaysRun = true) + public void teardown() + throws Exception + { + closeAll(server); + server = null; + } + + @BeforeMethod(alwaysRun = true) + public void clearUpLoggingSessions() + { + invalidateAllTokens(); + } + + @Test + public void testSuccessfulAuthenticationWithHttpGetOnlyRedirectHandler() + throws Exception + { + try (RedirectHandlerFixture ignore = RedirectHandlerFixture.withHandler(new HttpGetOnlyRedirectHandler()); + Connection connection = createConnection(); + Statement statement = connection.createStatement()) { + assertThat(statement.execute("SELECT 123")).isTrue(); + } + } + + /** + * Ignored due to lack of ui environment with web-browser on CI servers. + * Still this test is useful for local environments. + */ + @Test(enabled = false) + public void testSuccessfulAuthenticationWithDefaultBrowserRedirect() + throws Exception + { + try (Connection connection = createConnection(); + Statement statement = connection.createStatement()) { + assertThat(statement.execute("SELECT 123")).isTrue(); + } + } + + @Test + public void testAuthenticationFailsAfterUnfinishedRedirect() + throws Exception + { + try (RedirectHandlerFixture ignore = RedirectHandlerFixture.withHandler(new NoOpRedirectHandler()); + Connection connection = createConnection(); + Statement statement = connection.createStatement()) { + assertThatThrownBy(() -> statement.execute("SELECT 123")) + .isInstanceOf(SQLException.class); + } + } + + @Test + public void testAuthenticationFailsAfterRedirectException() + throws Exception + { + try (RedirectHandlerFixture ignore = RedirectHandlerFixture.withHandler(new FailingRedirectHandler()); + Connection connection = createConnection(); + Statement statement = connection.createStatement()) { + assertThatThrownBy(() -> statement.execute("SELECT 123")) + .isInstanceOf(SQLException.class) + .hasCauseExactlyInstanceOf(RedirectException.class); + } + } + + @Test + public void testAuthenticationFailsAfterServerAuthenticationFailure() + throws Exception + { + try (RedirectHandlerFixture ignore = RedirectHandlerFixture.withHandler(new HttpGetOnlyRedirectHandler()); + AutoCloseable ignore2 = TokenPollingErrorFixture.withPollingError("error occurred during token polling"); + Connection connection = createConnection(); + Statement statement = connection.createStatement()) { + assertThatThrownBy(() -> statement.execute("SELECT 123")) + .isInstanceOf(SQLException.class) + .hasMessage("error occurred during token polling"); + } + } + + @Test + public void testAuthenticationFailsAfterReceivingMalformedHeaderFromServer() + throws Exception + { + try (RedirectHandlerFixture ignore = RedirectHandlerFixture.withHandler(new HttpGetOnlyRedirectHandler()); + AutoCloseable ignored = WwwAuthenticateHeaderFixture.withWwwAuthenticate("Bearer no-valid-fields"); + Connection connection = createConnection(); + Statement statement = connection.createStatement()) { + assertThatThrownBy(() -> statement.execute("SELECT 123")) + .isInstanceOf(SQLException.class) + .hasCauseInstanceOf(ClientException.class) + .hasMessage("Authentication failed: Authentication required"); + } + } + + @Test + public void testAuthenticationReusesObtainedTokenPerConnection() + throws Exception + { + try (RedirectHandlerFixture ignore = RedirectHandlerFixture.withHandler(new HttpGetOnlyRedirectHandler()); + Connection connection = createConnection(); + Statement statement = connection.createStatement()) { + statement.execute("SELECT 123"); + statement.execute("SELECT 123"); + statement.execute("SELECT 123"); + + assertThat(countIssuedTokens()).isEqualTo(1); + } + } + + @Test + public void testAuthenticationAfterInitialTokenHasBeenInvalidated() + throws Exception + { + try (RedirectHandlerFixture ignore = RedirectHandlerFixture.withHandler(new HttpGetOnlyRedirectHandler()); + Connection connection = createConnection(); + Statement statement = connection.createStatement()) { + statement.execute("SELECT 123"); + + invalidateAllTokens(); + assertThat(countIssuedTokens()).isEqualTo(0); + + assertThat(statement.execute("SELECT 123")).isTrue(); + } + } + + private Connection createConnection() + throws Exception + { + String url = format("jdbc:presto://localhost:%s", server.getHttpsAddress().getPort()); + Properties properties = new Properties(); + properties.setProperty("SSL", "true"); + properties.setProperty("SSLTrustStorePath", new File(getResource("localhost.truststore").toURI()).getPath()); + properties.setProperty("SSLTrustStorePassword", "changeit"); + properties.setProperty("externalAuthentication", "true"); + properties.setProperty("externalAuthenticationTimeout", "2s"); + properties.setProperty("user", "test"); + return DriverManager.getConnection(url, properties); + } + + private static Multibinder authenticatorBinder(Binder binder) + { + return newSetBinder(binder, Authenticator.class); + } + + public static Module authenticatorModule(Class clazz, Module module) + { + Module authModule = binder -> authenticatorBinder(binder).addBinding().to(clazz).in(Scopes.SINGLETON); + return installModuleIf( + SecurityConfig.class, + config -> true, + combine(module, authModule)); + } + + private static class DummyExternalAuthModule + extends AbstractConfigurationAwareModule + { + private final IntSupplier port; + + public DummyExternalAuthModule(IntSupplier port) + { + this.port = requireNonNull(port, "port is null"); + } + + @Override + protected void setup(Binder ignored) + { + Module test = authenticatorModule(DummyAuthenticator.class, binder -> { + binder.bind(Authentications.class).in(SINGLETON); + binder.bind(IntSupplier.class).toInstance(port); + jaxrsBinder(binder).bind(DummyExternalAuthResources.class); + }); + + install(test); + } + } + + private static class Authentications + { + private final Map logginSessions = new ConcurrentHashMap<>(); + private final Set validTokens = ConcurrentHashMap.newKeySet(); + + public String startAuthentication() + { + String sessionId = UUID.randomUUID().toString(); + logginSessions.put(sessionId, ""); + return sessionId; + } + + public void logIn(String sessionId) + { + String token = sessionId + "_token"; + validTokens.add(token); + logginSessions.put(sessionId, token); + } + + public Optional getToken(String sessionId) + throws IllegalArgumentException + { + return Optional.ofNullable(logginSessions.get(sessionId)) + .filter(s -> !s.isEmpty()); + } + + public boolean verifyToken(String token) + { + return validTokens.contains(token); + } + + public void invalidateAllTokens() + { + validTokens.clear(); + } + + public int countValidTokens() + { + return validTokens.size(); + } + } + + private void invalidateAllTokens() + { + Authentications authentications = server.getInstance(Key.get(Authentications.class)); + authentications.invalidateAllTokens(); + } + + private int countIssuedTokens() + { + Authentications authentications = server.getInstance(Key.get(Authentications.class)); + return authentications.countValidTokens(); + } + + public static class DummyAuthenticator + implements Authenticator + { + private final IntSupplier port; + private final Authentications authentications; + + @Inject + public DummyAuthenticator(IntSupplier port, Authentications authentications) + { + this.port = requireNonNull(port, "port is null"); + this.authentications = requireNonNull(authentications, "authentications is null"); + } + + @Override + public Principal authenticate(HttpServletRequest request) + throws AuthenticationException + { + Optional authHeader = Optional.ofNullable(request.getHeader(AUTHORIZATION)); + List bearerHeaders = authHeader.isPresent() ? ImmutableList.of(authHeader.get()) : ImmutableList.of(); + if (bearerHeaders.stream() + .filter(header -> header.startsWith("Bearer ")) + .anyMatch(header -> authentications.verifyToken(header.substring("Bearer ".length())))) { + return new BasicPrincipal("user"); + } + + String sessionId = authentications.startAuthentication(); + + throw Optional.ofNullable(WwwAuthenticateHeaderFixture.HEADER.get()) + .map(header -> new AuthenticationException("Authentication required", header)) + .orElseGet(() -> new AuthenticationException( + "Authentication required", + format("Bearer x_redirect_server=\"http://localhost:%s/v1/authentications/dummy/logins/%s\", " + + "x_token_server=\"http://localhost:%s/v1/authentications/dummy/%s\"", + port.getAsInt(), sessionId, port.getAsInt(), sessionId))); + } + } + + @Path("/v1/authentications/dummy") + public static class DummyExternalAuthResources + { + private final Authentications authentications; + + @Inject + public DummyExternalAuthResources(Authentications authentications) + { + this.authentications = authentications; + } + + @GET + @Produces(TEXT_PLAIN) + @Path("logins/{sessionId}") + public String logInUser(@PathParam("sessionId") String sessionId) + { + authentications.logIn(sessionId); + return "User has been successfully logged in during " + sessionId + " session"; + } + + @GET + @Path("{sessionId}") + public Response getToken(@PathParam("sessionId") String sessionId, @Context HttpServletRequest request) + { + try { + return Optional.ofNullable(TokenPollingErrorFixture.ERROR.get()) + .map(error -> Response.ok(format("{ \"error\" : \"%s\"}", error), APPLICATION_JSON_TYPE).build()) + .orElseGet(() -> authentications.getToken(sessionId) + .map(token -> Response.ok(format("{ \"token\" : \"%s\"}", token), APPLICATION_JSON_TYPE).build()) + .orElseGet(() -> Response.ok(format("{ \"nextUri\" : \"%s\" }", request.getRequestURI()), APPLICATION_JSON_TYPE).build())); + } + catch (IllegalArgumentException ex) { + return Response.status(NOT_FOUND).build(); + } + } + } + + public static class HttpGetOnlyRedirectHandler + implements RedirectHandler + { + @Override + public void redirectTo(URI uri) + throws RedirectException + { + OkHttpClient client = new OkHttpClient(); + + Request request = new Request.Builder() + .url(HttpUrl.get(uri.toString())) + .build(); + + try (okhttp3.Response response = client.newCall(request).execute()) { + if (response.code() != HTTP_OK) { + throw new RedirectException("HTTP GET failed with status " + response.code()); + } + } + catch (IOException e) { + throw new RedirectException("Redirection failed", e); + } + } + } + + public static class NoOpRedirectHandler + implements RedirectHandler + { + @Override + public void redirectTo(URI uri) + throws RedirectException + {} + } + + public static class FailingRedirectHandler + implements RedirectHandler + { + @Override + public void redirectTo(URI uri) + throws RedirectException + { + throw new RedirectException("Redirect to uri has failed " + uri); + } + } + + static class RedirectHandlerFixture + implements AutoCloseable + { + private static final RedirectHandlerFixture INSTANCE = new RedirectHandlerFixture(); + + private RedirectHandlerFixture() {} + + public static RedirectHandlerFixture withHandler(RedirectHandler handler) + { + setRedirectHandler(handler); + return INSTANCE; + } + + @Override + public void close() + { + setRedirectHandler(new DesktopBrowserRedirectHandler()); + } + } + + static class TokenPollingErrorFixture + implements AutoCloseable + { + private static final AtomicReference ERROR = new AtomicReference<>(null); + + public static AutoCloseable withPollingError(String error) + { + if (ERROR.compareAndSet(null, error)) { + return new TokenPollingErrorFixture(); + } + throw new ConcurrentModificationException("polling errors can't be invoked in parallel"); + } + + @Override + public void close() + { + ERROR.set(null); + } + } + + static class WwwAuthenticateHeaderFixture + implements AutoCloseable + { + private static final AtomicReference HEADER = new AtomicReference<>(null); + + public static AutoCloseable withWwwAuthenticate(String header) + { + if (HEADER.compareAndSet(null, header)) { + return new WwwAuthenticateHeaderFixture(); + } + throw new ConcurrentModificationException("with WWW-Authenticate header can't be invoked in parallel"); + } + + @Override + public void close() + { + HEADER.set(null); + } + } +} diff --git a/presto-main-base/src/main/java/com/facebook/presto/server/security/SecurityConfig.java b/presto-main-base/src/main/java/com/facebook/presto/server/security/SecurityConfig.java index 9ac1cb678ba29..f190d86a35879 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/server/security/SecurityConfig.java +++ b/presto-main-base/src/main/java/com/facebook/presto/server/security/SecurityConfig.java @@ -40,7 +40,8 @@ public enum AuthenticationType KERBEROS, PASSWORD, JWT, - CUSTOM + CUSTOM, + OAUTH2 } @NotNull @@ -56,7 +57,7 @@ public SecurityConfig setAuthenticationTypes(List authentica } @Config("http-server.authentication.type") - @ConfigDescription("Authentication types (supported types: CERTIFICATE, KERBEROS, PASSWORD, JWT, CUSTOM)") + @ConfigDescription("Authentication types (supported types: CERTIFICATE, KERBEROS, PASSWORD, JWT, CUSTOM, OAUTH2)") public SecurityConfig setAuthenticationTypes(String types) { if (types == null) { From 498ba83378b9bb91d274a9688280edc79f3c6bdc Mon Sep 17 00:00:00 2001 From: Anant Aneja <1797669+aaneja@users.noreply.github.com> Date: Wed, 30 Apr 2025 10:51:39 +0530 Subject: [PATCH 105/113] CLI external authentication and composite redirect handler Partial cherry-pick but contains the following commits https://github.com/trinodb/trino/commit/e8a8b5ab https://github.com/trinodb/trino/commit/7b98764a Co-authored-by: Stephen Yugel Co-authored-by: Szymon Homa Co-authored-by: Mateusz Gajewski --- .../facebook/presto/cli/ClientOptions.java | 7 ++++ .../java/com/facebook/presto/cli/Console.java | 4 +- .../com/facebook/presto/cli/QueryRunner.java | 42 ++++++++++++++++++- .../facebook/presto/cli/AbstractCliTest.java | 4 +- 4 files changed, 54 insertions(+), 3 deletions(-) diff --git a/presto-cli/src/main/java/com/facebook/presto/cli/ClientOptions.java b/presto-cli/src/main/java/com/facebook/presto/cli/ClientOptions.java index f3f15e9324953..aa3d27b1502d6 100644 --- a/presto-cli/src/main/java/com/facebook/presto/cli/ClientOptions.java +++ b/presto-cli/src/main/java/com/facebook/presto/cli/ClientOptions.java @@ -15,6 +15,7 @@ import com.facebook.airlift.units.Duration; import com.facebook.presto.client.ClientSession; +import com.facebook.presto.client.auth.external.ExternalRedirectStrategy; import com.google.common.annotations.VisibleForTesting; import com.google.common.base.CharMatcher; import com.google.common.base.Splitter; @@ -162,6 +163,12 @@ public class ClientOptions @Option(name = "--disable-redirects", title = "disable redirects", description = "Disable client following redirects from server") public boolean disableRedirects; + @Option(name = "--external-authentication", title = "enable external authentication", description = "Enable external authentication") + public boolean externalAuthentication; + + @Option(name = "--external-authentication-redirect-handler", title = "external authentication redirect handler", description = "External authentication redirect handlers") + public List externalAuthenticationRedirectHandler = new ArrayList<>(); + public enum OutputFormat { ALIGNED, diff --git a/presto-cli/src/main/java/com/facebook/presto/cli/Console.java b/presto-cli/src/main/java/com/facebook/presto/cli/Console.java index 95a0dc76d4078..a94c67a5c1823 100644 --- a/presto-cli/src/main/java/com/facebook/presto/cli/Console.java +++ b/presto-cli/src/main/java/com/facebook/presto/cli/Console.java @@ -142,7 +142,9 @@ public boolean run() Optional.ofNullable(clientOptions.krb5KeytabPath), Optional.ofNullable(clientOptions.krb5CredentialCachePath), !clientOptions.krb5DisableRemoteServiceHostnameCanonicalization, - !clientOptions.disableRedirects)) { + !clientOptions.disableRedirects, + clientOptions.externalAuthentication, + clientOptions.externalAuthenticationRedirectHandler)) { if (hasQuery) { return executeCommand(queryRunner, query, clientOptions.outputFormat, clientOptions.ignoreErrors); } diff --git a/presto-cli/src/main/java/com/facebook/presto/cli/QueryRunner.java b/presto-cli/src/main/java/com/facebook/presto/cli/QueryRunner.java index 1ecffacc29bc2..9b88c5326d954 100644 --- a/presto-cli/src/main/java/com/facebook/presto/cli/QueryRunner.java +++ b/presto-cli/src/main/java/com/facebook/presto/cli/QueryRunner.java @@ -16,11 +16,21 @@ import com.facebook.presto.client.ClientSession; import com.facebook.presto.client.OkHttpUtil; import com.facebook.presto.client.StatementClient; +import com.facebook.presto.client.auth.external.CompositeRedirectHandler; +import com.facebook.presto.client.auth.external.ExternalAuthenticator; +import com.facebook.presto.client.auth.external.ExternalRedirectStrategy; +import com.facebook.presto.client.auth.external.HttpTokenPoller; +import com.facebook.presto.client.auth.external.KnownToken; +import com.facebook.presto.client.auth.external.RedirectHandler; +import com.facebook.presto.client.auth.external.TokenPoller; +import com.google.common.collect.ImmutableList; import com.google.common.net.HostAndPort; import okhttp3.OkHttpClient; import java.io.Closeable; import java.io.File; +import java.time.Duration; +import java.util.List; import java.util.Optional; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Consumer; @@ -73,7 +83,9 @@ public QueryRunner( Optional kerberosKeytabPath, Optional kerberosCredentialCachePath, boolean kerberosUseCanonicalHostname, - boolean followRedirects) + boolean followRedirects, + boolean externalAuthentication, + List externalRedirectHandlers) { this.session = new AtomicReference<>(requireNonNull(session, "session is null")); this.debug = debug; @@ -94,6 +106,7 @@ public QueryRunner( setupHttpProxy(builder, httpProxy); setupBasicAuth(builder, session, user, password); setupTokenAuth(builder, session, accessToken); + setupExternalAuth(builder, session, externalAuthentication, externalRedirectHandlers, sslSetup); if (kerberosRemoteServiceName.isPresent()) { checkArgument(session.getServer().getScheme().equalsIgnoreCase("https"), @@ -180,4 +193,31 @@ private static void setupTokenAuth( clientBuilder.addInterceptor(tokenAuth(accessToken.get())); } } + + private static void setupExternalAuth( + OkHttpClient.Builder builder, + ClientSession session, + boolean enabled, + List externalRedirectHandlers, + Consumer sslSetup) + { + if (!enabled) { + return; + } + + checkArgument(session.getServer().getScheme().equalsIgnoreCase("https"), + "Authentication using externalAuthentication requires HTTPS to be enabled"); + + List redirectStrategy = externalRedirectHandlers.isEmpty() ? ImmutableList.of(ExternalRedirectStrategy.ALL) : externalRedirectHandlers; + RedirectHandler redirectHandler = new CompositeRedirectHandler(redirectStrategy); + TokenPoller poller = new HttpTokenPoller(builder.build(), sslSetup); + ExternalAuthenticator authenticator = new ExternalAuthenticator( + redirectHandler, + poller, + KnownToken.local(), + Duration.ofMinutes(10)); + + builder.authenticator(authenticator); + builder.addInterceptor(authenticator); + } } diff --git a/presto-cli/src/test/java/com/facebook/presto/cli/AbstractCliTest.java b/presto-cli/src/test/java/com/facebook/presto/cli/AbstractCliTest.java index 97e8804910d04..dfa6fb0a3fbdb 100644 --- a/presto-cli/src/test/java/com/facebook/presto/cli/AbstractCliTest.java +++ b/presto-cli/src/test/java/com/facebook/presto/cli/AbstractCliTest.java @@ -150,7 +150,9 @@ protected static QueryRunner createQueryRunner(ClientSession clientSession, bool Optional.empty(), Optional.empty(), false, - true); + true, + false, + ImmutableList.of()); } protected static QueryRunner createQueryRunner(ClientSession clientSession) From 3e1ea50b06d0d9edbfc6e65a5782923791e5e792 Mon Sep 17 00:00:00 2001 From: Anant Aneja <1797669+aaneja@users.noreply.github.com> Date: Wed, 30 Apr 2025 11:40:28 +0530 Subject: [PATCH 106/113] Add support for OAuth2 authentication Partial cherry-pick of the following commits - https://github.com/trinodb/trino/commit/3879f455 https://github.com/trinodb/trino/commit/cd3da24c https://github.com/trinodb/trino/commit/dcb6f0bf https://github.com/trinodb/trino/commit/7cdd1336 https://github.com/trinodb/trino/commit/8b8b0bec https://github.com/trinodb/trino/commit/15e53ffd Co-authored-by: Stephen Yugel Co-authored-by: lukasz-walkiewicz Co-authored-by: Nik Hodgkinson --- pom.xml | 18 + .../main/sphinx/security/authorization.rst | 6 +- .../src/main/sphinx/security/oauth2.rst | 9 + presto-main-base/pom.xml | 3 + presto-main/pom.xml | 79 +++- .../facebook/presto/server/PrestoServer.java | 8 + .../facebook/presto/server/WebUiResource.java | 32 +- .../server/security/AuthenticationFilter.java | 88 +++- .../DefaultWebUiAuthenticationManager.java | 32 ++ .../server/security/ServerSecurityModule.java | 20 + .../security/WebUiAuthenticationManager.java | 27 ++ .../oauth2/ChallengeFailedException.java | 28 ++ .../server/security/oauth2/ForOAuth2.java | 31 ++ .../security/oauth2/ForRefreshTokens.java | 31 ++ .../security/oauth2/JweTokenSerializer.java | 179 +++++++ .../oauth2/JweTokenSerializerModule.java | 62 +++ .../server/security/oauth2/JwtUtil.java | 45 ++ .../oauth2/NimbusAirliftHttpClient.java | 137 ++++++ .../security/oauth2/NimbusHttpClient.java | 31 ++ .../security/oauth2/NimbusOAuth2Client.java | 437 ++++++++++++++++++ .../server/security/oauth2/NonceCookie.java | 90 ++++ .../OAuth2AuthenticationSupportModule.java | 33 ++ .../security/oauth2/OAuth2Authenticator.java | 146 ++++++ .../oauth2/OAuth2CallbackResource.java | 80 ++++ .../server/security/oauth2/OAuth2Client.java | 88 ++++ .../server/security/oauth2/OAuth2Config.java | 254 ++++++++++ .../security/oauth2/OAuth2ErrorCode.java | 46 ++ .../oauth2/OAuth2ServerConfigProvider.java | 67 +++ .../server/security/oauth2/OAuth2Service.java | 285 ++++++++++++ .../security/oauth2/OAuth2ServiceModule.java | 105 +++++ .../security/oauth2/OAuth2TokenExchange.java | 146 ++++++ .../oauth2/OAuth2TokenExchangeResource.java | 141 ++++++ .../security/oauth2/OAuth2TokenHandler.java | 21 + .../server/security/oauth2/OAuth2Utils.java | 79 ++++ .../security/oauth2/OAuthWebUiCookie.java | 61 +++ .../Oauth2WebUiAuthenticationManager.java | 119 +++++ .../server/security/oauth2/OidcDiscovery.java | 159 +++++++ .../security/oauth2/OidcDiscoveryConfig.java | 149 ++++++ .../security/oauth2/RefreshTokensConfig.java | 96 ++++ .../oauth2/StaticConfigurationProvider.java | 44 ++ .../StaticOAuth2ServerConfiguration.java | 105 +++++ .../security/oauth2/TokenPairSerializer.java | 92 ++++ .../security/oauth2/TokenRefresher.java | 68 +++ .../server/security/oauth2/ZstdCodec.java | 53 +++ .../server/testing/TestingPrestoServer.java | 5 + .../src/main/resources/oauth2/failure.html | 60 +++ .../src/main/resources/oauth2/success.html | 60 +++ .../presto/TestClientRequestFilterPlugin.java | 3 +- .../presto/server/MockHttpServletRequest.java | 31 +- .../oauth2/TestJweTokenSerializer.java | 171 +++++++ .../security/oauth2/TestOAuth2Config.java | 93 ++++ .../security/oauth2/TestOAuth2Utils.java | 55 +++ .../oauth2/TestOidcDiscoveryConfig.java | 67 +++ .../oauth2/TestRefreshTokensConfig.java | 74 +++ .../oauth2/TestingHydraIdentityProvider.java | 352 ++++++++++++++ .../oauth2/TokenEndpointAuthMethod.java | 33 ++ presto-native-execution/pom.xml | 7 +- presto-ui/src/static/ouath2/logout.html | 51 ++ 58 files changed, 4856 insertions(+), 36 deletions(-) create mode 100644 presto-docs/src/main/sphinx/security/oauth2.rst create mode 100644 presto-main/src/main/java/com/facebook/presto/server/security/DefaultWebUiAuthenticationManager.java create mode 100644 presto-main/src/main/java/com/facebook/presto/server/security/WebUiAuthenticationManager.java create mode 100644 presto-main/src/main/java/com/facebook/presto/server/security/oauth2/ChallengeFailedException.java create mode 100644 presto-main/src/main/java/com/facebook/presto/server/security/oauth2/ForOAuth2.java create mode 100644 presto-main/src/main/java/com/facebook/presto/server/security/oauth2/ForRefreshTokens.java create mode 100644 presto-main/src/main/java/com/facebook/presto/server/security/oauth2/JweTokenSerializer.java create mode 100644 presto-main/src/main/java/com/facebook/presto/server/security/oauth2/JweTokenSerializerModule.java create mode 100644 presto-main/src/main/java/com/facebook/presto/server/security/oauth2/JwtUtil.java create mode 100644 presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NimbusAirliftHttpClient.java create mode 100644 presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NimbusHttpClient.java create mode 100644 presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NimbusOAuth2Client.java create mode 100644 presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NonceCookie.java create mode 100644 presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2AuthenticationSupportModule.java create mode 100644 presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Authenticator.java create mode 100644 presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2CallbackResource.java create mode 100644 presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Client.java create mode 100644 presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Config.java create mode 100644 presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2ErrorCode.java create mode 100644 presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2ServerConfigProvider.java create mode 100644 presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Service.java create mode 100644 presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2ServiceModule.java create mode 100644 presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2TokenExchange.java create mode 100644 presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2TokenExchangeResource.java create mode 100644 presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2TokenHandler.java create mode 100644 presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Utils.java create mode 100644 presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuthWebUiCookie.java create mode 100644 presto-main/src/main/java/com/facebook/presto/server/security/oauth2/Oauth2WebUiAuthenticationManager.java create mode 100644 presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OidcDiscovery.java create mode 100644 presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OidcDiscoveryConfig.java create mode 100644 presto-main/src/main/java/com/facebook/presto/server/security/oauth2/RefreshTokensConfig.java create mode 100644 presto-main/src/main/java/com/facebook/presto/server/security/oauth2/StaticConfigurationProvider.java create mode 100644 presto-main/src/main/java/com/facebook/presto/server/security/oauth2/StaticOAuth2ServerConfiguration.java create mode 100644 presto-main/src/main/java/com/facebook/presto/server/security/oauth2/TokenPairSerializer.java create mode 100644 presto-main/src/main/java/com/facebook/presto/server/security/oauth2/TokenRefresher.java create mode 100644 presto-main/src/main/java/com/facebook/presto/server/security/oauth2/ZstdCodec.java create mode 100644 presto-main/src/main/resources/oauth2/failure.html create mode 100644 presto-main/src/main/resources/oauth2/success.html create mode 100644 presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestJweTokenSerializer.java create mode 100644 presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2Config.java create mode 100644 presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2Utils.java create mode 100644 presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOidcDiscoveryConfig.java create mode 100644 presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestRefreshTokensConfig.java create mode 100644 presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestingHydraIdentityProvider.java create mode 100644 presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TokenEndpointAuthMethod.java create mode 100644 presto-ui/src/static/ouath2/logout.html diff --git a/pom.xml b/pom.xml index 892935ac9e8f8..5107da0b0dd06 100644 --- a/pom.xml +++ b/pom.xml @@ -2700,6 +2700,24 @@ mariadb-java-client ${dep.mariadb.version} + + + com.nimbusds + nimbus-jose-jwt + 9.30.2 + + + + com.nimbusds + oauth2-oidc-sdk + 11.0 + + + org.aw2 + asm + + + diff --git a/presto-docs/src/main/sphinx/security/authorization.rst b/presto-docs/src/main/sphinx/security/authorization.rst index 6092a45ca6235..a354c72ed0401 100644 --- a/presto-docs/src/main/sphinx/security/authorization.rst +++ b/presto-docs/src/main/sphinx/security/authorization.rst @@ -45,9 +45,9 @@ so make sure you have authentication enabled. http-server.authentication.type=CERTIFICATE -- It is also possible to specify other authentication types such as - ``KERBEROS``, ``PASSWORD`` and ``JWT``. Additional configuration may be - needed. + - It is also possible to specify other authentication types such as + ``KERBEROS``, ``PASSWORD``, ``JWT``, and ``OAUTH2``. Additional configuration may be + needed. .. code-block:: none diff --git a/presto-docs/src/main/sphinx/security/oauth2.rst b/presto-docs/src/main/sphinx/security/oauth2.rst new file mode 100644 index 0000000000000..0858ebc23ed7b --- /dev/null +++ b/presto-docs/src/main/sphinx/security/oauth2.rst @@ -0,0 +1,9 @@ +============= +Oauth 2.0 Authentication +============= + +Presto server configuration +--------------------------- + +Refresh Tokens +--------------------------- diff --git a/presto-main-base/pom.xml b/presto-main-base/pom.xml index b9ff6288fe744..9c779dafd33c1 100644 --- a/presto-main-base/pom.xml +++ b/presto-main-base/pom.xml @@ -539,6 +539,9 @@ com.facebook.presto.testing.assertions + + com.facebook.presto.server.MockHttpServletRequest + diff --git a/presto-main/pom.xml b/presto-main/pom.xml index 194d70a0610c8..d2fb706da59eb 100644 --- a/presto-main/pom.xml +++ b/presto-main/pom.xml @@ -157,98 +157,122 @@ com.facebook.airlift.drift drift-codec + com.google.guava guava + com.google.errorprone error_prone_annotations + com.facebook.airlift.drift drift-client + com.facebook.airlift.drift drift-transport-netty + io.netty netty-buffer + com.facebook.airlift.drift drift-api + com.facebook.airlift http-client + joda-time joda-time + com.facebook.airlift concurrent + io.airlift slice + com.facebook.presto presto-analyzer + com.facebook.airlift units + com.facebook.presto presto-client + com.facebook.airlift.drift drift-protocol + com.facebook.airlift stats + com.facebook.presto presto-memory-context + jakarta.inject jakarta.inject-api + jakarta.servlet jakarta.servlet-api + com.facebook.presto presto-spi + com.facebook.airlift trace-token + com.facebook.airlift http-server + com.facebook.airlift event + com.facebook.airlift bootstrap + com.facebook.presto presto-parser @@ -283,7 +307,7 @@ io.projectreactor.netty reactor-netty-http - + io.micrometer micrometer-core @@ -349,6 +373,47 @@ runtime + + org.apache.commons + commons-lang3 + + + + com.nimbusds + nimbus-jose-jwt + + + + com.nimbusds + oauth2-oidc-sdk + + + org.ow2.asm + asm + + + + + + com.fasterxml.jackson.core + jackson-core + + + + com.squareup.okhttp3 + okhttp + + + + io.airlift + aircompressor + + + + net.jodah + failsafe + + com.facebook.presto @@ -415,6 +480,18 @@ + + + org.testcontainers + testcontainers + test + + + + org.testcontainers + postgresql + test + diff --git a/presto-main/src/main/java/com/facebook/presto/server/PrestoServer.java b/presto-main/src/main/java/com/facebook/presto/server/PrestoServer.java index ab6f45c8174ae..c8c01053e4dc7 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/PrestoServer.java +++ b/presto-main/src/main/java/com/facebook/presto/server/PrestoServer.java @@ -53,8 +53,10 @@ import com.facebook.presto.nodeManager.PluginNodeManager; import com.facebook.presto.security.AccessControlManager; import com.facebook.presto.security.AccessControlModule; +import com.facebook.presto.server.security.oauth2.OAuth2Client; import com.facebook.presto.server.security.PasswordAuthenticatorManager; import com.facebook.presto.server.security.PrestoAuthenticatorManager; +import com.facebook.presto.server.security.SecurityConfig; import com.facebook.presto.server.security.ServerSecurityModule; import com.facebook.presto.spi.function.SqlFunction; import com.facebook.presto.sql.analyzer.FeaturesConfig; @@ -88,6 +90,7 @@ import static com.facebook.airlift.json.JsonBinder.jsonBinder; import static com.facebook.presto.server.PrestoSystemRequirements.verifyJvmRequirements; import static com.facebook.presto.server.PrestoSystemRequirements.verifySystemTimeIsReasonable; +import static com.facebook.presto.server.security.SecurityConfig.AuthenticationType.OAUTH2; import static com.google.common.base.Strings.nullToEmpty; import static java.util.Objects.requireNonNull; @@ -214,6 +217,11 @@ public void run() startAssociatedProcesses(injector); + SecurityConfig securityConfig = injector.getInstance(SecurityConfig.class); + if (securityConfig.getAuthenticationTypes().contains(OAUTH2)) { + injector.getInstance(OAuth2Client.class).load(); + } + injector.getInstance(Announcer.class).start(); log.info("======== SERVER STARTED ========"); diff --git a/presto-main/src/main/java/com/facebook/presto/server/WebUiResource.java b/presto-main/src/main/java/com/facebook/presto/server/WebUiResource.java index 2eeb67b5f5d0b..a434eee5d1fab 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/WebUiResource.java +++ b/presto-main/src/main/java/com/facebook/presto/server/WebUiResource.java @@ -13,6 +13,8 @@ */ package com.facebook.presto.server; +import com.facebook.presto.server.security.oauth2.OAuthWebUiCookie; + import jakarta.annotation.security.RolesAllowed; import jakarta.ws.rs.GET; import jakarta.ws.rs.HeaderParam; @@ -21,7 +23,10 @@ import jakarta.ws.rs.core.Response; import jakarta.ws.rs.core.UriInfo; +import java.util.Optional; + import static com.facebook.presto.server.security.RoleType.ADMIN; +import static com.facebook.presto.server.security.oauth2.OAuth2Utils.getLastURLParameter; import static com.google.common.base.Strings.isNullOrEmpty; import static com.google.common.net.HttpHeaders.X_FORWARDED_PROTO; import static jakarta.ws.rs.core.Response.Status.MOVED_PERMANENTLY; @@ -30,6 +35,8 @@ @RolesAllowed(ADMIN) public class WebUiResource { + public static final String UI_ENDPOINT = "/"; + @GET public Response redirectIndexHtml( @HeaderParam(X_FORWARDED_PROTO) String proto, @@ -38,9 +45,30 @@ public Response redirectIndexHtml( if (isNullOrEmpty(proto)) { proto = uriInfo.getRequestUri().getScheme(); } + Optional lastURL = getLastURLParameter(uriInfo.getQueryParameters()); + if (lastURL.isPresent()) { + return Response + .seeOther(uriInfo.getRequestUriBuilder().scheme(proto).uri(lastURL.get()).build()) + .build(); + } + + return Response + .temporaryRedirect(uriInfo.getRequestUriBuilder().scheme(proto).path("/ui/").replaceQuery("").build()) + .build(); + } - return Response.status(MOVED_PERMANENTLY) - .location(uriInfo.getRequestUriBuilder().scheme(proto).path("/ui/").build()) + @GET + @Path("/logout") + public Response logout( + @HeaderParam(X_FORWARDED_PROTO) String proto, + @Context UriInfo uriInfo) + { + if (isNullOrEmpty(proto)) { + proto = uriInfo.getRequestUri().getScheme(); + } + return Response + .temporaryRedirect(uriInfo.getBaseUriBuilder().scheme(proto).path("/ui/logout.html").build()) + .cookie(OAuthWebUiCookie.delete()) .build(); } } diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/AuthenticationFilter.java b/presto-main/src/main/java/com/facebook/presto/server/security/AuthenticationFilter.java index 6e5907735c8a8..2043e5402dc86 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/security/AuthenticationFilter.java +++ b/presto-main/src/main/java/com/facebook/presto/server/security/AuthenticationFilter.java @@ -16,6 +16,7 @@ import com.facebook.airlift.http.server.AuthenticationException; import com.facebook.airlift.http.server.Authenticator; import com.facebook.presto.ClientRequestFilterManager; +import com.facebook.presto.server.security.oauth2.OAuth2Authenticator; import com.facebook.presto.spi.ClientRequestFilter; import com.facebook.presto.spi.PrestoException; import com.google.common.base.Joiner; @@ -45,7 +46,11 @@ import java.util.List; import java.util.Map; import java.util.Set; +import java.util.stream.Collectors; +import static com.facebook.presto.server.WebUiResource.UI_ENDPOINT; +import static com.facebook.presto.server.security.oauth2.OAuth2CallbackResource.CALLBACK_ENDPOINT; +import static com.facebook.presto.server.security.oauth2.OAuth2TokenExchangeResource.TOKEN_ENDPOINT; import static com.facebook.presto.spi.StandardErrorCode.HEADER_MODIFICATION_ATTEMPT; import static com.google.common.io.ByteStreams.copy; import static com.google.common.io.ByteStreams.nullOutputStream; @@ -64,15 +69,53 @@ public class AuthenticationFilter private final boolean allowForwardedHttps; private final ClientRequestFilterManager clientRequestFilterManager; private final List headersBlockList = ImmutableList.of("X-Presto-Transaction-Id", "X-Presto-Started-Transaction-Id", "X-Presto-Clear-Transaction-Id", "X-Presto-Trace-Token"); + private final WebUiAuthenticationManager webUiAuthenticationManager; + private final boolean isOauth2Enabled; @Inject - public AuthenticationFilter(List authenticators, SecurityConfig securityConfig, ClientRequestFilterManager clientRequestFilterManager) + public AuthenticationFilter(List authenticators, SecurityConfig securityConfig, ClientRequestFilterManager clientRequestFilterManager, WebUiAuthenticationManager webUiAuthenticationManager) { this.authenticators = ImmutableList.copyOf(requireNonNull(authenticators, "authenticators is null")); + this.webUiAuthenticationManager = requireNonNull(webUiAuthenticationManager, "webUiAuthenticationManager is null"); + this.isOauth2Enabled = this.authenticators.stream() + .anyMatch(a -> a.getClass().equals(OAuth2Authenticator.class)); this.allowForwardedHttps = requireNonNull(securityConfig, "securityConfig is null").getAllowForwardedHttps(); this.clientRequestFilterManager = requireNonNull(clientRequestFilterManager, "clientRequestFilterManager is null"); } + private static void skipRequestBody(HttpServletRequest request) + throws IOException + { + // If we send the challenge without consuming the body of the request, + // the server will close the connection after sending the response. + // The client may interpret this as a failed request and not resend the + // request with the authentication header. We can avoid this behavior + // in the client by reading and discarding the entire body of the + // unauthenticated request before sending the response. + try (InputStream inputStream = request.getInputStream()) { + copy(inputStream, nullOutputStream()); + } + } + + public static ServletRequest withPrincipal(HttpServletRequest request, Principal principal) + { + requireNonNull(principal, "principal is null"); + return new HttpServletRequestWrapper(request) + { + @Override + public Principal getUserPrincipal() + { + return principal; + } + }; + } + + public static boolean isPublic(HttpServletRequest request) + { + return request.getPathInfo().startsWith(TOKEN_ENDPOINT) + || request.getPathInfo().startsWith(CALLBACK_ENDPOINT); + } + @Override public void init(FilterConfig filterConfig) {} @@ -86,6 +129,13 @@ public void doFilter(ServletRequest servletRequest, ServletResponse servletRespo HttpServletRequest request = (HttpServletRequest) servletRequest; HttpServletResponse response = (HttpServletResponse) servletResponse; + // Check if it's a request going to the web UI side. + if (isWebUiRequest(request) && isOauth2Enabled) { + // call web authenticator + this.webUiAuthenticationManager.handleRequest(request, response, nextFilter); + return; + } + // skip authentication if non-secure or not configured if (!doesRequestSupportAuthentication(request)) { nextFilter.doFilter(request, response); @@ -108,6 +158,7 @@ public void doFilter(ServletRequest servletRequest, ServletResponse servletRespo e.getAuthenticateHeader().ifPresent(authenticateHeaders::add); continue; } + // authentication succeeded HttpServletRequest wrappedRequest = mergeExtraHeaders(request, principal); nextFilter.doFilter(withPrincipal(wrappedRequest, principal), response); @@ -117,6 +168,11 @@ public void doFilter(ServletRequest servletRequest, ServletResponse servletRespo // authentication failed skipRequestBody(request); + // Browsers have special handling for the BASIC challenge authenticate header so we need to filter them out if the WebUI Oauth Token is present. + if (isOauth2Enabled && OAuth2Authenticator.extractTokenFromCookie(request).isPresent()) { + authenticateHeaders = authenticateHeaders.stream().filter(value -> value.contains("x_token_server")).collect(Collectors.toSet()); + } + for (String value : authenticateHeaders) { response.addHeader(WWW_AUTHENTICATE, value); } @@ -182,6 +238,9 @@ public HttpServletRequest mergeExtraHeaders(HttpServletRequest request, Principa private boolean doesRequestSupportAuthentication(HttpServletRequest request) { + if (isPublic(request)) { + return false; + } if (authenticators.isEmpty()) { return false; } @@ -194,31 +253,10 @@ private boolean doesRequestSupportAuthentication(HttpServletRequest request) return false; } - private static ServletRequest withPrincipal(HttpServletRequest request, Principal principal) - { - requireNonNull(principal, "principal is null"); - return new HttpServletRequestWrapper(request) - { - @Override - public Principal getUserPrincipal() - { - return principal; - } - }; - } - - private static void skipRequestBody(HttpServletRequest request) - throws IOException + private boolean isWebUiRequest(HttpServletRequest request) { - // If we send the challenge without consuming the body of the request, - // the server will close the connection after sending the response. - // The client may interpret this as a failed request and not resend the - // request with the authentication header. We can avoid this behavior - // in the client by reading and discarding the entire body of the - // unauthenticated request before sending the response. - try (InputStream inputStream = request.getInputStream()) { - copy(inputStream, nullOutputStream()); - } + String pathInfo = request.getPathInfo(); + return pathInfo == null || pathInfo.equals(UI_ENDPOINT) || pathInfo.startsWith("/ui"); } public static class ModifiedHttpServletRequest diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/DefaultWebUiAuthenticationManager.java b/presto-main/src/main/java/com/facebook/presto/server/security/DefaultWebUiAuthenticationManager.java new file mode 100644 index 0000000000000..dd0691f88e9a1 --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/server/security/DefaultWebUiAuthenticationManager.java @@ -0,0 +1,32 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security; + +import javax.servlet.FilterChain; +import javax.servlet.ServletException; +import javax.servlet.http.HttpServletRequest; +import javax.servlet.http.HttpServletResponse; + +import java.io.IOException; + +public class DefaultWebUiAuthenticationManager + implements WebUiAuthenticationManager +{ + @Override + public void handleRequest(HttpServletRequest request, HttpServletResponse response, FilterChain nextFilter) + throws IOException, ServletException + { + nextFilter.doFilter(request, response); + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/ServerSecurityModule.java b/presto-main/src/main/java/com/facebook/presto/server/security/ServerSecurityModule.java index cd7943e40b313..48e076d7dd405 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/security/ServerSecurityModule.java +++ b/presto-main/src/main/java/com/facebook/presto/server/security/ServerSecurityModule.java @@ -18,11 +18,18 @@ import com.facebook.airlift.http.server.CertificateAuthenticator; import com.facebook.airlift.http.server.KerberosAuthenticator; import com.facebook.airlift.http.server.KerberosConfig; +import com.facebook.airlift.http.server.TheServlet; import com.facebook.presto.server.security.SecurityConfig.AuthenticationType; +import com.facebook.presto.server.security.oauth2.OAuth2AuthenticationSupportModule; +import com.facebook.presto.server.security.oauth2.OAuth2Authenticator; +import com.facebook.presto.server.security.oauth2.OAuth2Config; +import com.facebook.presto.server.security.oauth2.Oauth2WebUiAuthenticationManager; import com.google.inject.Binder; import com.google.inject.Scopes; import com.google.inject.multibindings.Multibinder; +import javax.servlet.Filter; + import java.util.List; import static com.facebook.airlift.configuration.ConfigBinder.configBinder; @@ -30,8 +37,10 @@ import static com.facebook.presto.server.security.SecurityConfig.AuthenticationType.CUSTOM; import static com.facebook.presto.server.security.SecurityConfig.AuthenticationType.JWT; import static com.facebook.presto.server.security.SecurityConfig.AuthenticationType.KERBEROS; +import static com.facebook.presto.server.security.SecurityConfig.AuthenticationType.OAUTH2; import static com.facebook.presto.server.security.SecurityConfig.AuthenticationType.PASSWORD; import static com.google.inject.multibindings.Multibinder.newSetBinder; +import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder; public class ServerSecurityModule extends AbstractConfigurationAwareModule @@ -39,6 +48,10 @@ public class ServerSecurityModule @Override protected void setup(Binder binder) { + newOptionalBinder(binder, WebUiAuthenticationManager.class).setDefault().to(DefaultWebUiAuthenticationManager.class).in(Scopes.SINGLETON); + newSetBinder(binder, Filter.class, TheServlet.class).addBinding() + .to(AuthenticationFilter.class).in(Scopes.SINGLETON); + binder.bind(PasswordAuthenticatorManager.class).in(Scopes.SINGLETON); binder.bind(PrestoAuthenticatorManager.class).in(Scopes.SINGLETON); @@ -60,6 +73,13 @@ else if (authType == JWT) { configBinder(binder).bindConfig(JsonWebTokenConfig.class); authBinder.addBinding().to(JsonWebTokenAuthenticator.class).in(Scopes.SINGLETON); } + else if (authType == OAUTH2) { + newOptionalBinder(binder, WebUiAuthenticationManager.class).setBinding().to(Oauth2WebUiAuthenticationManager.class).in(Scopes.SINGLETON); + install(new OAuth2AuthenticationSupportModule()); + binder.bind(OAuth2Authenticator.class).in(Scopes.SINGLETON); + configBinder(binder).bindConfig(OAuth2Config.class); + authBinder.addBinding().to(OAuth2Authenticator.class).in(Scopes.SINGLETON); + } else if (authType == CUSTOM) { authBinder.addBinding().to(CustomPrestoAuthenticator.class).in(Scopes.SINGLETON); } diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/WebUiAuthenticationManager.java b/presto-main/src/main/java/com/facebook/presto/server/security/WebUiAuthenticationManager.java new file mode 100644 index 0000000000000..e34a495463c37 --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/server/security/WebUiAuthenticationManager.java @@ -0,0 +1,27 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security; + +import javax.servlet.FilterChain; +import javax.servlet.ServletException; +import javax.servlet.http.HttpServletRequest; +import javax.servlet.http.HttpServletResponse; + +import java.io.IOException; + +public interface WebUiAuthenticationManager +{ + void handleRequest(HttpServletRequest request, HttpServletResponse response, FilterChain nextFilter) + throws IOException, ServletException; +} diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/ChallengeFailedException.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/ChallengeFailedException.java new file mode 100644 index 0000000000000..f72a7c7c537ea --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/ChallengeFailedException.java @@ -0,0 +1,28 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +public class ChallengeFailedException + extends Exception +{ + public ChallengeFailedException(String message) + { + super(message); + } + + public ChallengeFailedException(String message, Throwable cause) + { + super(message, cause); + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/ForOAuth2.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/ForOAuth2.java new file mode 100644 index 0000000000000..ee66543301512 --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/ForOAuth2.java @@ -0,0 +1,31 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import com.google.inject.BindingAnnotation; + +import java.lang.annotation.Retention; +import java.lang.annotation.Target; + +import static java.lang.annotation.ElementType.FIELD; +import static java.lang.annotation.ElementType.METHOD; +import static java.lang.annotation.ElementType.PARAMETER; +import static java.lang.annotation.RetentionPolicy.RUNTIME; + +@Retention(RUNTIME) +@Target({FIELD, PARAMETER, METHOD}) +@BindingAnnotation +public @interface ForOAuth2 +{ +} diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/ForRefreshTokens.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/ForRefreshTokens.java new file mode 100644 index 0000000000000..d4d3939c61320 --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/ForRefreshTokens.java @@ -0,0 +1,31 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import com.google.inject.BindingAnnotation; + +import java.lang.annotation.Retention; +import java.lang.annotation.Target; + +import static java.lang.annotation.ElementType.FIELD; +import static java.lang.annotation.ElementType.METHOD; +import static java.lang.annotation.ElementType.PARAMETER; +import static java.lang.annotation.RetentionPolicy.RUNTIME; + +@Retention(RUNTIME) +@Target({FIELD, PARAMETER, METHOD}) +@BindingAnnotation +public @interface ForRefreshTokens +{ +} diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/JweTokenSerializer.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/JweTokenSerializer.java new file mode 100644 index 0000000000000..30b8196f6c57f --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/JweTokenSerializer.java @@ -0,0 +1,179 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import com.nimbusds.jose.EncryptionMethod; +import com.nimbusds.jose.JOSEException; +import com.nimbusds.jose.JWEAlgorithm; +import com.nimbusds.jose.JWEHeader; +import com.nimbusds.jose.JWEObject; +import com.nimbusds.jose.KeyLengthException; +import com.nimbusds.jose.Payload; +import com.nimbusds.jose.crypto.AESDecrypter; +import com.nimbusds.jose.crypto.AESEncrypter; +import io.airlift.units.Duration; +import io.jsonwebtoken.Claims; +import io.jsonwebtoken.CompressionCodec; +import io.jsonwebtoken.CompressionException; +import io.jsonwebtoken.Header; +import io.jsonwebtoken.JwtBuilder; +import io.jsonwebtoken.JwtParser; + +import javax.crypto.KeyGenerator; +import javax.crypto.SecretKey; + +import java.security.NoSuchAlgorithmException; +import java.text.ParseException; +import java.time.Clock; +import java.util.Date; +import java.util.Map; +import java.util.Optional; + +import static com.facebook.presto.server.security.oauth2.JwtUtil.newJwtBuilder; +import static com.facebook.presto.server.security.oauth2.JwtUtil.newJwtParserBuilder; +import static com.google.common.base.Preconditions.checkState; +import static java.lang.String.format; +import static java.util.Objects.requireNonNull; + +public class JweTokenSerializer + implements TokenPairSerializer +{ + private static final JWEAlgorithm ALGORITHM = JWEAlgorithm.A256KW; + private static final EncryptionMethod ENCRYPTION_METHOD = EncryptionMethod.A256CBC_HS512; + private static final CompressionCodec COMPRESSION_CODEC = new ZstdCodec(); + private static final String ACCESS_TOKEN_KEY = "access_token"; + private static final String EXPIRATION_TIME_KEY = "expiration_time"; + private static final String REFRESH_TOKEN_KEY = "refresh_token"; + private final OAuth2Client client; + private final Clock clock; + private final String issuer; + private final String audience; + private final Duration tokenExpiration; + private final JwtParser parser; + private final AESEncrypter jweEncrypter; + private final AESDecrypter jweDecrypter; + private final String principalField; + + public JweTokenSerializer( + RefreshTokensConfig config, + OAuth2Client client, + String issuer, + String audience, + String principalField, + Clock clock, + Duration tokenExpiration) + throws KeyLengthException, NoSuchAlgorithmException + { + SecretKey secretKey = createKey(requireNonNull(config, "config is null")); + this.jweEncrypter = new AESEncrypter(secretKey); + this.jweDecrypter = new AESDecrypter(secretKey); + this.client = requireNonNull(client, "client is null"); + this.issuer = requireNonNull(issuer, "issuer is null"); + this.principalField = requireNonNull(principalField, "principalField is null"); + this.audience = requireNonNull(audience, "issuer is null"); + this.clock = requireNonNull(clock, "clock is null"); + this.tokenExpiration = requireNonNull(tokenExpiration, "tokenExpiration is null"); + + this.parser = newJwtParserBuilder() + .setClock(() -> Date.from(clock.instant())) + .requireIssuer(this.issuer) + .requireAudience(this.audience) + .setCompressionCodecResolver(JweTokenSerializer::resolveCompressionCodec) + .build(); + } + + @Override + public TokenPair deserialize(String token) + { + requireNonNull(token, "token is null"); + + try { + JWEObject jwe = JWEObject.parse(token); + jwe.decrypt(jweDecrypter); + Claims claims = parser.parseClaimsJwt(jwe.getPayload().toString()).getBody(); + return TokenPair.accessAndRefreshTokens( + claims.get(ACCESS_TOKEN_KEY, String.class), + claims.get(EXPIRATION_TIME_KEY, Date.class), + claims.get(REFRESH_TOKEN_KEY, String.class)); + } + catch (ParseException ex) { + throw new IllegalArgumentException("Malformed jwt token", ex); + } + catch (JOSEException ex) { + throw new IllegalArgumentException("Decryption failed", ex); + } + } + + @Override + public String serialize(TokenPair tokenPair) + { + requireNonNull(tokenPair, "tokenPair is null"); + + Optional> accessTokenClaims = client.getClaims(tokenPair.getAccessToken()); + if (!accessTokenClaims.isPresent()) { + throw new IllegalArgumentException("Claims are missing"); + } + Map claims = accessTokenClaims.get(); + if (!claims.containsKey(principalField)) { + throw new IllegalArgumentException(format("%s field is missing", principalField)); + } + JwtBuilder jwt = newJwtBuilder() + .setExpiration(Date.from(clock.instant().plusMillis(tokenExpiration.toMillis()))) + .claim(principalField, claims.get(principalField).toString()) + .setAudience(audience) + .setIssuer(issuer) + .claim(ACCESS_TOKEN_KEY, tokenPair.getAccessToken()) + .claim(EXPIRATION_TIME_KEY, tokenPair.getExpiration()) + .claim(REFRESH_TOKEN_KEY, tokenPair.getRefreshToken().orElseThrow(JweTokenSerializer::throwExceptionForNonExistingRefreshToken)) + .compressWith(COMPRESSION_CODEC); + + try { + JWEObject jwe = new JWEObject( + new JWEHeader(ALGORITHM, ENCRYPTION_METHOD), + new Payload(jwt.compact())); + jwe.encrypt(jweEncrypter); + return jwe.serialize(); + } + catch (JOSEException ex) { + throw new IllegalStateException("Encryption failed", ex); + } + } + + private static SecretKey createKey(RefreshTokensConfig config) + throws NoSuchAlgorithmException + { + SecretKey signingKey = config.getSecretKey(); + if (signingKey == null) { + KeyGenerator generator = KeyGenerator.getInstance("AES"); + generator.init(256); + return generator.generateKey(); + } + return signingKey; + } + + private static RuntimeException throwExceptionForNonExistingRefreshToken() + { + throw new IllegalStateException("Expected refresh token to be present. Please check your identity provider setup, or disable refresh tokens"); + } + + private static CompressionCodec resolveCompressionCodec(Header header) + throws CompressionException + { + if (header.getCompressionAlgorithm() != null) { + checkState(header.getCompressionAlgorithm().equals(ZstdCodec.CODEC_NAME), "Unknown codec '%s' used for token compression", header.getCompressionAlgorithm()); + return COMPRESSION_CODEC; + } + return null; + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/JweTokenSerializerModule.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/JweTokenSerializerModule.java new file mode 100644 index 0000000000000..1819f79a33137 --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/JweTokenSerializerModule.java @@ -0,0 +1,62 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import com.facebook.airlift.configuration.AbstractConfigurationAwareModule; +import com.facebook.presto.client.NodeVersion; +import com.google.inject.Binder; +import com.google.inject.Inject; +import com.google.inject.Key; +import com.google.inject.Provides; +import com.google.inject.Singleton; +import com.nimbusds.jose.KeyLengthException; + +import java.security.NoSuchAlgorithmException; +import java.time.Clock; +import java.time.Duration; + +import static com.facebook.airlift.configuration.ConfigBinder.configBinder; +import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder; + +public class JweTokenSerializerModule + extends AbstractConfigurationAwareModule +{ + @Override + protected void setup(Binder binder) + { + configBinder(binder).bindConfig(RefreshTokensConfig.class); + RefreshTokensConfig config = buildConfigObject(RefreshTokensConfig.class); + newOptionalBinder(binder, Key.get(Duration.class, ForRefreshTokens.class)).setBinding().toInstance(Duration.ofMillis(config.getTokenExpiration().toMillis())); + } + + @Provides + @Singleton + @Inject + public TokenPairSerializer getTokenPairSerializer( + OAuth2Client client, + NodeVersion nodeVersion, + RefreshTokensConfig config, + OAuth2Config oAuth2Config) + throws KeyLengthException, NoSuchAlgorithmException + { + return new JweTokenSerializer( + config, + client, + config.getIssuer() + "_" + nodeVersion.getVersion(), + config.getAudience(), + oAuth2Config.getPrincipalField(), + Clock.systemUTC(), + config.getTokenExpiration()); + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/JwtUtil.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/JwtUtil.java new file mode 100644 index 0000000000000..c2ea7114fd322 --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/JwtUtil.java @@ -0,0 +1,45 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import io.jsonwebtoken.JwtBuilder; +import io.jsonwebtoken.JwtParserBuilder; +import io.jsonwebtoken.Jwts; +import io.jsonwebtoken.io.Deserializer; +import io.jsonwebtoken.io.Serializer; +import io.jsonwebtoken.jackson.io.JacksonDeserializer; +import io.jsonwebtoken.jackson.io.JacksonSerializer; + +import java.util.Map; + +// avoid reflection and services lookup +public final class JwtUtil +{ + private static final Serializer> JWT_SERIALIZER = new JacksonSerializer<>(); + private static final Deserializer> JWT_DESERIALIZER = new JacksonDeserializer<>(); + + private JwtUtil() {} + + public static JwtBuilder newJwtBuilder() + { + return Jwts.builder() + .serializeToJsonWith(JWT_SERIALIZER); + } + + public static JwtParserBuilder newJwtParserBuilder() + { + return Jwts.parserBuilder() + .deserializeJsonWith(JWT_DESERIALIZER); + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NimbusAirliftHttpClient.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NimbusAirliftHttpClient.java new file mode 100644 index 0000000000000..04d3eb17aad01 --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NimbusAirliftHttpClient.java @@ -0,0 +1,137 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import com.facebook.airlift.http.client.HttpClient; +import com.facebook.airlift.http.client.Request; +import com.facebook.airlift.http.client.Response; +import com.facebook.airlift.http.client.ResponseHandler; +import com.facebook.airlift.http.client.ResponseHandlerUtils; +import com.facebook.airlift.http.client.StringResponseHandler; +import com.google.common.collect.ImmutableMultimap; +import com.nimbusds.jose.util.Resource; +import com.nimbusds.oauth2.sdk.ParseException; +import com.nimbusds.oauth2.sdk.http.HTTPRequest; +import com.nimbusds.oauth2.sdk.http.HTTPResponse; + +import javax.inject.Inject; +import javax.ws.rs.core.UriBuilder; + +import java.io.IOException; +import java.net.URISyntaxException; +import java.net.URL; + +import static com.facebook.airlift.http.client.Request.Builder.prepareGet; +import static com.facebook.airlift.http.client.StaticBodyGenerator.createStaticBodyGenerator; +import static com.facebook.airlift.http.client.StringResponseHandler.createStringResponseHandler; +import static com.google.common.net.HttpHeaders.CONTENT_TYPE; +import static com.nimbusds.oauth2.sdk.http.HTTPRequest.Method.DELETE; +import static com.nimbusds.oauth2.sdk.http.HTTPRequest.Method.GET; +import static com.nimbusds.oauth2.sdk.http.HTTPRequest.Method.POST; +import static com.nimbusds.oauth2.sdk.http.HTTPRequest.Method.PUT; +import static java.lang.String.format; +import static java.nio.charset.StandardCharsets.UTF_8; +import static java.util.Objects.requireNonNull; + +public class NimbusAirliftHttpClient + implements NimbusHttpClient +{ + private final HttpClient httpClient; + + @Inject + public NimbusAirliftHttpClient(@ForOAuth2 HttpClient httpClient) + { + this.httpClient = requireNonNull(httpClient, "httpClient is null"); + } + + @Override + public Resource retrieveResource(URL url) + throws IOException + { + try { + StringResponseHandler.StringResponse response = httpClient.execute( + prepareGet().setUri(url.toURI()).build(), + createStringResponseHandler()); + return new Resource(response.getBody(), response.getHeader(CONTENT_TYPE)); + } + catch (URISyntaxException e) { + throw new RuntimeException(e); + } + } + + @Override + public T execute(com.nimbusds.oauth2.sdk.Request nimbusRequest, Parser parser) + { + HTTPRequest httpRequest = nimbusRequest.toHTTPRequest(); + HTTPRequest.Method method = httpRequest.getMethod(); + + Request.Builder request = new Request.Builder() + .setMethod(method.name()) + .setFollowRedirects(httpRequest.getFollowRedirects()); + + UriBuilder url = UriBuilder.fromUri(httpRequest.getURI()); + if (method.equals(GET) || method.equals(DELETE)) { + httpRequest.getQueryParameters().forEach((key, value) -> url.queryParam(key, value.toArray())); + } + + url.fragment(httpRequest.getFragment()); + + request.setUri(url.build()); + + ImmutableMultimap.Builder headers = ImmutableMultimap.builder(); + httpRequest.getHeaderMap().forEach(headers::putAll); + request.addHeaders(headers.build()); + + if (method.equals(POST) || method.equals(PUT)) { + String query = httpRequest.getQuery(); + if (query != null) { + request.setBodyGenerator(createStaticBodyGenerator(httpRequest.getQuery(), UTF_8)); + } + } + return httpClient.execute(request.build(), new NimbusResponseHandler<>(parser)); + } + + public static class NimbusResponseHandler + implements ResponseHandler + { + private final StringResponseHandler handler = createStringResponseHandler(); + private final Parser parser; + + public NimbusResponseHandler(Parser parser) + { + this.parser = requireNonNull(parser, "parser is null"); + } + + @Override + public T handleException(Request request, Exception exception) + { + throw ResponseHandlerUtils.propagate(request, exception); + } + + @Override + public T handle(Request request, Response response) + { + StringResponseHandler.StringResponse stringResponse = handler.handle(request, response); + HTTPResponse nimbusResponse = new HTTPResponse(response.getStatusCode()); + response.getHeaders().asMap().forEach((name, values) -> nimbusResponse.setHeader(name.toString(), values.toArray(new String[0]))); + nimbusResponse.setContent(stringResponse.getBody()); + try { + return parser.parse(nimbusResponse); + } + catch (ParseException e) { + throw new RuntimeException(format("Unable to parse response status=[%d], body=[%s]", stringResponse.getStatusCode(), stringResponse.getBody()), e); + } + } + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NimbusHttpClient.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NimbusHttpClient.java new file mode 100644 index 0000000000000..b8c0e491dec4b --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NimbusHttpClient.java @@ -0,0 +1,31 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import com.nimbusds.jose.util.ResourceRetriever; +import com.nimbusds.oauth2.sdk.ParseException; +import com.nimbusds.oauth2.sdk.Request; +import com.nimbusds.oauth2.sdk.http.HTTPResponse; + +public interface NimbusHttpClient + extends ResourceRetriever +{ + T execute(Request request, Parser parser); + + interface Parser + { + T parse(HTTPResponse response) + throws ParseException; + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NimbusOAuth2Client.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NimbusOAuth2Client.java new file mode 100644 index 0000000000000..ff24e9e477c21 --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NimbusOAuth2Client.java @@ -0,0 +1,437 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import com.facebook.airlift.log.Logger; +import com.facebook.presto.server.security.oauth2.OAuth2ServerConfigProvider.OAuth2ServerConfig; +import com.google.common.collect.ImmutableSet; +import com.google.common.collect.Ordering; +import com.nimbusds.jose.JOSEException; +import com.nimbusds.jose.JWSAlgorithm; +import com.nimbusds.jose.JWSHeader; +import com.nimbusds.jose.jwk.source.RemoteJWKSet; +import com.nimbusds.jose.proc.BadJOSEException; +import com.nimbusds.jose.proc.JWSKeySelector; +import com.nimbusds.jose.proc.JWSVerificationKeySelector; +import com.nimbusds.jose.proc.SecurityContext; +import com.nimbusds.jwt.JWTClaimsSet; +import com.nimbusds.jwt.proc.DefaultJWTClaimsVerifier; +import com.nimbusds.jwt.proc.DefaultJWTProcessor; +import com.nimbusds.jwt.proc.JWTProcessor; +import com.nimbusds.oauth2.sdk.AccessTokenResponse; +import com.nimbusds.oauth2.sdk.AuthorizationCode; +import com.nimbusds.oauth2.sdk.AuthorizationCodeGrant; +import com.nimbusds.oauth2.sdk.AuthorizationGrant; +import com.nimbusds.oauth2.sdk.AuthorizationRequest; +import com.nimbusds.oauth2.sdk.ParseException; +import com.nimbusds.oauth2.sdk.RefreshTokenGrant; +import com.nimbusds.oauth2.sdk.Scope; +import com.nimbusds.oauth2.sdk.TokenRequest; +import com.nimbusds.oauth2.sdk.auth.ClientSecretBasic; +import com.nimbusds.oauth2.sdk.auth.Secret; +import com.nimbusds.oauth2.sdk.id.ClientID; +import com.nimbusds.oauth2.sdk.id.Issuer; +import com.nimbusds.oauth2.sdk.id.State; +import com.nimbusds.oauth2.sdk.token.AccessToken; +import com.nimbusds.oauth2.sdk.token.BearerAccessToken; +import com.nimbusds.oauth2.sdk.token.RefreshToken; +import com.nimbusds.oauth2.sdk.token.Tokens; +import com.nimbusds.openid.connect.sdk.AuthenticationRequest; +import com.nimbusds.openid.connect.sdk.Nonce; +import com.nimbusds.openid.connect.sdk.OIDCTokenResponse; +import com.nimbusds.openid.connect.sdk.UserInfoRequest; +import com.nimbusds.openid.connect.sdk.UserInfoResponse; +import com.nimbusds.openid.connect.sdk.claims.AccessTokenHash; +import com.nimbusds.openid.connect.sdk.claims.IDTokenClaimsSet; +import com.nimbusds.openid.connect.sdk.token.OIDCTokens; +import com.nimbusds.openid.connect.sdk.validators.AccessTokenValidator; +import com.nimbusds.openid.connect.sdk.validators.IDTokenValidator; +import com.nimbusds.openid.connect.sdk.validators.InvalidHashException; +import io.airlift.units.Duration; + +import javax.inject.Inject; + +import java.net.MalformedURLException; +import java.net.URI; +import java.time.Instant; +import java.util.Date; +import java.util.HashSet; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.stream.Stream; + +import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Preconditions.checkState; +import static com.google.common.collect.ImmutableSet.toImmutableSet; +import static com.google.common.hash.Hashing.sha256; +import static com.nimbusds.oauth2.sdk.ResponseType.CODE; +import static com.nimbusds.openid.connect.sdk.OIDCScopeValue.OPENID; +import static java.nio.charset.StandardCharsets.UTF_8; +import static java.util.Objects.requireNonNull; +import static java.util.concurrent.TimeUnit.SECONDS; + +public class NimbusOAuth2Client + implements OAuth2Client +{ + private static final Logger LOG = Logger.get(NimbusAirliftHttpClient.class); + + private final Issuer issuer; + private final ClientID clientId; + private final ClientSecretBasic clientAuth; + private final Scope scope; + private final String principalField; + private final Set accessTokenAudiences; + private final Duration maxClockSkew; + private final NimbusHttpClient httpClient; + private final OAuth2ServerConfigProvider serverConfigurationProvider; + private volatile boolean loaded; + private URI authUrl; + private URI tokenUrl; + private Optional userinfoUrl; + private JWSKeySelector jwsKeySelector; + private JWTProcessor accessTokenProcessor; + private AuthorizationCodeFlow flow; + + @Inject + public NimbusOAuth2Client(OAuth2Config oauthConfig, OAuth2ServerConfigProvider serverConfigurationProvider, NimbusHttpClient httpClient) + { + requireNonNull(oauthConfig, "oauthConfig is null"); + issuer = new Issuer(oauthConfig.getIssuer()); + clientId = new ClientID(oauthConfig.getClientId()); + clientAuth = new ClientSecretBasic(clientId, new Secret(oauthConfig.getClientSecret())); + scope = Scope.parse(oauthConfig.getScopes()); + principalField = oauthConfig.getPrincipalField(); + maxClockSkew = oauthConfig.getMaxClockSkew(); + + accessTokenAudiences = new HashSet<>(oauthConfig.getAdditionalAudiences()); + accessTokenAudiences.add(clientId.getValue()); + accessTokenAudiences.add(null); // A null value in the set allows JWTs with no audience + + this.serverConfigurationProvider = requireNonNull(serverConfigurationProvider, "serverConfigurationProvider is null"); + this.httpClient = requireNonNull(httpClient, "httpClient is null"); + } + + @Override + public void load() + { + OAuth2ServerConfig config = serverConfigurationProvider.get(); + this.authUrl = config.getAuthUrl(); + this.tokenUrl = config.getTokenUrl(); + this.userinfoUrl = config.getUserinfoUrl(); + try { + jwsKeySelector = new JWSVerificationKeySelector<>( + Stream.concat(JWSAlgorithm.Family.RSA.stream(), JWSAlgorithm.Family.EC.stream()).collect(toImmutableSet()), + new RemoteJWKSet<>(config.getJwksUrl().toURL(), httpClient)); + } + catch (MalformedURLException e) { + throw new RuntimeException(e); + } + + DefaultJWTProcessor processor = new DefaultJWTProcessor<>(); + processor.setJWSKeySelector(jwsKeySelector); + DefaultJWTClaimsVerifier accessTokenVerifier = new DefaultJWTClaimsVerifier<>( + accessTokenAudiences, + new JWTClaimsSet.Builder() + .issuer(config.getAccessTokenIssuer().orElse(issuer.getValue())) + .build(), + ImmutableSet.of(principalField), + ImmutableSet.of()); + accessTokenVerifier.setMaxClockSkew((int) maxClockSkew.roundTo(SECONDS)); + processor.setJWTClaimsSetVerifier(accessTokenVerifier); + accessTokenProcessor = processor; + flow = scope.contains(OPENID) ? new OAuth2WithOidcExtensionsCodeFlow() : new OAuth2AuthorizationCodeFlow(); + loaded = true; + } + + @Override + public Request createAuthorizationRequest(String state, URI callbackUri) + { + checkState(loaded, "OAuth2 client not initialized"); + return flow.createAuthorizationRequest(state, callbackUri); + } + + @Override + public Response getOAuth2Response(String code, URI callbackUri, Optional nonce) + throws ChallengeFailedException + { + checkState(loaded, "OAuth2 client not initialized"); + return flow.getOAuth2Response(code, callbackUri, nonce); + } + + @Override + public Optional> getClaims(String accessToken) + { + checkState(loaded, "OAuth2 client not initialized"); + return getJWTClaimsSet(accessToken).map(JWTClaimsSet::getClaims); + } + + @Override + public Response refreshTokens(String refreshToken) + throws ChallengeFailedException + { + checkState(loaded, "OAuth2 client not initialized"); + return flow.refreshTokens(refreshToken); + } + + private interface AuthorizationCodeFlow + { + Request createAuthorizationRequest(String state, URI callbackUri); + + Response getOAuth2Response(String code, URI callbackUri, Optional nonce) + throws ChallengeFailedException; + + Response refreshTokens(String refreshToken) + throws ChallengeFailedException; + } + + private class OAuth2AuthorizationCodeFlow + implements AuthorizationCodeFlow + { + @Override + public Request createAuthorizationRequest(String state, URI callbackUri) + { + return new Request( + new AuthorizationRequest.Builder(CODE, clientId) + .redirectionURI(callbackUri) + .scope(scope) + .endpointURI(authUrl) + .state(new State(state)) + .build() + .toURI(), + Optional.empty()); + } + + @Override + public Response getOAuth2Response(String code, URI callbackUri, Optional nonce) + throws ChallengeFailedException + { + checkArgument(!nonce.isPresent(), "Unexpected nonce provided"); + AccessTokenResponse tokenResponse = getTokenResponse(code, callbackUri, AccessTokenResponse::parse); + Tokens tokens = tokenResponse.toSuccessResponse().getTokens(); + return toResponse(tokens, Optional.empty()); + } + + @Override + public Response refreshTokens(String refreshToken) + throws ChallengeFailedException + { + requireNonNull(refreshToken, "refreshToken is null"); + AccessTokenResponse tokenResponse = getTokenResponse(refreshToken, AccessTokenResponse::parse); + return toResponse(tokenResponse.toSuccessResponse().getTokens(), Optional.of(refreshToken)); + } + + private Response toResponse(Tokens tokens, Optional existingRefreshToken) + throws ChallengeFailedException + { + AccessToken accessToken = tokens.getAccessToken(); + RefreshToken refreshToken = tokens.getRefreshToken(); + JWTClaimsSet claims = getJWTClaimsSet(accessToken.getValue()).orElseThrow(() -> new ChallengeFailedException("invalid access token")); + return new Response( + accessToken.getValue(), + determineExpiration(getExpiration(accessToken), claims.getExpirationTime()), + buildRefreshToken(refreshToken, existingRefreshToken)); + } + } + + private class OAuth2WithOidcExtensionsCodeFlow + implements AuthorizationCodeFlow + { + private final IDTokenValidator idTokenValidator; + + public OAuth2WithOidcExtensionsCodeFlow() + { + idTokenValidator = new IDTokenValidator(issuer, clientId, jwsKeySelector, null); + idTokenValidator.setMaxClockSkew((int) maxClockSkew.roundTo(SECONDS)); + } + + @Override + public Request createAuthorizationRequest(String state, URI callbackUri) + { + String nonce = new Nonce().getValue(); + return new Request( + new AuthenticationRequest.Builder(CODE, scope, clientId, callbackUri) + .endpointURI(authUrl) + .state(new State(state)) + .nonce(new Nonce(hashNonce(nonce))) + .build() + .toURI(), + Optional.of(nonce)); + } + + @Override + public Response getOAuth2Response(String code, URI callbackUri, Optional nonce) + throws ChallengeFailedException + { + if (!nonce.isPresent()) { + throw new ChallengeFailedException("Missing nonce"); + } + + OIDCTokenResponse tokenResponse = getTokenResponse(code, callbackUri, OIDCTokenResponse::parse); + OIDCTokens tokens = tokenResponse.getOIDCTokens(); + validateTokens(tokens, nonce); + return toResponse(tokens, Optional.empty()); + } + + @Override + public Response refreshTokens(String refreshToken) + throws ChallengeFailedException + { + OIDCTokenResponse tokenResponse = getTokenResponse(refreshToken, OIDCTokenResponse::parse); + OIDCTokens tokens = tokenResponse.getOIDCTokens(); + validateTokens(tokens); + return toResponse(tokens, Optional.of(refreshToken)); + } + + private Response toResponse(OIDCTokens tokens, Optional existingRefreshToken) + throws ChallengeFailedException + { + AccessToken accessToken = tokens.getAccessToken(); + RefreshToken refreshToken = tokens.getRefreshToken(); + JWTClaimsSet claims = getJWTClaimsSet(accessToken.getValue()).orElseThrow(() -> new ChallengeFailedException("invalid access token")); + return new Response( + accessToken.getValue(), + determineExpiration(getExpiration(accessToken), claims.getExpirationTime()), + buildRefreshToken(refreshToken, existingRefreshToken)); + } + + private void validateTokens(OIDCTokens tokens, Optional nonce) + throws ChallengeFailedException + { + try { + IDTokenClaimsSet idToken = idTokenValidator.validate( + tokens.getIDToken(), + nonce.map(this::hashNonce) + .map(Nonce::new) + .orElse(null)); + AccessTokenHash accessTokenHash = idToken.getAccessTokenHash(); + if (accessTokenHash != null) { + AccessTokenValidator.validate(tokens.getAccessToken(), ((JWSHeader) tokens.getIDToken().getHeader()).getAlgorithm(), accessTokenHash); + } + } + catch (BadJOSEException | JOSEException | InvalidHashException e) { + throw new ChallengeFailedException("Cannot validate tokens", e); + } + } + + private void validateTokens(OIDCTokens tokens) + throws ChallengeFailedException + { + validateTokens(tokens, Optional.empty()); + } + + private String hashNonce(String nonce) + { + return sha256() + .hashString(nonce, UTF_8) + .toString(); + } + } + + private T getTokenResponse(String code, URI callbackUri, NimbusAirliftHttpClient.Parser parser) + throws ChallengeFailedException + { + return getTokenResponse(new AuthorizationCodeGrant(new AuthorizationCode(code), callbackUri), parser); + } + + private T getTokenResponse(String refreshToken, NimbusAirliftHttpClient.Parser parser) + throws ChallengeFailedException + { + return getTokenResponse(new RefreshTokenGrant(new RefreshToken(refreshToken)), parser); + } + + private T getTokenResponse(AuthorizationGrant authorizationGrant, NimbusAirliftHttpClient.Parser parser) + throws ChallengeFailedException + { + T tokenResponse = httpClient.execute(new TokenRequest(tokenUrl, clientAuth, authorizationGrant, scope), parser); + if (!tokenResponse.indicatesSuccess()) { + throw new ChallengeFailedException("Error while fetching access token: " + tokenResponse.toErrorResponse().toJSONObject()); + } + return tokenResponse; + } + + private Optional getJWTClaimsSet(String accessToken) + { + if (userinfoUrl.isPresent()) { + return queryUserInfo(accessToken); + } + return parseAccessToken(accessToken); + } + + private Optional queryUserInfo(String accessToken) + { + try { + UserInfoResponse response = httpClient.execute(new UserInfoRequest(userinfoUrl.get(), new BearerAccessToken(accessToken)), UserInfoResponse::parse); + if (!response.indicatesSuccess()) { + LOG.error("Received bad response from userinfo endpoint: " + response.toErrorResponse().getErrorObject()); + return Optional.empty(); + } + return Optional.of(response.toSuccessResponse().getUserInfo().toJWTClaimsSet()); + } + catch (ParseException | RuntimeException e) { + LOG.error(e, "Received bad response from userinfo endpoint"); + return Optional.empty(); + } + } + + private Optional parseAccessToken(String accessToken) + { + try { + return Optional.of(accessTokenProcessor.process(accessToken, null)); + } + catch (java.text.ParseException | BadJOSEException | JOSEException e) { + LOG.error(e, "Failed to parse JWT access token"); + return Optional.empty(); + } + } + + private static Instant determineExpiration(Optional validUntil, Date expiration) + throws ChallengeFailedException + { + if (validUntil.isPresent()) { + if (expiration != null) { + return Ordering.natural().min(validUntil.get(), expiration.toInstant()); + } + + return validUntil.get(); + } + + if (expiration != null) { + return expiration.toInstant(); + } + + throw new ChallengeFailedException("no valid expiration date"); + } + + private Optional buildRefreshToken(RefreshToken refreshToken, Optional existingRefreshToken) + { + Optional firstOption = Optional.ofNullable(refreshToken) + .map(RefreshToken::getValue); + + if (firstOption.isPresent()) { + return firstOption; + } + else if (existingRefreshToken.isPresent()) { + return existingRefreshToken; + } + else { + return Optional.empty(); + } + } + + private static Optional getExpiration(AccessToken accessToken) + { + return accessToken.getLifetime() != 0 ? Optional.of(Instant.now().plusSeconds(accessToken.getLifetime())) : Optional.empty(); + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NonceCookie.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NonceCookie.java new file mode 100644 index 0000000000000..7475d0c46cd9f --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NonceCookie.java @@ -0,0 +1,90 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import org.apache.commons.lang3.StringUtils; + +import javax.ws.rs.core.Cookie; +import javax.ws.rs.core.NewCookie; + +import java.time.Instant; +import java.util.Date; +import java.util.Optional; + +import static com.facebook.presto.server.security.oauth2.OAuth2CallbackResource.CALLBACK_ENDPOINT; +import static com.google.common.base.Predicates.not; +import static javax.ws.rs.core.Cookie.DEFAULT_VERSION; +import static javax.ws.rs.core.NewCookie.DEFAULT_MAX_AGE; + +public final class NonceCookie +{ + // prefix according to: https://tools.ietf.org/html/draft-ietf-httpbis-rfc6265bis-05#section-4.1.3.1 + public static final String NONCE_COOKIE = "__Secure-Presto-Nonce"; + + private NonceCookie() {} + + public static NewCookie create(String nonce, Instant tokenExpiration) + { + return new NewCookie( + NONCE_COOKIE, + nonce, + CALLBACK_ENDPOINT, + null, + DEFAULT_VERSION, + null, + DEFAULT_MAX_AGE, + Date.from(tokenExpiration), + true, + true); + } + + public static javax.servlet.http.Cookie createServletCookie(String nonce, Instant tokenExpiration) + { + return toServletCookie(create(nonce, tokenExpiration)); + } + + public static javax.servlet.http.Cookie toServletCookie(NewCookie cookie) + { + javax.servlet.http.Cookie servletCookie = new javax.servlet.http.Cookie(cookie.getName(), cookie.getValue()); + servletCookie.setPath(cookie.getPath()); + servletCookie.setVersion(cookie.getVersion()); + servletCookie.setMaxAge(cookie.getMaxAge()); + servletCookie.setSecure(cookie.isSecure()); + servletCookie.setHttpOnly(cookie.isHttpOnly()); + + return servletCookie; + } + + public static Optional read(Cookie cookie) + { + return Optional.ofNullable(cookie) + .map(Cookie::getValue) + .filter(not(StringUtils::isBlank)); + } + + public static NewCookie delete() + { + return new NewCookie( + NONCE_COOKIE, + "delete", + CALLBACK_ENDPOINT, + null, + DEFAULT_VERSION, + null, + 0, + null, + true, + true); + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2AuthenticationSupportModule.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2AuthenticationSupportModule.java new file mode 100644 index 0000000000000..88cd2b5e0f23d --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2AuthenticationSupportModule.java @@ -0,0 +1,33 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import com.facebook.airlift.configuration.AbstractConfigurationAwareModule; +import com.google.inject.Binder; +import com.google.inject.Scopes; + +import static com.facebook.airlift.jaxrs.JaxrsBinder.jaxrsBinder; + +public class OAuth2AuthenticationSupportModule + extends AbstractConfigurationAwareModule +{ + @Override + protected void setup(Binder binder) + { + binder.bind(OAuth2TokenExchange.class).in(Scopes.SINGLETON); + binder.bind(OAuth2TokenHandler.class).to(OAuth2TokenExchange.class).in(Scopes.SINGLETON); + jaxrsBinder(binder).bind(OAuth2TokenExchangeResource.class); + install(new OAuth2ServiceModule()); + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Authenticator.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Authenticator.java new file mode 100644 index 0000000000000..3fe980f061364 --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Authenticator.java @@ -0,0 +1,146 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import com.facebook.airlift.http.server.AuthenticationException; +import com.facebook.airlift.http.server.Authenticator; +import com.facebook.airlift.http.server.BasicPrincipal; +import com.facebook.airlift.log.Logger; +import org.apache.commons.lang3.StringUtils; + +import javax.inject.Inject; +import javax.servlet.http.Cookie; +import javax.servlet.http.HttpServletRequest; + +import java.net.URI; +import java.security.Principal; +import java.sql.Date; +import java.time.Instant; +import java.util.Arrays; +import java.util.Map; +import java.util.Optional; +import java.util.UUID; + +import static com.facebook.presto.server.security.oauth2.OAuth2TokenExchangeResource.getInitiateUri; +import static com.facebook.presto.server.security.oauth2.OAuth2TokenExchangeResource.getTokenUri; +import static com.facebook.presto.server.security.oauth2.OAuth2Utils.getSchemeUriBuilder; +import static com.facebook.presto.server.security.oauth2.OAuthWebUiCookie.OAUTH2_COOKIE; +import static com.google.common.base.Strings.nullToEmpty; +import static com.google.common.net.HttpHeaders.AUTHORIZATION; +import static java.lang.String.format; +import static java.util.Objects.requireNonNull; + +public class OAuth2Authenticator + implements Authenticator +{ + private static final Logger Log = Logger.get(OAuth2Authenticator.class); + private final String principalField; + private final OAuth2Client client; + private final TokenPairSerializer tokenPairSerializer; + private final TokenRefresher tokenRefresher; + + @Inject + public OAuth2Authenticator(OAuth2Client client, OAuth2Config config, TokenRefresher tokenRefresher, TokenPairSerializer tokenPairSerializer) + { + this.client = requireNonNull(client, "service is null"); + this.principalField = config.getPrincipalField(); + requireNonNull(config, "oauth2Config is null"); + this.tokenRefresher = requireNonNull(tokenRefresher, "tokenRefresher is null"); + this.tokenPairSerializer = requireNonNull(tokenPairSerializer, "tokenPairSerializer is null"); + } + + public Principal authenticate(HttpServletRequest request) throws AuthenticationException + { + String token = extractToken(request); + TokenPairSerializer.TokenPair tokenPair; + try { + tokenPair = tokenPairSerializer.deserialize(token); + } + catch (IllegalArgumentException e) { + Log.error(e, "Failed to deserialize the OAuth token"); + throw needAuthentication(request, Optional.empty(), "Invalid Credentials"); + } + + if (tokenPair.getExpiration().before(Date.from(Instant.now()))) { + throw needAuthentication(request, Optional.of(token), "Invalid Credentials"); + } + Optional> claims = client.getClaims(tokenPair.getAccessToken()); + + if (!claims.isPresent()) { + throw needAuthentication(request, Optional.ofNullable(token), "Invalid Credentials"); + } + String principal = (String) claims.get().get(principalField); + if (StringUtils.isEmpty(principal)) { + Log.warn("The subject is not present we need to authenticate"); + needAuthentication(request, Optional.empty(), "Invalid Credentials"); + } + + return new BasicPrincipal(principal); + } + + public String extractToken(HttpServletRequest request) throws AuthenticationException + { + Optional cookieToken = this.extractTokenFromCookie(request); + Optional headerToken = this.extractTokenFromHeader(request); + + if (!cookieToken.isPresent() && !headerToken.isPresent()) { + throw needAuthentication(request, Optional.empty(), "Invalid Credentials"); + } + + return cookieToken.orElseGet(() -> headerToken.get()); + } + + public Optional extractTokenFromHeader(HttpServletRequest request) + { + String authHeader = nullToEmpty(request.getHeader(AUTHORIZATION)); + int space = authHeader.indexOf(' '); + if ((space < 0) || !authHeader.substring(0, space).equalsIgnoreCase("bearer")) { + return Optional.empty(); + } + + return Optional.ofNullable(authHeader.substring(space + 1).trim()) + .filter(t -> !t.isEmpty()); + } + + public static Optional extractTokenFromCookie(HttpServletRequest request) + { + Cookie[] cookies = Optional.ofNullable(request.getCookies()).orElse(new Cookie[0]); + return Optional.ofNullable(Arrays.stream(cookies) + .filter(cookie -> cookie.getName().equals(OAUTH2_COOKIE)) + .findFirst() + .map(c -> c.getValue()) + .orElse(null)); + } + + private AuthenticationException needAuthentication(HttpServletRequest request, Optional currentToken, String message) + { + URI baseUri = getSchemeUriBuilder(request).build(); + return currentToken + .map(tokenPairSerializer::deserialize) + .flatMap(tokenRefresher::refreshToken) + .map(refreshId -> baseUri.resolve(getTokenUri(refreshId))) + .map(tokenUri -> new AuthenticationException(message, format("Bearer x_token_server=\"%s\"", tokenUri))) + .orElseGet(() -> buildNeedAuthentication(request, message)); + } + + private AuthenticationException buildNeedAuthentication(HttpServletRequest request, String message) + { + UUID authId = UUID.randomUUID(); + URI baseUri = getSchemeUriBuilder(request).build(); + URI initiateUri = baseUri.resolve(getInitiateUri(authId)); + URI tokenUri = baseUri.resolve(getTokenUri(authId)); + + return new AuthenticationException(message, format("Bearer x_redirect_server=\"%s\", x_token_server=\"%s\"", initiateUri, tokenUri)); + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2CallbackResource.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2CallbackResource.java new file mode 100644 index 0000000000000..ef5b6d6330d38 --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2CallbackResource.java @@ -0,0 +1,80 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import com.facebook.airlift.log.Logger; + +import javax.inject.Inject; +import javax.servlet.http.HttpServletRequest; +import javax.ws.rs.CookieParam; +import javax.ws.rs.GET; +import javax.ws.rs.Path; +import javax.ws.rs.Produces; +import javax.ws.rs.QueryParam; +import javax.ws.rs.core.Context; +import javax.ws.rs.core.Cookie; +import javax.ws.rs.core.Response; +import javax.ws.rs.core.UriBuilder; + +import static com.facebook.presto.server.security.oauth2.NonceCookie.NONCE_COOKIE; +import static com.facebook.presto.server.security.oauth2.OAuth2Utils.getSchemeUriBuilder; +import static java.util.Objects.requireNonNull; +import static javax.ws.rs.core.MediaType.TEXT_HTML; +import static javax.ws.rs.core.Response.Status.BAD_REQUEST; + +@Path(OAuth2CallbackResource.CALLBACK_ENDPOINT) +public class OAuth2CallbackResource +{ + private static final Logger LOG = Logger.get(OAuth2CallbackResource.class); + + public static final String CALLBACK_ENDPOINT = "/oauth2/callback"; + + private final OAuth2Service service; + + @Inject + public OAuth2CallbackResource(OAuth2Service service) + { + this.service = requireNonNull(service, "service is null"); + } + + @GET + @Produces(TEXT_HTML) + public Response callback( + @QueryParam("state") String state, + @QueryParam("code") String code, + @QueryParam("error") String error, + @QueryParam("error_description") String errorDescription, + @QueryParam("error_uri") String errorUri, + @CookieParam(NONCE_COOKIE) Cookie nonce, + @Context HttpServletRequest request) + { + if (error != null) { + return service.handleOAuth2Error(state, error, errorDescription, errorUri); + } + + try { + requireNonNull(state, "state is null"); + requireNonNull(code, "code is null"); + UriBuilder builder = getSchemeUriBuilder(request); + return service.finishOAuth2Challenge(state, code, builder.build().resolve(CALLBACK_ENDPOINT), NonceCookie.read(nonce), request); + } + catch (RuntimeException e) { + LOG.error(e, "Authentication response could not be verified: state=%s", state); + return Response.status(BAD_REQUEST) + .cookie(NonceCookie.delete()) + .entity(service.getInternalFailureHtml("Authentication response could not be verified")) + .build(); + } + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Client.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Client.java new file mode 100644 index 0000000000000..fd68a2a5a2e06 --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Client.java @@ -0,0 +1,88 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import java.net.URI; +import java.time.Instant; +import java.util.Map; +import java.util.Optional; + +import static java.util.Objects.requireNonNull; + +public interface OAuth2Client +{ + void load(); + + Request createAuthorizationRequest(String state, URI callbackUri); + + Response getOAuth2Response(String code, URI callbackUri, Optional nonce) + throws ChallengeFailedException; + + Optional> getClaims(String accessToken); + + Response refreshTokens(String refreshToken) + throws ChallengeFailedException; + + class Request + { + private final URI authorizationUri; + private final Optional nonce; + + public Request(URI authorizationUri, Optional nonce) + { + this.authorizationUri = requireNonNull(authorizationUri, "authorizationUri is null"); + this.nonce = requireNonNull(nonce, "nonce is null"); + } + + public URI getAuthorizationUri() + { + return authorizationUri; + } + + public Optional getNonce() + { + return nonce; + } + } + + class Response + { + private final String accessToken; + private final Instant expiration; + + private final Optional refreshToken; + + public Response(String accessToken, Instant expiration, Optional refreshToken) + { + this.accessToken = requireNonNull(accessToken, "accessToken is null"); + this.expiration = requireNonNull(expiration, "expiration is null"); + this.refreshToken = requireNonNull(refreshToken, "refreshToken is null"); + } + + public String getAccessToken() + { + return accessToken; + } + + public Instant getExpiration() + { + return expiration; + } + + public Optional getRefreshToken() + { + return refreshToken; + } + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Config.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Config.java new file mode 100644 index 0000000000000..8de3e4442138b --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Config.java @@ -0,0 +1,254 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import com.facebook.airlift.configuration.Config; +import com.facebook.airlift.configuration.ConfigDescription; +import com.facebook.airlift.configuration.ConfigSecuritySensitive; +import com.google.common.base.Splitter; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; +import io.airlift.units.Duration; +import io.airlift.units.MinDuration; + +import javax.validation.constraints.NotNull; + +import java.io.File; +import java.util.Collections; +import java.util.List; +import java.util.Optional; +import java.util.Set; +import java.util.concurrent.TimeUnit; + +import static com.facebook.presto.server.security.oauth2.OAuth2Service.OPENID_SCOPE; +import static com.google.common.base.Strings.nullToEmpty; + +public class OAuth2Config +{ + private Optional stateKey = Optional.empty(); + private String issuer; + private String clientId; + private String clientSecret; + private Set scopes = ImmutableSet.of(OPENID_SCOPE); + private String principalField = "sub"; + private Optional groupsField = Optional.empty(); + private List additionalAudiences = Collections.emptyList(); + private Duration challengeTimeout = new Duration(15, TimeUnit.MINUTES); + private Duration maxClockSkew = new Duration(1, TimeUnit.MINUTES); + private Optional userMappingPattern = Optional.empty(); + private Optional userMappingFile = Optional.empty(); + private boolean enableRefreshTokens; + private boolean enableDiscovery = true; + + public Optional getStateKey() + { + return stateKey; + } + + @Config("http-server.authentication.oauth2.state-key") + @ConfigDescription("A secret key used by HMAC algorithm to sign the state parameter") + public OAuth2Config setStateKey(String stateKey) + { + this.stateKey = Optional.ofNullable(stateKey); + return this; + } + + @NotNull + public String getIssuer() + { + return issuer; + } + + @Config("http-server.authentication.oauth2.issuer") + @ConfigDescription("The required issuer of a token") + public OAuth2Config setIssuer(String issuer) + { + this.issuer = issuer; + return this; + } + + @NotNull + public String getClientId() + { + return clientId; + } + + @Config("http-server.authentication.oauth2.client-id") + @ConfigDescription("Client ID") + public OAuth2Config setClientId(String clientId) + { + this.clientId = clientId; + return this; + } + + @NotNull + public String getClientSecret() + { + return clientSecret; + } + + @Config("http-server.authentication.oauth2.client-secret") + @ConfigSecuritySensitive + @ConfigDescription("Client secret") + public OAuth2Config setClientSecret(String clientSecret) + { + this.clientSecret = clientSecret; + return this; + } + + @NotNull + public List getAdditionalAudiences() + { + return additionalAudiences; + } + + public OAuth2Config setAdditionalAudiences(List additionalAudiences) + { + this.additionalAudiences = ImmutableList.copyOf(additionalAudiences); + return this; + } + + @Config("http-server.authentication.oauth2.additional-audiences") + @ConfigDescription("Additional audiences to trust in addition to the Client ID") + public OAuth2Config setAdditionalAudiences(String additionalAudiences) + { + Splitter splitter = Splitter.on(',').trimResults().omitEmptyStrings(); + this.additionalAudiences = ImmutableList.copyOf(splitter.split(nullToEmpty(additionalAudiences))); + return this; + } + + @NotNull + public Set getScopes() + { + return scopes; + } + + @Config("http-server.authentication.oauth2.scopes") + @ConfigDescription("Scopes requested by the server during OAuth2 authorization challenge") + public OAuth2Config setScopes(String scopes) + { + Splitter splitter = Splitter.on(',').trimResults().omitEmptyStrings(); + this.scopes = ImmutableSet.copyOf(splitter.split(nullToEmpty(scopes))); + return this; + } + + @NotNull + public String getPrincipalField() + { + return principalField; + } + + @Config("http-server.authentication.oauth2.principal-field") + @ConfigDescription("The claim to use as the principal") + public OAuth2Config setPrincipalField(String principalField) + { + this.principalField = principalField; + return this; + } + + public Optional getGroupsField() + { + return groupsField; + } + + @Config("http-server.authentication.oauth2.groups-field") + @ConfigDescription("Groups field in the claim") + public OAuth2Config setGroupsField(String groupsField) + { + this.groupsField = Optional.ofNullable(groupsField); + return this; + } + + @MinDuration("1ms") + @NotNull + public Duration getChallengeTimeout() + { + return challengeTimeout; + } + + @Config("http-server.authentication.oauth2.challenge-timeout") + @ConfigDescription("Maximum duration of OAuth2 authorization challenge") + public OAuth2Config setChallengeTimeout(Duration challengeTimeout) + { + this.challengeTimeout = challengeTimeout; + return this; + } + + @MinDuration("0s") + @NotNull + public Duration getMaxClockSkew() + { + return maxClockSkew; + } + + @Config("http-server.authentication.oauth2.max-clock-skew") + @ConfigDescription("Max clock skew between the Authorization Server and the coordinator") + public OAuth2Config setMaxClockSkew(Duration maxClockSkew) + { + this.maxClockSkew = maxClockSkew; + return this; + } + + public Optional getUserMappingPattern() + { + return userMappingPattern; + } + + @Config("http-server.authentication.oauth2.user-mapping.pattern") + @ConfigDescription("Regex to match against user name") + public OAuth2Config setUserMappingPattern(String userMappingPattern) + { + this.userMappingPattern = Optional.ofNullable(userMappingPattern); + return this; + } + + public Optional getUserMappingFile() + { + return userMappingFile; + } + + @Config("http-server.authentication.oauth2.user-mapping.file") + @ConfigDescription("File containing rules for mapping user") + public OAuth2Config setUserMappingFile(File userMappingFile) + { + this.userMappingFile = Optional.ofNullable(userMappingFile); + return this; + } + + public boolean isEnableRefreshTokens() + { + return enableRefreshTokens; + } + + @Config("http-server.authentication.oauth2.refresh-tokens") + @ConfigDescription("Enables OpenID refresh tokens usage") + public OAuth2Config setEnableRefreshTokens(boolean enableRefreshTokens) + { + this.enableRefreshTokens = enableRefreshTokens; + return this; + } + + public boolean isEnableDiscovery() + { + return enableDiscovery; + } + + @Config("http-server.authentication.oauth2.oidc.discovery") + @ConfigDescription("Enable OpenID Provider Issuer discovery") + public OAuth2Config setEnableDiscovery(boolean enableDiscovery) + { + this.enableDiscovery = enableDiscovery; + return this; + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2ErrorCode.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2ErrorCode.java new file mode 100644 index 0000000000000..058ce8be92b1e --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2ErrorCode.java @@ -0,0 +1,46 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import java.util.Arrays; + +public enum OAuth2ErrorCode +{ + ACCESS_DENIED("access_denied", "OAuth2 server denied the login"), + UNAUTHORIZED_CLIENT("unauthorized_client", "OAuth2 server does not allow request from this Presto server"), + SERVER_ERROR("server_error", "OAuth2 server had a failure"), + TEMPORARILY_UNAVAILABLE("temporarily_unavailable", "OAuth2 server is temporarily unavailable"); + + private final String code; + private final String message; + + OAuth2ErrorCode(String code, String message) + { + this.code = code; + this.message = message; + } + + public static OAuth2ErrorCode fromString(String codeStr) + { + return Arrays.stream(OAuth2ErrorCode.values()) + .filter(value -> codeStr.equalsIgnoreCase(value.code)) + .findFirst() + .orElseThrow(() -> new IllegalArgumentException("No enum constant " + OAuth2ErrorCode.class.getCanonicalName() + "." + codeStr)); + } + + public String getMessage() + { + return this.message; + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2ServerConfigProvider.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2ServerConfigProvider.java new file mode 100644 index 0000000000000..9455e5622605d --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2ServerConfigProvider.java @@ -0,0 +1,67 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import java.net.URI; +import java.util.Optional; + +import static java.util.Objects.requireNonNull; + +public interface OAuth2ServerConfigProvider +{ + OAuth2ServerConfig get(); + + class OAuth2ServerConfig + { + private final Optional accessTokenIssuer; + private final URI authUrl; + private final URI tokenUrl; + private final URI jwksUrl; + private final Optional userinfoUrl; + + public OAuth2ServerConfig(Optional accessTokenIssuer, URI authUrl, URI tokenUrl, URI jwksUrl, Optional userinfoUrl) + { + this.accessTokenIssuer = requireNonNull(accessTokenIssuer, "accessTokenIssuer is null"); + this.authUrl = requireNonNull(authUrl, "authUrl is null"); + this.tokenUrl = requireNonNull(tokenUrl, "tokenUrl is null"); + this.jwksUrl = requireNonNull(jwksUrl, "jwksUrl is null"); + this.userinfoUrl = requireNonNull(userinfoUrl, "userinfoUrl is null"); + } + + public Optional getAccessTokenIssuer() + { + return accessTokenIssuer; + } + + public URI getAuthUrl() + { + return authUrl; + } + + public URI getTokenUrl() + { + return tokenUrl; + } + + public URI getJwksUrl() + { + return jwksUrl; + } + + public Optional getUserinfoUrl() + { + return userinfoUrl; + } + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Service.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Service.java new file mode 100644 index 0000000000000..40a97bb1b82ca --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Service.java @@ -0,0 +1,285 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import com.facebook.airlift.log.Logger; +import com.google.common.io.Resources; +import io.jsonwebtoken.Claims; +import io.jsonwebtoken.JwtParser; + +import javax.inject.Inject; +import javax.servlet.http.HttpServletRequest; +import javax.servlet.http.HttpServletResponse; +import javax.ws.rs.core.HttpHeaders; +import javax.ws.rs.core.Response; +import javax.ws.rs.core.UriBuilder; + +import java.io.IOException; +import java.net.URI; +import java.security.Key; +import java.security.SecureRandom; +import java.time.Duration; +import java.time.Instant; +import java.time.temporal.TemporalAmount; +import java.util.Date; +import java.util.Optional; +import java.util.Random; + +import static com.facebook.presto.server.security.oauth2.JwtUtil.newJwtBuilder; +import static com.facebook.presto.server.security.oauth2.JwtUtil.newJwtParserBuilder; +import static com.facebook.presto.server.security.oauth2.OAuth2Utils.getSchemeUriBuilder; +import static com.facebook.presto.server.security.oauth2.TokenPairSerializer.TokenPair.fromOAuth2Response; +import static com.google.common.base.Strings.nullToEmpty; +import static com.google.common.base.Verify.verify; +import static com.google.common.hash.Hashing.sha256; +import static io.jsonwebtoken.security.Keys.hmacShaKeyFor; +import static java.lang.String.format; +import static java.nio.charset.StandardCharsets.UTF_8; +import static java.time.Instant.now; +import static java.util.Objects.requireNonNull; +import static javax.ws.rs.core.Response.Status.BAD_REQUEST; +import static javax.ws.rs.core.Response.Status.FORBIDDEN; + +public class OAuth2Service +{ + private static final Logger logger = Logger.get(OAuth2Service.class); + + public static final String OPENID_SCOPE = "openid"; + + private static final String STATE_AUDIENCE_UI = "presto_oauth_ui"; + private static final String FAILURE_REPLACEMENT_TEXT = ""; + private static final Random SECURE_RANDOM = new SecureRandom(); + public static final String HANDLER_STATE_CLAIM = "handler_state"; + + private final OAuth2Client client; + private final Optional tokenExpiration; + private final TokenPairSerializer tokenPairSerializer; + + private final String successHtml; + private final String failureHtml; + + private final TemporalAmount challengeTimeout; + private final Key stateHmac; + private final JwtParser jwtParser; + + private final OAuth2TokenHandler tokenHandler; + + @Inject + public OAuth2Service( + OAuth2Client client, + OAuth2Config oauth2Config, + OAuth2TokenHandler tokenHandler, + TokenPairSerializer tokenPairSerializer, + @ForRefreshTokens Optional tokenExpiration) + throws IOException + { + this.client = requireNonNull(client, "client is null"); + requireNonNull(oauth2Config, "oauth2Config is null"); + + this.successHtml = Resources.toString(Resources.getResource(getClass(), "/oauth2/success.html"), UTF_8); + this.failureHtml = Resources.toString(Resources.getResource(getClass(), "/oauth2/failure.html"), UTF_8); + verify(failureHtml.contains(FAILURE_REPLACEMENT_TEXT), "failure.html does not contain the replacement text"); + + this.challengeTimeout = Duration.ofMillis(oauth2Config.getChallengeTimeout().toMillis()); + this.stateHmac = hmacShaKeyFor(oauth2Config.getStateKey() + .map(key -> sha256().hashString(key, UTF_8).asBytes()) + .orElseGet(() -> secureRandomBytes(32))); + this.jwtParser = newJwtParserBuilder() + .setSigningKey(stateHmac) + .requireAudience(STATE_AUDIENCE_UI) + .build(); + + this.tokenHandler = requireNonNull(tokenHandler, "tokenHandler is null"); + this.tokenPairSerializer = requireNonNull(tokenPairSerializer, "tokenPairSerializer is null"); + + this.tokenExpiration = requireNonNull(tokenExpiration, "tokenExpiration is null"); + } + + public Response startOAuth2Challenge(URI callbackUri, Optional handlerState) + { + Instant challengeExpiration = now().plus(challengeTimeout); + String state = newJwtBuilder() + .signWith(stateHmac) + .setAudience(STATE_AUDIENCE_UI) + .claim(HANDLER_STATE_CLAIM, handlerState.orElse(null)) + .setExpiration(Date.from(challengeExpiration)) + .compact(); + + OAuth2Client.Request request = client.createAuthorizationRequest(state, callbackUri); + Response.ResponseBuilder response = Response.seeOther(request.getAuthorizationUri()); + request.getNonce().ifPresent(nce -> response.cookie(NonceCookie.create(nce, challengeExpiration))); + return response.build(); + } + + public void startOAuth2Challenge(URI callbackUri, Optional handlerState, HttpServletResponse servletResponse) + throws IOException + { + Instant challengeExpiration = now().plus(challengeTimeout); + + OAuth2Client.Request challengeRequest = this.startChallenge(callbackUri, handlerState); + challengeRequest.getNonce().ifPresent(nce -> servletResponse.addCookie(NonceCookie.createServletCookie(nce, challengeExpiration))); + servletResponseSeeOther(challengeRequest.getAuthorizationUri().toString(), servletResponse); + } + + public void servletResponseSeeOther(String location, HttpServletResponse servletResponse) + throws IOException + { + // 303 is preferred over a 302 when this response is received by a POST/PUT/DELETE and the redirect should be done via a GET instead of original method + servletResponse.addHeader(HttpHeaders.LOCATION, location); + servletResponse.sendError(HttpServletResponse.SC_SEE_OTHER); + } + + private OAuth2Client.Request startChallenge(URI callbackUri, Optional handlerState) + { + Instant challengeExpiration = now().plus(challengeTimeout); + String state = newJwtBuilder() + .signWith(stateHmac) + .setAudience(STATE_AUDIENCE_UI) + .claim(HANDLER_STATE_CLAIM, handlerState.orElse(null)) + .setExpiration(Date.from(challengeExpiration)) + .compact(); + + return client.createAuthorizationRequest(state, callbackUri); + } + + public Response handleOAuth2Error(String state, String error, String errorDescription, String errorUri) + { + try { + Claims stateClaims = parseState(state); + Optional.ofNullable(stateClaims.get(HANDLER_STATE_CLAIM, String.class)) + .ifPresent(value -> + tokenHandler.setTokenExchangeError(value, + format("Authentication response could not be verified: error=%s, errorDescription=%s, errorUri=%s", + error, errorDescription, errorDescription))); + } + catch (ChallengeFailedException | RuntimeException e) { + logger.error(e, "Authentication response could not be verified invalid state: state=%s", state); + return Response.status(FORBIDDEN) + .entity(getInternalFailureHtml("Authentication response could not be verified")) + .cookie(NonceCookie.delete()) + .build(); + } + + logger.error("OAuth server returned an error: error=%s, error_description=%s, error_uri=%s, state=%s", error, errorDescription, errorUri, state); + return Response.ok() + .entity(getCallbackErrorHtml(error)) + .cookie(NonceCookie.delete()) + .build(); + } + + public Response finishOAuth2Challenge(String state, String code, URI callbackUri, Optional nonce, HttpServletRequest request) + { + Optional handlerState; + try { + Claims stateClaims = parseState(state); + handlerState = Optional.ofNullable(stateClaims.get(HANDLER_STATE_CLAIM, String.class)); + } + catch (ChallengeFailedException | RuntimeException e) { + logger.error(e, "Authentication response could not be verified invalid state: state=%s", state); + return Response.status(BAD_REQUEST) + .entity(getInternalFailureHtml("Authentication response could not be verified")) + .cookie(NonceCookie.delete()) + .build(); + } + + // Note: the Web UI may be disabled, so REST requests can not redirect to a success or error page inside the Web UI + try { + // fetch access token + OAuth2Client.Response oauth2Response = client.getOAuth2Response(code, callbackUri, nonce); + + if (!handlerState.isPresent()) { + UriBuilder uriBuilder = getSchemeUriBuilder(request); + return Response + .seeOther(uriBuilder.build().resolve("/ui/")) + .cookie( + OAuthWebUiCookie.create( + tokenPairSerializer.serialize( + fromOAuth2Response(oauth2Response)), + tokenExpiration + .map(expiration -> Instant.now().plus(expiration)) + .orElse(oauth2Response.getExpiration())), + NonceCookie.delete()) + .build(); + } + + tokenHandler.setAccessToken(handlerState.get(), tokenPairSerializer.serialize(fromOAuth2Response(oauth2Response))); + + Response.ResponseBuilder builder = Response.ok(getSuccessHtml()); + builder.cookie( + OAuthWebUiCookie.create( + tokenPairSerializer.serialize(fromOAuth2Response(oauth2Response)), + tokenExpiration.map(expiration -> Instant.now().plus(expiration)) + .orElse(oauth2Response.getExpiration()))); + + return builder.cookie(NonceCookie.delete()).build(); + } + catch (ChallengeFailedException | RuntimeException e) { + logger.error(e, "Authentication response could not be verified: state=%s", state); + + handlerState.ifPresent(value -> + tokenHandler.setTokenExchangeError(value, format("Authentication response could not be verified: state=%s", value))); + return Response.status(BAD_REQUEST) + .cookie(NonceCookie.delete()) + .entity(getInternalFailureHtml("Authentication response could not be verified")) + .build(); + } + } + + private Claims parseState(String state) + throws ChallengeFailedException + { + try { + return jwtParser + .parseClaimsJws(state) + .getBody(); + } + catch (RuntimeException e) { + throw new ChallengeFailedException("State validation failed", e); + } + } + + public String getSuccessHtml() + { + return successHtml; + } + + public String getCallbackErrorHtml(String errorCode) + { + return failureHtml.replace(FAILURE_REPLACEMENT_TEXT, getOAuth2ErrorMessage(errorCode)); + } + + public String getInternalFailureHtml(String errorMessage) + { + return failureHtml.replace(FAILURE_REPLACEMENT_TEXT, nullToEmpty(errorMessage)); + } + + private static byte[] secureRandomBytes(int count) + { + byte[] bytes = new byte[count]; + SECURE_RANDOM.nextBytes(bytes); + return bytes; + } + + private static String getOAuth2ErrorMessage(String errorCode) + { + try { + OAuth2ErrorCode code = OAuth2ErrorCode.fromString(errorCode); + return code.getMessage(); + } + catch (IllegalArgumentException e) { + logger.error(e, "Unknown error code received code=%s", errorCode); + return "OAuth2 unknown error code: " + errorCode; + } + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2ServiceModule.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2ServiceModule.java new file mode 100644 index 0000000000000..31b8740edbaf4 --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2ServiceModule.java @@ -0,0 +1,105 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import com.facebook.airlift.configuration.AbstractConfigurationAwareModule; +import com.google.inject.Binder; +import com.google.inject.Inject; +import com.google.inject.Key; +import com.google.inject.Provides; +import com.google.inject.Scopes; +import com.google.inject.Singleton; + +import java.time.Duration; + +import static com.facebook.airlift.configuration.ConditionalModule.installModuleIf; +import static com.facebook.airlift.configuration.ConfigBinder.configBinder; +import static com.facebook.airlift.http.client.HttpClientBinder.httpClientBinder; +import static com.facebook.airlift.jaxrs.JaxrsBinder.jaxrsBinder; +import static com.facebook.presto.server.security.oauth2.TokenPairSerializer.ACCESS_TOKEN_ONLY_SERIALIZER; +import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder; + +public class OAuth2ServiceModule + extends AbstractConfigurationAwareModule +{ + @Override + protected void setup(Binder binder) + { + jaxrsBinder(binder).bind(OAuth2CallbackResource.class); + configBinder(binder).bindConfig(OAuth2Config.class); + binder.bind(OAuth2Service.class).in(Scopes.SINGLETON); + binder.bind(OAuth2TokenHandler.class).to(OAuth2TokenExchange.class).in(Scopes.SINGLETON); + binder.bind(NimbusHttpClient.class).to(NimbusAirliftHttpClient.class).in(Scopes.SINGLETON); + newOptionalBinder(binder, OAuth2Client.class) + .setDefault() + .to(NimbusOAuth2Client.class) + .in(Scopes.SINGLETON); + install(installModuleIf(OAuth2Config.class, OAuth2Config::isEnableDiscovery, this::bindOidcDiscovery, this::bindStaticConfiguration)); + install(installModuleIf(OAuth2Config.class, OAuth2Config::isEnableRefreshTokens, this::enableRefreshTokens, this::disableRefreshTokens)); + httpClientBinder(binder) + .bindHttpClient("oauth2-jwk", ForOAuth2.class) + // Reset to defaults to override InternalCommunicationModule changes to this client default configuration. + // Setting a keystore and/or a truststore for internal communication changes the default SSL configuration + // for all clients in this guice context. This does not make sense for this client which will very rarely + // use the same SSL configuration, so using the system default truststore makes more sense. + .withConfigDefaults(config -> config + .setKeyStorePath(null) + .setKeyStorePassword(null) + .setTrustStorePath(null) + .setTrustStorePassword(null)); + } + + private void enableRefreshTokens(Binder binder) + { + install(new JweTokenSerializerModule()); + } + + private void disableRefreshTokens(Binder binder) + { + binder.bind(TokenPairSerializer.class).toInstance(ACCESS_TOKEN_ONLY_SERIALIZER); + newOptionalBinder(binder, Key.get(Duration.class, ForRefreshTokens.class)); + } + + @Singleton + @Provides + @Inject + public TokenRefresher getTokenRefresher(TokenPairSerializer tokenAssembler, OAuth2TokenHandler tokenHandler, OAuth2Client oAuth2Client) + { + return new TokenRefresher(tokenAssembler, tokenHandler, oAuth2Client); + } + + private void bindStaticConfiguration(Binder binder) + { + configBinder(binder).bindConfig(StaticOAuth2ServerConfiguration.class); + binder.bind(OAuth2ServerConfigProvider.class).to(StaticConfigurationProvider.class).in(Scopes.SINGLETON); + } + + private void bindOidcDiscovery(Binder binder) + { + configBinder(binder).bindConfig(OidcDiscoveryConfig.class); + binder.bind(OAuth2ServerConfigProvider.class).to(OidcDiscovery.class).in(Scopes.SINGLETON); + } + + @Override + public int hashCode() + { + return OAuth2ServiceModule.class.hashCode(); + } + + @Override + public boolean equals(Object obj) + { + return obj instanceof OAuth2ServiceModule; + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2TokenExchange.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2TokenExchange.java new file mode 100644 index 0000000000000..ff06a4bfa9c4c --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2TokenExchange.java @@ -0,0 +1,146 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import com.google.common.cache.CacheBuilder; +import com.google.common.cache.CacheLoader; +import com.google.common.cache.LoadingCache; +import com.google.common.hash.Hashing; +import com.google.common.util.concurrent.ListenableFuture; +import com.google.common.util.concurrent.SettableFuture; +import io.airlift.units.Duration; + +import javax.annotation.PreDestroy; +import javax.inject.Inject; + +import java.nio.charset.StandardCharsets; +import java.util.Optional; +import java.util.UUID; +import java.util.concurrent.Future; +import java.util.concurrent.ScheduledExecutorService; + +import static com.facebook.airlift.concurrent.Threads.daemonThreadsNamed; +import static com.google.common.util.concurrent.Futures.nonCancellationPropagating; +import static java.util.Objects.requireNonNull; +import static java.util.concurrent.Executors.newSingleThreadScheduledExecutor; +import static java.util.concurrent.TimeUnit.MILLISECONDS; +import static java.util.concurrent.TimeUnit.SECONDS; + +public class OAuth2TokenExchange + implements OAuth2TokenHandler +{ + public static final Duration MAX_POLL_TIME = new Duration(10, SECONDS); + private static final TokenPoll TOKEN_POLL_TIMED_OUT = TokenPoll.error("Authentication has timed out"); + private static final TokenPoll TOKEN_POLL_DROPPED = TokenPoll.error("Authentication has been finished by the client"); + + private final LoadingCache> cache; + private final ScheduledExecutorService executor = newSingleThreadScheduledExecutor(daemonThreadsNamed("oauth2-token-exchange")); + + @Inject + public OAuth2TokenExchange(OAuth2Config config) + { + long challengeTimeoutMillis = config.getChallengeTimeout().toMillis(); + this.cache = buildUnsafeCache( + CacheBuilder.newBuilder() + .expireAfterWrite(challengeTimeoutMillis + (MAX_POLL_TIME.toMillis() * 10), MILLISECONDS) + .removalListener(notification -> notification.getValue().set(TOKEN_POLL_TIMED_OUT)), + new CacheLoader>() + { + @Override + public SettableFuture load(String authIdHash) + { + SettableFuture future = SettableFuture.create(); + Future timeout = executor.schedule(() -> future.set(TOKEN_POLL_TIMED_OUT), challengeTimeoutMillis, MILLISECONDS); + future.addListener(() -> timeout.cancel(true), executor); + return future; + } + }); + } + + private static LoadingCache buildUnsafeCache(CacheBuilder cacheBuilder, CacheLoader cacheLoader) + { + return cacheBuilder.build(cacheLoader); + } + + @PreDestroy + public void stop() + { + executor.shutdownNow(); + } + + @Override + public void setAccessToken(String authIdHash, String accessToken) + { + cache.getUnchecked(authIdHash).set(TokenPoll.token(accessToken)); + } + + @Override + public void setTokenExchangeError(String authIdHash, String message) + { + cache.getUnchecked(authIdHash).set(TokenPoll.error(message)); + } + + public ListenableFuture getTokenPoll(UUID authId) + { + return nonCancellationPropagating(cache.getUnchecked(hashAuthId(authId))); + } + + public void dropToken(UUID authId) + { + cache.getUnchecked(hashAuthId(authId)).set(TOKEN_POLL_DROPPED); + } + + public static String hashAuthId(UUID authId) + { + return Hashing.sha256() + .hashString(authId.toString(), StandardCharsets.UTF_8) + .toString(); + } + + public static class TokenPoll + { + private final Optional token; + private final Optional error; + + private TokenPoll(String token, String error) + { + this.token = Optional.ofNullable(token); + this.error = Optional.ofNullable(error); + } + + static TokenPoll token(String token) + { + requireNonNull(token, "token is null"); + + return new TokenPoll(token, null); + } + + static TokenPoll error(String error) + { + requireNonNull(error, "error is null"); + + return new TokenPoll(null, error); + } + + public Optional getToken() + { + return token; + } + + public Optional getError() + { + return error; + } + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2TokenExchangeResource.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2TokenExchangeResource.java new file mode 100644 index 0000000000000..36d27a0fe1252 --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2TokenExchangeResource.java @@ -0,0 +1,141 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import com.facebook.airlift.json.JsonCodec; +import com.facebook.airlift.json.JsonCodecFactory; +import com.facebook.presto.dispatcher.DispatchExecutor; +import com.facebook.presto.server.security.oauth2.OAuth2TokenExchange.TokenPoll; +import com.google.common.base.VerifyException; +import com.google.common.collect.ImmutableMap; +import com.google.common.util.concurrent.Futures; +import com.google.common.util.concurrent.ListenableFuture; +import com.google.common.util.concurrent.ListeningExecutorService; + +import javax.inject.Inject; +import javax.servlet.http.HttpServletRequest; +import javax.ws.rs.BadRequestException; +import javax.ws.rs.DELETE; +import javax.ws.rs.GET; +import javax.ws.rs.Path; +import javax.ws.rs.PathParam; +import javax.ws.rs.Produces; +import javax.ws.rs.container.AsyncResponse; +import javax.ws.rs.container.Suspended; +import javax.ws.rs.core.Context; +import javax.ws.rs.core.MediaType; +import javax.ws.rs.core.Response; +import javax.ws.rs.core.UriBuilder; + +import java.util.Map; +import java.util.Optional; +import java.util.UUID; + +import static com.facebook.airlift.http.server.AsyncResponseHandler.bindAsyncResponse; +import static com.facebook.presto.server.security.oauth2.OAuth2CallbackResource.CALLBACK_ENDPOINT; +import static com.facebook.presto.server.security.oauth2.OAuth2TokenExchange.MAX_POLL_TIME; +import static com.facebook.presto.server.security.oauth2.OAuth2TokenExchange.hashAuthId; +import static com.facebook.presto.server.security.oauth2.OAuth2Utils.getSchemeUriBuilder; +import static java.util.Objects.requireNonNull; +import static javax.ws.rs.core.MediaType.APPLICATION_JSON_TYPE; + +@Path(OAuth2TokenExchangeResource.TOKEN_ENDPOINT) +public class OAuth2TokenExchangeResource +{ + public static final String TOKEN_ENDPOINT = "/oauth2/token/"; + private static final JsonCodec> MAP_CODEC = new JsonCodecFactory().mapJsonCodec(String.class, Object.class); + private final OAuth2TokenExchange tokenExchange; + private final OAuth2Service service; + private final ListeningExecutorService responseExecutor; + + @Inject + public OAuth2TokenExchangeResource(OAuth2TokenExchange tokenExchange, OAuth2Service service, DispatchExecutor executor) + { + this.tokenExchange = requireNonNull(tokenExchange, "tokenExchange is null"); + this.service = requireNonNull(service, "service is null"); + this.responseExecutor = requireNonNull(executor, "executor is null").getExecutor(); + } + + @Path("initiate/{authIdHash}") + @GET + @Produces(MediaType.APPLICATION_JSON) + public Response initiateTokenExchange(@PathParam("authIdHash") String authIdHash, @Context HttpServletRequest request) + { + UriBuilder builder = getSchemeUriBuilder(request); + return service.startOAuth2Challenge(builder.build().resolve(CALLBACK_ENDPOINT), Optional.ofNullable(authIdHash)); + } + + @Path("{authId}") + @GET + @Produces(MediaType.APPLICATION_JSON) + public void getAuthenticationToken(@PathParam("authId") UUID authId, @Suspended AsyncResponse asyncResponse, @Context HttpServletRequest request) + { + if (authId == null) { + throw new BadRequestException(); + } + + // Do not drop the response from the cache on failure, as this would result in a + // hang if the client retries the request. The response will timeout eventually. + ListenableFuture tokenFuture = tokenExchange.getTokenPoll(authId); + ListenableFuture responseFuture = Futures.transform(tokenFuture, OAuth2TokenExchangeResource::toResponse, responseExecutor); + bindAsyncResponse(asyncResponse, responseFuture, responseExecutor) + .withTimeout(MAX_POLL_TIME, pendingResponse(request)); + } + + private static Response toResponse(TokenPoll poll) + { + if (poll.getError().isPresent()) { + return Response.ok(jsonMap("error", poll.getError().get()), APPLICATION_JSON_TYPE).build(); + } + if (poll.getToken().isPresent()) { + return Response.ok(jsonMap("token", poll.getToken().get()), APPLICATION_JSON_TYPE).build(); + } + throw new VerifyException("invalid TokenPoll state"); + } + + private static Response pendingResponse(HttpServletRequest request) + { + UriBuilder builder = getSchemeUriBuilder(request); + return Response.ok(jsonMap("nextUri", builder.build()), APPLICATION_JSON_TYPE).build(); + } + + @DELETE + @Path("{authId}") + public Response deleteAuthenticationToken(@PathParam("authId") UUID authId) + { + if (authId == null) { + throw new BadRequestException(); + } + + tokenExchange.dropToken(authId); + return Response + .ok() + .build(); + } + + public static String getTokenUri(UUID authId) + { + return TOKEN_ENDPOINT + authId; + } + + public static String getInitiateUri(UUID authId) + { + return TOKEN_ENDPOINT + "initiate/" + hashAuthId(authId); + } + + private static String jsonMap(String key, Object value) + { + return MAP_CODEC.toJson(ImmutableMap.of(key, value)); + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2TokenHandler.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2TokenHandler.java new file mode 100644 index 0000000000000..027d3c27d90f4 --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2TokenHandler.java @@ -0,0 +1,21 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +public interface OAuth2TokenHandler +{ + void setAccessToken(String hashedState, String accessToken); + + void setTokenExchangeError(String hashedState, String errorMessage); +} diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Utils.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Utils.java new file mode 100644 index 0000000000000..4d9dba90ec2f1 --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Utils.java @@ -0,0 +1,79 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import javax.servlet.http.HttpServletRequest; +import javax.ws.rs.core.MultivaluedMap; +import javax.ws.rs.core.UriBuilder; + +import java.util.List; +import java.util.Map; +import java.util.Optional; + +import static com.google.common.net.HttpHeaders.X_FORWARDED_PROTO; + +public final class OAuth2Utils +{ + private OAuth2Utils() {} + + /** + * Returns a UriBuilder with the scheme set based upon the X_FORWARDED_PROTO header. + * If the header exists on the request we set the scheme to what is in that header. i.e. https. + * If the header is not set then we use the scheme on the request. + * + * Ex: If you are using a load balancer to handle ssl forwarding for Presto. You must set the + * X_FORWARDED_PROTO header in the load balancer to 'https'. For any callback or redirect url's + * for the OAUTH2 Login flow must use the scheme of https. + * + * @param request HttpServletRequest + * @return a new instance of UriBuilder with the scheme set. + */ + public static UriBuilder getSchemeUriBuilder(HttpServletRequest request) + { + Optional forwardedProto = Optional.ofNullable(request.getHeader(X_FORWARDED_PROTO)); + + UriBuilder builder = UriBuilder.fromUri(getFullRequestURL(request)); + if (forwardedProto.isPresent()) { + builder.scheme(forwardedProto.get()); + } + else { + builder.scheme(request.getScheme()); + } + + return builder; + } + + /** + * Finds the lastURL query parameter in the request. + * + * @return Optional the value of the lastURL parameter + */ + public static Optional getLastURLParameter(MultivaluedMap queryParams) + { + Optional>> lastUrl = queryParams.entrySet().stream().filter(qp -> qp.getKey().equals("lastURL")).findFirst(); + if (lastUrl.isPresent() && lastUrl.get().getValue().size() > 0) { + return Optional.ofNullable(lastUrl.get().getValue().get(0)); + } + + return Optional.empty(); + } + + public static String getFullRequestURL(HttpServletRequest request) + { + StringBuilder requestURL = new StringBuilder(request.getRequestURL()); + String queryString = request.getQueryString(); + + return queryString == null ? requestURL.toString() : requestURL.append("?").append(queryString).toString(); + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuthWebUiCookie.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuthWebUiCookie.java new file mode 100644 index 0000000000000..477c0d0d3581d --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuthWebUiCookie.java @@ -0,0 +1,61 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import javax.ws.rs.core.NewCookie; + +import java.time.Instant; +import java.util.Date; + +import static javax.ws.rs.core.Cookie.DEFAULT_VERSION; +import static javax.ws.rs.core.NewCookie.DEFAULT_MAX_AGE; + +public final class OAuthWebUiCookie +{ + // prefix according to: https://tools.ietf.org/html/draft-ietf-httpbis-rfc6265bis-05#section-4.1.3.1 + public static final String OAUTH2_COOKIE = "__Secure-Presto-OAuth2-Token"; + + public static final String API_PATH = "/"; + + private OAuthWebUiCookie() {} + + public static NewCookie create(String token, Instant tokenExpiration) + { + return new NewCookie( + OAUTH2_COOKIE, + token, + API_PATH, + null, + DEFAULT_VERSION, + null, + DEFAULT_MAX_AGE, + Date.from(tokenExpiration), + true, + true); + } + public static NewCookie delete() + { + return new NewCookie( + OAUTH2_COOKIE, + "delete", + API_PATH, + null, + DEFAULT_VERSION, + null, + 0, + null, + true, + true); + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/Oauth2WebUiAuthenticationManager.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/Oauth2WebUiAuthenticationManager.java new file mode 100644 index 0000000000000..d31b993e5b437 --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/Oauth2WebUiAuthenticationManager.java @@ -0,0 +1,119 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import com.facebook.airlift.http.server.AuthenticationException; +import com.facebook.airlift.log.Logger; +import com.facebook.presto.server.security.WebUiAuthenticationManager; +import com.facebook.presto.server.security.oauth2.TokenPairSerializer.TokenPair; + +import javax.inject.Inject; +import javax.servlet.FilterChain; +import javax.servlet.ServletException; +import javax.servlet.http.Cookie; +import javax.servlet.http.HttpServletRequest; +import javax.servlet.http.HttpServletResponse; +import javax.ws.rs.core.UriBuilder; + +import java.io.IOException; +import java.security.Principal; +import java.time.Duration; +import java.time.Instant; +import java.util.Optional; + +import static com.facebook.presto.server.security.AuthenticationFilter.withPrincipal; +import static com.facebook.presto.server.security.oauth2.OAuth2CallbackResource.CALLBACK_ENDPOINT; +import static com.facebook.presto.server.security.oauth2.OAuth2Utils.getSchemeUriBuilder; +import static java.util.Objects.requireNonNull; + +public class Oauth2WebUiAuthenticationManager + implements WebUiAuthenticationManager +{ + private static final Logger logger = Logger.get(Oauth2WebUiAuthenticationManager.class); + private final OAuth2Service oAuth2Service; + private final OAuth2Authenticator oAuth2Authenticator; + private final TokenPairSerializer tokenPairSerializer; + private final OAuth2Client client; + private final Optional tokenExpiration; + + @Inject + public Oauth2WebUiAuthenticationManager(OAuth2Service oAuth2Service, OAuth2Authenticator oAuth2Authenticator, TokenPairSerializer tokenPairSerializer, OAuth2Client client, @ForRefreshTokens Optional tokenExpiration) + { + this.oAuth2Service = requireNonNull(oAuth2Service, "oauth2Service is null"); + this.oAuth2Authenticator = requireNonNull(oAuth2Authenticator, "oauth2Authenticator is null"); + this.tokenPairSerializer = requireNonNull(tokenPairSerializer, "tokenPairSerializer is null"); + this.client = requireNonNull(client, "oauth2Client is null"); + this.tokenExpiration = requireNonNull(tokenExpiration, "tokenExpiration is null"); + } + + public void handleRequest(HttpServletRequest request, HttpServletResponse response, FilterChain nextFilter) + throws IOException, ServletException + { + try { + Principal principal = this.oAuth2Authenticator.authenticate(request); + nextFilter.doFilter(withPrincipal(request, principal), response); + } + catch (AuthenticationException e) { + needAuthentication(request, response); + } + } + + private Optional getTokenPair(HttpServletRequest request) + { + try { + Optional token = this.oAuth2Authenticator.extractTokenFromCookie(request); + if (token.isPresent()) { + return Optional.ofNullable(tokenPairSerializer.deserialize(token.get())); + } + else { + return Optional.empty(); + } + } + catch (Exception e) { + logger.error(e, "Exception occurred during token pair deserialization"); + return Optional.empty(); + } + } + + private void needAuthentication(HttpServletRequest request, HttpServletResponse response) + throws IOException + { + Optional tokenPair = getTokenPair(request); + Optional refreshToken = tokenPair.flatMap(TokenPair::getRefreshToken); + if (refreshToken.isPresent()) { + try { + OAuth2Client.Response refreshRes = client.refreshTokens(refreshToken.get()); + String serializeToken = tokenPairSerializer.serialize(TokenPair.fromOAuth2Response(refreshRes)); + UriBuilder builder = getSchemeUriBuilder(request); + Cookie newCookie = NonceCookie.toServletCookie(OAuthWebUiCookie.create(serializeToken, tokenExpiration.map(expiration -> Instant.now().plus(expiration)).orElse(refreshRes.getExpiration()))); + response.addCookie(newCookie); + response.sendRedirect(builder.build().toString()); + } + catch (ChallengeFailedException e) { + logger.error(e, "Token refresh challenge has failed"); + this.startOauth2Challenge(request, response); + } + } + else { + this.startOauth2Challenge(request, response); + } + } + + private void startOauth2Challenge(HttpServletRequest request, HttpServletResponse response) + throws IOException + { + UriBuilder builder = getSchemeUriBuilder(request); + this.oAuth2Service.startOAuth2Challenge(builder.build().resolve(CALLBACK_ENDPOINT), Optional.empty(), response); + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OidcDiscovery.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OidcDiscovery.java new file mode 100644 index 0000000000000..1b6bae30cc8c5 --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OidcDiscovery.java @@ -0,0 +1,159 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import com.facebook.airlift.json.JsonObjectMapperProvider; +import com.facebook.airlift.log.Logger; +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.nimbusds.oauth2.sdk.ParseException; +import com.nimbusds.oauth2.sdk.http.HTTPResponse; +import com.nimbusds.oauth2.sdk.id.Issuer; +import com.nimbusds.openid.connect.sdk.op.OIDCProviderConfigurationRequest; +import com.nimbusds.openid.connect.sdk.op.OIDCProviderMetadata; +import net.jodah.failsafe.Failsafe; +import net.jodah.failsafe.RetryPolicy; + +import javax.inject.Inject; + +import java.net.URI; +import java.time.Duration; +import java.util.Optional; + +import static com.facebook.airlift.http.client.HttpStatus.OK; +import static com.facebook.airlift.http.client.HttpStatus.REQUEST_TIMEOUT; +import static com.facebook.airlift.http.client.HttpStatus.TOO_MANY_REQUESTS; +import static com.facebook.presto.server.security.oauth2.StaticOAuth2ServerConfiguration.ACCESS_TOKEN_ISSUER; +import static com.facebook.presto.server.security.oauth2.StaticOAuth2ServerConfiguration.AUTH_URL; +import static com.facebook.presto.server.security.oauth2.StaticOAuth2ServerConfiguration.JWKS_URL; +import static com.facebook.presto.server.security.oauth2.StaticOAuth2ServerConfiguration.TOKEN_URL; +import static com.facebook.presto.server.security.oauth2.StaticOAuth2ServerConfiguration.USERINFO_URL; +import static com.google.common.base.Preconditions.checkState; +import static java.lang.String.format; +import static java.util.Objects.requireNonNull; + +public class OidcDiscovery + implements OAuth2ServerConfigProvider +{ + private static final Logger LOG = Logger.get(OidcDiscovery.class); + + private static final ObjectMapper OBJECT_MAPPER = new JsonObjectMapperProvider().get(); + private final Issuer issuer; + private final Duration discoveryTimeout; + private final boolean userinfoEndpointEnabled; + private final Optional accessTokenIssuer; + private final Optional authUrl; + private final Optional tokenUrl; + private final Optional jwksUrl; + private final Optional userinfoUrl; + private final NimbusHttpClient httpClient; + + @Inject + public OidcDiscovery(OAuth2Config oauthConfig, OidcDiscoveryConfig oidcConfig, NimbusHttpClient httpClient) + { + requireNonNull(oauthConfig, "oauthConfig is null"); + issuer = new Issuer(requireNonNull(oauthConfig.getIssuer(), "issuer is null")); + requireNonNull(oidcConfig, "oidcConfig is null"); + userinfoEndpointEnabled = oidcConfig.isUserinfoEndpointEnabled(); + discoveryTimeout = Duration.ofMillis(requireNonNull(oidcConfig.getDiscoveryTimeout(), "discoveryTimeout is null").toMillis()); + accessTokenIssuer = requireNonNull(oidcConfig.getAccessTokenIssuer(), "accessTokenIssuer is null"); + authUrl = requireNonNull(oidcConfig.getAuthUrl(), "authUrl is null"); + tokenUrl = requireNonNull(oidcConfig.getTokenUrl(), "tokenUrl is null"); + jwksUrl = requireNonNull(oidcConfig.getJwksUrl(), "jwksUrl is null"); + userinfoUrl = requireNonNull(oidcConfig.getUserinfoUrl(), "userinfoUrl is null"); + this.httpClient = requireNonNull(httpClient, "httpClient is null"); + } + + @Override + public OAuth2ServerConfig get() + { + return Failsafe.with(new RetryPolicy<>() + .withMaxAttempts(-1) + .withMaxDuration(discoveryTimeout) + .withDelay(Duration.ofSeconds(1)) + .abortOn(IllegalStateException.class) + .onFailedAttempt(attempt -> LOG.debug("OpenID Connect Metadata read failed: %s", attempt.getLastFailure()))) + .get(() -> httpClient.execute(new OIDCProviderConfigurationRequest(issuer), this::parseConfigurationResponse)); + } + + private OAuth2ServerConfig parseConfigurationResponse(HTTPResponse response) + throws ParseException + { + int statusCode = response.getStatusCode(); + if (statusCode != OK.code()) { + // stop on any client errors other than REQUEST_TIMEOUT and TOO_MANY_REQUESTS + if (statusCode < 400 || statusCode >= 500 || statusCode == REQUEST_TIMEOUT.code() || statusCode == TOO_MANY_REQUESTS.code()) { + throw new RuntimeException("Invalid response from OpenID Metadata endpoint: " + statusCode); + } + else { + throw new IllegalStateException(format("Invalid response from OpenID Metadata endpoint. Expected response code to be %s, but was %s", OK.code(), statusCode)); + } + } + return readConfiguration(response.getContent()); + } + + private OAuth2ServerConfig readConfiguration(String body) + throws ParseException + { + OIDCProviderMetadata metadata = OIDCProviderMetadata.parse(body); + checkMetadataState(issuer.equals(metadata.getIssuer()), "The value of the \"issuer\" claim in Metadata document different than the Issuer URL used for the Configuration Request."); + try { + JsonNode metadataJson = OBJECT_MAPPER.readTree(body); + Optional userinfoEndpoint; + if (userinfoEndpointEnabled) { + userinfoEndpoint = getOptionalField("userinfo_endpoint", Optional.ofNullable(metadata.getUserInfoEndpointURI()).map(URI::toString), USERINFO_URL, userinfoUrl); + } + else { + userinfoEndpoint = Optional.empty(); + } + return new OAuth2ServerConfig( + // AD FS server can include "access_token_issuer" field in OpenID Provider Metadata. + // It's not a part of the OIDC standard thus have to be handled separately. + // see: https://docs.microsoft.com/en-us/openspecs/windows_protocols/ms-oidce/f629647a-4825-465b-80bb-32c7e9cec2c8 + getOptionalField("access_token_issuer", Optional.ofNullable(metadataJson.get("access_token_issuer")).map(JsonNode::textValue), ACCESS_TOKEN_ISSUER, accessTokenIssuer), + getRequiredField("authorization_endpoint", metadata.getAuthorizationEndpointURI(), AUTH_URL, authUrl), + getRequiredField("token_endpoint", metadata.getTokenEndpointURI(), TOKEN_URL, tokenUrl), + getRequiredField("jwks_uri", metadata.getJWKSetURI(), JWKS_URL, jwksUrl), + userinfoEndpoint.map(URI::create)); + } + catch (JsonProcessingException e) { + throw new ParseException("Invalid JSON value", e); + } + } + + private static URI getRequiredField(String metadataField, URI metadataValue, String configurationField, Optional configurationValue) + { + Optional uri = getOptionalField(metadataField, Optional.ofNullable(metadataValue).map(URI::toString), configurationField, configurationValue); + checkMetadataState(uri.isPresent(), "Missing required \"%s\" property.", metadataField); + return URI.create(uri.get()); + } + + private static Optional getOptionalField(String metadataField, Optional metadataValue, String configurationField, Optional configurationValue) + { + if (configurationValue.isPresent()) { + if (!configurationValue.equals(metadataValue)) { + LOG.warn("Overriding \"%s=%s\" from OpenID metadata document with value \"%s=%s\" defined in configuration", + metadataField, metadataValue.orElse(""), configurationField, configurationValue.orElse("")); + } + return configurationValue; + } + return metadataValue; + } + + private static void checkMetadataState(boolean expression, String additionalMessage, String... additionalMessageArgs) + { + checkState(expression, "Invalid response from OpenID Metadata endpoint. " + additionalMessage, (Object[]) additionalMessageArgs); + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OidcDiscoveryConfig.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OidcDiscoveryConfig.java new file mode 100644 index 0000000000000..126a5308b7f94 --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OidcDiscoveryConfig.java @@ -0,0 +1,149 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import com.facebook.airlift.configuration.Config; +import com.facebook.airlift.configuration.ConfigDescription; +import io.airlift.units.Duration; + +import javax.validation.constraints.NotNull; + +import java.util.Optional; + +import static com.facebook.presto.server.security.oauth2.StaticOAuth2ServerConfiguration.ACCESS_TOKEN_ISSUER; +import static com.facebook.presto.server.security.oauth2.StaticOAuth2ServerConfiguration.AUTH_URL; +import static com.facebook.presto.server.security.oauth2.StaticOAuth2ServerConfiguration.JWKS_URL; +import static com.facebook.presto.server.security.oauth2.StaticOAuth2ServerConfiguration.TOKEN_URL; +import static com.facebook.presto.server.security.oauth2.StaticOAuth2ServerConfiguration.USERINFO_URL; +import static java.util.concurrent.TimeUnit.SECONDS; + +public class OidcDiscoveryConfig +{ + private Duration discoveryTimeout = new Duration(30, SECONDS); + private boolean userinfoEndpointEnabled = true; + + //TODO Left for backward compatibility, remove after the next release/a couple of releases + private Optional accessTokenIssuer = Optional.empty(); + private Optional authUrl = Optional.empty(); + private Optional tokenUrl = Optional.empty(); + private Optional jwksUrl = Optional.empty(); + private Optional userinfoUrl = Optional.empty(); + + @NotNull + public Duration getDiscoveryTimeout() + { + return discoveryTimeout; + } + + @Config("http-server.authentication.oauth2.oidc.discovery.timeout") + @ConfigDescription("OpenID Connect discovery timeout") + public OidcDiscoveryConfig setDiscoveryTimeout(Duration discoveryTimeout) + { + this.discoveryTimeout = discoveryTimeout; + return this; + } + + public boolean isUserinfoEndpointEnabled() + { + return userinfoEndpointEnabled; + } + + @Config("http-server.authentication.oauth2.oidc.use-userinfo-endpoint") + @ConfigDescription("Use userinfo endpoint from OpenID connect metadata document") + public OidcDiscoveryConfig setUserinfoEndpointEnabled(boolean userinfoEndpointEnabled) + { + this.userinfoEndpointEnabled = userinfoEndpointEnabled; + return this; + } + + @NotNull + @Deprecated + public Optional getAccessTokenIssuer() + { + return accessTokenIssuer; + } + + @Config(ACCESS_TOKEN_ISSUER) + @ConfigDescription("The required issuer for access tokens") + @Deprecated + public OidcDiscoveryConfig setAccessTokenIssuer(String accessTokenIssuer) + { + this.accessTokenIssuer = Optional.ofNullable(accessTokenIssuer); + return this; + } + + @NotNull + @Deprecated + public Optional getAuthUrl() + { + return authUrl; + } + + @Config(AUTH_URL) + @ConfigDescription("URL of the authorization server's authorization endpoint") + @Deprecated + public OidcDiscoveryConfig setAuthUrl(String authUrl) + { + this.authUrl = Optional.ofNullable(authUrl); + return this; + } + + @NotNull + @Deprecated + public Optional getTokenUrl() + { + return tokenUrl; + } + + @Config(TOKEN_URL) + @ConfigDescription("URL of the authorization server's token endpoint") + @Deprecated + public OidcDiscoveryConfig setTokenUrl(String tokenUrl) + { + this.tokenUrl = Optional.ofNullable(tokenUrl); + return this; + } + + @NotNull + @Deprecated + public Optional getJwksUrl() + { + return jwksUrl; + } + + @Config(JWKS_URL) + @ConfigDescription("URL of the authorization server's JWKS (JSON Web Key Set) endpoint") + @Deprecated + public OidcDiscoveryConfig setJwksUrl(String jwksUrl) + { + this.jwksUrl = Optional.ofNullable(jwksUrl); + return this; + } + + @NotNull + @Deprecated + public Optional getUserinfoUrl() + { + return userinfoUrl; + } + + @Config(USERINFO_URL) + @ConfigDescription("URL of the userinfo endpoint") + @Deprecated + public OidcDiscoveryConfig setUserinfoUrl(String userinfoUrl) + { + this.userinfoUrl = Optional.ofNullable(userinfoUrl); + return this; + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/RefreshTokensConfig.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/RefreshTokensConfig.java new file mode 100644 index 0000000000000..69bcd717e7451 --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/RefreshTokensConfig.java @@ -0,0 +1,96 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import com.facebook.airlift.configuration.Config; +import com.facebook.airlift.configuration.ConfigDescription; +import com.facebook.airlift.configuration.ConfigSecuritySensitive; +import io.airlift.units.Duration; +import io.jsonwebtoken.io.Decoders; +import io.jsonwebtoken.security.Keys; + +import javax.crypto.SecretKey; +import javax.validation.constraints.NotEmpty; + +import static com.google.common.base.Strings.isNullOrEmpty; +import static java.util.concurrent.TimeUnit.HOURS; + +public class RefreshTokensConfig +{ + private Duration tokenExpiration = Duration.succinctDuration(1, HOURS); + private static final String coordinator = "Presto_coordinator"; + private String issuer = coordinator; + private String audience = coordinator; + + private SecretKey secretKey; + + public Duration getTokenExpiration() + { + return tokenExpiration; + } + + @Config("http-server.authentication.oauth2.refresh-tokens.issued-token.timeout") + @ConfigDescription("Expiration time for issued token. It needs to be equal or lower than duration of refresh token issued by IdP") + public RefreshTokensConfig setTokenExpiration(Duration tokenExpiration) + { + this.tokenExpiration = tokenExpiration; + return this; + } + + @NotEmpty + public String getIssuer() + { + return issuer; + } + + @Config("http-server.authentication.oauth2.refresh-tokens.issued-token.issuer") + @ConfigDescription("Issuer representing this coordinator instance, that will be used in issued token. In addition current Version will be added to it") + public RefreshTokensConfig setIssuer(String issuer) + { + this.issuer = issuer; + return this; + } + + @NotEmpty + public String getAudience() + { + return audience; + } + + @Config("http-server.authentication.oauth2.refresh-tokens.issued-token.audience") + @ConfigDescription("Audience representing this coordinator instance, that will be used in issued token") + public RefreshTokensConfig setAudience(String audience) + { + this.audience = audience; + return this; + } + + @Config("http-server.authentication.oauth2.refresh-tokens.secret-key") + @ConfigDescription("Base64 encoded secret key used to encrypt generated token") + @ConfigSecuritySensitive + public RefreshTokensConfig setSecretKey(String key) + { + if (isNullOrEmpty(key)) { + return this; + } + + secretKey = Keys.hmacShaKeyFor(Decoders.BASE64.decode(key)); + return this; + } + + public SecretKey getSecretKey() + { + return secretKey; + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/StaticConfigurationProvider.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/StaticConfigurationProvider.java new file mode 100644 index 0000000000000..627c42b3a5814 --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/StaticConfigurationProvider.java @@ -0,0 +1,44 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import javax.inject.Inject; + +import java.net.URI; + +import static java.util.Objects.requireNonNull; + +public class StaticConfigurationProvider + implements OAuth2ServerConfigProvider +{ + private final OAuth2ServerConfig config; + + @Inject + StaticConfigurationProvider(StaticOAuth2ServerConfiguration config) + { + requireNonNull(config, "config is null"); + this.config = new OAuth2ServerConfig( + config.getAccessTokenIssuer(), + URI.create(config.getAuthUrl()), + URI.create(config.getTokenUrl()), + URI.create(config.getJwksUrl()), + config.getUserinfoUrl().map(URI::create)); + } + + @Override + public OAuth2ServerConfig get() + { + return config; + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/StaticOAuth2ServerConfiguration.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/StaticOAuth2ServerConfiguration.java new file mode 100644 index 0000000000000..2c497f05a27cb --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/StaticOAuth2ServerConfiguration.java @@ -0,0 +1,105 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import com.facebook.airlift.configuration.Config; +import com.facebook.airlift.configuration.ConfigDescription; + +import javax.validation.constraints.NotNull; + +import java.util.Optional; + +public class StaticOAuth2ServerConfiguration +{ + public static final String ACCESS_TOKEN_ISSUER = "http-server.authentication.oauth2.access-token-issuer"; + public static final String AUTH_URL = "http-server.authentication.oauth2.auth-url"; + public static final String TOKEN_URL = "http-server.authentication.oauth2.token-url"; + public static final String JWKS_URL = "http-server.authentication.oauth2.jwks-url"; + public static final String USERINFO_URL = "http-server.authentication.oauth2.userinfo-url"; + + private Optional accessTokenIssuer = Optional.empty(); + private String authUrl; + private String tokenUrl; + private String jwksUrl; + private Optional userinfoUrl = Optional.empty(); + + @NotNull + public Optional getAccessTokenIssuer() + { + return accessTokenIssuer; + } + + @Config(ACCESS_TOKEN_ISSUER) + @ConfigDescription("The required issuer for access tokens") + public StaticOAuth2ServerConfiguration setAccessTokenIssuer(String accessTokenIssuer) + { + this.accessTokenIssuer = Optional.ofNullable(accessTokenIssuer); + return this; + } + + @NotNull + public String getAuthUrl() + { + return authUrl; + } + + @Config(AUTH_URL) + @ConfigDescription("URL of the authorization server's authorization endpoint") + public StaticOAuth2ServerConfiguration setAuthUrl(String authUrl) + { + this.authUrl = authUrl; + return this; + } + + @NotNull + public String getTokenUrl() + { + return tokenUrl; + } + + @Config(TOKEN_URL) + @ConfigDescription("URL of the authorization server's token endpoint") + public StaticOAuth2ServerConfiguration setTokenUrl(String tokenUrl) + { + this.tokenUrl = tokenUrl; + return this; + } + + @NotNull + public String getJwksUrl() + { + return jwksUrl; + } + + @Config(JWKS_URL) + @ConfigDescription("URL of the authorization server's JWKS (JSON Web Key Set) endpoint") + public StaticOAuth2ServerConfiguration setJwksUrl(String jwksUrl) + { + this.jwksUrl = jwksUrl; + return this; + } + + public Optional getUserinfoUrl() + { + return userinfoUrl; + } + + @Config(USERINFO_URL) + @ConfigDescription("URL of the userinfo endpoint") + public StaticOAuth2ServerConfiguration setUserinfoUrl(String userinfoUrl) + { + this.userinfoUrl = Optional.ofNullable(userinfoUrl); + return this; + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/TokenPairSerializer.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/TokenPairSerializer.java new file mode 100644 index 0000000000000..715374cc8e61d --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/TokenPairSerializer.java @@ -0,0 +1,92 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.facebook.presto.server.security.oauth2; + +import com.facebook.presto.server.security.oauth2.OAuth2Client.Response; + +import javax.annotation.Nullable; + +import java.util.Date; +import java.util.Optional; + +import static java.lang.Long.MAX_VALUE; +import static java.util.Objects.requireNonNull; + +public interface TokenPairSerializer +{ + TokenPairSerializer ACCESS_TOKEN_ONLY_SERIALIZER = new TokenPairSerializer() + { + @Override + public TokenPair deserialize(String token) + { + return TokenPair.accessToken(token); + } + + @Override + public String serialize(TokenPair tokenPair) + { + return tokenPair.getAccessToken(); + } + }; + + TokenPair deserialize(String token); + + String serialize(TokenPair tokenPair); + + class TokenPair + { + private final String accessToken; + private final Date expiration; + private final Optional refreshToken; + + private TokenPair(String accessToken, Date expiration, Optional refreshToken) + { + this.accessToken = requireNonNull(accessToken, "accessToken is nul"); + this.expiration = requireNonNull(expiration, "expiration is null"); + this.refreshToken = requireNonNull(refreshToken, "refreshToken is null"); + } + + public static TokenPair accessToken(String accessToken) + { + return new TokenPair(accessToken, new Date(MAX_VALUE), Optional.empty()); + } + + public static TokenPair fromOAuth2Response(Response tokens) + { + requireNonNull(tokens, "tokens is null"); + return new TokenPair(tokens.getAccessToken(), Date.from(tokens.getExpiration()), tokens.getRefreshToken()); + } + + public static TokenPair accessAndRefreshTokens(String accessToken, Date expiration, @Nullable String refreshToken) + { + return new TokenPair(accessToken, expiration, Optional.ofNullable(refreshToken)); + } + + public String getAccessToken() + { + return accessToken; + } + + public Date getExpiration() + { + return expiration; + } + + public Optional getRefreshToken() + { + return refreshToken; + } + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/TokenRefresher.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/TokenRefresher.java new file mode 100644 index 0000000000000..411103f6f8724 --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/TokenRefresher.java @@ -0,0 +1,68 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import com.facebook.presto.server.security.oauth2.OAuth2Client.Response; +import com.facebook.presto.server.security.oauth2.TokenPairSerializer.TokenPair; + +import java.util.Optional; +import java.util.UUID; + +import static com.facebook.presto.server.security.oauth2.OAuth2TokenExchange.hashAuthId; +import static java.util.Objects.requireNonNull; + +public class TokenRefresher +{ + private final TokenPairSerializer tokenAssembler; + private final OAuth2TokenHandler tokenHandler; + private final OAuth2Client client; + + public TokenRefresher(TokenPairSerializer tokenAssembler, OAuth2TokenHandler tokenHandler, OAuth2Client client) + { + this.tokenAssembler = requireNonNull(tokenAssembler, "tokenAssembler is null"); + this.tokenHandler = requireNonNull(tokenHandler, "tokenHandler is null"); + this.client = requireNonNull(client, "oAuth2Client is null"); + } + + public Optional refreshToken(TokenPair tokenPair) + { + requireNonNull(tokenPair, "tokenPair is null"); + + Optional refreshToken = tokenPair.getRefreshToken(); + if (refreshToken.isPresent()) { + UUID refreshingId = UUID.randomUUID(); + try { + refreshToken(refreshToken.get(), refreshingId); + return Optional.of(refreshingId); + } + // If Refresh token has expired then restart the flow + catch (RuntimeException exception) { + return Optional.empty(); + } + } + return Optional.empty(); + } + + private void refreshToken(String refreshToken, UUID refreshingId) + { + try { + Response response = client.refreshTokens(refreshToken); + String serializedToken = tokenAssembler.serialize(TokenPair.fromOAuth2Response(response)); + tokenHandler.setAccessToken(hashAuthId(refreshingId), serializedToken); + } + catch (ChallengeFailedException e) { + tokenHandler.setTokenExchangeError(hashAuthId(refreshingId), "Token refreshing has failed: " + e.getMessage()); + } + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/ZstdCodec.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/ZstdCodec.java new file mode 100644 index 0000000000000..1065d90416240 --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/ZstdCodec.java @@ -0,0 +1,53 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import io.airlift.compress.zstd.ZstdCompressor; +import io.airlift.compress.zstd.ZstdDecompressor; +import io.jsonwebtoken.CompressionCodec; +import io.jsonwebtoken.CompressionException; + +import static java.lang.Math.toIntExact; +import static java.util.Arrays.copyOfRange; + +public class ZstdCodec + implements CompressionCodec +{ + public static final String CODEC_NAME = "ZSTD"; + + @Override + public String getAlgorithmName() + { + return CODEC_NAME; + } + + @Override + public byte[] compress(byte[] bytes) + throws CompressionException + { + ZstdCompressor compressor = new ZstdCompressor(); + byte[] compressed = new byte[compressor.maxCompressedLength(bytes.length)]; + int outputSize = compressor.compress(bytes, 0, bytes.length, compressed, 0, compressed.length); + return copyOfRange(compressed, 0, outputSize); + } + + @Override + public byte[] decompress(byte[] bytes) + throws CompressionException + { + byte[] output = new byte[toIntExact(ZstdDecompressor.getDecompressedSize(bytes, 0, bytes.length))]; + new ZstdDecompressor().decompress(bytes, 0, bytes.length, output, 0, output.length); + return output; + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/server/testing/TestingPrestoServer.java b/presto-main/src/main/java/com/facebook/presto/server/testing/TestingPrestoServer.java index e58509da40263..c46f3a1df23a7 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/testing/TestingPrestoServer.java +++ b/presto-main/src/main/java/com/facebook/presto/server/testing/TestingPrestoServer.java @@ -228,6 +228,11 @@ public TestingPrestoServer(List additionalModules) this(true, ImmutableMap.of(), null, null, new SqlParserOptions(), additionalModules); } + public TestingPrestoServer(Map properties) throws Exception + { + this(true, properties, null, null, new SqlParserOptions(), ImmutableList.of()); + } + public TestingPrestoServer( boolean coordinator, Map properties, diff --git a/presto-main/src/main/resources/oauth2/failure.html b/presto-main/src/main/resources/oauth2/failure.html new file mode 100644 index 0000000000000..1f4dc659ad49d --- /dev/null +++ b/presto-main/src/main/resources/oauth2/failure.html @@ -0,0 +1,60 @@ + + + + + + + + + Failure + + + + + + + + + + + + + + + + + + + + + + + + + + + +

OAuth2 authentication failed

+ +

+ + + + + + diff --git a/presto-main/src/main/resources/oauth2/success.html b/presto-main/src/main/resources/oauth2/success.html new file mode 100644 index 0000000000000..dac401a1d9e6d --- /dev/null +++ b/presto-main/src/main/resources/oauth2/success.html @@ -0,0 +1,60 @@ + + + + + + + + + Success + + + + + + + + + + + + + + + + + + + + + + + + + + + +

OAuth2 authentication succeeded

+ +

This browser window can be closed

+ + + + + + diff --git a/presto-main/src/test/java/com/facebook/presto/TestClientRequestFilterPlugin.java b/presto-main/src/test/java/com/facebook/presto/TestClientRequestFilterPlugin.java index d073f8ed96fee..eb295a80fbade 100644 --- a/presto-main/src/test/java/com/facebook/presto/TestClientRequestFilterPlugin.java +++ b/presto-main/src/test/java/com/facebook/presto/TestClientRequestFilterPlugin.java @@ -16,6 +16,7 @@ import com.facebook.airlift.http.server.Authenticator; import com.facebook.presto.server.MockHttpServletRequest; import com.facebook.presto.server.security.AuthenticationFilter; +import com.facebook.presto.server.security.DefaultWebUiAuthenticationManager; import com.facebook.presto.server.security.SecurityConfig; import com.facebook.presto.server.testing.TestingPrestoServer; import com.facebook.presto.spi.ClientRequestFilter; @@ -111,7 +112,7 @@ private AuthenticationFilter setupAuthenticationFilter(List authenticators = createAuthenticators(); SecurityConfig securityConfig = createSecurityConfig(); - return new AuthenticationFilter(authenticators, securityConfig, clientRequestFilterManager); + return new AuthenticationFilter(authenticators, securityConfig, clientRequestFilterManager, new DefaultWebUiAuthenticationManager()); } } diff --git a/presto-main/src/test/java/com/facebook/presto/server/MockHttpServletRequest.java b/presto-main/src/test/java/com/facebook/presto/server/MockHttpServletRequest.java index b1fe03962f481..0c85c083070e9 100644 --- a/presto-main/src/test/java/com/facebook/presto/server/MockHttpServletRequest.java +++ b/presto-main/src/test/java/com/facebook/presto/server/MockHttpServletRequest.java @@ -30,8 +30,10 @@ import jakarta.servlet.http.HttpSession; import jakarta.servlet.http.HttpUpgradeHandler; import jakarta.servlet.http.Part; +import jakarts.ws.rs.core.UriBuilder; import java.io.BufferedReader; +import java.net.URI; import java.security.Principal; import java.util.Collection; import java.util.Enumeration; @@ -50,12 +52,14 @@ public class MockHttpServletRequest private final ListMultimap headers; private final String remoteAddress; private final Map attributes; + private final String requestUrl; public MockHttpServletRequest(ListMultimap headers, String remoteAddress, Map attributes) { this.headers = ImmutableListMultimap.copyOf(requireNonNull(headers, "headers is null")); this.remoteAddress = requireNonNull(remoteAddress, "remoteAddress is null"); this.attributes = new HashMap<>(requireNonNull(attributes, "attributes is null")); + this.requestUrl = null; } public MockHttpServletRequest(ListMultimap headers) @@ -64,6 +68,14 @@ public MockHttpServletRequest(ListMultimap headers) this(headers, DEFAULT_ADDRESS, ImmutableMap.of()); } + public MockHttpServletRequest(ListMultimap headers, String remoteAddress, String requestUrl) + { + this.headers = ImmutableListMultimap.copyOf(requireNonNull(headers, "headers is null")); + this.remoteAddress = requireNonNull(remoteAddress, "remoteAddress is null"); + this.requestUrl = requireNonNull(requestUrl, "requestUrl is null"); + this.attributes = ImmutableMap.of(); + } + @Override public String getAuthType() { @@ -145,7 +157,12 @@ public String getContextPath() @Override public String getQueryString() { - throw new UnsupportedOperationException(); + if (this.requestUrl == null) { + throw new UnsupportedOperationException(); + } + URI uri = UriBuilder.fromUri(this.requestUrl).build(); + + return uri.getQuery(); } @Override @@ -181,7 +198,10 @@ public String getRequestURI() @Override public StringBuffer getRequestURL() { - throw new UnsupportedOperationException(); + if (this.requestUrl == null) { + throw new UnsupportedOperationException(); + } + return new StringBuffer(this.requestUrl); } @Override @@ -337,7 +357,12 @@ public String getProtocol() @Override public String getScheme() { - throw new UnsupportedOperationException(); + if (this.requestUrl == null) { + throw new UnsupportedOperationException(); + } + URI uri = UriBuilder.fromUri(this.requestUrl).build(); + + return uri.getScheme(); } @Override diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestJweTokenSerializer.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestJweTokenSerializer.java new file mode 100644 index 0000000000000..37dae99191556 --- /dev/null +++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestJweTokenSerializer.java @@ -0,0 +1,171 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import com.facebook.presto.server.security.oauth2.TokenPairSerializer.TokenPair; +import com.nimbusds.jose.KeyLengthException; +import io.airlift.units.Duration; +import io.jsonwebtoken.ExpiredJwtException; +import io.jsonwebtoken.Jwts; +import org.testng.annotations.Test; + +import java.net.URI; +import java.security.GeneralSecurityException; +import java.time.Clock; +import java.time.Instant; +import java.time.ZoneId; +import java.time.ZonedDateTime; +import java.util.Calendar; +import java.util.Date; +import java.util.Map; +import java.util.Optional; + +import static com.facebook.presto.server.security.oauth2.TokenPairSerializer.TokenPair.accessAndRefreshTokens; +import static io.airlift.units.Duration.succinctDuration; +import static java.time.temporal.ChronoUnit.MILLIS; +import static java.util.concurrent.TimeUnit.MINUTES; +import static java.util.concurrent.TimeUnit.SECONDS; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +public class TestJweTokenSerializer +{ + @Test + public void testSerialization() + throws Exception + { + JweTokenSerializer serializer = tokenSerializer(Clock.systemUTC(), succinctDuration(5, SECONDS)); + + Date expiration = new Calendar.Builder().setDate(2022, 6, 22).build().getTime(); + String serializedTokenPair = serializer.serialize(accessAndRefreshTokens("access_token", expiration, "refresh_token")); + TokenPair deserializedTokenPair = serializer.deserialize(serializedTokenPair); + + assertThat(deserializedTokenPair.getAccessToken()).isEqualTo("access_token"); + assertThat(deserializedTokenPair.getExpiration()).isEqualTo(expiration); + assertThat(deserializedTokenPair.getRefreshToken()).isEqualTo(Optional.of("refresh_token")); + } + + @Test + public void testTokenDeserializationAfterTimeoutButBeforeExpirationExtension() + throws Exception + { + TestingClock clock = new TestingClock(); + JweTokenSerializer serializer = tokenSerializer( + clock, + succinctDuration(12, MINUTES)); + Date expiration = new Calendar.Builder().setDate(2022, 6, 22).build().getTime(); + String serializedTokenPair = serializer.serialize(accessAndRefreshTokens("access_token", expiration, "refresh_token")); + clock.advanceBy(succinctDuration(10, MINUTES)); + TokenPair deserializedTokenPair = serializer.deserialize(serializedTokenPair); + + assertThat(deserializedTokenPair.getAccessToken()).isEqualTo("access_token"); + assertThat(deserializedTokenPair.getExpiration()).isEqualTo(expiration); + assertThat(deserializedTokenPair.getRefreshToken()).isEqualTo(Optional.of("refresh_token")); + } + + @Test + public void testTokenDeserializationAfterTimeoutAndExpirationExtension() + throws Exception + { + TestingClock clock = new TestingClock(); + + JweTokenSerializer serializer = tokenSerializer( + clock, + succinctDuration(12, MINUTES)); + Date expiration = new Calendar.Builder().setDate(2022, 6, 22).build().getTime(); + String serializedTokenPair = serializer.serialize(accessAndRefreshTokens("access_token", expiration, "refresh_token")); + + clock.advanceBy(succinctDuration(20, MINUTES)); + assertThatThrownBy(() -> serializer.deserialize(serializedTokenPair)) + .isExactlyInstanceOf(ExpiredJwtException.class); + } + + private JweTokenSerializer tokenSerializer(Clock clock, Duration tokenExpiration) + throws GeneralSecurityException, KeyLengthException + { + return new JweTokenSerializer( + new RefreshTokensConfig(), + new Oauth2ClientStub(), + "presto_coordinator_test_version", + "presto_coordinator", + "sub", + clock, + tokenExpiration); + } + + static class Oauth2ClientStub + implements OAuth2Client + { + private final Map claims = Jwts.claims() + .setSubject("user"); + + @Override + public void load() + { + } + + @Override + public Request createAuthorizationRequest(String state, URI callbackUri) + { + throw new UnsupportedOperationException("operation is not yet supported"); + } + + @Override + public Response getOAuth2Response(String code, URI callbackUri, Optional nonce) + { + throw new UnsupportedOperationException("operation is not yet supported"); + } + + @Override + public Optional> getClaims(String accessToken) + { + return Optional.of(claims); + } + + @Override + public Response refreshTokens(String refreshToken) + { + throw new UnsupportedOperationException("operation is not yet supported"); + } + } + + private static class TestingClock + extends Clock + { + private Instant currentTime = ZonedDateTime.of(2022, 5, 6, 10, 15, 0, 0, ZoneId.systemDefault()).toInstant(); + + @Override + public ZoneId getZone() + { + return ZoneId.systemDefault(); + } + + @Override + public Clock withZone(ZoneId zone) + { + return this; + } + + @Override + public Instant instant() + { + return currentTime; + } + + public void advanceBy(Duration currentTimeDelta) + { + this.currentTime = currentTime.plus(currentTimeDelta.toMillis(), MILLIS); + } + } +} diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2Config.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2Config.java new file mode 100644 index 0000000000000..e8dca428ab98b --- /dev/null +++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2Config.java @@ -0,0 +1,93 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import com.google.common.collect.ImmutableMap; +import io.airlift.units.Duration; +import org.testng.annotations.Test; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Map; + +import static com.facebook.airlift.configuration.testing.ConfigAssertions.assertFullMapping; +import static com.facebook.airlift.configuration.testing.ConfigAssertions.assertRecordedDefaults; +import static com.facebook.airlift.configuration.testing.ConfigAssertions.recordDefaults; +import static java.util.concurrent.TimeUnit.MINUTES; +import static java.util.concurrent.TimeUnit.SECONDS; + +public class TestOAuth2Config +{ + @Test + public void testDefaults() + { + assertRecordedDefaults(recordDefaults(OAuth2Config.class) + .setStateKey(null) + .setIssuer(null) + .setClientId(null) + .setClientSecret(null) + .setScopes("openid") + .setChallengeTimeout(new Duration(15, MINUTES)) + .setPrincipalField("sub") + .setGroupsField(null) + .setAdditionalAudiences("") + .setMaxClockSkew(new Duration(1, MINUTES)) + .setUserMappingPattern(null) + .setUserMappingFile(null) + .setEnableRefreshTokens(false) + .setEnableDiscovery(true)); + } + + @Test + public void testExplicitPropertyMappings() + throws IOException + { + Path userMappingFile = Files.createTempFile(null, null); + Map properties = ImmutableMap.builder() + .put("http-server.authentication.oauth2.state-key", "key-secret") + .put("http-server.authentication.oauth2.issuer", "http://127.0.0.1:9000/oauth2") + .put("http-server.authentication.oauth2.client-id", "another-consumer") + .put("http-server.authentication.oauth2.client-secret", "consumer-secret") + .put("http-server.authentication.oauth2.scopes", "email,offline") + .put("http-server.authentication.oauth2.principal-field", "some-field") + .put("http-server.authentication.oauth2.groups-field", "groups") + .put("http-server.authentication.oauth2.additional-audiences", "test-aud1,test-aud2") + .put("http-server.authentication.oauth2.challenge-timeout", "90s") + .put("http-server.authentication.oauth2.max-clock-skew", "15s") + .put("http-server.authentication.oauth2.user-mapping.pattern", "(.*)@something") + .put("http-server.authentication.oauth2.user-mapping.file", userMappingFile.toString()) + .put("http-server.authentication.oauth2.refresh-tokens", "true") + .put("http-server.authentication.oauth2.oidc.discovery", "false") + .build(); + + OAuth2Config expected = new OAuth2Config() + .setStateKey("key-secret") + .setIssuer("http://127.0.0.1:9000/oauth2") + .setClientId("another-consumer") + .setClientSecret("consumer-secret") + .setScopes("email, offline") + .setPrincipalField("some-field") + .setGroupsField("groups") + .setAdditionalAudiences("test-aud1,test-aud2") + .setChallengeTimeout(new Duration(90, SECONDS)) + .setMaxClockSkew(new Duration(15, SECONDS)) + .setUserMappingPattern("(.*)@something") + .setUserMappingFile(userMappingFile.toFile()) + .setEnableRefreshTokens(true) + .setEnableDiscovery(false); + + assertFullMapping(properties, expected); + } +} diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2Utils.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2Utils.java new file mode 100644 index 0000000000000..0d21db65f4ec7 --- /dev/null +++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2Utils.java @@ -0,0 +1,55 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import com.facebook.presto.server.MockHttpServletRequest; +import com.google.common.collect.ImmutableListMultimap; +import org.testng.annotations.Test; + +import javax.servlet.http.HttpServletRequest; +import javax.ws.rs.core.UriBuilder; + +import static com.facebook.presto.server.security.oauth2.OAuth2Utils.getSchemeUriBuilder; +import static com.google.common.net.HttpHeaders.X_FORWARDED_PROTO; +import static org.testng.Assert.assertEquals; + +public class TestOAuth2Utils +{ + @Test + public void testGetSchemeUriBuilderNoProtoHeader() + { + HttpServletRequest request = new MockHttpServletRequest( + ImmutableListMultimap.builder() + .build(), + "testRemote", + "http://www.example.com"); + + UriBuilder builder = getSchemeUriBuilder(request); + assertEquals(builder.build().getScheme(), "http"); + } + + @Test + public void testGetSchemeUriBuilderProtoHeader() + { + HttpServletRequest request = new MockHttpServletRequest( + ImmutableListMultimap.builder() + .put(X_FORWARDED_PROTO, "https") + .build(), + "testRemote", + "http://www.example.com"); + + UriBuilder builder = getSchemeUriBuilder(request); + assertEquals(builder.build().getScheme(), "https"); + } +} diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOidcDiscoveryConfig.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOidcDiscoveryConfig.java new file mode 100644 index 0000000000000..5144b3f934ea9 --- /dev/null +++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOidcDiscoveryConfig.java @@ -0,0 +1,67 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import com.google.common.collect.ImmutableMap; +import io.airlift.units.Duration; +import org.testng.annotations.Test; + +import java.util.Map; + +import static com.facebook.airlift.configuration.testing.ConfigAssertions.assertFullMapping; +import static com.facebook.airlift.configuration.testing.ConfigAssertions.assertRecordedDefaults; +import static com.facebook.airlift.configuration.testing.ConfigAssertions.recordDefaults; +import static java.util.concurrent.TimeUnit.MINUTES; +import static java.util.concurrent.TimeUnit.SECONDS; + +public class TestOidcDiscoveryConfig +{ + @Test + public void testDefaults() + { + assertRecordedDefaults(recordDefaults(OidcDiscoveryConfig.class) + .setDiscoveryTimeout(new Duration(30, SECONDS)) + .setUserinfoEndpointEnabled(true) + .setAccessTokenIssuer(null) + .setAuthUrl(null) + .setTokenUrl(null) + .setJwksUrl(null) + .setUserinfoUrl(null)); + } + + @Test + public void testExplicitPropertyMapping() + { + Map properties = ImmutableMap.builder() + .put("http-server.authentication.oauth2.oidc.discovery.timeout", "1m") + .put("http-server.authentication.oauth2.oidc.use-userinfo-endpoint", "false") + .put("http-server.authentication.oauth2.access-token-issuer", "https://issuer.com/at") + .put("http-server.authentication.oauth2.auth-url", "https://issuer.com/auth") + .put("http-server.authentication.oauth2.token-url", "https://issuer.com/token") + .put("http-server.authentication.oauth2.jwks-url", "https://issuer.com/jwks.json") + .put("http-server.authentication.oauth2.userinfo-url", "https://issuer.com/user") + .build(); + + OidcDiscoveryConfig expected = new OidcDiscoveryConfig() + .setDiscoveryTimeout(new Duration(1, MINUTES)) + .setUserinfoEndpointEnabled(false) + .setAccessTokenIssuer("https://issuer.com/at") + .setAuthUrl("https://issuer.com/auth") + .setTokenUrl("https://issuer.com/token") + .setJwksUrl("https://issuer.com/jwks.json") + .setUserinfoUrl("https://issuer.com/user"); + + assertFullMapping(properties, expected); + } +} diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestRefreshTokensConfig.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestRefreshTokensConfig.java new file mode 100644 index 0000000000000..d365a513573de --- /dev/null +++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestRefreshTokensConfig.java @@ -0,0 +1,74 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import com.google.common.collect.ImmutableMap; +import org.testng.annotations.Test; + +import javax.crypto.KeyGenerator; +import javax.crypto.SecretKey; + +import java.security.NoSuchAlgorithmException; +import java.util.Map; + +import static com.facebook.airlift.configuration.testing.ConfigAssertions.assertFullMapping; +import static com.facebook.airlift.configuration.testing.ConfigAssertions.assertRecordedDefaults; +import static com.facebook.airlift.configuration.testing.ConfigAssertions.recordDefaults; +import static io.airlift.units.Duration.succinctDuration; +import static io.jsonwebtoken.io.Encoders.BASE64; +import static java.util.concurrent.TimeUnit.HOURS; + +public class TestRefreshTokensConfig +{ + @Test + public void testDefaults() + { + assertRecordedDefaults(recordDefaults(RefreshTokensConfig.class) + .setTokenExpiration(succinctDuration(1, HOURS)) + .setIssuer("Presto_coordinator") + .setAudience("Presto_coordinator") + .setSecretKey(null)); + } + + @Test + public void testExplicitPropertyMappings() + throws Exception + { + String encodedBase64SecretKey = BASE64.encode(generateKey()); + + Map properties = ImmutableMap.builder() + .put("http-server.authentication.oauth2.refresh-tokens.issued-token.timeout", "24h") + .put("http-server.authentication.oauth2.refresh-tokens.issued-token.issuer", "issuer") + .put("http-server.authentication.oauth2.refresh-tokens.issued-token.audience", "audience") + .put("http-server.authentication.oauth2.refresh-tokens.secret-key", encodedBase64SecretKey) + .build(); + + RefreshTokensConfig expected = new RefreshTokensConfig() + .setTokenExpiration(succinctDuration(24, HOURS)) + .setIssuer("issuer") + .setAudience("audience") + .setSecretKey(encodedBase64SecretKey); + + assertFullMapping(properties, expected); + } + + private byte[] generateKey() + throws NoSuchAlgorithmException + { + KeyGenerator generator = KeyGenerator.getInstance("AES"); + generator.init(256); + SecretKey key = generator.generateKey(); + return key.getEncoded(); + } +} diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestingHydraIdentityProvider.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestingHydraIdentityProvider.java new file mode 100644 index 0000000000000..f0b77820c717c --- /dev/null +++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestingHydraIdentityProvider.java @@ -0,0 +1,352 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import com.facebook.airlift.http.server.HttpServerConfig; +import com.facebook.airlift.http.server.HttpServerInfo; +import com.facebook.airlift.http.server.testing.TestingHttpServer; +import com.facebook.airlift.log.Level; +import com.facebook.airlift.log.Logging; +import com.facebook.airlift.node.NodeInfo; +import com.facebook.presto.server.testing.TestingPrestoServer; +import com.facebook.presto.util.AutoCloseableCloser; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.io.Resources; +import com.google.inject.Key; +import okhttp3.HttpUrl; +import okhttp3.MediaType; +import okhttp3.OkHttpClient; +import okhttp3.Request; +import okhttp3.RequestBody; +import okhttp3.Response; +import org.testcontainers.containers.FixedHostPortGenericContainer; +import org.testcontainers.containers.GenericContainer; +import org.testcontainers.containers.Network; +import org.testcontainers.containers.PostgreSQLContainer; +import org.testcontainers.containers.startupcheck.OneShotStartupCheckStrategy; +import org.testcontainers.containers.wait.strategy.Wait; +import org.testcontainers.containers.wait.strategy.WaitAllStrategy; +import org.testcontainers.utility.MountableFile; + +import javax.servlet.http.HttpServlet; +import javax.servlet.http.HttpServletRequest; +import javax.servlet.http.HttpServletResponse; + +import java.io.Closeable; +import java.io.IOException; +import java.net.URI; +import java.time.Duration; +import java.util.List; +import java.util.Optional; + +import static com.facebook.presto.client.OkHttpUtil.setupInsecureSsl; +import static com.facebook.presto.server.security.oauth2.TokenEndpointAuthMethod.CLIENT_SECRET_BASIC; +import static com.google.common.base.Throwables.throwIfUnchecked; +import static java.util.Objects.requireNonNull; +import static javax.servlet.http.HttpServletResponse.SC_NOT_FOUND; +import static javax.ws.rs.core.MediaType.APPLICATION_JSON; + +public class TestingHydraIdentityProvider + implements Closeable +{ + private static final String HYDRA_IMAGE = "oryd/hydra:v1.10.6"; + private static final String ISSUER = "https://localhost:4444/"; + private static final String DSN = "postgres://hydra:mysecretpassword@database:5432/hydra?sslmode=disable"; + + private final Network network = Network.newNetwork(); + + private final PostgreSQLContainer databaseContainer = new PostgreSQLContainer<>() + .withNetwork(network) + .withNetworkAliases("database") + .withUsername("hydra") + .withPassword("mysecretpassword") + .withDatabaseName("hydra"); + + private final GenericContainer migrationContainer = createHydraContainer() + .withCommand("migrate", "sql", "--yes", DSN) + .withStartupCheckStrategy(new OneShotStartupCheckStrategy().withTimeout(Duration.ofMinutes(5))); + + private final AutoCloseableCloser closer = AutoCloseableCloser.create(); + private final ObjectMapper mapper = new ObjectMapper(); + private final Duration ttlAccessToken; + private final boolean useJwt; + private final boolean exposeFixedPorts; + private final OkHttpClient httpClient; + private FixedHostPortGenericContainer hydraContainer; + + public TestingHydraIdentityProvider(Duration ttlAccessToken, boolean useJwt, boolean exposeFixedPorts) + { + this.ttlAccessToken = requireNonNull(ttlAccessToken, "ttlAccessToken is null"); + this.useJwt = useJwt; + this.exposeFixedPorts = exposeFixedPorts; + OkHttpClient.Builder httpClientBuilder = new OkHttpClient.Builder(); + setupInsecureSsl(httpClientBuilder); + httpClientBuilder.followRedirects(false); + httpClient = httpClientBuilder.build(); + closer.register(network); + closer.register(databaseContainer); + closer.register(migrationContainer); + } + + public void start() + throws Exception + { + databaseContainer.start(); + migrationContainer.start(); + TestingHttpServer loginAndConsentServer = createTestingLoginAndConsentServer(); + closer.register(loginAndConsentServer::stop); + loginAndConsentServer.start(); + URI loginAndConsentBaseUrl = loginAndConsentServer.getBaseUrl(); + + hydraContainer = createHydraContainer() + .withNetworkAliases("hydra") + .withExposedPorts(4444, 4445) + .withEnv("DSN", DSN) + .withEnv("URLS_SELF_ISSUER", ISSUER) + .withEnv("URLS_CONSENT", loginAndConsentBaseUrl + "/consent") + .withEnv("URLS_LOGIN", loginAndConsentBaseUrl + "/login") + .withEnv("SERVE_TLS_KEY_PATH", "/tmp/certs/localhost.pem") + .withEnv("SERVE_TLS_CERT_PATH", "/tmp/certs/localhost.pem") + .withEnv("TTL_ACCESS_TOKEN", ttlAccessToken.getSeconds() + "s") + .withEnv("STRATEGIES_ACCESS_TOKEN", useJwt ? "jwt" : null) + .withEnv("LOG_LEAK_SENSITIVE_VALUES", "true") + .withCommand("serve", "all") + .withCopyFileToContainer(MountableFile.forClasspathResource("/cert"), "/tmp/certs") + .waitingFor(new WaitAllStrategy() + .withStrategy(Wait.forLogMessage(".*Setting up http server on :4444.*", 1)) + .withStrategy(Wait.forLogMessage(".*Setting up http server on :4445.*", 1))); + if (exposeFixedPorts) { + hydraContainer = hydraContainer + .withFixedExposedPort(4444, 4444) + .withFixedExposedPort(4445, 4445); + } + closer.register(hydraContainer); + hydraContainer.start(); + } + + public FixedHostPortGenericContainer createHydraContainer() + { + return new FixedHostPortGenericContainer<>(HYDRA_IMAGE).withNetwork(network); + } + + public void createClient( + String clientId, + String clientSecret, + TokenEndpointAuthMethod tokenEndpointAuthMethod, + List audiences, + String callbackUrl) + { + createHydraContainer() + .withCommand("clients", "create", + "--endpoint", "https://hydra:4445", + "--skip-tls-verify", + "--id", clientId, + "--secret", clientSecret, + "--audience", String.join(",", audiences), + "--grant-types", "authorization_code,refresh_token,client_credentials", + "--response-types", "token,code,id_token", + "--scope", "openid,offline", + "--token-endpoint-auth-method", tokenEndpointAuthMethod.getValue(), + "--callbacks", callbackUrl) + .withStartupCheckStrategy(new OneShotStartupCheckStrategy().withTimeout(Duration.ofSeconds(30))) + .start(); + } + + public int getAuthPort() + { + return hydraContainer.getMappedPort(4444); + } + + public int getAdminPort() + { + return hydraContainer.getMappedPort(4445); + } + + @Override + public void close() + throws IOException + { + try { + closer.close(); + } + catch (Exception e) { + throwIfUnchecked(e); + throw new RuntimeException(e); + } + } + + private TestingHttpServer createTestingLoginAndConsentServer() + throws IOException + { + NodeInfo nodeInfo = new NodeInfo("test"); + HttpServerConfig config = new HttpServerConfig().setHttpPort(0); + HttpServerInfo httpServerInfo = new HttpServerInfo(config, nodeInfo); + return new TestingHttpServer(httpServerInfo, nodeInfo, config, new AcceptAllLoginsAndConsentsServlet(), ImmutableMap.of(), ImmutableMap.of(), Optional.empty()); + } + + private class AcceptAllLoginsAndConsentsServlet + extends HttpServlet + { + private final ObjectMapper mapper = new ObjectMapper(); + private final OkHttpClient httpClient; + + public AcceptAllLoginsAndConsentsServlet() + { + OkHttpClient.Builder httpClientBuilder = new OkHttpClient.Builder(); + setupInsecureSsl(httpClientBuilder); + httpClient = httpClientBuilder.build(); + } + + @Override + protected void doGet(HttpServletRequest request, HttpServletResponse response) + throws IOException + { + if (request.getPathInfo().equals("/login")) { + acceptLogin(request, response); + return; + } + if (request.getPathInfo().contains("/consent")) { + acceptConsent(request, response); + return; + } + response.setStatus(SC_NOT_FOUND); + } + + private void acceptLogin(HttpServletRequest request, HttpServletResponse response) + throws IOException + { + String loginChallenge = request.getParameter("login_challenge"); + try (Response loginAcceptResponse = acceptLogin(loginChallenge)) { + sendRedirect(loginAcceptResponse, response); + } + } + + private void acceptConsent(HttpServletRequest request, HttpServletResponse response) + throws IOException + { + String consentChallenge = request.getParameter("consent_challenge"); + JsonNode consentRequest = getConsentRequest(consentChallenge); + try (Response acceptConsentResponse = acceptConsent(consentChallenge, consentRequest)) { + sendRedirect(acceptConsentResponse, response); + } + } + + private Response acceptLogin(String loginChallenge) + throws IOException + { + return httpClient.newCall( + new Request.Builder() + .url("https://localhost:" + getAdminPort() + "/oauth2/auth/requests/login/accept?login_challenge=" + loginChallenge) + .put(RequestBody.create( + MediaType.parse(APPLICATION_JSON), + mapper.writeValueAsString(mapper.createObjectNode().put("subject", "foo@bar.com")))) + .build()) + .execute(); + } + + private JsonNode getConsentRequest(String consentChallenge) + throws IOException + { + try (Response response = httpClient.newCall( + new Request.Builder() + .url("https://localhost:" + getAdminPort() + "/oauth2/auth/requests/consent?consent_challenge=" + consentChallenge) + .get() + .build()) + .execute()) { + requireNonNull(response.body()); + return mapper.readTree(response.body().byteStream()); + } + } + + private Response acceptConsent(String consentChallenge, JsonNode consentRequest) + throws IOException + { + return httpClient.newCall( + new Request.Builder() + .url("https://localhost:" + getAdminPort() + "/oauth2/auth/requests/consent/accept?consent_challenge=" + consentChallenge) + .put(RequestBody.create( + MediaType.parse(APPLICATION_JSON), + mapper.writeValueAsString(mapper.createObjectNode() + .set("grant_scope", consentRequest.get("requested_scope")) + .set("grant_access_token_audience", consentRequest.get("requested_access_token_audience"))))) + .build()) + .execute(); + } + + private void sendRedirect(Response redirectResponse, HttpServletResponse response) + throws IOException + { + requireNonNull(redirectResponse.body()); + response.sendRedirect( + toHostUrl(mapper.readTree(redirectResponse.body().byteStream()) + .get("redirect_to") + .textValue())); + } + + private String toHostUrl(String url) + { + return HttpUrl.get(URI.create(url)) + .newBuilder() + .port(getAuthPort()) + .toString(); + } + } + + private static void runTestServer(boolean useJwt) + throws Exception + { + try (TestingHydraIdentityProvider service = new TestingHydraIdentityProvider(Duration.ofMinutes(30), useJwt, true)) { + service.start(); + service.createClient( + "presto-client", + "presto-secret", + CLIENT_SECRET_BASIC, + ImmutableList.of("https://localhost:8443/ui"), + "https://localhost:8443/oauth2/callback"); + ImmutableMap.Builder config = ImmutableMap.builder() + .put("http-server.https.port", "8443") + .put("http-server.https.enabled", "true") + .put("http-server.https.keystore.path", Resources.getResource("cert/localhost.pem").getPath()) + .put("http-server.https.keystore.key", "") + .put("http-server.authentication.type", "OAUTH2") + .put("http-server.authentication.oauth2.issuer", ISSUER) + .put("http-server.authentication.oauth2.client-id", "presto-client") + .put("http-server.authentication.oauth2.client-secret", "presto-secret") + .put("http-server.authentication.oauth2.user-mapping.pattern", "(.*)@.*") + .put("http-server.authentication.oauth2.oidc.use-userinfo-endpoint", String.valueOf(!useJwt)) + .put("oauth2-jwk.http-client.trust-store-path", Resources.getResource("cert/localhost.pem").getPath()); + try (TestingPrestoServer server = new TestingPrestoServer(config.build())) { + server.getInstance(Key.get(OAuth2Client.class)).load(); + Thread.sleep(Long.MAX_VALUE); + } + } + } + + public static void main(String[] args) + throws Exception + { + Logging logging = Logging.initialize(); + try { + logging.setLevel(OAuth2Service.class.getName(), Level.DEBUG); + runTestServer(false); + } + finally { + logging.setLevel(OAuth2Service.class.getName(), Level.INFO); + } + } +} diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TokenEndpointAuthMethod.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TokenEndpointAuthMethod.java new file mode 100644 index 0000000000000..c05e636e9b11c --- /dev/null +++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TokenEndpointAuthMethod.java @@ -0,0 +1,33 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import static java.util.Objects.requireNonNull; + +public enum TokenEndpointAuthMethod +{ + CLIENT_SECRET_BASIC("client_secret_basic"); + + private final String value; + + TokenEndpointAuthMethod(String value) + { + this.value = requireNonNull(value, "value is null"); + } + + public String getValue() + { + return value; + } +} diff --git a/presto-native-execution/pom.xml b/presto-native-execution/pom.xml index f9150f63a4daf..3dab4aa3b4e23 100644 --- a/presto-native-execution/pom.xml +++ b/presto-native-execution/pom.xml @@ -100,7 +100,12 @@ com.facebook.presto presto-main - test + + + org.apache.commons + commons-lang3 + + diff --git a/presto-ui/src/static/ouath2/logout.html b/presto-ui/src/static/ouath2/logout.html new file mode 100644 index 0000000000000..821687918d81c --- /dev/null +++ b/presto-ui/src/static/ouath2/logout.html @@ -0,0 +1,51 @@ + + + + + + + + + Logout + + + + + + + + + + + + + + + + + + +

OAuth2 logout succeeded

+ +

This browser window can be closed

+ + + + + + From 398056cde974766aa4b64a9fca03ab8cab151ac6 Mon Sep 17 00:00:00 2001 From: Anant Aneja <1797669+aaneja@users.noreply.github.com> Date: Wed, 30 Apr 2025 12:57:51 +0530 Subject: [PATCH 107/113] Tests for the OAuth2 Authentication and OIDC discovery Co-authored-by: Stephen Yugel Co-authored-by: lukasz-walkiewicz --- presto-main/pom.xml | 5 + .../server/testing/TestingPrestoServer.java | 5 + .../BaseOAuth2AuthenticationFilterTest.java | 400 ++++++++++++++++++ .../security/oauth2/SimpleProxyServer.java | 197 +++++++++ ...TestDualAuthenticationFilterWithOAuth.java | 252 +++++++++++ ...TestOAuth2AuthenticationFilterWithJwt.java | 51 +++ ...tOAuth2AuthenticationFilterWithOpaque.java | 54 +++ .../security/oauth2/TestOidcDiscovery.java | 365 ++++++++++++++++ .../oauth2/TestingHydraIdentityProvider.java | 33 ++ .../src/test/resources/cert/generate.sh | 7 + .../src/test/resources/cert/localhost.conf | 20 + .../src/test/resources/cert/localhost.pem | 83 ++++ .../openid-configuration-invalid-issuer.json | 106 +++++ ...onfiguration-with-access-token-issuer.json | 87 ++++ ...openid-configuration-without-userinfo.json | 92 ++++ .../resources/oidc/openid-configuration.json | 106 +++++ 16 files changed, 1863 insertions(+) create mode 100644 presto-main/src/test/java/com/facebook/presto/server/security/oauth2/BaseOAuth2AuthenticationFilterTest.java create mode 100644 presto-main/src/test/java/com/facebook/presto/server/security/oauth2/SimpleProxyServer.java create mode 100644 presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestDualAuthenticationFilterWithOAuth.java create mode 100644 presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2AuthenticationFilterWithJwt.java create mode 100644 presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2AuthenticationFilterWithOpaque.java create mode 100644 presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOidcDiscovery.java create mode 100755 presto-main/src/test/resources/cert/generate.sh create mode 100644 presto-main/src/test/resources/cert/localhost.conf create mode 100644 presto-main/src/test/resources/cert/localhost.pem create mode 100644 presto-main/src/test/resources/oidc/openid-configuration-invalid-issuer.json create mode 100644 presto-main/src/test/resources/oidc/openid-configuration-with-access-token-issuer.json create mode 100644 presto-main/src/test/resources/oidc/openid-configuration-without-userinfo.json create mode 100644 presto-main/src/test/resources/oidc/openid-configuration.json diff --git a/presto-main/pom.xml b/presto-main/pom.xml index d2fb706da59eb..73ff20c4cb1d3 100644 --- a/presto-main/pom.xml +++ b/presto-main/pom.xml @@ -414,6 +414,11 @@ failsafe
+ + com.squareup.okhttp3 + okhttp-urlconnection + + com.facebook.presto diff --git a/presto-main/src/main/java/com/facebook/presto/server/testing/TestingPrestoServer.java b/presto-main/src/main/java/com/facebook/presto/server/testing/TestingPrestoServer.java index c46f3a1df23a7..980ba54f389b8 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/testing/TestingPrestoServer.java +++ b/presto-main/src/main/java/com/facebook/presto/server/testing/TestingPrestoServer.java @@ -624,6 +624,11 @@ public HostAndPort getHttpsAddress() return HostAndPort.fromParts(httpsUri.getHost(), httpsUri.getPort()); } + public URI getHttpBaseUrl() + { + return server.getHttpServerInfo().getHttpUri(); + } + public CatalogManager getCatalogManager() { return catalogManager; diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/BaseOAuth2AuthenticationFilterTest.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/BaseOAuth2AuthenticationFilterTest.java new file mode 100644 index 0000000000000..2f3b6a41053ab --- /dev/null +++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/BaseOAuth2AuthenticationFilterTest.java @@ -0,0 +1,400 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import com.facebook.airlift.json.JsonCodec; +import com.facebook.airlift.log.Level; +import com.facebook.airlift.log.Logging; +import com.facebook.airlift.testing.Closeables; +import com.facebook.presto.server.testing.TestingPrestoServer; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.inject.Key; +import io.jsonwebtoken.impl.DefaultClaims; +import okhttp3.Cookie; +import okhttp3.CookieJar; +import okhttp3.HttpUrl; +import okhttp3.JavaNetCookieJar; +import okhttp3.OkHttpClient; +import okhttp3.Request; +import okhttp3.Response; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + +import java.io.IOException; +import java.net.CookieManager; +import java.net.CookieStore; +import java.net.HttpCookie; +import java.net.MalformedURLException; +import java.net.URI; +import java.net.URL; +import java.security.KeyPairGenerator; +import java.security.NoSuchAlgorithmException; +import java.time.Duration; +import java.time.Instant; +import java.util.List; +import java.util.Optional; +import java.util.UUID; + +import static com.facebook.airlift.testing.Assertions.assertLessThan; +import static com.facebook.presto.client.OkHttpUtil.setupInsecureSsl; +import static com.facebook.presto.server.security.oauth2.JwtUtil.newJwtBuilder; +import static com.facebook.presto.server.security.oauth2.OAuthWebUiCookie.OAUTH2_COOKIE; +import static com.facebook.presto.server.security.oauth2.TokenEndpointAuthMethod.CLIENT_SECRET_BASIC; +import static io.airlift.units.Duration.nanosSince; +import static java.util.concurrent.TimeUnit.MILLISECONDS; +import static java.util.concurrent.TimeUnit.SECONDS; +import static javax.servlet.http.HttpServletResponse.SC_OK; +import static javax.ws.rs.core.HttpHeaders.AUTHORIZATION; +import static javax.ws.rs.core.HttpHeaders.LOCATION; +import static javax.ws.rs.core.Response.Status.OK; +import static javax.ws.rs.core.Response.Status.SEE_OTHER; +import static javax.ws.rs.core.Response.Status.UNAUTHORIZED; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.fail; +import static org.testng.Assert.assertEquals; + +public abstract class BaseOAuth2AuthenticationFilterTest +{ + protected static final Duration TTL_ACCESS_TOKEN_IN_SECONDS = Duration.ofSeconds(5); + + protected static final String PRESTO_CLIENT_ID = "presto-client"; + protected static final String PRESTO_CLIENT_SECRET = "presto-secret"; + private static final String PRESTO_AUDIENCE = PRESTO_CLIENT_ID; + private static final String ADDITIONAL_AUDIENCE = "https://external-service.com"; + protected static final String TRUSTED_CLIENT_ID = "trusted-client"; + protected static final String TRUSTED_CLIENT_SECRET = "trusted-secret"; + private static final String UNTRUSTED_CLIENT_ID = "untrusted-client"; + private static final String UNTRUSTED_CLIENT_SECRET = "untrusted-secret"; + private static final String UNTRUSTED_CLIENT_AUDIENCE = "https://untrusted.com"; + + private final Logging logging = Logging.initialize(); + protected final OkHttpClient httpClient; + protected TestingHydraIdentityProvider hydraIdP; + + private TestingPrestoServer server; + + private SimpleProxyServer simpleProxy; + private URI uiUri; + + private URI proxyURI; + + protected BaseOAuth2AuthenticationFilterTest() + { + OkHttpClient.Builder httpClientBuilder = new OkHttpClient.Builder(); + setupInsecureSsl(httpClientBuilder); + httpClientBuilder.followRedirects(false); + httpClient = httpClientBuilder.build(); + } + + static void waitForNodeRefresh(TestingPrestoServer server) + throws InterruptedException + { + long start = System.nanoTime(); + while (server.refreshNodes().getActiveNodes().size() < 1) { + assertLessThan(nanosSince(start), new io.airlift.units.Duration(10, SECONDS)); + MILLISECONDS.sleep(10); + } + } + + @BeforeClass + public void setup() + throws Exception + { + logging.setLevel(OAuth2Service.class.getName(), Level.DEBUG); + + hydraIdP = getHydraIdp(); + String idpUrl = "https://localhost:" + hydraIdP.getAuthPort(); + server = new TestingPrestoServer(getOAuth2Config(idpUrl)); + server.getInstance(Key.get(OAuth2Client.class)).load(); + waitForNodeRefresh(server); + // Due to problems with the Presto OSS project related to the AuthenticationFilter we have to run Presto behind a Proxy and terminate SSL at the proxy. + simpleProxy = new SimpleProxyServer(server.getHttpBaseUrl()); + MILLISECONDS.sleep(1000); + proxyURI = simpleProxy.getHttpsBaseUrl(); + uiUri = proxyURI.resolve("/"); + + hydraIdP.createClient( + PRESTO_CLIENT_ID, + PRESTO_CLIENT_SECRET, + CLIENT_SECRET_BASIC, + ImmutableList.of(PRESTO_AUDIENCE, ADDITIONAL_AUDIENCE), + simpleProxy.getHttpsBaseUrl() + "/oauth2/callback"); + hydraIdP.createClient( + TRUSTED_CLIENT_ID, + TRUSTED_CLIENT_SECRET, + CLIENT_SECRET_BASIC, + ImmutableList.of(TRUSTED_CLIENT_ID), + simpleProxy.getHttpsBaseUrl() + "/oauth2/callback"); + hydraIdP.createClient( + UNTRUSTED_CLIENT_ID, + UNTRUSTED_CLIENT_SECRET, + CLIENT_SECRET_BASIC, + ImmutableList.of(UNTRUSTED_CLIENT_AUDIENCE), + "https://untrusted.com/callback"); + } + + protected abstract ImmutableMap getOAuth2Config(String idpUrl); + + protected abstract TestingHydraIdentityProvider getHydraIdp() + throws Exception; + + @AfterClass(alwaysRun = true) + public void tearDown() + throws Exception + { + Closeables.closeAll(server, hydraIdP, simpleProxy); + } + + @Test + public void testUnauthorizedApiCall() + throws IOException + { + try (Response response = httpClient + .newCall(apiCall().build()) + .execute()) { + assertUnauthorizedResponse(response); + } + } + + @Test + public void testUnauthorizedUICall() + throws IOException + { + try (Response response = httpClient + .newCall(uiCall().build()) + .execute()) { + assertRedirectResponse(response); + } + } + + @Test + public void testUnsignedToken() + throws NoSuchAlgorithmException, IOException + { + KeyPairGenerator keyGenerator = KeyPairGenerator.getInstance("RSA"); + keyGenerator.initialize(4096); + long now = Instant.now().getEpochSecond(); + String token = newJwtBuilder() + .setHeaderParam("alg", "RS256") + .setHeaderParam("kid", "public:f467aa08-1c1b-4cde-ba45-84b0ef5d2ba8") + .setHeaderParam("typ", "JWT") + .setClaims( + new DefaultClaims( + ImmutableMap.builder() + .put("aud", ImmutableList.of()) + .put("client_id", PRESTO_CLIENT_ID) + .put("exp", now + 60L) + .put("iat", now) + .put("iss", "https://hydra:4444/") + .put("jti", UUID.randomUUID()) + .put("nbf", now) + .put("scp", ImmutableList.of("openid")) + .put("sub", "foo@bar.com") + .build())) + .signWith(keyGenerator.generateKeyPair().getPrivate()) + .compact(); + try (Response response = httpClientWithOAuth2Cookie(token, false) + .newCall(apiCall().build()) + .execute()) { + assertUnauthorizedResponse(response); + } + } + + @Test + public void testTokenWithInvalidAudience() + throws IOException + { + String token = hydraIdP.getToken(UNTRUSTED_CLIENT_ID, UNTRUSTED_CLIENT_SECRET, ImmutableList.of(UNTRUSTED_CLIENT_AUDIENCE)); + try (Response response = httpClientWithOAuth2Cookie(token, false) + .newCall(apiCall().build()) + .execute()) { + assertUnauthorizedResponse(response); + } + } + + @Test + public void testTokenFromTrustedClient() + throws IOException + { + String token = hydraIdP.getToken(TRUSTED_CLIENT_ID, TRUSTED_CLIENT_SECRET, ImmutableList.of(TRUSTED_CLIENT_ID)); + assertUICallWithCookie(token); + } + + @Test + public void testTokenWithMultipleAudiences() + throws IOException + { + String token = hydraIdP.getToken(PRESTO_CLIENT_ID, PRESTO_CLIENT_SECRET, ImmutableList.of(PRESTO_AUDIENCE, ADDITIONAL_AUDIENCE)); + assertUICallWithCookie(token); + } + + @Test + public void testSuccessfulFlow() + throws Exception + { + // create a new HttpClient which follows redirects and give access to cookies + CookieManager cookieManager = new CookieManager(); + CookieStore cookieStore = cookieManager.getCookieStore(); + OkHttpClient.Builder httpClientBuilder = new OkHttpClient.Builder(); + setupInsecureSsl(httpClientBuilder); + OkHttpClient httpClient = httpClientBuilder + .followRedirects(true) + .cookieJar(new JavaNetCookieJar(cookieManager)) + .build(); + + assertThat(cookieStore.get(uiUri)).isEmpty(); + + // access UI and follow redirects in order to get OAuth2 cookie + Response response = httpClient.newCall( + new Request.Builder() + .url(uiUri.toURL()) + .get() + .build()) + .execute(); + + assertEquals(response.code(), SC_OK); + + Optional oauth2Cookie = cookieStore.get(uiUri) + .stream() + .filter(cookie -> cookie.getName().equals(OAUTH2_COOKIE)) + .findFirst(); + assertThat(oauth2Cookie).isNotEmpty(); + assertOAuth2Cookie(oauth2Cookie.get()); + assertUICallWithCookie(oauth2Cookie.get().getValue()); + } + + @Test + public void testExpiredAccessToken() + throws Exception + { + String token = hydraIdP.getToken(PRESTO_CLIENT_ID, PRESTO_CLIENT_SECRET, ImmutableList.of(PRESTO_AUDIENCE)); + assertUICallWithCookie(token); + Thread.sleep(TTL_ACCESS_TOKEN_IN_SECONDS.plusSeconds(1).toMillis()); // wait for the token expiration = ttl of access token + 1 sec + try (Response response = httpClientWithOAuth2Cookie(token, false).newCall(apiCall().build()).execute()) { + assertUnauthorizedResponse(response); + } + } + + private Request.Builder uiCall() + { + return new Request.Builder() + .url(proxyURI.resolve("/").toString()) + .get(); + } + + private Request.Builder apiCall() + { + return new Request.Builder() + .url(proxyURI.resolve("/v1/cluster").toString()) + .get(); + } + + private void assertOAuth2Cookie(HttpCookie cookie) + { + assertThat(cookie.getName()).isEqualTo(OAUTH2_COOKIE); + assertThat(cookie.getDomain()).isIn(proxyURI.getHost()); + assertThat(cookie.getPath()).isEqualTo("/"); + assertThat(cookie.getSecure()).isTrue(); + assertThat(cookie.isHttpOnly()).isTrue(); + assertThat(cookie.getMaxAge()).isLessThanOrEqualTo(TTL_ACCESS_TOKEN_IN_SECONDS.getSeconds()); + validateAccessToken(cookie.getValue()); + } + protected void validateAccessToken(String cookieValue) + { + Request request = new Request.Builder().url("https://localhost:" + hydraIdP.getAuthPort() + "/userinfo").addHeader(AUTHORIZATION, "Bearer " + cookieValue).build(); + try (Response response = httpClient.newCall(request).execute()) { + assertThat(response.body()).isNotNull(); + DefaultClaims claims = new DefaultClaims(JsonCodec.mapJsonCodec(String.class, Object.class).fromJson(response.body().bytes())); + assertThat(claims.getSubject()).isEqualTo("foo@bar.com"); + } + catch (IOException e) { + fail("Exception while calling /userinfo", e); + } + } + + private void assertUICallWithCookie(String cookieValue) + throws IOException + { + OkHttpClient httpClient = httpClientWithOAuth2Cookie(cookieValue, true); + // pass access token in Presto UI cookie + try (Response response = httpClient.newCall(uiCall().build()) + .execute()) { + assertThat(response.code()).isEqualTo(OK.getStatusCode()); + } + } + + @SuppressWarnings("NullableProblems") + private OkHttpClient httpClientWithOAuth2Cookie(String cookieValue, boolean followRedirects) + { + OkHttpClient.Builder httpClientBuilder = new OkHttpClient.Builder(); + setupInsecureSsl(httpClientBuilder); + httpClientBuilder.followRedirects(followRedirects); + httpClientBuilder.cookieJar(new CookieJar() + { + @Override + public void saveFromResponse(HttpUrl url, List cookies) + { + } + + @Override + public List loadForRequest(HttpUrl url) + { + return ImmutableList.of(new Cookie.Builder() + .domain(proxyURI.getHost()) + .path("/") + .name(OAUTH2_COOKIE) + .value(cookieValue) + .secure() + .build()); + } + }); + return httpClientBuilder.build(); + } + + private void assertRedirectResponse(Response response) + throws MalformedURLException + { + assertThat(response.code()).isEqualTo(SEE_OTHER.getStatusCode()); + assertRedirectUrl(response.header(LOCATION)); + } + + private void assertUnauthorizedResponse(Response response) + throws IOException + { + assertThat(response.code()).isEqualTo(UNAUTHORIZED.getStatusCode()); + assertThat(response.body()).isNotNull(); + // NOTE that our errors come in looking like an HTML page since we don't do anything special on the server side so it just is like that. + assertThat(response.body().string()).contains("Invalid Credentials"); + } + + private void assertRedirectUrl(String redirectUrl) + throws MalformedURLException + { + assertThat(redirectUrl).isNotNull(); + URL location = new URL(redirectUrl); + HttpUrl url = HttpUrl.parse(redirectUrl); + assertThat(url).isNotNull(); + assertThat(location.getProtocol()).isEqualTo("https"); + assertThat(location.getHost()).isEqualTo("localhost"); + assertThat(location.getPort()).isEqualTo(hydraIdP.getAuthPort()); + assertThat(location.getPath()).isEqualTo("/oauth2/auth"); + assertThat(url.queryParameterValues("response_type")).isEqualTo(ImmutableList.of("code")); + assertThat(url.queryParameterValues("scope")).isEqualTo(ImmutableList.of("openid")); + assertThat(url.queryParameterValues("redirect_uri")).isEqualTo(ImmutableList.of(proxyURI + "/oauth2/callback")); + assertThat(url.queryParameterValues("client_id")).isEqualTo(ImmutableList.of(PRESTO_CLIENT_ID)); + assertThat(url.queryParameterValues("state")).isNotNull(); + } +} diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/SimpleProxyServer.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/SimpleProxyServer.java new file mode 100644 index 0000000000000..36e9921930fca --- /dev/null +++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/SimpleProxyServer.java @@ -0,0 +1,197 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import com.facebook.airlift.http.server.HttpServerConfig; +import com.facebook.airlift.http.server.HttpServerInfo; +import com.facebook.airlift.http.server.testing.TestingHttpServer; +import com.facebook.airlift.log.Logger; +import com.facebook.airlift.node.NodeInfo; +import com.google.common.collect.ImmutableMap; +import com.google.common.io.Resources; +import okhttp3.Headers; +import okhttp3.OkHttpClient; +import okhttp3.Request; +import okhttp3.Response; + +import javax.servlet.ServletException; +import javax.servlet.http.Cookie; +import javax.servlet.http.HttpServlet; +import javax.servlet.http.HttpServletRequest; +import javax.servlet.http.HttpServletResponse; +import javax.ws.rs.core.UriBuilder; + +import java.io.Closeable; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.net.URI; +import java.util.List; +import java.util.Optional; + +import static com.facebook.presto.client.OkHttpUtil.setupInsecureSsl; +import static com.facebook.presto.server.security.oauth2.OAuth2Utils.getFullRequestURL; +import static com.google.common.base.Throwables.throwIfUnchecked; +import static com.google.common.net.HttpHeaders.HOST; +import static com.google.common.net.HttpHeaders.X_FORWARDED_FOR; +import static com.google.common.net.HttpHeaders.X_FORWARDED_PROTO; +import static java.lang.String.format; +import static java.util.concurrent.TimeUnit.SECONDS; + +public class SimpleProxyServer + implements Closeable +{ + private final TestingHttpServer server; + + public SimpleProxyServer(URI forwardBaseURI) + throws Exception + { + server = createSimpleProxyServer(forwardBaseURI); + server.start(); + } + + @Override + public void close() + throws IOException + { + try { + server.stop(); + } + catch (Exception e) { + throwIfUnchecked(e); + throw new RuntimeException(e); + } + } + + public URI getHttpsBaseUrl() + { + return server.getHttpServerInfo().getHttpsUri(); + } + + private TestingHttpServer createSimpleProxyServer(URI forwardBaseURI) + throws IOException + { + NodeInfo nodeInfo = new NodeInfo("test"); + HttpServerConfig config = new HttpServerConfig() + .setHttpPort(0) + .setHttpsEnabled(true) + .setHttpsPort(0) + .setKeystorePath(Resources.getResource("cert/localhost.pem").getPath()); + HttpServerInfo httpServerInfo = new HttpServerInfo(config, nodeInfo); + return new TestingHttpServer(httpServerInfo, nodeInfo, config, new SimpleProxy(forwardBaseURI), ImmutableMap.of(), ImmutableMap.of(), Optional.empty()); + } + + private class SimpleProxy + extends HttpServlet + { + private final OkHttpClient httpClient; + private final URI forwardBaseURI; + + private final Logger logger = Logger.get(SimpleProxy.class); + + public SimpleProxy(URI forwardBaseURI) + { + this.forwardBaseURI = forwardBaseURI; + OkHttpClient.Builder httpClientBuilder = new OkHttpClient.Builder(); + setupInsecureSsl(httpClientBuilder); + httpClient = httpClientBuilder + .followRedirects(false) + .connectTimeout(100, SECONDS) + .writeTimeout(100, SECONDS) + .readTimeout(100, SECONDS) + .build(); + } + + @Override + protected void service(HttpServletRequest request, HttpServletResponse servletResponse) + throws ServletException, IOException + { + UriBuilder requestUriBuilder = UriBuilder.fromUri(getFullRequestURL(request)); + requestUriBuilder + .scheme("http") + .host(forwardBaseURI.getHost()) + .port(forwardBaseURI.getPort()); + + String hostHeader = new StringBuilder().append(request.getRemoteHost()).append(":").append(request.getLocalPort()).toString(); + Cookie[] cookies = Optional.ofNullable(request.getCookies()).orElse(new Cookie[0]); + String requestUri = requestUriBuilder.build().toString(); + Request.Builder reqBuilder = new Request.Builder() + .url(requestUri) + .addHeader(X_FORWARDED_PROTO, "https") + .addHeader(X_FORWARDED_FOR, request.getRemoteAddr()) + .addHeader(HOST, hostHeader) + .get(); + + if (cookies.length > 0) { + for (Cookie cookie : cookies) { + reqBuilder.addHeader("Cookie", cookie.getName() + "=" + cookie.getValue()); + } + } + Response response; + try { + response = httpClient.newCall(reqBuilder.build()).execute(); + servletResponse.setStatus(response.code()); + + Headers responseHeaders = response.headers(); + responseHeaders.names().stream().forEach(headerName -> { + // Headers can have multiple values + List headerValues = responseHeaders.values(headerName); + headerValues.forEach(headerValue -> { + servletResponse.addHeader(headerName, headerValue); + }); + }); + + //copy the response body to the servlet response. + InputStream is = response.body().byteStream(); + OutputStream os = servletResponse.getOutputStream(); + byte[] buffer = new byte[10 * 1024]; + int read; + while ((read = is.read(buffer)) != -1) { + os.write(buffer, 0, read); + } + } + catch (Exception e) { + logger.error(format("Encountered an error while proxying request to %s", requestUri), e); + servletResponse.sendError(HttpServletResponse.SC_INTERNAL_SERVER_ERROR); + } + } + } + + // This is just to help iterate changes that might need to be made in the future to the Simple Proxy Server for test purposes. + // Waiting for the tests to run can be really slow so having this helper here is nice if you want quicker feedback. + private static void runTestServer() + throws Exception + { + SimpleProxyServer test = new SimpleProxyServer(new URI("")); + OkHttpClient.Builder httpClientBuilder = new OkHttpClient.Builder(); + setupInsecureSsl(httpClientBuilder); + OkHttpClient client = httpClientBuilder.build(); + Request.Builder req = new Request.Builder().url(test.getHttpsBaseUrl().resolve("/v1/query").toString()).get(); + Logger logger = Logger.get("Run Test Server Debug Helper"); + try { + client.newCall(req.build()).execute(); + } + catch (Exception e) { + logger.error(e); + } + + test.close(); + } + + public static void main(String[] args) + throws Exception + { + runTestServer(); + } +} diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestDualAuthenticationFilterWithOAuth.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestDualAuthenticationFilterWithOAuth.java new file mode 100644 index 0000000000000..465a930dc30ae --- /dev/null +++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestDualAuthenticationFilterWithOAuth.java @@ -0,0 +1,252 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.facebook.presto.server.security.oauth2; + +import com.facebook.airlift.log.Level; +import com.facebook.airlift.log.Logging; +import com.facebook.presto.server.testing.TestingPrestoServer; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.io.Resources; +import com.google.inject.Key; +import okhttp3.Cookie; +import okhttp3.CookieJar; +import okhttp3.HttpUrl; +import okhttp3.OkHttpClient; +import okhttp3.Request; +import okhttp3.Response; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + +import java.io.IOException; +import java.net.URI; +import java.time.Duration; +import java.util.Base64; +import java.util.List; + +import static com.facebook.airlift.testing.Assertions.assertLessThan; +import static com.facebook.presto.client.OkHttpUtil.setupInsecureSsl; +import static com.facebook.presto.server.security.oauth2.OAuthWebUiCookie.OAUTH2_COOKIE; +import static com.facebook.presto.server.security.oauth2.TokenEndpointAuthMethod.CLIENT_SECRET_BASIC; +import static com.google.common.net.HttpHeaders.AUTHORIZATION; +import static com.google.common.net.HttpHeaders.WWW_AUTHENTICATE; +import static io.airlift.units.Duration.nanosSince; +import static java.util.concurrent.TimeUnit.MILLISECONDS; +import static java.util.concurrent.TimeUnit.SECONDS; +import static javax.ws.rs.core.Response.Status.OK; +import static javax.ws.rs.core.Response.Status.UNAUTHORIZED; +import static org.assertj.core.api.Assertions.assertThat; + +public class TestDualAuthenticationFilterWithOAuth +{ + protected static final Duration TTL_ACCESS_TOKEN_IN_SECONDS = Duration.ofSeconds(5); + protected static final String PRESTO_CLIENT_ID = "presto-client"; + protected static final String PRESTO_CLIENT_SECRET = "presto-secret"; + private static final String PRESTO_AUDIENCE = PRESTO_CLIENT_ID; + private static final String ADDITIONAL_AUDIENCE = "https://external-service.com"; + protected static final String TRUSTED_CLIENT_ID = "trusted-client"; + protected static final String TRUSTED_CLIENT_SECRET = "trusted-secret"; + private static final String UNTRUSTED_CLIENT_ID = "untrusted-client"; + private static final String UNTRUSTED_CLIENT_SECRET = "untrusted-secret"; + private static final String UNTRUSTED_CLIENT_AUDIENCE = "https://untrusted.com"; + + private final Logging logging = Logging.initialize(); + protected final OkHttpClient httpClient; + protected TestingHydraIdentityProvider hydraIdP; + + private TestingPrestoServer server; + + private SimpleProxyServer simpleProxy; + private URI uiUri; + + private URI proxyURI; + + protected TestDualAuthenticationFilterWithOAuth() + { + OkHttpClient.Builder httpClientBuilder = new OkHttpClient.Builder(); + setupInsecureSsl(httpClientBuilder); + httpClientBuilder.followRedirects(false); + httpClient = httpClientBuilder.build(); + } + + static void waitForNodeRefresh(TestingPrestoServer server) + throws InterruptedException + { + long start = System.nanoTime(); + while (server.refreshNodes().getActiveNodes().size() < 1) { + assertLessThan(nanosSince(start), new io.airlift.units.Duration(10, SECONDS)); + MILLISECONDS.sleep(10); + } + } + + protected ImmutableMap getConfig(String idpUrl) + { + return ImmutableMap.builder() + .put("http-server.authentication.allow-forwarded-https", "true") + .put("http-server.authentication.type", "OAUTH2,PASSWORD") + .put("http-server.authentication.oauth2.issuer", "https://localhost:4444/") + .put("http-server.authentication.oauth2.auth-url", idpUrl + "/oauth2/auth") + .put("http-server.authentication.oauth2.token-url", idpUrl + "/oauth2/token") + .put("http-server.authentication.oauth2.jwks-url", idpUrl + "/.well-known/jwks.json") + .put("http-server.authentication.oauth2.client-id", PRESTO_CLIENT_ID) + .put("http-server.authentication.oauth2.client-secret", PRESTO_CLIENT_SECRET) + .put("http-server.authentication.oauth2.additional-audiences", TRUSTED_CLIENT_ID) + .put("http-server.authentication.oauth2.max-clock-skew", "0s") + .put("http-server.authentication.oauth2.user-mapping.pattern", "(.*)(@.*)?") + .put("http-server.authentication.oauth2.oidc.discovery", "false") + .put("oauth2-jwk.http-client.trust-store-path", Resources.getResource("cert/localhost.pem").getPath()) + .build(); + } + + @BeforeClass + public void setup() + throws Exception + { + logging.setLevel(OAuth2Service.class.getName(), Level.DEBUG); + hydraIdP = new TestingHydraIdentityProvider(TTL_ACCESS_TOKEN_IN_SECONDS, true, false); + hydraIdP.start(); + String idpUrl = "https://localhost:" + hydraIdP.getAuthPort(); + server = new TestingPrestoServer(getConfig(idpUrl)); + server.getInstance(Key.get(OAuth2Client.class)).load(); + waitForNodeRefresh(server); + // Due to problems with the Presto OSS project related to the AuthenticationFilter we have to run Presto behind a Proxy and terminate SSL at the proxy. + simpleProxy = new SimpleProxyServer(server.getHttpBaseUrl()); + MILLISECONDS.sleep(1000); + proxyURI = simpleProxy.getHttpsBaseUrl(); + uiUri = proxyURI.resolve("/"); + + hydraIdP.createClient( + PRESTO_CLIENT_ID, + PRESTO_CLIENT_SECRET, + CLIENT_SECRET_BASIC, + ImmutableList.of(PRESTO_AUDIENCE, ADDITIONAL_AUDIENCE), + simpleProxy.getHttpsBaseUrl() + "/oauth2/callback"); + } + + @Test + public void testExpiredOAuthToken() + throws Exception + { + String token = hydraIdP.getToken(PRESTO_CLIENT_ID, PRESTO_CLIENT_SECRET, ImmutableList.of(PRESTO_AUDIENCE)); + assertUICallWithCookie(token); + Thread.sleep(TTL_ACCESS_TOKEN_IN_SECONDS.plusSeconds(1).toMillis()); // wait for the token expiration = ttl of access token + 1 sec + try (Response response = httpClientWithOAuth2Cookie(token, false).newCall(apiCall().build()).execute()) { + assertUnauthorizedOAuthOnlyHeaders(response); + } + } + + @Test + public void testNoAuth() + throws Exception + { + try (Response response = httpClient + .newCall(apiCall().build()) + .execute()) { + assertAllUnauthorizedHeaders(response); + } + } + + @Test + public void testInvalidBasicAuth() + throws Exception + { + String userPass = "test:password"; + String basicAuth = "Basic " + Base64.getEncoder().encodeToString(userPass.getBytes()); + try (Response response = httpClient + .newCall(apiCall().addHeader(AUTHORIZATION, basicAuth).build()) + .execute()) { + assertAllUnauthorizedHeaders(response); + } + } + + private Request.Builder apiCall() + { + return new Request.Builder() + .url(proxyURI.resolve("/v1/cluster").toString()) + .get(); + } + + private void assertUnauthorizedOAuthOnlyHeaders(Response response) + throws IOException + { + String redirectServer = "x_redirect_server=\"" + proxyURI.resolve("/oauth2/token/initiate/"); + String tokenServer = "x_token_server=\"" + proxyURI.resolve("/oauth2/token/"); + assertUnauthorizedResponse(response); + List headers = response.headers(WWW_AUTHENTICATE); + assertThat(headers.size()).isEqualTo(1); + assertThat(headers.get(0)).contains(tokenServer, redirectServer); + } + + private void assertAllUnauthorizedHeaders(Response response) + throws IOException + { + String redirectServer = "x_redirect_server=\"" + proxyURI.resolve("/oauth2/token/initiate/").toString(); + String tokenServer = "x_token_server=\"" + proxyURI.resolve("/oauth2/token/"); + assertUnauthorizedResponse(response); + List headers = response.headers(WWW_AUTHENTICATE); + assertThat(headers.size()).isEqualTo(2); + assertThat(headers.stream().allMatch(h -> + h.contains("Basic realm=\"Presto\"") || + (h.contains(redirectServer) && h.contains(tokenServer)) + )).isTrue(); + } + + private void assertUICallWithCookie(String cookieValue) + throws IOException + { + OkHttpClient httpClient = httpClientWithOAuth2Cookie(cookieValue, true); + // pass access token in Presto UI cookie + try (Response response = httpClient.newCall(apiCall().build()) + .execute()) { + assertThat(response.code()).isEqualTo(OK.getStatusCode()); + } + } + + private OkHttpClient httpClientWithOAuth2Cookie(String cookieValue, boolean followRedirects) + { + OkHttpClient.Builder httpClientBuilder = new OkHttpClient.Builder(); + setupInsecureSsl(httpClientBuilder); + httpClientBuilder.followRedirects(followRedirects); + httpClientBuilder.cookieJar(new CookieJar() + { + @Override + public void saveFromResponse(HttpUrl url, List cookies) + { + } + + @Override + public List loadForRequest(HttpUrl url) + { + return ImmutableList.of(new Cookie.Builder() + .domain(proxyURI.getHost()) + .path("/") + .name(OAUTH2_COOKIE) + .value(cookieValue) + .secure() + .build()); + } + }); + return httpClientBuilder.build(); + } + + private void assertUnauthorizedResponse(Response response) + throws IOException + { + assertThat(response.code()).isEqualTo(UNAUTHORIZED.getStatusCode()); + assertThat(response.body()).isNotNull(); + // NOTE that our errors come in looking like an HTML page since we don't do anything special on the server side so it just is like that. + assertThat(response.body().string()).contains("Invalid Credentials"); + } +} diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2AuthenticationFilterWithJwt.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2AuthenticationFilterWithJwt.java new file mode 100644 index 0000000000000..a4eb2f6069038 --- /dev/null +++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2AuthenticationFilterWithJwt.java @@ -0,0 +1,51 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import com.google.common.collect.ImmutableMap; +import com.google.common.io.Resources; + +public class TestOAuth2AuthenticationFilterWithJwt + extends BaseOAuth2AuthenticationFilterTest +{ + @Override + protected ImmutableMap getOAuth2Config(String idpUrl) + { + return ImmutableMap.builder() + .put("http-server.authentication.allow-forwarded-https", "true") + .put("http-server.authentication.type", "OAUTH2") + .put("http-server.authentication.oauth2.issuer", "https://localhost:4444/") + .put("http-server.authentication.oauth2.auth-url", idpUrl + "/oauth2/auth") + .put("http-server.authentication.oauth2.token-url", idpUrl + "/oauth2/token") + .put("http-server.authentication.oauth2.jwks-url", idpUrl + "/.well-known/jwks.json") + .put("http-server.authentication.oauth2.client-id", PRESTO_CLIENT_ID) + .put("http-server.authentication.oauth2.client-secret", PRESTO_CLIENT_SECRET) + .put("http-server.authentication.oauth2.additional-audiences", TRUSTED_CLIENT_ID) + .put("http-server.authentication.oauth2.max-clock-skew", "0s") + .put("http-server.authentication.oauth2.user-mapping.pattern", "(.*)(@.*)?") + .put("http-server.authentication.oauth2.oidc.discovery", "false") + .put("oauth2-jwk.http-client.trust-store-path", Resources.getResource("cert/localhost.pem").getPath()) + .build(); + } + + @Override + protected TestingHydraIdentityProvider getHydraIdp() + throws Exception + { + TestingHydraIdentityProvider hydraIdP = new TestingHydraIdentityProvider(TTL_ACCESS_TOKEN_IN_SECONDS, true, false); + hydraIdP.start(); + + return hydraIdP; + } +} diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2AuthenticationFilterWithOpaque.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2AuthenticationFilterWithOpaque.java new file mode 100644 index 0000000000000..b36453cccfef5 --- /dev/null +++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2AuthenticationFilterWithOpaque.java @@ -0,0 +1,54 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import com.google.common.collect.ImmutableMap; +import com.google.common.io.Resources; + +public class TestOAuth2AuthenticationFilterWithOpaque + extends BaseOAuth2AuthenticationFilterTest +{ + @Override + protected ImmutableMap getOAuth2Config(String idpUrl) + { + return ImmutableMap.builder() + .put("http-server.authentication.allow-forwarded-https", "true") + .put("http-server.authentication.type", "OAUTH2") + .put("http-server.authentication.oauth2.issuer", "https://localhost:4444/") + .put("http-server.authentication.oauth2.auth-url", idpUrl + "/oauth2/auth") + .put("http-server.authentication.oauth2.token-url", idpUrl + "/oauth2/token") + .put("http-server.authentication.oauth2.jwks-url", idpUrl + "/.well-known/jwks.json") + .put("http-server.authentication.oauth2.userinfo-url", idpUrl + "/userinfo") + .put("http-server.authentication.oauth2.client-id", PRESTO_CLIENT_ID) + .put("http-server.authentication.oauth2.client-secret", PRESTO_CLIENT_SECRET) + // This is necessary as Hydra does not return `sub` from `/userinfo` for client credential grants. + .put("http-server.authentication.oauth2.principal-field", "iss") + .put("http-server.authentication.oauth2.additional-audiences", TRUSTED_CLIENT_ID) + .put("http-server.authentication.oauth2.max-clock-skew", "0s") + .put("http-server.authentication.oauth2.user-mapping.pattern", "(.*)(@.*)?") + .put("http-server.authentication.oauth2.oidc.discovery", "false") + .put("oauth2-jwk.http-client.trust-store-path", Resources.getResource("cert/localhost.pem").getPath()) + .build(); + } + + @Override + protected TestingHydraIdentityProvider getHydraIdp() + throws Exception + { + TestingHydraIdentityProvider hydraIdP = new TestingHydraIdentityProvider(TTL_ACCESS_TOKEN_IN_SECONDS, false, false); + hydraIdP.start(); + + return hydraIdP; + } +} diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOidcDiscovery.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOidcDiscovery.java new file mode 100644 index 0000000000000..d6e65f9d73c4a --- /dev/null +++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOidcDiscovery.java @@ -0,0 +1,365 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.server.security.oauth2; + +import com.facebook.airlift.http.server.Authenticator; +import com.facebook.airlift.http.server.HttpServerConfig; +import com.facebook.airlift.http.server.HttpServerInfo; +import com.facebook.airlift.http.server.testing.TestingHttpServer; +import com.facebook.airlift.node.NodeInfo; +import com.facebook.presto.server.security.oauth2.OAuth2ServerConfigProvider.OAuth2ServerConfig; +import com.facebook.presto.server.testing.TestingPrestoServer; +import com.google.common.collect.ImmutableMap; +import com.google.common.io.Resources; +import com.google.inject.Key; +import com.google.inject.TypeLiteral; +import org.testng.annotations.DataProvider; +import org.testng.annotations.Test; + +import javax.servlet.http.HttpServlet; +import javax.servlet.http.HttpServletRequest; +import javax.servlet.http.HttpServletResponse; + +import java.io.IOException; +import java.net.URI; +import java.time.Instant; +import java.util.List; +import java.util.Map; +import java.util.Optional; + +import static com.facebook.airlift.http.client.HttpStatus.TOO_MANY_REQUESTS; +import static com.facebook.presto.server.security.oauth2.BaseOAuth2AuthenticationFilterTest.PRESTO_CLIENT_ID; +import static com.facebook.presto.server.security.oauth2.BaseOAuth2AuthenticationFilterTest.PRESTO_CLIENT_SECRET; +import static java.nio.charset.StandardCharsets.UTF_8; +import static java.util.Objects.requireNonNull; +import static javax.ws.rs.core.HttpHeaders.CONTENT_TYPE; +import static javax.ws.rs.core.MediaType.APPLICATION_JSON; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +public class TestOidcDiscovery +{ + @Test(dataProvider = "staticConfiguration") + public void testStaticConfiguration(Optional accessTokenPath, Optional userinfoPath) + throws Exception + { + try (MetadataServer metadataServer = new MetadataServer(ImmutableMap.of("/jwks.json", "jwk/jwk-public.json"))) { + URI issuer = metadataServer.getBaseUrl(); + Optional accessTokenIssuer = accessTokenPath.map(issuer::resolve); + Optional userinfoUrl = userinfoPath.map(issuer::resolve); + ImmutableMap.Builder properties = ImmutableMap.builder() + .put("http-server.authentication.oauth2.issuer", metadataServer.getBaseUrl().toString()) + .put("http-server.authentication.oauth2.oidc.discovery", "false") + .put("http-server.authentication.oauth2.auth-url", issuer.resolve("/connect/authorize").toString()) + .put("http-server.authentication.oauth2.token-url", issuer.resolve("/connect/token").toString()) + .put("http-server.authentication.oauth2.jwks-url", issuer.resolve("/jwks.json").toString()); + accessTokenIssuer.map(URI::toString).ifPresent(uri -> properties.put("http-server.authentication.oauth2.access-token-issuer", uri)); + userinfoUrl.map(URI::toString).ifPresent(uri -> properties.put("http-server.authentication.oauth2.userinfo-url", uri)); + try (TestingPrestoServer server = createServer(properties.build())) { + assertConfiguration(server, issuer, accessTokenIssuer.map(issuer::resolve), userinfoUrl.map(issuer::resolve)); + } + } + } + + @DataProvider(name = "staticConfiguration") + public static Object[][] staticConfiguration() + { + return new Object[][] { + {Optional.empty(), Optional.empty()}, + {Optional.of("/access-token-issuer"), Optional.of("/userinfo")}, + }; + } + + @Test(dataProvider = "oidcDiscovery") + public void testOidcDiscovery(String configuration, Optional accessTokenIssuer, Optional userinfoUrl) + throws Exception + { + try (MetadataServer metadataServer = new MetadataServer( + ImmutableMap.builder() + .put("/.well-known/openid-configuration", "oidc/" + configuration) + .put("/jwks.json", "jwk/jwk-public.json") + .build()); + TestingPrestoServer server = createServer( + ImmutableMap.builder() + .put("http-server.authentication.oauth2.issuer", metadataServer.getBaseUrl().toString()) + .put("http-server.authentication.oauth2.oidc.discovery", "true") + .build())) { + URI issuer = metadataServer.getBaseUrl(); + assertConfiguration(server, issuer, accessTokenIssuer.map(issuer::resolve), userinfoUrl.map(issuer::resolve)); + } + } + + @DataProvider(name = "oidcDiscovery") + public static Object[][] oidcDiscovery() + { + return new Object[][] { + {"openid-configuration.json", Optional.empty(), Optional.of("/connect/userinfo")}, + {"openid-configuration-without-userinfo.json", Optional.empty(), Optional.empty()}, + {"openid-configuration-with-access-token-issuer.json", Optional.of("http://access-token-issuer.com/adfs/services/trust"), Optional.of("/connect/userinfo")}, + }; + } + + @Test + public void testIssuerCheck() + { + assertThatThrownBy(() -> { + try (MetadataServer metadataServer = new MetadataServer( + ImmutableMap.builder() + .put("/.well-known/openid-configuration", "oidc/openid-configuration-invalid-issuer.json") + .put("/jwks.json", "jwk/jwk-public.json") + .build()); + TestingPrestoServer server = createServer( + ImmutableMap.builder() + .put("http-server.authentication.oauth2.issuer", metadataServer.getBaseUrl().toString()) + .put("http-server.authentication.oauth2.oidc.discovery", "true") + .build())) { + // should throw an exception + server.getInstance(Key.get(OAuth2ServerConfigProvider.class)).get(); + } + }).hasMessageContaining( + "Invalid response from OpenID Metadata endpoint. " + + "The value of the \"issuer\" claim in Metadata document different than the Issuer URL used for the Configuration Request."); + } + + @Test + public void testStopOnClientError() + { + assertThatThrownBy(() -> { + try (MetadataServer metadataServer = new MetadataServer(ImmutableMap.of()); + TestingPrestoServer server = createServer( + ImmutableMap.builder() + .put("http-server.authentication.oauth2.issuer", metadataServer.getBaseUrl().toString()) + .put("http-server.authentication.oauth2.oidc.discovery", "true") + .build())) { + // should throw an exception + server.getInstance(Key.get(OAuth2ServerConfigProvider.class)).get(); + } + }).hasMessageContaining("Invalid response from OpenID Metadata endpoint. Expected response code to be 200, but was 404"); + } + + @Test + public void testOidcDiscoveryRetrying() + throws Exception + { + try (MetadataServer metadataServer = new MetadataServer(new MetadataServletWithStartup( + ImmutableMap.builder() + .put("/.well-known/openid-configuration", "oidc/openid-configuration.json") + .put("/jwks.json", "jwk/jwk-public.json") + .build(), 5)); + TestingPrestoServer server = createServer( + ImmutableMap.builder() + .put("http-server.authentication.oauth2.issuer", metadataServer.getBaseUrl().toString()) + .put("http-server.authentication.oauth2.oidc.discovery", "true") + .put("http-server.authentication.oauth2.oidc.discovery.timeout", "10s") + .build())) { + URI issuer = metadataServer.getBaseUrl(); + assertConfiguration(server, issuer, Optional.empty(), Optional.of(issuer.resolve("/connect/userinfo"))); + } + } + + @Test + public void testOidcDiscoveryTimesOut() + { + assertThatThrownBy(() -> { + try (MetadataServer metadataServer = new MetadataServer(new MetadataServletWithStartup( + ImmutableMap.builder() + .put("/.well-known/openid-configuration", "oidc/openid-configuration.json") + .put("/jwks.json", "jwk/jwk-public.json") + .build(), 10)); + TestingPrestoServer server = createServer( + ImmutableMap.builder() + .put("http-server.authentication.oauth2.issuer", metadataServer.getBaseUrl().toString()) + .put("http-server.authentication.oauth2.oidc.discovery", "true") + .put("http-server.authentication.oauth2.oidc.discovery.timeout", "5s") + .build())) { + // should throw an exception + server.getInstance(Key.get(OAuth2ServerConfigProvider.class)).get(); + } + }).hasMessageContaining("Invalid response from OpenID Metadata endpoint: 429"); + } + + @Test + public void testIgnoringUserinfoUrl() + throws Exception + { + try (MetadataServer metadataServer = new MetadataServer( + ImmutableMap.builder() + .put("/.well-known/openid-configuration", "oidc/openid-configuration.json") + .put("/jwks.json", "jwk/jwk-public.json") + .build()); + TestingPrestoServer server = createServer( + ImmutableMap.builder() + .put("http-server.authentication.oauth2.issuer", metadataServer.getBaseUrl().toString()) + .put("http-server.authentication.oauth2.oidc.discovery", "true") + .put("http-server.authentication.oauth2.oidc.use-userinfo-endpoint", "false") + .build())) { + URI issuer = metadataServer.getBaseUrl(); + assertConfiguration(server, issuer, Optional.empty(), Optional.empty()); + } + } + + @Test + public void testBackwardCompatibility() + throws Exception + { + try (MetadataServer metadataServer = new MetadataServer( + ImmutableMap.builder() + .put("/.well-known/openid-configuration", "oidc/openid-configuration-with-access-token-issuer.json") + .put("/jwks.json", "jwk/jwk-public.json") + .build())) { + URI issuer = metadataServer.getBaseUrl(); + URI authUrl = issuer.resolve("/custom-authorize"); + URI tokenUrl = issuer.resolve("/custom-token"); + URI jwksUrl = issuer.resolve("/custom-jwks.json"); + String accessTokenIssuer = issuer.resolve("/custom-access-token-issuer").toString(); + URI userinfoUrl = issuer.resolve("/custom-userinfo-url"); + try (TestingPrestoServer server = createServer( + ImmutableMap.builder() + .put("http-server.authentication.oauth2.issuer", issuer.toString()) + .put("http-server.authentication.oauth2.oidc.discovery", "true") + .put("http-server.authentication.oauth2.auth-url", authUrl.toString()) + .put("http-server.authentication.oauth2.token-url", tokenUrl.toString()) + .put("http-server.authentication.oauth2.jwks-url", jwksUrl.toString()) + .put("http-server.authentication.oauth2.access-token-issuer", accessTokenIssuer) + .put("http-server.authentication.oauth2.userinfo-url", userinfoUrl.toString()) + .build())) { + assertComponents(server); + OAuth2ServerConfig config = server.getInstance(Key.get(OAuth2ServerConfigProvider.class)).get(); + assertThat(config.getAccessTokenIssuer()).isEqualTo(Optional.of(accessTokenIssuer)); + assertThat(config.getAuthUrl()).isEqualTo(authUrl); + assertThat(config.getTokenUrl()).isEqualTo(tokenUrl); + assertThat(config.getJwksUrl()).isEqualTo(jwksUrl); + assertThat(config.getUserinfoUrl()).isEqualTo(Optional.of(userinfoUrl)); + } + } + } + + private static void assertConfiguration(TestingPrestoServer server, URI issuer, Optional accessTokenIssuer, Optional userinfoUrl) + { + assertComponents(server); + OAuth2ServerConfig config = server.getInstance(Key.get(OAuth2ServerConfigProvider.class)).get(); + assertThat(config.getAccessTokenIssuer()).isEqualTo(accessTokenIssuer.map(URI::toString)); + assertThat(config.getAuthUrl()).isEqualTo(issuer.resolve("/connect/authorize")); + assertThat(config.getTokenUrl()).isEqualTo(issuer.resolve("/connect/token")); + assertThat(config.getJwksUrl()).isEqualTo(issuer.resolve("/jwks.json")); + assertThat(config.getUserinfoUrl()).isEqualTo(userinfoUrl); + } + + private static void assertComponents(TestingPrestoServer server) + { + List authenticators = server.getInstance(Key.get(new TypeLiteral>() {})); + assertThat(authenticators).hasSize(1); + assertThat(authenticators.get(0)).isInstanceOf(OAuth2Authenticator.class); +// assertThat(server.getInstance(Key.get(WebUiAuthenticationFilter.class))).isInstanceOf(OAuth2WebUiAuthenticationFilter.class); + // does not throw an exception + server.getInstance(Key.get(OAuth2Client.class)).load(); + } + + private static TestingPrestoServer createServer(Map configuration) + throws Exception + { + ImmutableMap config = ImmutableMap.builder() + .put("http-server.authentication.allow-forwarded-https", "true") + .put("http-server.authentication.type", "OAUTH2") + .put("http-server.authentication.oauth2.client-id", PRESTO_CLIENT_ID) + .put("http-server.authentication.oauth2.client-secret", PRESTO_CLIENT_SECRET) + .putAll(configuration) + .build(); + + return new TestingPrestoServer(config); + } + + public static class MetadataServer + implements AutoCloseable + { + private final TestingHttpServer httpServer; + + public MetadataServer(Map responseMapping) + throws Exception + { + this(new MetadataServlet(responseMapping)); + } + + public MetadataServer(HttpServlet servlet) + throws Exception + { + NodeInfo nodeInfo = new NodeInfo("test"); + HttpServerConfig config = new HttpServerConfig().setHttpPort(0); + HttpServerInfo httpServerInfo = new HttpServerInfo(config, nodeInfo); + httpServer = new TestingHttpServer(httpServerInfo, nodeInfo, config, servlet, ImmutableMap.of(), ImmutableMap.of(), Optional.empty()); + httpServer.start(); + } + + public URI getBaseUrl() + { + return httpServer.getBaseUrl(); + } + + @Override + public void close() + throws Exception + { + httpServer.stop(); + } + } + + public static class MetadataServlet + extends HttpServlet + { + private final Map responseMapping; + + public MetadataServlet(Map responseMapping) + { + this.responseMapping = requireNonNull(responseMapping, "responseMapping is null"); + } + + @Override + protected void doGet(HttpServletRequest request, HttpServletResponse response) + throws IOException + { + String fileName = responseMapping.get(request.getPathInfo()); + if (fileName == null) { + response.setStatus(404); + return; + } + response.setHeader(CONTENT_TYPE, APPLICATION_JSON); + String body = Resources.toString(Resources.getResource(fileName), UTF_8); + body = body.replaceAll("https://issuer.com", request.getRequestURL().toString().replace("/.well-known/openid-configuration", "")); + response.getWriter().write(body); + } + } + + public static class MetadataServletWithStartup + extends MetadataServlet + { + private final Instant startTime; + + public MetadataServletWithStartup(Map responseMapping, int startupInSeconds) + { + super(responseMapping); + startTime = Instant.now().plusSeconds(startupInSeconds); + } + + @Override + protected void doGet(HttpServletRequest request, HttpServletResponse response) + throws IOException + { + if (Instant.now().isBefore(startTime)) { + response.setStatus(TOO_MANY_REQUESTS.code()); + return; + } + super.doGet(request, response); + } + } +} diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestingHydraIdentityProvider.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestingHydraIdentityProvider.java index f0b77820c717c..ee2cb16482900 100644 --- a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestingHydraIdentityProvider.java +++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestingHydraIdentityProvider.java @@ -28,6 +28,9 @@ import com.google.common.collect.ImmutableMap; import com.google.common.io.Resources; import com.google.inject.Key; +import com.nimbusds.oauth2.sdk.GrantType; +import okhttp3.Credentials; +import okhttp3.FormBody; import okhttp3.HttpUrl; import okhttp3.MediaType; import okhttp3.OkHttpClient; @@ -46,6 +49,7 @@ import javax.servlet.http.HttpServlet; import javax.servlet.http.HttpServletRequest; import javax.servlet.http.HttpServletResponse; +import javax.ws.rs.core.HttpHeaders; import java.io.Closeable; import java.io.IOException; @@ -56,9 +60,11 @@ import static com.facebook.presto.client.OkHttpUtil.setupInsecureSsl; import static com.facebook.presto.server.security.oauth2.TokenEndpointAuthMethod.CLIENT_SECRET_BASIC; +import static com.google.common.base.Preconditions.checkState; import static com.google.common.base.Throwables.throwIfUnchecked; import static java.util.Objects.requireNonNull; import static javax.servlet.http.HttpServletResponse.SC_NOT_FOUND; +import static javax.servlet.http.HttpServletResponse.SC_OK; import static javax.ws.rs.core.MediaType.APPLICATION_JSON; public class TestingHydraIdentityProvider @@ -89,6 +95,11 @@ public class TestingHydraIdentityProvider private final OkHttpClient httpClient; private FixedHostPortGenericContainer hydraContainer; + public TestingHydraIdentityProvider() + { + this(Duration.ofMinutes(30), true, false); + } + public TestingHydraIdentityProvider(Duration ttlAccessToken, boolean useJwt, boolean exposeFixedPorts) { this.ttlAccessToken = requireNonNull(ttlAccessToken, "ttlAccessToken is null"); @@ -167,6 +178,28 @@ public void createClient( .start(); } + public String getToken(String clientId, String clientSecret, List audiences) + throws IOException + { + try (Response response = httpClient + .newCall( + new Request.Builder() + .url("https://localhost:" + getAuthPort() + "/oauth2/token") + .addHeader(HttpHeaders.AUTHORIZATION, Credentials.basic(clientId, clientSecret)) + .post(new FormBody.Builder() + .add("grant_type", GrantType.CLIENT_CREDENTIALS.getValue()) + .add("audience", String.join(" ", audiences)) + .build()) + .build()) + .execute()) { + checkState(response.code() == SC_OK); + requireNonNull(response.body()); + return mapper.readTree(response.body().byteStream()) + .get("access_token") + .textValue(); + } + } + public int getAuthPort() { return hydraContainer.getMappedPort(4444); diff --git a/presto-main/src/test/resources/cert/generate.sh b/presto-main/src/test/resources/cert/generate.sh new file mode 100755 index 0000000000000..a84749045a5e4 --- /dev/null +++ b/presto-main/src/test/resources/cert/generate.sh @@ -0,0 +1,7 @@ +#!/bin/sh + +set -eux + +openssl req -new -x509 -newkey rsa:4096 -sha256 -nodes -keyout localhost.key -days 3560 -out localhost.crt -config localhost.conf +cat localhost.crt localhost.key > localhost.pem + diff --git a/presto-main/src/test/resources/cert/localhost.conf b/presto-main/src/test/resources/cert/localhost.conf new file mode 100644 index 0000000000000..560e1a454c8d9 --- /dev/null +++ b/presto-main/src/test/resources/cert/localhost.conf @@ -0,0 +1,20 @@ +[req] +default_bits = 4096 +prompt = no +default_md = sha256 +x509_extensions = v3_req +distinguished_name = dn + +[dn] +C = US +ST = California +L = Palo Alto +O = PrestoTest +CN = PrestoTest + +[v3_req] +subjectAltName = @alt_names + +[alt_names] +IP.1 = 127.0.0.1 +DNS.1 = localhost diff --git a/presto-main/src/test/resources/cert/localhost.pem b/presto-main/src/test/resources/cert/localhost.pem new file mode 100644 index 0000000000000..2d8a8fb289876 --- /dev/null +++ b/presto-main/src/test/resources/cert/localhost.pem @@ -0,0 +1,83 @@ +-----BEGIN CERTIFICATE----- +MIIFYTCCA0mgAwIBAgIJAKUofzuCtcfnMA0GCSqGSIb3DQEBCwUAMGAxCzAJBgNV +BAYTAlVTMRMwEQYDVQQIDApDYWxpZm9ybmlhMRIwEAYDVQQHDAlQYWxvIEFsdG8x +EzARBgNVBAoMClByZXN0b1Rlc3QxEzARBgNVBAMMClByZXN0b1Rlc3QwHhcNMjIw +OTA2MTgyOTUxWhcNMzIwNjA1MTgyOTUxWjBgMQswCQYDVQQGEwJVUzETMBEGA1UE +CAwKQ2FsaWZvcm5pYTESMBAGA1UEBwwJUGFsbyBBbHRvMRMwEQYDVQQKDApQcmVz +dG9UZXN0MRMwEQYDVQQDDApQcmVzdG9UZXN0MIICIjANBgkqhkiG9w0BAQEFAAOC +Ag8AMIICCgKCAgEA28P/OFPTMWu5AUt2YF3IvfFtZ2FRioB02+FIE21KtYM1r8w2 +2GRyvLuT8vBaYoh8bNSGI2x2R1NtfOalaUCUfr9XRyZPcP4FEE5x0QRK2SYYOfzr +URx5gv3SSlhahmSjsFAojpG7lUBsKpopFcjZb0wSq3hFDVHQN57Xzmt1YHbTZrEt +5yFyqt2AYRVHz8XxJbsUOy514/YGQfLLZqukSLYk055qFIclzFqXU+/cg6UVpl7U +hlLTo0GApBQ2eLGCBDZqXhkCf2U1lMGGVLsJFNmGaumLV88yZmRYQC9MJWlCPCGG +ZcyKNxjZq70SbmjIA6s0FVcXYZ0z6xQqDpVBichLebrtR8ShKU29u1ITL2kaF1iK +gZi8FzEKwnzxLlTZACeBfGeELl9HKUmUvOwU9LHp0UX4fZLlcj5IYXk0IMExFEYa +qfKYxThmdo1Gmpl1orW3mnx+BK8/VtMn7RquHTgQof/dry3EAMo+7/N9v3tI4m+5 +99g+DoHyOdpWYOsTtGmkYcWLG8/ka2lzcaLx347VBRBgNa/afKquPK3ogi04wYYl +K8wBciyhg+J7MkC0k1Q5ek25qynotHLkIGC+LsmEOuw8kApc8/ZFvA5yavn4Dsm5 +x94CmVhJQqL5Dqr6IEmzxvWjFWCC4mD2os52Ff1RlEGuEB5B9CAeaWCSuGECAwEA +AaMeMBwwGgYDVR0RBBMwEYcEfwAAAYIJbG9jYWxob3N0MA0GCSqGSIb3DQEBCwUA +A4ICAQCnGqgmdRRIuEij5bAHHF+CkbtkHbOO0fTxmd5/Cd5vkiCUR0uhwBg9JUbC +5zl/wBTx8guc0b7Wjvr3gwBUWQhGR2LKQMfatFZSy8nPcfSkZUGjY4s7KYI8mjPY +1Lri/EE1gu2p+iB9Dw4EnHW3QSneyy4yrkpLcjywbkF93SsThgHQ27gEK3/HeGzY +dFQY17z6zRNnlkMU2JVh0VptE6xtdR3WAMhscVx4dpEjFz9FSKnAYE2svTAy0OTD +8+W795a4/eVCxdHw/3PqR4XSO8isdSimeSTtdpRqsDrW4jx4IIGVFMQPin5XABAl +Wbbs8VMZPB1OvLSpmPtV79o/EklWS9x0MtbXF6iT5VBIbP1JoQLHQ9O0+V4PT0CH +8f8+Wc2UqLskqJZOyDzV3Y81mcdmKYcWNN8LJBk6PA1Pt/RDAg58QW1KfUVleey8 +eaAMw7d9BoIM/nUe0Q4TPll0WJnFMWHLUItgs7YuN/39kOnxGQi3iKQCMV2qwRip +tHqTvw3fHXQxEbZZhLxInC04+pOF3ZpqdzkeaZsXpwklV+uDw1u+5rzlheVrErwA +BzAXENpJKZLs4mjHko+Z4loNzGsJzRci/YRkwSKW6stW04RWWEKQkc9YhFsDYDmt +I70j8+R7et3nx2DqaQ18WgQ/5xVsEoXg5sgFjj1cRVFDQdrexA== +-----END CERTIFICATE----- +-----BEGIN PRIVATE KEY----- +MIIJQgIBADANBgkqhkiG9w0BAQEFAASCCSwwggkoAgEAAoICAQDbw/84U9Mxa7kB +S3ZgXci98W1nYVGKgHTb4UgTbUq1gzWvzDbYZHK8u5Py8FpiiHxs1IYjbHZHU218 +5qVpQJR+v1dHJk9w/gUQTnHRBErZJhg5/OtRHHmC/dJKWFqGZKOwUCiOkbuVQGwq +mikVyNlvTBKreEUNUdA3ntfOa3VgdtNmsS3nIXKq3YBhFUfPxfEluxQ7LnXj9gZB +8stmq6RItiTTnmoUhyXMWpdT79yDpRWmXtSGUtOjQYCkFDZ4sYIENmpeGQJ/ZTWU +wYZUuwkU2YZq6YtXzzJmZFhAL0wlaUI8IYZlzIo3GNmrvRJuaMgDqzQVVxdhnTPr +FCoOlUGJyEt5uu1HxKEpTb27UhMvaRoXWIqBmLwXMQrCfPEuVNkAJ4F8Z4QuX0cp +SZS87BT0senRRfh9kuVyPkhheTQgwTEURhqp8pjFOGZ2jUaamXWitbeafH4Erz9W +0yftGq4dOBCh/92vLcQAyj7v832/e0jib7n32D4OgfI52lZg6xO0aaRhxYsbz+Rr +aXNxovHfjtUFEGA1r9p8qq48reiCLTjBhiUrzAFyLKGD4nsyQLSTVDl6TbmrKei0 +cuQgYL4uyYQ67DyQClzz9kW8DnJq+fgOybnH3gKZWElCovkOqvogSbPG9aMVYILi +YPaiznYV/VGUQa4QHkH0IB5pYJK4YQIDAQABAoICAFSVcz2yxa5Xz7T33m/oqMKy +kXEgu8ma919JrfwMLJ0AC0HGT7Wps5+AcskmSSNzdLBOe/JWZI+/RHy2KSQBfyXp +byYrUJgkrL5B8vyHsmcxilGHTurBEuOf3bhPmUfwpC/QKkv1O0WOrhMXkoiX7Vgw +516nw6wEuScvM9B2+45NLcBwoUI8VW3+ItM65ZDKlq32+ypsD2PV5UKsuCykE28I +69OnPRz5h0rH80aTI0Rn3ZVTGmk4p8xGAcUlInIBoBEPAJGG/rcZtS2z7ofeFPi5 +YEr16HO7g6J1LKJHkf80LBIItTmpJ+lc3yqCcv2bxp/i3QD5rD4dy0XHVQiX2cj/ +KQSh64MAszRtXZE+GOIAT6UOxyEZqYrPNopXZWEROpM4W/c20HsA1ffpXu0abOcY +rkZhWmKk5mJNBAe+EE6aoxGJ9Th6x7WO7WSjiAmXIxjoAk0YClOwMf6y8+5fuqg8 +aEtG4GOz2bfUksbE3AzaXMcu2o4tt+4+paZA0pkhB9Za1ySP5N4pCQ1kvKhdstZi +e9laJtBS3YSS+rm2pe3g/rA9x5OrTDpms8S3LvFflG1KdPzJeBX7zRg6AHFBhAPI +OpFvsdeVDHOmGe/MYjC0gpLkoVBJLY7cqiicqUUWYaFm0JMoma8VS2h+Y4+cV39N +erXXLflO+Zn1EhiWrWYBAoIBAQD8K3evijZepPil3wtLff/dHgWzOsIubkBWEALu +M+Ab9zm0ahZ1muEej3p0U+Eq8TUbOAqaOWWndFMiqTcjco3zluSIh7wYajY827uG +5Fuyw3zIUm1g5HicZyVsqBJxRcSxkNqzFMMnvhFqUNyrsKHL+/GsCnWiuZhz8z8h +j9l616DF6LpefZrDL1nne4xfliZuE99BWNBi7n1n48q/yxpZ+PyNjAZFUu2wcOFE ++/z5wuuNwywk0gsWEPiqt1LAIbLwHYaTfIwK/OO19ScfFjzw8ZRNLcNEhFbFyZfR +aVkL22uqU7vtBzA1ld2AtogDTEdMQ1ruFd/ZIoOiHK6OpUhJAoIBAQDfGoeTi51A +Fh6gDh4/eiXKpggUvxBbLOkyfcb3j9VCbPmm+09YYR+b99FD7jhM2gKY1QXscvm5 +LFwMiXBNOv110PL1H0lGaNnv5sRFgQwAd7BpBNSAxYQCjF7RZz1GBHq3iLaAYtY4 +jRVPGL+n55E84jzi2Ip856Nf4TaTf98rsslSPbr6zwFvsGALFOki7hJivcelS1Qm +3U6UuJ3CmWv6GA3BY6BCj+5pXzR8+5bisYuIFUm1uW6sSn+1aaogWtOTRIFqf3mT +WSewoOcJpXHs5kBHkRVMycTIUaijbMdoANQzsZ8d6dfcyxpfrncxEaf+VJWKxEVp +OVSTUbpn/d9ZAoIBAEBF0/up/rGg2r9sWjSjqNNzE4DbOSMcdsl4y0ZrcnOuT8bs +Q002bKqdZ1i/CGUplZ+aaRlmB8Lmo0nyV1txlzy++QDTl92hNLHOT73R9o1ZxjRI +zhgkI5m5sJBBRnIYlkmr4hJC+HrotweiFJyuKI8VaEOxZspTA8iJ901WnNfyncfT +yazL1uZo60FU/DJg0uq7pevB91s/7jbMmKDJ462LCNQLHI4O1QZjvwcWMyR1yhQX +6uh3oNu+96KLl0vhSvpojCSLWiZyzpdSJOaHhIDlEieZwmt0T6mZ+FgnwcqD4q1H +Kl7/tgnyaMKlw4UTrBiEEmkcqjFt2p83MEarWgECggEBAKzFnrhkJiK6/nx0cng1 +345PhXKLg98XqH+xZ6PPfxcxzSPC+m82x4PBJg21LWRWcCxqy2uQnlMIR0BuLsmg +JShX585rrBMan6toyhYJGYJDLhol42rViqVujv8bNBhE38PB25MQ91RT7WyTfdhJ +O/AqQ3xotNaFi790aQ9Qt0Lf8Yf+xg30wOf9bmMmjmS+eP5+eV1IOKLgPzpsvb81 +kKjcd8qLnE/vpnFziPJA41gqpiN8WNiiAVLrXnremSD1NWOWaaJPlZbGNDZUZJbT +yKXsqVrCv/v3RKzcj/v/AW1JNwvRQaeor8IMhyARu7wEMFSErEoKNLaH7zcm03Q0 +5gECggEAcaeUlvCqcQiU/i42f9cHu2wx7kDKil6JS8ZYsJ0dGl5g9u2+VyeUBJvI +jLTrnwaGyn+Gx+PlGaC0QK5i9WjbblxrPuWFJk63pGDzLW35jrM55QC9kiI4C1Q7 +TyWOfjGBGa+Gi+v5bo+B3osgAHx7F2SQaoZ2C6Qbo+CMvDCihYFsPcOVngNzf+7+ +QkPOOO5ixlopHyekKe6dcMvb1PzovSEK3DZoh0XH4M5hz3LNk3Z57seNyN9pxcLI +U4qSYvcgYwbYZWjXymU6LVsx5lKJ4RrJGS56K0Q3ebC9nz7pNoVGKxvFAo4OI179 +pTI3VZB66Rqurj0RWq452LJE1Onb5A== +-----END PRIVATE KEY----- diff --git a/presto-main/src/test/resources/oidc/openid-configuration-invalid-issuer.json b/presto-main/src/test/resources/oidc/openid-configuration-invalid-issuer.json new file mode 100644 index 0000000000000..2f23e8ac98154 --- /dev/null +++ b/presto-main/src/test/resources/oidc/openid-configuration-invalid-issuer.json @@ -0,0 +1,106 @@ +{ + "issuer": "https://invalid-issuer.com", + "authorization_endpoint": "https://invalid-issuer.com/connect/authorize", + "token_endpoint": "https://invalid-issuer.com/connect/token", + "token_endpoint_auth_methods_supported": [ + "client_secret_basic", + "private_key_jwt" + ], + "token_endpoint_auth_signing_alg_values_supported": [ + "RS256", + "ES256" + ], + "userinfo_endpoint": "https://invalid-issuer.com/connect/userinfo", + "check_session_iframe": "https://invalid-issuer.com/connect/check_session", + "end_session_endpoint": "https://invalid-issuer.com/connect/end_session", + "jwks_uri": "https://invalid-issuer.com/jwks.json", + "registration_endpoint": "https://invalid-issuer.com/connect/register", + "scopes_supported": [ + "openid", + "profile", + "email", + "address", + "phone", + "offline_access" + ], + "response_types_supported": [ + "code", + "code id_token", + "id_token", + "token id_token" + ], + "acr_values_supported": [ + "urn:mace:incommon:iap:silver", + "urn:mace:incommon:iap:bronze" + ], + "subject_types_supported": [ + "public", + "pairwise" + ], + "userinfo_signing_alg_values_supported": [ + "RS256", + "ES256", + "HS256" + ], + "userinfo_encryption_alg_values_supported": [ + "RSA1_5", + "A128KW" + ], + "userinfo_encryption_enc_values_supported": [ + "A128CBC-HS256", + "A128GCM" + ], + "id_token_signing_alg_values_supported": [ + "RS256", + "ES256", + "HS256" + ], + "id_token_encryption_alg_values_supported": [ + "RSA1_5", + "A128KW" + ], + "id_token_encryption_enc_values_supported": [ + "A128CBC-HS256", + "A128GCM" + ], + "request_object_signing_alg_values_supported": [ + "none", + "RS256", + "ES256" + ], + "display_values_supported": [ + "page", + "popup" + ], + "claim_types_supported": [ + "normal", + "distributed" + ], + "claims_supported": [ + "sub", + "iss", + "auth_time", + "acr", + "name", + "given_name", + "family_name", + "nickname", + "profile", + "picture", + "website", + "email", + "email_verified", + "locale", + "zoneinfo", + "http://example.info/claims/groups" + ], + "claims_parameter_supported": true, + "service_documentation": "http://invalid-issuer.com/connect/service_documentation.html", + "ui_locales_supported": [ + "en-US", + "en-GB", + "en-CA", + "fr-FR", + "fr-CA" + ] +} diff --git a/presto-main/src/test/resources/oidc/openid-configuration-with-access-token-issuer.json b/presto-main/src/test/resources/oidc/openid-configuration-with-access-token-issuer.json new file mode 100644 index 0000000000000..905032ce3df1d --- /dev/null +++ b/presto-main/src/test/resources/oidc/openid-configuration-with-access-token-issuer.json @@ -0,0 +1,87 @@ +{ + "issuer": "https://issuer.com", + "authorization_endpoint": "https://issuer.com/connect/authorize", + "token_endpoint": "https://issuer.com/connect/token", + "jwks_uri": "https://issuer.com/jwks.json", + "token_endpoint_auth_methods_supported": [ + "client_secret_post", + "client_secret_basic", + "private_key_jwt", + "windows_client_authentication" + ], + "response_types_supported": [ + "code", + "id_token", + "code id_token", + "id_token token", + "code token", + "code id_token token" + ], + "response_modes_supported": [ + "query", + "fragment", + "form_post" + ], + "grant_types_supported": [ + "authorization_code", + "refresh_token", + "client_credentials", + "urn:ietf:params:oauth:grant-type:jwt-bearer", + "implicit", + "password", + "srv_challenge", + "urn:ietf:params:oauth:grant-type:device_code", + "device_code" + ], + "subject_types_supported": [ + "pairwise" + ], + "scopes_supported": [ + "user_impersonation", + "vpn_cert", + "email", + "openid", + "profile", + "allatclaims", + "logon_cert", + "aza", + "winhello_cert" + ], + "id_token_signing_alg_values_supported": [ + "RS256" + ], + "token_endpoint_auth_signing_alg_values_supported": [ + "RS256" + ], + "access_token_issuer": "http://access-token-issuer.com/adfs/services/trust", + "claims_supported": [ + "aud", + "iss", + "iat", + "exp", + "auth_time", + "nonce", + "at_hash", + "c_hash", + "sub", + "upn", + "unique_name", + "pwd_url", + "pwd_exp", + "mfa_auth_time", + "sid", + "nbf" + ], + "microsoft_multi_refresh_token": true, + "userinfo_endpoint": "https://issuer.com/connect/userinfo", + "capabilities": [], + "end_session_endpoint": "https://issuer.com/adfs/oauth2/logout", + "as_access_token_token_binding_supported": true, + "as_refresh_token_token_binding_supported": true, + "resource_access_token_token_binding_supported": true, + "op_id_token_token_binding_supported": true, + "rp_id_token_token_binding_supported": true, + "frontchannel_logout_supported": true, + "frontchannel_logout_session_supported": true, + "device_authorization_endpoint": "https://issuer.com/adfs/oauth2/devicecode" +} diff --git a/presto-main/src/test/resources/oidc/openid-configuration-without-userinfo.json b/presto-main/src/test/resources/oidc/openid-configuration-without-userinfo.json new file mode 100644 index 0000000000000..96f8c6048acf6 --- /dev/null +++ b/presto-main/src/test/resources/oidc/openid-configuration-without-userinfo.json @@ -0,0 +1,92 @@ +{ + "issuer": "https://issuer.com", + "authorization_endpoint": "https://issuer.com/connect/authorize", + "token_endpoint": "https://issuer.com/connect/token", + "token_endpoint_auth_methods_supported": [ + "client_secret_basic", + "private_key_jwt" + ], + "token_endpoint_auth_signing_alg_values_supported": [ + "RS256", + "ES256" + ], + "check_session_iframe": "https://issuer.com/connect/check_session", + "end_session_endpoint": "https://issuer.com/connect/end_session", + "jwks_uri": "https://issuer.com/jwks.json", + "registration_endpoint": "https://issuer.com/connect/register", + "scopes_supported": [ + "openid", + "profile", + "email", + "address", + "phone", + "offline_access" + ], + "response_types_supported": [ + "code", + "code id_token", + "id_token", + "token id_token" + ], + "acr_values_supported": [ + "urn:mace:incommon:iap:silver", + "urn:mace:incommon:iap:bronze" + ], + "subject_types_supported": [ + "public", + "pairwise" + ], + "id_token_signing_alg_values_supported": [ + "RS256", + "ES256", + "HS256" + ], + "id_token_encryption_alg_values_supported": [ + "RSA1_5", + "A128KW" + ], + "id_token_encryption_enc_values_supported": [ + "A128CBC-HS256", + "A128GCM" + ], + "request_object_signing_alg_values_supported": [ + "none", + "RS256", + "ES256" + ], + "display_values_supported": [ + "page", + "popup" + ], + "claim_types_supported": [ + "normal", + "distributed" + ], + "claims_supported": [ + "sub", + "iss", + "auth_time", + "acr", + "name", + "given_name", + "family_name", + "nickname", + "profile", + "picture", + "website", + "email", + "email_verified", + "locale", + "zoneinfo", + "http://example.info/claims/groups" + ], + "claims_parameter_supported": true, + "service_documentation": "http://issuer.com/connect/service_documentation.html", + "ui_locales_supported": [ + "en-US", + "en-GB", + "en-CA", + "fr-FR", + "fr-CA" + ] +} diff --git a/presto-main/src/test/resources/oidc/openid-configuration.json b/presto-main/src/test/resources/oidc/openid-configuration.json new file mode 100644 index 0000000000000..39c456fca85b6 --- /dev/null +++ b/presto-main/src/test/resources/oidc/openid-configuration.json @@ -0,0 +1,106 @@ +{ + "issuer": "https://issuer.com", + "authorization_endpoint": "https://issuer.com/connect/authorize", + "token_endpoint": "https://issuer.com/connect/token", + "token_endpoint_auth_methods_supported": [ + "client_secret_basic", + "private_key_jwt" + ], + "token_endpoint_auth_signing_alg_values_supported": [ + "RS256", + "ES256" + ], + "userinfo_endpoint": "https://issuer.com/connect/userinfo", + "check_session_iframe": "https://issuer.com/connect/check_session", + "end_session_endpoint": "https://issuer.com/connect/end_session", + "jwks_uri": "https://issuer.com/jwks.json", + "registration_endpoint": "https://issuer.com/connect/register", + "scopes_supported": [ + "openid", + "profile", + "email", + "address", + "phone", + "offline_access" + ], + "response_types_supported": [ + "code", + "code id_token", + "id_token", + "token id_token" + ], + "acr_values_supported": [ + "urn:mace:incommon:iap:silver", + "urn:mace:incommon:iap:bronze" + ], + "subject_types_supported": [ + "public", + "pairwise" + ], + "userinfo_signing_alg_values_supported": [ + "RS256", + "ES256", + "HS256" + ], + "userinfo_encryption_alg_values_supported": [ + "RSA1_5", + "A128KW" + ], + "userinfo_encryption_enc_values_supported": [ + "A128CBC-HS256", + "A128GCM" + ], + "id_token_signing_alg_values_supported": [ + "RS256", + "ES256", + "HS256" + ], + "id_token_encryption_alg_values_supported": [ + "RSA1_5", + "A128KW" + ], + "id_token_encryption_enc_values_supported": [ + "A128CBC-HS256", + "A128GCM" + ], + "request_object_signing_alg_values_supported": [ + "none", + "RS256", + "ES256" + ], + "display_values_supported": [ + "page", + "popup" + ], + "claim_types_supported": [ + "normal", + "distributed" + ], + "claims_supported": [ + "sub", + "iss", + "auth_time", + "acr", + "name", + "given_name", + "family_name", + "nickname", + "profile", + "picture", + "website", + "email", + "email_verified", + "locale", + "zoneinfo", + "http://example.info/claims/groups" + ], + "claims_parameter_supported": true, + "service_documentation": "http://issuer.com/connect/service_documentation.html", + "ui_locales_supported": [ + "en-US", + "en-GB", + "en-CA", + "fr-FR", + "fr-CA" + ] +} From 6247d6c5518d01b3a9d4d9b74d2d32cdcf6a3e5f Mon Sep 17 00:00:00 2001 From: auden-woolfson Date: Tue, 10 Jun 2025 15:45:29 -0700 Subject: [PATCH 108/113] Fix build and test failures from oauth2 w/o CLI --- pom.xml | 2 +- .../facebook/presto/cli/ClientOptions.java | 7 --- .../java/com/facebook/presto/cli/Console.java | 4 +- .../com/facebook/presto/cli/QueryRunner.java | 42 +------------- .../facebook/presto/cli/AbstractCliTest.java | 4 +- presto-client/pom.xml | 17 ++++-- .../auth/external/ExternalAuthenticator.java | 3 +- .../auth/external/MemoryCachedKnownToken.java | 2 +- presto-jdbc/pom.xml | 19 ++++--- .../presto/jdbc/ConnectionProperties.java | 2 +- .../jdbc/TestJdbcExternalAuthentication.java | 23 ++++---- .../server/security/SecurityConfig.java | 3 +- presto-main/etc/regex-map.txt | 3 + presto-main/pom.xml | 15 ++--- .../facebook/presto/server/PrestoServer.java | 2 +- .../facebook/presto/server/WebUiResource.java | 2 - .../server/security/AuthenticationFilter.java | 32 +++++------ .../DefaultWebUiAuthenticationManager.java | 8 +-- .../server/security/ServerSecurityModule.java | 20 +++++-- .../security/WebUiAuthenticationManager.java | 8 +-- .../security/oauth2/JweTokenSerializer.java | 2 +- .../oauth2/NimbusAirliftHttpClient.java | 2 +- .../security/oauth2/NimbusOAuth2Client.java | 56 ++++++++++++++++++- .../server/security/oauth2/NonceCookie.java | 16 +++--- .../security/oauth2/OAuth2Authenticator.java | 4 +- .../oauth2/OAuth2CallbackResource.java | 24 ++++---- .../server/security/oauth2/OAuth2Config.java | 7 +-- .../server/security/oauth2/OAuth2Service.java | 18 +++--- .../security/oauth2/OAuth2TokenExchange.java | 4 +- .../oauth2/OAuth2TokenExchangeResource.java | 28 +++++----- .../server/security/oauth2/OAuth2Utils.java | 6 +- ... => OAuth2WebUiAuthenticationManager.java} | 21 +++---- .../security/oauth2/OAuthWebUiCookie.java | 6 +- .../security/oauth2/OidcDiscoveryConfig.java | 5 +- .../security/oauth2/RefreshTokensConfig.java | 4 +- .../StaticOAuth2ServerConfiguration.java | 3 +- .../security/oauth2/TokenPairSerializer.java | 3 +- .../presto/server/MockHttpServletRequest.java | 2 +- .../BaseOAuth2AuthenticationFilterTest.java | 31 +++++----- .../security/oauth2/SimpleProxyServer.java | 13 ++--- ...TestDualAuthenticationFilterWithOAuth.java | 19 +++---- .../oauth2/TestJweTokenSerializer.java | 4 +- ...TestOAuth2AuthenticationFilterWithJwt.java | 6 ++ ...tOAuth2AuthenticationFilterWithOpaque.java | 6 ++ .../security/oauth2/TestOAuth2Config.java | 2 +- .../security/oauth2/TestOAuth2Utils.java | 5 +- .../security/oauth2/TestOidcDiscovery.java | 22 +++++--- .../oauth2/TestOidcDiscoveryConfig.java | 2 +- .../oauth2/TestRefreshTokensConfig.java | 2 +- .../oauth2/TestingHydraIdentityProvider.java | 24 +++++--- presto-native-execution/pom.xml | 7 +-- .../facebook/presto/router/RouterModule.java | 2 + .../src/static}/oauth2/failure.html | 0 .../src/static/{ouath2 => oauth2}/logout.html | 0 .../src/static}/oauth2/success.html | 0 55 files changed, 304 insertions(+), 270 deletions(-) create mode 100644 presto-main/etc/regex-map.txt rename presto-main/src/main/java/com/facebook/presto/server/security/oauth2/{Oauth2WebUiAuthenticationManager.java => OAuth2WebUiAuthenticationManager.java} (89%) rename {presto-main/src/main/resources => presto-ui/src/static}/oauth2/failure.html (100%) rename presto-ui/src/static/{ouath2 => oauth2}/logout.html (100%) rename {presto-main/src/main/resources => presto-ui/src/static}/oauth2/success.html (100%) diff --git a/pom.xml b/pom.xml index 5107da0b0dd06..186bcaecac7b0 100644 --- a/pom.xml +++ b/pom.xml @@ -2704,7 +2704,7 @@ com.nimbusds nimbus-jose-jwt - 9.30.2 + 10.0.2 diff --git a/presto-cli/src/main/java/com/facebook/presto/cli/ClientOptions.java b/presto-cli/src/main/java/com/facebook/presto/cli/ClientOptions.java index aa3d27b1502d6..f3f15e9324953 100644 --- a/presto-cli/src/main/java/com/facebook/presto/cli/ClientOptions.java +++ b/presto-cli/src/main/java/com/facebook/presto/cli/ClientOptions.java @@ -15,7 +15,6 @@ import com.facebook.airlift.units.Duration; import com.facebook.presto.client.ClientSession; -import com.facebook.presto.client.auth.external.ExternalRedirectStrategy; import com.google.common.annotations.VisibleForTesting; import com.google.common.base.CharMatcher; import com.google.common.base.Splitter; @@ -163,12 +162,6 @@ public class ClientOptions @Option(name = "--disable-redirects", title = "disable redirects", description = "Disable client following redirects from server") public boolean disableRedirects; - @Option(name = "--external-authentication", title = "enable external authentication", description = "Enable external authentication") - public boolean externalAuthentication; - - @Option(name = "--external-authentication-redirect-handler", title = "external authentication redirect handler", description = "External authentication redirect handlers") - public List externalAuthenticationRedirectHandler = new ArrayList<>(); - public enum OutputFormat { ALIGNED, diff --git a/presto-cli/src/main/java/com/facebook/presto/cli/Console.java b/presto-cli/src/main/java/com/facebook/presto/cli/Console.java index a94c67a5c1823..95a0dc76d4078 100644 --- a/presto-cli/src/main/java/com/facebook/presto/cli/Console.java +++ b/presto-cli/src/main/java/com/facebook/presto/cli/Console.java @@ -142,9 +142,7 @@ public boolean run() Optional.ofNullable(clientOptions.krb5KeytabPath), Optional.ofNullable(clientOptions.krb5CredentialCachePath), !clientOptions.krb5DisableRemoteServiceHostnameCanonicalization, - !clientOptions.disableRedirects, - clientOptions.externalAuthentication, - clientOptions.externalAuthenticationRedirectHandler)) { + !clientOptions.disableRedirects)) { if (hasQuery) { return executeCommand(queryRunner, query, clientOptions.outputFormat, clientOptions.ignoreErrors); } diff --git a/presto-cli/src/main/java/com/facebook/presto/cli/QueryRunner.java b/presto-cli/src/main/java/com/facebook/presto/cli/QueryRunner.java index 9b88c5326d954..1ecffacc29bc2 100644 --- a/presto-cli/src/main/java/com/facebook/presto/cli/QueryRunner.java +++ b/presto-cli/src/main/java/com/facebook/presto/cli/QueryRunner.java @@ -16,21 +16,11 @@ import com.facebook.presto.client.ClientSession; import com.facebook.presto.client.OkHttpUtil; import com.facebook.presto.client.StatementClient; -import com.facebook.presto.client.auth.external.CompositeRedirectHandler; -import com.facebook.presto.client.auth.external.ExternalAuthenticator; -import com.facebook.presto.client.auth.external.ExternalRedirectStrategy; -import com.facebook.presto.client.auth.external.HttpTokenPoller; -import com.facebook.presto.client.auth.external.KnownToken; -import com.facebook.presto.client.auth.external.RedirectHandler; -import com.facebook.presto.client.auth.external.TokenPoller; -import com.google.common.collect.ImmutableList; import com.google.common.net.HostAndPort; import okhttp3.OkHttpClient; import java.io.Closeable; import java.io.File; -import java.time.Duration; -import java.util.List; import java.util.Optional; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Consumer; @@ -83,9 +73,7 @@ public QueryRunner( Optional kerberosKeytabPath, Optional kerberosCredentialCachePath, boolean kerberosUseCanonicalHostname, - boolean followRedirects, - boolean externalAuthentication, - List externalRedirectHandlers) + boolean followRedirects) { this.session = new AtomicReference<>(requireNonNull(session, "session is null")); this.debug = debug; @@ -106,7 +94,6 @@ public QueryRunner( setupHttpProxy(builder, httpProxy); setupBasicAuth(builder, session, user, password); setupTokenAuth(builder, session, accessToken); - setupExternalAuth(builder, session, externalAuthentication, externalRedirectHandlers, sslSetup); if (kerberosRemoteServiceName.isPresent()) { checkArgument(session.getServer().getScheme().equalsIgnoreCase("https"), @@ -193,31 +180,4 @@ private static void setupTokenAuth( clientBuilder.addInterceptor(tokenAuth(accessToken.get())); } } - - private static void setupExternalAuth( - OkHttpClient.Builder builder, - ClientSession session, - boolean enabled, - List externalRedirectHandlers, - Consumer sslSetup) - { - if (!enabled) { - return; - } - - checkArgument(session.getServer().getScheme().equalsIgnoreCase("https"), - "Authentication using externalAuthentication requires HTTPS to be enabled"); - - List redirectStrategy = externalRedirectHandlers.isEmpty() ? ImmutableList.of(ExternalRedirectStrategy.ALL) : externalRedirectHandlers; - RedirectHandler redirectHandler = new CompositeRedirectHandler(redirectStrategy); - TokenPoller poller = new HttpTokenPoller(builder.build(), sslSetup); - ExternalAuthenticator authenticator = new ExternalAuthenticator( - redirectHandler, - poller, - KnownToken.local(), - Duration.ofMinutes(10)); - - builder.authenticator(authenticator); - builder.addInterceptor(authenticator); - } } diff --git a/presto-cli/src/test/java/com/facebook/presto/cli/AbstractCliTest.java b/presto-cli/src/test/java/com/facebook/presto/cli/AbstractCliTest.java index dfa6fb0a3fbdb..97e8804910d04 100644 --- a/presto-cli/src/test/java/com/facebook/presto/cli/AbstractCliTest.java +++ b/presto-cli/src/test/java/com/facebook/presto/cli/AbstractCliTest.java @@ -150,9 +150,7 @@ protected static QueryRunner createQueryRunner(ClientSession clientSession, bool Optional.empty(), Optional.empty(), false, - true, - false, - ImmutableList.of()); + true); } protected static QueryRunner createQueryRunner(ClientSession clientSession) diff --git a/presto-client/pom.xml b/presto-client/pom.xml index 3547e6c17c001..6435a2cda2551 100644 --- a/presto-client/pom.xml +++ b/presto-client/pom.xml @@ -21,7 +21,6 @@ com.facebook.presto presto-spi - @@ -56,11 +55,6 @@ jackson-databind - - com.facebook.airlift - concurrent - - com.facebook.airlift json @@ -106,6 +100,11 @@ okio-jvm + + com.google.code.findbugs + jsr305 + + org.jetbrains.kotlin @@ -118,6 +117,12 @@ + + com.facebook.airlift + concurrent + test + + com.facebook.presto presto-testng-services diff --git a/presto-client/src/main/java/com/facebook/presto/client/auth/external/ExternalAuthenticator.java b/presto-client/src/main/java/com/facebook/presto/client/auth/external/ExternalAuthenticator.java index 69bd0faf72e0f..60f568dc7c173 100644 --- a/presto-client/src/main/java/com/facebook/presto/client/auth/external/ExternalAuthenticator.java +++ b/presto-client/src/main/java/com/facebook/presto/client/auth/external/ExternalAuthenticator.java @@ -15,6 +15,7 @@ import com.facebook.presto.client.ClientException; import com.google.common.annotations.VisibleForTesting; +import jakarta.annotation.Nullable; import okhttp3.Authenticator; import okhttp3.Challenge; import okhttp3.Interceptor; @@ -22,8 +23,6 @@ import okhttp3.Response; import okhttp3.Route; -import javax.annotation.Nullable; - import java.io.IOException; import java.net.URI; import java.net.URISyntaxException; diff --git a/presto-client/src/main/java/com/facebook/presto/client/auth/external/MemoryCachedKnownToken.java b/presto-client/src/main/java/com/facebook/presto/client/auth/external/MemoryCachedKnownToken.java index 104841ac2dd14..0b1253a6deaa6 100644 --- a/presto-client/src/main/java/com/facebook/presto/client/auth/external/MemoryCachedKnownToken.java +++ b/presto-client/src/main/java/com/facebook/presto/client/auth/external/MemoryCachedKnownToken.java @@ -67,7 +67,7 @@ public Optional getToken() public void setupToken(Supplier> tokenSource) { // Try to lock and generate new token. If some other thread (Connection) has - // already obtained writeLock and is generating new token, then skipp this + // already obtained writeLock and is generating new token, then skip this // to block on getToken() if (writeLock.tryLock()) { try { diff --git a/presto-jdbc/pom.xml b/presto-jdbc/pom.xml index 071330ea73ed4..d5515e0d71b7c 100644 --- a/presto-jdbc/pom.xml +++ b/presto-jdbc/pom.xml @@ -57,11 +57,6 @@ okhttp - - com.facebook.airlift - concurrent - - com.facebook.airlift security @@ -122,6 +117,12 @@ + + com.facebook.airlift + concurrent + test + + com.facebook.airlift configuration @@ -201,14 +202,14 @@ - javax.servlet - javax.servlet-api + jakarta.servlet + jakarta.servlet-api test - javax.ws.rs - javax.ws.rs-api + jakarta.ws.rs + jakarta.ws.rs-api test diff --git a/presto-jdbc/src/main/java/com/facebook/presto/jdbc/ConnectionProperties.java b/presto-jdbc/src/main/java/com/facebook/presto/jdbc/ConnectionProperties.java index 823eccdae7c01..931761723d9e5 100644 --- a/presto-jdbc/src/main/java/com/facebook/presto/jdbc/ConnectionProperties.java +++ b/presto-jdbc/src/main/java/com/facebook/presto/jdbc/ConnectionProperties.java @@ -13,12 +13,12 @@ */ package com.facebook.presto.jdbc; +import com.facebook.airlift.units.Duration; import com.facebook.presto.client.auth.external.ExternalRedirectStrategy; import com.google.common.base.Splitter; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import com.google.common.net.HostAndPort; -import io.airlift.units.Duration; import okhttp3.Protocol; import java.io.File; diff --git a/presto-jdbc/src/test/java/com/facebook/presto/jdbc/TestJdbcExternalAuthentication.java b/presto-jdbc/src/test/java/com/facebook/presto/jdbc/TestJdbcExternalAuthentication.java index 408ef01386164..ede9a534cf2a3 100644 --- a/presto-jdbc/src/test/java/com/facebook/presto/jdbc/TestJdbcExternalAuthentication.java +++ b/presto-jdbc/src/test/java/com/facebook/presto/jdbc/TestJdbcExternalAuthentication.java @@ -34,6 +34,13 @@ import com.google.inject.Module; import com.google.inject.Scopes; import com.google.inject.multibindings.Multibinder; +import jakarta.servlet.http.HttpServletRequest; +import jakarta.ws.rs.GET; +import jakarta.ws.rs.Path; +import jakarta.ws.rs.PathParam; +import jakarta.ws.rs.Produces; +import jakarta.ws.rs.core.Context; +import jakarta.ws.rs.core.Response; import okhttp3.HttpUrl; import okhttp3.OkHttpClient; import okhttp3.Request; @@ -42,14 +49,6 @@ import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; -import javax.servlet.http.HttpServletRequest; -import javax.ws.rs.GET; -import javax.ws.rs.Path; -import javax.ws.rs.PathParam; -import javax.ws.rs.Produces; -import javax.ws.rs.core.Context; -import javax.ws.rs.core.Response; - import java.io.File; import java.io.IOException; import java.net.URI; @@ -78,13 +77,13 @@ import static com.google.inject.Scopes.SINGLETON; import static com.google.inject.multibindings.Multibinder.newSetBinder; import static com.google.inject.util.Modules.combine; +import static jakarta.ws.rs.core.HttpHeaders.AUTHORIZATION; +import static jakarta.ws.rs.core.MediaType.APPLICATION_JSON_TYPE; +import static jakarta.ws.rs.core.MediaType.TEXT_PLAIN; +import static jakarta.ws.rs.core.Response.Status.NOT_FOUND; import static java.lang.String.format; import static java.net.HttpURLConnection.HTTP_OK; import static java.util.Objects.requireNonNull; -import static javax.ws.rs.core.HttpHeaders.AUTHORIZATION; -import static javax.ws.rs.core.MediaType.APPLICATION_JSON_TYPE; -import static javax.ws.rs.core.MediaType.TEXT_PLAIN; -import static javax.ws.rs.core.Response.Status.NOT_FOUND; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; diff --git a/presto-main-base/src/main/java/com/facebook/presto/server/security/SecurityConfig.java b/presto-main-base/src/main/java/com/facebook/presto/server/security/SecurityConfig.java index f190d86a35879..52d70ce60ac49 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/server/security/SecurityConfig.java +++ b/presto-main-base/src/main/java/com/facebook/presto/server/security/SecurityConfig.java @@ -41,6 +41,7 @@ public enum AuthenticationType PASSWORD, JWT, CUSTOM, + TEST_EXTERNAL, OAUTH2 } @@ -57,7 +58,7 @@ public SecurityConfig setAuthenticationTypes(List authentica } @Config("http-server.authentication.type") - @ConfigDescription("Authentication types (supported types: CERTIFICATE, KERBEROS, PASSWORD, JWT, CUSTOM, OAUTH2)") + @ConfigDescription("Authentication types (supported types: CERTIFICATE, KERBEROS, PASSWORD, JWT, CUSTOM, OAUTH2, TEST_EXTERNAL)") public SecurityConfig setAuthenticationTypes(String types) { if (types == null) { diff --git a/presto-main/etc/regex-map.txt b/presto-main/etc/regex-map.txt new file mode 100644 index 0000000000000..b61daa86bcc42 --- /dev/null +++ b/presto-main/etc/regex-map.txt @@ -0,0 +1,3 @@ +user=.* +internal=coordinator +admin=su.* \ No newline at end of file diff --git a/presto-main/pom.xml b/presto-main/pom.xml index 73ff20c4cb1d3..5f737a7ab3bd6 100644 --- a/presto-main/pom.xml +++ b/presto-main/pom.xml @@ -399,11 +399,6 @@ jackson-core - - com.squareup.okhttp3 - okhttp - - io.airlift aircompressor @@ -415,8 +410,8 @@ - com.squareup.okhttp3 - okhttp-urlconnection + net.minidev + json-smart @@ -532,7 +527,13 @@ com.facebook.airlift.drift:drift-protocol io.netty:netty-buffer javax.inject:javax.inject + com.squareup.okhttp3:okhttp + com.squareup.okhttp3:okhttp-urlconnection + + com.squareup.okhttp3:okhttp + com.squareup.okhttp3:okhttp-urlconnection + diff --git a/presto-main/src/main/java/com/facebook/presto/server/PrestoServer.java b/presto-main/src/main/java/com/facebook/presto/server/PrestoServer.java index c8c01053e4dc7..5c9f7bbe6099d 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/PrestoServer.java +++ b/presto-main/src/main/java/com/facebook/presto/server/PrestoServer.java @@ -53,11 +53,11 @@ import com.facebook.presto.nodeManager.PluginNodeManager; import com.facebook.presto.security.AccessControlManager; import com.facebook.presto.security.AccessControlModule; -import com.facebook.presto.server.security.oauth2.OAuth2Client; import com.facebook.presto.server.security.PasswordAuthenticatorManager; import com.facebook.presto.server.security.PrestoAuthenticatorManager; import com.facebook.presto.server.security.SecurityConfig; import com.facebook.presto.server.security.ServerSecurityModule; +import com.facebook.presto.server.security.oauth2.OAuth2Client; import com.facebook.presto.spi.function.SqlFunction; import com.facebook.presto.sql.analyzer.FeaturesConfig; import com.facebook.presto.sql.expressions.ExpressionOptimizerManager; diff --git a/presto-main/src/main/java/com/facebook/presto/server/WebUiResource.java b/presto-main/src/main/java/com/facebook/presto/server/WebUiResource.java index a434eee5d1fab..a2e908d47aea7 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/WebUiResource.java +++ b/presto-main/src/main/java/com/facebook/presto/server/WebUiResource.java @@ -14,7 +14,6 @@ package com.facebook.presto.server; import com.facebook.presto.server.security.oauth2.OAuthWebUiCookie; - import jakarta.annotation.security.RolesAllowed; import jakarta.ws.rs.GET; import jakarta.ws.rs.HeaderParam; @@ -29,7 +28,6 @@ import static com.facebook.presto.server.security.oauth2.OAuth2Utils.getLastURLParameter; import static com.google.common.base.Strings.isNullOrEmpty; import static com.google.common.net.HttpHeaders.X_FORWARDED_PROTO; -import static jakarta.ws.rs.core.Response.Status.MOVED_PERMANENTLY; @Path("/") @RolesAllowed(ADMIN) diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/AuthenticationFilter.java b/presto-main/src/main/java/com/facebook/presto/server/security/AuthenticationFilter.java index 2043e5402dc86..17d462cb1b822 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/security/AuthenticationFilter.java +++ b/presto-main/src/main/java/com/facebook/presto/server/security/AuthenticationFilter.java @@ -83,20 +83,6 @@ public AuthenticationFilter(List authenticators, SecurityConfig s this.clientRequestFilterManager = requireNonNull(clientRequestFilterManager, "clientRequestFilterManager is null"); } - private static void skipRequestBody(HttpServletRequest request) - throws IOException - { - // If we send the challenge without consuming the body of the request, - // the server will close the connection after sending the response. - // The client may interpret this as a failed request and not resend the - // request with the authentication header. We can avoid this behavior - // in the client by reading and discarding the entire body of the - // unauthenticated request before sending the response. - try (InputStream inputStream = request.getInputStream()) { - copy(inputStream, nullOutputStream()); - } - } - public static ServletRequest withPrincipal(HttpServletRequest request, Principal principal) { requireNonNull(principal, "principal is null"); @@ -110,7 +96,7 @@ public Principal getUserPrincipal() }; } - public static boolean isPublic(HttpServletRequest request) + private static boolean isRequestToOAuthEndpoint(HttpServletRequest request) { return request.getPathInfo().startsWith(TOKEN_ENDPOINT) || request.getPathInfo().startsWith(CALLBACK_ENDPOINT); @@ -238,7 +224,7 @@ public HttpServletRequest mergeExtraHeaders(HttpServletRequest request, Principa private boolean doesRequestSupportAuthentication(HttpServletRequest request) { - if (isPublic(request)) { + if (isRequestToOAuthEndpoint(request)) { return false; } if (authenticators.isEmpty()) { @@ -253,6 +239,20 @@ private boolean doesRequestSupportAuthentication(HttpServletRequest request) return false; } + private static void skipRequestBody(HttpServletRequest request) + throws IOException + { + // If we send the challenge without consuming the body of the request, + // the server will close the connection after sending the response. + // The client may interpret this as a failed request and not resend the + // request with the authentication header. We can avoid this behavior + // in the client by reading and discarding the entire body of the + // unauthenticated request before sending the response. + try (InputStream inputStream = request.getInputStream()) { + copy(inputStream, nullOutputStream()); + } + } + private boolean isWebUiRequest(HttpServletRequest request) { String pathInfo = request.getPathInfo(); diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/DefaultWebUiAuthenticationManager.java b/presto-main/src/main/java/com/facebook/presto/server/security/DefaultWebUiAuthenticationManager.java index dd0691f88e9a1..75ff4d072d711 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/security/DefaultWebUiAuthenticationManager.java +++ b/presto-main/src/main/java/com/facebook/presto/server/security/DefaultWebUiAuthenticationManager.java @@ -13,10 +13,10 @@ */ package com.facebook.presto.server.security; -import javax.servlet.FilterChain; -import javax.servlet.ServletException; -import javax.servlet.http.HttpServletRequest; -import javax.servlet.http.HttpServletResponse; +import jakarta.servlet.FilterChain; +import jakarta.servlet.ServletException; +import jakarta.servlet.http.HttpServletRequest; +import jakarta.servlet.http.HttpServletResponse; import java.io.IOException; diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/ServerSecurityModule.java b/presto-main/src/main/java/com/facebook/presto/server/security/ServerSecurityModule.java index 48e076d7dd405..6fe27fa7b7e96 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/security/ServerSecurityModule.java +++ b/presto-main/src/main/java/com/facebook/presto/server/security/ServerSecurityModule.java @@ -15,7 +15,10 @@ import com.facebook.airlift.configuration.AbstractConfigurationAwareModule; import com.facebook.airlift.http.server.Authenticator; +import com.facebook.airlift.http.server.Authorizer; import com.facebook.airlift.http.server.CertificateAuthenticator; +import com.facebook.airlift.http.server.ConfigurationBasedAuthorizer; +import com.facebook.airlift.http.server.ConfigurationBasedAuthorizerConfig; import com.facebook.airlift.http.server.KerberosAuthenticator; import com.facebook.airlift.http.server.KerberosConfig; import com.facebook.airlift.http.server.TheServlet; @@ -23,12 +26,11 @@ import com.facebook.presto.server.security.oauth2.OAuth2AuthenticationSupportModule; import com.facebook.presto.server.security.oauth2.OAuth2Authenticator; import com.facebook.presto.server.security.oauth2.OAuth2Config; -import com.facebook.presto.server.security.oauth2.Oauth2WebUiAuthenticationManager; +import com.facebook.presto.server.security.oauth2.OAuth2WebUiAuthenticationManager; import com.google.inject.Binder; import com.google.inject.Scopes; import com.google.inject.multibindings.Multibinder; - -import javax.servlet.Filter; +import jakarta.servlet.Filter; import java.util.List; @@ -39,6 +41,7 @@ import static com.facebook.presto.server.security.SecurityConfig.AuthenticationType.KERBEROS; import static com.facebook.presto.server.security.SecurityConfig.AuthenticationType.OAUTH2; import static com.facebook.presto.server.security.SecurityConfig.AuthenticationType.PASSWORD; +import static com.facebook.presto.server.security.SecurityConfig.AuthenticationType.TEST_EXTERNAL; import static com.google.inject.multibindings.Multibinder.newSetBinder; import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder; @@ -74,17 +77,24 @@ else if (authType == JWT) { authBinder.addBinding().to(JsonWebTokenAuthenticator.class).in(Scopes.SINGLETON); } else if (authType == OAUTH2) { - newOptionalBinder(binder, WebUiAuthenticationManager.class).setBinding().to(Oauth2WebUiAuthenticationManager.class).in(Scopes.SINGLETON); + newOptionalBinder(binder, WebUiAuthenticationManager.class).setBinding().to(OAuth2WebUiAuthenticationManager.class).in(Scopes.SINGLETON); install(new OAuth2AuthenticationSupportModule()); binder.bind(OAuth2Authenticator.class).in(Scopes.SINGLETON); configBinder(binder).bindConfig(OAuth2Config.class); authBinder.addBinding().to(OAuth2Authenticator.class).in(Scopes.SINGLETON); + + configBinder(binder).bindConfig(ConfigurationBasedAuthorizerConfig.class); + binder.bind(Authorizer.class).to(ConfigurationBasedAuthorizer.class).in(Scopes.SINGLETON); } else if (authType == CUSTOM) { authBinder.addBinding().to(CustomPrestoAuthenticator.class).in(Scopes.SINGLETON); } else { - throw new AssertionError("Unhandled auth type: " + authType); + // TEST_EXTERNAL is an authentication type used for testing the external auth flow for the JDBC driver. + // This is here as a guard since it's not a real authenticator but if I exclude it from the checks then teh error is thrown. + if (authType != TEST_EXTERNAL) { + throw new AssertionError("Unhandled auth type: " + authType); + } } } } diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/WebUiAuthenticationManager.java b/presto-main/src/main/java/com/facebook/presto/server/security/WebUiAuthenticationManager.java index e34a495463c37..322d416168942 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/security/WebUiAuthenticationManager.java +++ b/presto-main/src/main/java/com/facebook/presto/server/security/WebUiAuthenticationManager.java @@ -13,10 +13,10 @@ */ package com.facebook.presto.server.security; -import javax.servlet.FilterChain; -import javax.servlet.ServletException; -import javax.servlet.http.HttpServletRequest; -import javax.servlet.http.HttpServletResponse; +import jakarta.servlet.FilterChain; +import jakarta.servlet.ServletException; +import jakarta.servlet.http.HttpServletRequest; +import jakarta.servlet.http.HttpServletResponse; import java.io.IOException; diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/JweTokenSerializer.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/JweTokenSerializer.java index 30b8196f6c57f..b10ea12a8e8a2 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/JweTokenSerializer.java +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/JweTokenSerializer.java @@ -13,6 +13,7 @@ */ package com.facebook.presto.server.security.oauth2; +import com.facebook.airlift.units.Duration; import com.nimbusds.jose.EncryptionMethod; import com.nimbusds.jose.JOSEException; import com.nimbusds.jose.JWEAlgorithm; @@ -22,7 +23,6 @@ import com.nimbusds.jose.Payload; import com.nimbusds.jose.crypto.AESDecrypter; import com.nimbusds.jose.crypto.AESEncrypter; -import io.airlift.units.Duration; import io.jsonwebtoken.Claims; import io.jsonwebtoken.CompressionCodec; import io.jsonwebtoken.CompressionException; diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NimbusAirliftHttpClient.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NimbusAirliftHttpClient.java index 04d3eb17aad01..3aeeb4046f216 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NimbusAirliftHttpClient.java +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NimbusAirliftHttpClient.java @@ -24,9 +24,9 @@ import com.nimbusds.oauth2.sdk.ParseException; import com.nimbusds.oauth2.sdk.http.HTTPRequest; import com.nimbusds.oauth2.sdk.http.HTTPResponse; +import jakarta.ws.rs.core.UriBuilder; import javax.inject.Inject; -import javax.ws.rs.core.UriBuilder; import java.io.IOException; import java.net.URISyntaxException; diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NimbusOAuth2Client.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NimbusOAuth2Client.java index ff24e9e477c21..86b2fb96b7d8f 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NimbusOAuth2Client.java +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NimbusOAuth2Client.java @@ -14,6 +14,7 @@ package com.facebook.presto.server.security.oauth2; import com.facebook.airlift.log.Logger; +import com.facebook.airlift.units.Duration; import com.facebook.presto.server.security.oauth2.OAuth2ServerConfigProvider.OAuth2ServerConfig; import com.google.common.collect.ImmutableSet; import com.google.common.collect.Ordering; @@ -40,6 +41,7 @@ import com.nimbusds.oauth2.sdk.TokenRequest; import com.nimbusds.oauth2.sdk.auth.ClientSecretBasic; import com.nimbusds.oauth2.sdk.auth.Secret; +import com.nimbusds.oauth2.sdk.http.HTTPResponse; import com.nimbusds.oauth2.sdk.id.ClientID; import com.nimbusds.oauth2.sdk.id.Issuer; import com.nimbusds.oauth2.sdk.id.State; @@ -50,23 +52,27 @@ import com.nimbusds.openid.connect.sdk.AuthenticationRequest; import com.nimbusds.openid.connect.sdk.Nonce; import com.nimbusds.openid.connect.sdk.OIDCTokenResponse; +import com.nimbusds.openid.connect.sdk.UserInfoErrorResponse; import com.nimbusds.openid.connect.sdk.UserInfoRequest; import com.nimbusds.openid.connect.sdk.UserInfoResponse; +import com.nimbusds.openid.connect.sdk.UserInfoSuccessResponse; import com.nimbusds.openid.connect.sdk.claims.AccessTokenHash; import com.nimbusds.openid.connect.sdk.claims.IDTokenClaimsSet; import com.nimbusds.openid.connect.sdk.token.OIDCTokens; import com.nimbusds.openid.connect.sdk.validators.AccessTokenValidator; import com.nimbusds.openid.connect.sdk.validators.IDTokenValidator; import com.nimbusds.openid.connect.sdk.validators.InvalidHashException; -import io.airlift.units.Duration; +import net.minidev.json.JSONObject; import javax.inject.Inject; import java.net.MalformedURLException; import java.net.URI; import java.time.Instant; +import java.util.Collections; import java.util.Date; import java.util.HashSet; +import java.util.List; import java.util.Map; import java.util.Optional; import java.util.Set; @@ -74,6 +80,7 @@ import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Preconditions.checkState; +import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.collect.ImmutableSet.toImmutableSet; import static com.google.common.hash.Hashing.sha256; import static com.nimbusds.oauth2.sdk.ResponseType.CODE; @@ -372,7 +379,7 @@ private Optional getJWTClaimsSet(String accessToken) private Optional queryUserInfo(String accessToken) { try { - UserInfoResponse response = httpClient.execute(new UserInfoRequest(userinfoUrl.get(), new BearerAccessToken(accessToken)), UserInfoResponse::parse); + UserInfoResponse response = httpClient.execute(new UserInfoRequest(userinfoUrl.get(), new BearerAccessToken(accessToken)), this::parse); if (!response.indicatesSuccess()) { LOG.error("Received bad response from userinfo endpoint: " + response.toErrorResponse().getErrorObject()); return Optional.empty(); @@ -385,6 +392,51 @@ private Optional queryUserInfo(String accessToken) } } + // Using this parsing method for our /userinfo response from the IdP in order to allow for different principal + // fields as defined, and in the absence of the `sub` claim. This is a "hack" solution to alter the claims + // present in the response before calling the parser provided by the oidc sdk, which fails hard if the + // `sub` claim is missing. Note we also have to offload audience verification to this method since it + // is not handled in the library + public UserInfoResponse parse(HTTPResponse httpResponse) + throws ParseException + { + JSONObject body = httpResponse.getContentAsJSONObject(); + + String principal = (String) body.get(principalField); + if (principal == null) { + throw new ParseException(String.format("/userinfo response missing principal field %s", principalField)); + } + + if (!principalField.equals("sub") && body.get("sub") == null) { + body.put("sub", principal); + httpResponse.setBody(body.toJSONString()); + } + + Object audClaim = body.get("aud"); + List audiences; + + if (audClaim instanceof String) { + audiences = List.of((String) audClaim); + } + else if (audClaim instanceof List) { + audiences = ((List) audClaim).stream() + .filter(String.class::isInstance) + .map(String.class::cast) + .collect(toImmutableList()); + } + else { + throw new ParseException("Unsupported or missing 'aud' claim type in /userinfo response"); + } + + if (!(audiences.contains(clientId.getValue()) || !Collections.disjoint(audiences, accessTokenAudiences))) { + throw new ParseException("Invalid audience in /userinfo response"); + } + + return (httpResponse.getStatusCode() == 200) + ? UserInfoSuccessResponse.parse(httpResponse) + : UserInfoErrorResponse.parse(httpResponse); + } + private Optional parseAccessToken(String accessToken) { try { diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NonceCookie.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NonceCookie.java index 7475d0c46cd9f..d80e2cacc6941 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NonceCookie.java +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/NonceCookie.java @@ -13,19 +13,18 @@ */ package com.facebook.presto.server.security.oauth2; +import jakarta.ws.rs.core.Cookie; +import jakarta.ws.rs.core.NewCookie; import org.apache.commons.lang3.StringUtils; -import javax.ws.rs.core.Cookie; -import javax.ws.rs.core.NewCookie; - import java.time.Instant; import java.util.Date; import java.util.Optional; import static com.facebook.presto.server.security.oauth2.OAuth2CallbackResource.CALLBACK_ENDPOINT; import static com.google.common.base.Predicates.not; -import static javax.ws.rs.core.Cookie.DEFAULT_VERSION; -import static javax.ws.rs.core.NewCookie.DEFAULT_MAX_AGE; +import static jakarta.ws.rs.core.Cookie.DEFAULT_VERSION; +import static jakarta.ws.rs.core.NewCookie.DEFAULT_MAX_AGE; public final class NonceCookie { @@ -49,16 +48,15 @@ public static NewCookie create(String nonce, Instant tokenExpiration) true); } - public static javax.servlet.http.Cookie createServletCookie(String nonce, Instant tokenExpiration) + public static jakarta.servlet.http.Cookie createServletCookie(String nonce, Instant tokenExpiration) { return toServletCookie(create(nonce, tokenExpiration)); } - public static javax.servlet.http.Cookie toServletCookie(NewCookie cookie) + public static jakarta.servlet.http.Cookie toServletCookie(NewCookie cookie) { - javax.servlet.http.Cookie servletCookie = new javax.servlet.http.Cookie(cookie.getName(), cookie.getValue()); + jakarta.servlet.http.Cookie servletCookie = new jakarta.servlet.http.Cookie(cookie.getName(), cookie.getValue()); servletCookie.setPath(cookie.getPath()); - servletCookie.setVersion(cookie.getVersion()); servletCookie.setMaxAge(cookie.getMaxAge()); servletCookie.setSecure(cookie.isSecure()); servletCookie.setHttpOnly(cookie.isHttpOnly()); diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Authenticator.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Authenticator.java index 3fe980f061364..83c2de302c287 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Authenticator.java +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Authenticator.java @@ -17,11 +17,11 @@ import com.facebook.airlift.http.server.Authenticator; import com.facebook.airlift.http.server.BasicPrincipal; import com.facebook.airlift.log.Logger; +import jakarta.servlet.http.Cookie; +import jakarta.servlet.http.HttpServletRequest; import org.apache.commons.lang3.StringUtils; import javax.inject.Inject; -import javax.servlet.http.Cookie; -import javax.servlet.http.HttpServletRequest; import java.net.URI; import java.security.Principal; diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2CallbackResource.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2CallbackResource.java index ef5b6d6330d38..42d52450837d2 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2CallbackResource.java +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2CallbackResource.java @@ -14,24 +14,24 @@ package com.facebook.presto.server.security.oauth2; import com.facebook.airlift.log.Logger; +import jakarta.servlet.http.HttpServletRequest; +import jakarta.ws.rs.CookieParam; +import jakarta.ws.rs.GET; +import jakarta.ws.rs.Path; +import jakarta.ws.rs.Produces; +import jakarta.ws.rs.QueryParam; +import jakarta.ws.rs.core.Context; +import jakarta.ws.rs.core.Cookie; +import jakarta.ws.rs.core.Response; +import jakarta.ws.rs.core.UriBuilder; import javax.inject.Inject; -import javax.servlet.http.HttpServletRequest; -import javax.ws.rs.CookieParam; -import javax.ws.rs.GET; -import javax.ws.rs.Path; -import javax.ws.rs.Produces; -import javax.ws.rs.QueryParam; -import javax.ws.rs.core.Context; -import javax.ws.rs.core.Cookie; -import javax.ws.rs.core.Response; -import javax.ws.rs.core.UriBuilder; import static com.facebook.presto.server.security.oauth2.NonceCookie.NONCE_COOKIE; import static com.facebook.presto.server.security.oauth2.OAuth2Utils.getSchemeUriBuilder; +import static jakarta.ws.rs.core.MediaType.TEXT_HTML; +import static jakarta.ws.rs.core.Response.Status.BAD_REQUEST; import static java.util.Objects.requireNonNull; -import static javax.ws.rs.core.MediaType.TEXT_HTML; -import static javax.ws.rs.core.Response.Status.BAD_REQUEST; @Path(OAuth2CallbackResource.CALLBACK_ENDPOINT) public class OAuth2CallbackResource diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Config.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Config.java index 8de3e4442138b..b1b0f9513b2f7 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Config.java +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Config.java @@ -16,13 +16,12 @@ import com.facebook.airlift.configuration.Config; import com.facebook.airlift.configuration.ConfigDescription; import com.facebook.airlift.configuration.ConfigSecuritySensitive; +import com.facebook.airlift.units.Duration; +import com.facebook.airlift.units.MinDuration; import com.google.common.base.Splitter; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; -import io.airlift.units.Duration; -import io.airlift.units.MinDuration; - -import javax.validation.constraints.NotNull; +import jakarta.validation.constraints.NotNull; import java.io.File; import java.util.Collections; diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Service.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Service.java index 40a97bb1b82ca..78616a7d8cad5 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Service.java +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Service.java @@ -17,13 +17,13 @@ import com.google.common.io.Resources; import io.jsonwebtoken.Claims; import io.jsonwebtoken.JwtParser; +import jakarta.servlet.http.HttpServletRequest; +import jakarta.servlet.http.HttpServletResponse; +import jakarta.ws.rs.core.HttpHeaders; +import jakarta.ws.rs.core.Response; +import jakarta.ws.rs.core.UriBuilder; import javax.inject.Inject; -import javax.servlet.http.HttpServletRequest; -import javax.servlet.http.HttpServletResponse; -import javax.ws.rs.core.HttpHeaders; -import javax.ws.rs.core.Response; -import javax.ws.rs.core.UriBuilder; import java.io.IOException; import java.net.URI; @@ -44,12 +44,12 @@ import static com.google.common.base.Verify.verify; import static com.google.common.hash.Hashing.sha256; import static io.jsonwebtoken.security.Keys.hmacShaKeyFor; +import static jakarta.ws.rs.core.Response.Status.BAD_REQUEST; +import static jakarta.ws.rs.core.Response.Status.FORBIDDEN; import static java.lang.String.format; import static java.nio.charset.StandardCharsets.UTF_8; import static java.time.Instant.now; import static java.util.Objects.requireNonNull; -import static javax.ws.rs.core.Response.Status.BAD_REQUEST; -import static javax.ws.rs.core.Response.Status.FORBIDDEN; public class OAuth2Service { @@ -87,8 +87,8 @@ public OAuth2Service( this.client = requireNonNull(client, "client is null"); requireNonNull(oauth2Config, "oauth2Config is null"); - this.successHtml = Resources.toString(Resources.getResource(getClass(), "/oauth2/success.html"), UTF_8); - this.failureHtml = Resources.toString(Resources.getResource(getClass(), "/oauth2/failure.html"), UTF_8); + this.successHtml = Resources.toString(Resources.getResource(getClass(), "/webapp/oauth2/success.html"), UTF_8); + this.failureHtml = Resources.toString(Resources.getResource(getClass(), "/webapp/oauth2/failure.html"), UTF_8); verify(failureHtml.contains(FAILURE_REPLACEMENT_TEXT), "failure.html does not contain the replacement text"); this.challengeTimeout = Duration.ofMillis(oauth2Config.getChallengeTimeout().toMillis()); diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2TokenExchange.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2TokenExchange.java index ff06a4bfa9c4c..95e0bdc25eb01 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2TokenExchange.java +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2TokenExchange.java @@ -13,15 +13,15 @@ */ package com.facebook.presto.server.security.oauth2; +import com.facebook.airlift.units.Duration; import com.google.common.cache.CacheBuilder; import com.google.common.cache.CacheLoader; import com.google.common.cache.LoadingCache; import com.google.common.hash.Hashing; import com.google.common.util.concurrent.ListenableFuture; import com.google.common.util.concurrent.SettableFuture; -import io.airlift.units.Duration; +import jakarta.annotation.PreDestroy; -import javax.annotation.PreDestroy; import javax.inject.Inject; import java.nio.charset.StandardCharsets; diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2TokenExchangeResource.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2TokenExchangeResource.java index 36d27a0fe1252..5335712abe75e 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2TokenExchangeResource.java +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2TokenExchangeResource.java @@ -22,21 +22,21 @@ import com.google.common.util.concurrent.Futures; import com.google.common.util.concurrent.ListenableFuture; import com.google.common.util.concurrent.ListeningExecutorService; +import jakarta.servlet.http.HttpServletRequest; +import jakarta.ws.rs.BadRequestException; +import jakarta.ws.rs.DELETE; +import jakarta.ws.rs.GET; +import jakarta.ws.rs.Path; +import jakarta.ws.rs.PathParam; +import jakarta.ws.rs.Produces; +import jakarta.ws.rs.container.AsyncResponse; +import jakarta.ws.rs.container.Suspended; +import jakarta.ws.rs.core.Context; +import jakarta.ws.rs.core.MediaType; +import jakarta.ws.rs.core.Response; +import jakarta.ws.rs.core.UriBuilder; import javax.inject.Inject; -import javax.servlet.http.HttpServletRequest; -import javax.ws.rs.BadRequestException; -import javax.ws.rs.DELETE; -import javax.ws.rs.GET; -import javax.ws.rs.Path; -import javax.ws.rs.PathParam; -import javax.ws.rs.Produces; -import javax.ws.rs.container.AsyncResponse; -import javax.ws.rs.container.Suspended; -import javax.ws.rs.core.Context; -import javax.ws.rs.core.MediaType; -import javax.ws.rs.core.Response; -import javax.ws.rs.core.UriBuilder; import java.util.Map; import java.util.Optional; @@ -47,8 +47,8 @@ import static com.facebook.presto.server.security.oauth2.OAuth2TokenExchange.MAX_POLL_TIME; import static com.facebook.presto.server.security.oauth2.OAuth2TokenExchange.hashAuthId; import static com.facebook.presto.server.security.oauth2.OAuth2Utils.getSchemeUriBuilder; +import static jakarta.ws.rs.core.MediaType.APPLICATION_JSON_TYPE; import static java.util.Objects.requireNonNull; -import static javax.ws.rs.core.MediaType.APPLICATION_JSON_TYPE; @Path(OAuth2TokenExchangeResource.TOKEN_ENDPOINT) public class OAuth2TokenExchangeResource diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Utils.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Utils.java index 4d9dba90ec2f1..0b5f8f62db545 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Utils.java +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2Utils.java @@ -13,9 +13,9 @@ */ package com.facebook.presto.server.security.oauth2; -import javax.servlet.http.HttpServletRequest; -import javax.ws.rs.core.MultivaluedMap; -import javax.ws.rs.core.UriBuilder; +import jakarta.servlet.http.HttpServletRequest; +import jakarta.ws.rs.core.MultivaluedMap; +import jakarta.ws.rs.core.UriBuilder; import java.util.List; import java.util.Map; diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/Oauth2WebUiAuthenticationManager.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2WebUiAuthenticationManager.java similarity index 89% rename from presto-main/src/main/java/com/facebook/presto/server/security/oauth2/Oauth2WebUiAuthenticationManager.java rename to presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2WebUiAuthenticationManager.java index d31b993e5b437..925019f5ac442 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/Oauth2WebUiAuthenticationManager.java +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuth2WebUiAuthenticationManager.java @@ -17,14 +17,14 @@ import com.facebook.airlift.log.Logger; import com.facebook.presto.server.security.WebUiAuthenticationManager; import com.facebook.presto.server.security.oauth2.TokenPairSerializer.TokenPair; +import jakarta.servlet.FilterChain; +import jakarta.servlet.ServletException; +import jakarta.servlet.http.Cookie; +import jakarta.servlet.http.HttpServletRequest; +import jakarta.servlet.http.HttpServletResponse; +import jakarta.ws.rs.core.UriBuilder; import javax.inject.Inject; -import javax.servlet.FilterChain; -import javax.servlet.ServletException; -import javax.servlet.http.Cookie; -import javax.servlet.http.HttpServletRequest; -import javax.servlet.http.HttpServletResponse; -import javax.ws.rs.core.UriBuilder; import java.io.IOException; import java.security.Principal; @@ -33,14 +33,15 @@ import java.util.Optional; import static com.facebook.presto.server.security.AuthenticationFilter.withPrincipal; +import static com.facebook.presto.server.security.oauth2.OAuth2Authenticator.extractTokenFromCookie; import static com.facebook.presto.server.security.oauth2.OAuth2CallbackResource.CALLBACK_ENDPOINT; import static com.facebook.presto.server.security.oauth2.OAuth2Utils.getSchemeUriBuilder; import static java.util.Objects.requireNonNull; -public class Oauth2WebUiAuthenticationManager +public class OAuth2WebUiAuthenticationManager implements WebUiAuthenticationManager { - private static final Logger logger = Logger.get(Oauth2WebUiAuthenticationManager.class); + private static final Logger logger = Logger.get(OAuth2WebUiAuthenticationManager.class); private final OAuth2Service oAuth2Service; private final OAuth2Authenticator oAuth2Authenticator; private final TokenPairSerializer tokenPairSerializer; @@ -48,7 +49,7 @@ public class Oauth2WebUiAuthenticationManager private final Optional tokenExpiration; @Inject - public Oauth2WebUiAuthenticationManager(OAuth2Service oAuth2Service, OAuth2Authenticator oAuth2Authenticator, TokenPairSerializer tokenPairSerializer, OAuth2Client client, @ForRefreshTokens Optional tokenExpiration) + public OAuth2WebUiAuthenticationManager(OAuth2Service oAuth2Service, OAuth2Authenticator oAuth2Authenticator, TokenPairSerializer tokenPairSerializer, OAuth2Client client, @ForRefreshTokens Optional tokenExpiration) { this.oAuth2Service = requireNonNull(oAuth2Service, "oauth2Service is null"); this.oAuth2Authenticator = requireNonNull(oAuth2Authenticator, "oauth2Authenticator is null"); @@ -72,7 +73,7 @@ public void handleRequest(HttpServletRequest request, HttpServletResponse respon private Optional getTokenPair(HttpServletRequest request) { try { - Optional token = this.oAuth2Authenticator.extractTokenFromCookie(request); + Optional token = extractTokenFromCookie(request); if (token.isPresent()) { return Optional.ofNullable(tokenPairSerializer.deserialize(token.get())); } diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuthWebUiCookie.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuthWebUiCookie.java index 477c0d0d3581d..867e7b60c3dd8 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuthWebUiCookie.java +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OAuthWebUiCookie.java @@ -13,13 +13,13 @@ */ package com.facebook.presto.server.security.oauth2; -import javax.ws.rs.core.NewCookie; +import jakarta.ws.rs.core.NewCookie; import java.time.Instant; import java.util.Date; -import static javax.ws.rs.core.Cookie.DEFAULT_VERSION; -import static javax.ws.rs.core.NewCookie.DEFAULT_MAX_AGE; +import static jakarta.ws.rs.core.Cookie.DEFAULT_VERSION; +import static jakarta.ws.rs.core.NewCookie.DEFAULT_MAX_AGE; public final class OAuthWebUiCookie { diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OidcDiscoveryConfig.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OidcDiscoveryConfig.java index 126a5308b7f94..d5d29589fc5a4 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OidcDiscoveryConfig.java +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/OidcDiscoveryConfig.java @@ -15,9 +15,8 @@ import com.facebook.airlift.configuration.Config; import com.facebook.airlift.configuration.ConfigDescription; -import io.airlift.units.Duration; - -import javax.validation.constraints.NotNull; +import com.facebook.airlift.units.Duration; +import jakarta.validation.constraints.NotNull; import java.util.Optional; diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/RefreshTokensConfig.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/RefreshTokensConfig.java index 69bcd717e7451..7d25c142f0332 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/RefreshTokensConfig.java +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/RefreshTokensConfig.java @@ -16,12 +16,12 @@ import com.facebook.airlift.configuration.Config; import com.facebook.airlift.configuration.ConfigDescription; import com.facebook.airlift.configuration.ConfigSecuritySensitive; -import io.airlift.units.Duration; +import com.facebook.airlift.units.Duration; import io.jsonwebtoken.io.Decoders; import io.jsonwebtoken.security.Keys; +import jakarta.validation.constraints.NotEmpty; import javax.crypto.SecretKey; -import javax.validation.constraints.NotEmpty; import static com.google.common.base.Strings.isNullOrEmpty; import static java.util.concurrent.TimeUnit.HOURS; diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/StaticOAuth2ServerConfiguration.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/StaticOAuth2ServerConfiguration.java index 2c497f05a27cb..51f116abe09b5 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/StaticOAuth2ServerConfiguration.java +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/StaticOAuth2ServerConfiguration.java @@ -15,8 +15,7 @@ import com.facebook.airlift.configuration.Config; import com.facebook.airlift.configuration.ConfigDescription; - -import javax.validation.constraints.NotNull; +import jakarta.validation.constraints.NotNull; import java.util.Optional; diff --git a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/TokenPairSerializer.java b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/TokenPairSerializer.java index 715374cc8e61d..f178a4048adc8 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/TokenPairSerializer.java +++ b/presto-main/src/main/java/com/facebook/presto/server/security/oauth2/TokenPairSerializer.java @@ -15,8 +15,7 @@ package com.facebook.presto.server.security.oauth2; import com.facebook.presto.server.security.oauth2.OAuth2Client.Response; - -import javax.annotation.Nullable; +import jakarta.annotation.Nullable; import java.util.Date; import java.util.Optional; diff --git a/presto-main/src/test/java/com/facebook/presto/server/MockHttpServletRequest.java b/presto-main/src/test/java/com/facebook/presto/server/MockHttpServletRequest.java index 0c85c083070e9..2033b62f88262 100644 --- a/presto-main/src/test/java/com/facebook/presto/server/MockHttpServletRequest.java +++ b/presto-main/src/test/java/com/facebook/presto/server/MockHttpServletRequest.java @@ -30,7 +30,7 @@ import jakarta.servlet.http.HttpSession; import jakarta.servlet.http.HttpUpgradeHandler; import jakarta.servlet.http.Part; -import jakarts.ws.rs.core.UriBuilder; +import jakarta.ws.rs.core.UriBuilder; import java.io.BufferedReader; import java.net.URI; diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/BaseOAuth2AuthenticationFilterTest.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/BaseOAuth2AuthenticationFilterTest.java index 2f3b6a41053ab..4540c655cb655 100644 --- a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/BaseOAuth2AuthenticationFilterTest.java +++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/BaseOAuth2AuthenticationFilterTest.java @@ -49,19 +49,19 @@ import java.util.UUID; import static com.facebook.airlift.testing.Assertions.assertLessThan; +import static com.facebook.airlift.units.Duration.nanosSince; import static com.facebook.presto.client.OkHttpUtil.setupInsecureSsl; import static com.facebook.presto.server.security.oauth2.JwtUtil.newJwtBuilder; import static com.facebook.presto.server.security.oauth2.OAuthWebUiCookie.OAUTH2_COOKIE; import static com.facebook.presto.server.security.oauth2.TokenEndpointAuthMethod.CLIENT_SECRET_BASIC; -import static io.airlift.units.Duration.nanosSince; +import static jakarta.servlet.http.HttpServletResponse.SC_OK; +import static jakarta.ws.rs.core.HttpHeaders.AUTHORIZATION; +import static jakarta.ws.rs.core.HttpHeaders.LOCATION; +import static jakarta.ws.rs.core.Response.Status.OK; +import static jakarta.ws.rs.core.Response.Status.SEE_OTHER; +import static jakarta.ws.rs.core.Response.Status.UNAUTHORIZED; import static java.util.concurrent.TimeUnit.MILLISECONDS; import static java.util.concurrent.TimeUnit.SECONDS; -import static javax.servlet.http.HttpServletResponse.SC_OK; -import static javax.ws.rs.core.HttpHeaders.AUTHORIZATION; -import static javax.ws.rs.core.HttpHeaders.LOCATION; -import static javax.ws.rs.core.Response.Status.OK; -import static javax.ws.rs.core.Response.Status.SEE_OTHER; -import static javax.ws.rs.core.Response.Status.UNAUTHORIZED; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.fail; import static org.testng.Assert.assertEquals; @@ -104,7 +104,7 @@ static void waitForNodeRefresh(TestingPrestoServer server) { long start = System.nanoTime(); while (server.refreshNodes().getActiveNodes().size() < 1) { - assertLessThan(nanosSince(start), new io.airlift.units.Duration(10, SECONDS)); + assertLessThan(nanosSince(start), new com.facebook.airlift.units.Duration(10, SECONDS)); MILLISECONDS.sleep(10); } } @@ -123,7 +123,7 @@ public void setup() // Due to problems with the Presto OSS project related to the AuthenticationFilter we have to run Presto behind a Proxy and terminate SSL at the proxy. simpleProxy = new SimpleProxyServer(server.getHttpBaseUrl()); MILLISECONDS.sleep(1000); - proxyURI = simpleProxy.getHttpsBaseUrl(); + proxyURI = URI.create("https://127.0.0.1:" + simpleProxy.getHttpsBaseUrl().getPort()); uiUri = proxyURI.resolve("/"); hydraIdP.createClient( @@ -131,13 +131,13 @@ public void setup() PRESTO_CLIENT_SECRET, CLIENT_SECRET_BASIC, ImmutableList.of(PRESTO_AUDIENCE, ADDITIONAL_AUDIENCE), - simpleProxy.getHttpsBaseUrl() + "/oauth2/callback"); + proxyURI + "/oauth2/callback"); hydraIdP.createClient( TRUSTED_CLIENT_ID, TRUSTED_CLIENT_SECRET, CLIENT_SECRET_BASIC, ImmutableList.of(TRUSTED_CLIENT_ID), - simpleProxy.getHttpsBaseUrl() + "/oauth2/callback"); + proxyURI + "/oauth2/callback"); hydraIdP.createClient( UNTRUSTED_CLIENT_ID, UNTRUSTED_CLIENT_SECRET, @@ -146,7 +146,8 @@ public void setup() "https://untrusted.com/callback"); } - protected abstract ImmutableMap getOAuth2Config(String idpUrl); + protected abstract ImmutableMap getOAuth2Config(String idpUrl) + throws IOException; protected abstract TestingHydraIdentityProvider getHydraIdp() throws Exception; @@ -312,6 +313,7 @@ private void assertOAuth2Cookie(HttpCookie cookie) assertThat(cookie.getMaxAge()).isLessThanOrEqualTo(TTL_ACCESS_TOKEN_IN_SECONDS.getSeconds()); validateAccessToken(cookie.getValue()); } + protected void validateAccessToken(String cookieValue) { Request request = new Request.Builder().url("https://localhost:" + hydraIdP.getAuthPort() + "/userinfo").addHeader(AUTHORIZATION, "Bearer " + cookieValue).build(); @@ -352,13 +354,14 @@ public void saveFromResponse(HttpUrl url, List cookies) @Override public List loadForRequest(HttpUrl url) { - return ImmutableList.of(new Cookie.Builder() + Cookie cookie = new Cookie.Builder() .domain(proxyURI.getHost()) .path("/") .name(OAUTH2_COOKIE) .value(cookieValue) .secure() - .build()); + .build(); + return ImmutableList.of(cookie); } }); return httpClientBuilder.build(); diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/SimpleProxyServer.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/SimpleProxyServer.java index 36e9921930fca..99e8bb43c4b8d 100644 --- a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/SimpleProxyServer.java +++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/SimpleProxyServer.java @@ -20,18 +20,17 @@ import com.facebook.airlift.node.NodeInfo; import com.google.common.collect.ImmutableMap; import com.google.common.io.Resources; +import jakarta.servlet.ServletException; +import jakarta.servlet.http.Cookie; +import jakarta.servlet.http.HttpServlet; +import jakarta.servlet.http.HttpServletRequest; +import jakarta.servlet.http.HttpServletResponse; +import jakarta.ws.rs.core.UriBuilder; import okhttp3.Headers; import okhttp3.OkHttpClient; import okhttp3.Request; import okhttp3.Response; -import javax.servlet.ServletException; -import javax.servlet.http.Cookie; -import javax.servlet.http.HttpServlet; -import javax.servlet.http.HttpServletRequest; -import javax.servlet.http.HttpServletResponse; -import javax.ws.rs.core.UriBuilder; - import java.io.Closeable; import java.io.IOException; import java.io.InputStream; diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestDualAuthenticationFilterWithOAuth.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestDualAuthenticationFilterWithOAuth.java index 465a930dc30ae..089c6b9c80aaf 100644 --- a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestDualAuthenticationFilterWithOAuth.java +++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestDualAuthenticationFilterWithOAuth.java @@ -37,16 +37,17 @@ import java.util.List; import static com.facebook.airlift.testing.Assertions.assertLessThan; +import static com.facebook.airlift.units.Duration.nanosSince; import static com.facebook.presto.client.OkHttpUtil.setupInsecureSsl; import static com.facebook.presto.server.security.oauth2.OAuthWebUiCookie.OAUTH2_COOKIE; import static com.facebook.presto.server.security.oauth2.TokenEndpointAuthMethod.CLIENT_SECRET_BASIC; import static com.google.common.net.HttpHeaders.AUTHORIZATION; import static com.google.common.net.HttpHeaders.WWW_AUTHENTICATE; -import static io.airlift.units.Duration.nanosSince; +import static jakarta.ws.rs.core.Response.Status.OK; +import static jakarta.ws.rs.core.Response.Status.UNAUTHORIZED; +import static java.io.File.createTempFile; import static java.util.concurrent.TimeUnit.MILLISECONDS; import static java.util.concurrent.TimeUnit.SECONDS; -import static javax.ws.rs.core.Response.Status.OK; -import static javax.ws.rs.core.Response.Status.UNAUTHORIZED; import static org.assertj.core.api.Assertions.assertThat; public class TestDualAuthenticationFilterWithOAuth @@ -57,10 +58,6 @@ public class TestDualAuthenticationFilterWithOAuth private static final String PRESTO_AUDIENCE = PRESTO_CLIENT_ID; private static final String ADDITIONAL_AUDIENCE = "https://external-service.com"; protected static final String TRUSTED_CLIENT_ID = "trusted-client"; - protected static final String TRUSTED_CLIENT_SECRET = "trusted-secret"; - private static final String UNTRUSTED_CLIENT_ID = "untrusted-client"; - private static final String UNTRUSTED_CLIENT_SECRET = "untrusted-secret"; - private static final String UNTRUSTED_CLIENT_AUDIENCE = "https://untrusted.com"; private final Logging logging = Logging.initialize(); protected final OkHttpClient httpClient; @@ -69,7 +66,6 @@ public class TestDualAuthenticationFilterWithOAuth private TestingPrestoServer server; private SimpleProxyServer simpleProxy; - private URI uiUri; private URI proxyURI; @@ -86,12 +82,13 @@ static void waitForNodeRefresh(TestingPrestoServer server) { long start = System.nanoTime(); while (server.refreshNodes().getActiveNodes().size() < 1) { - assertLessThan(nanosSince(start), new io.airlift.units.Duration(10, SECONDS)); + assertLessThan(nanosSince(start), new com.facebook.airlift.units.Duration(10, SECONDS)); MILLISECONDS.sleep(10); } } protected ImmutableMap getConfig(String idpUrl) + throws IOException { return ImmutableMap.builder() .put("http-server.authentication.allow-forwarded-https", "true") @@ -106,6 +103,7 @@ protected ImmutableMap getConfig(String idpUrl) .put("http-server.authentication.oauth2.max-clock-skew", "0s") .put("http-server.authentication.oauth2.user-mapping.pattern", "(.*)(@.*)?") .put("http-server.authentication.oauth2.oidc.discovery", "false") + .put("configuration-based-authorizer.role-regex-map.file-path", createTempFile("regex-map", null).getAbsolutePath().toString()) .put("oauth2-jwk.http-client.trust-store-path", Resources.getResource("cert/localhost.pem").getPath()) .build(); } @@ -124,8 +122,7 @@ public void setup() // Due to problems with the Presto OSS project related to the AuthenticationFilter we have to run Presto behind a Proxy and terminate SSL at the proxy. simpleProxy = new SimpleProxyServer(server.getHttpBaseUrl()); MILLISECONDS.sleep(1000); - proxyURI = simpleProxy.getHttpsBaseUrl(); - uiUri = proxyURI.resolve("/"); + proxyURI = URI.create("https://127.0.0.1:" + simpleProxy.getHttpsBaseUrl().getPort()); hydraIdP.createClient( PRESTO_CLIENT_ID, diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestJweTokenSerializer.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestJweTokenSerializer.java index 37dae99191556..7a41c0fdaeeea 100644 --- a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestJweTokenSerializer.java +++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestJweTokenSerializer.java @@ -13,9 +13,9 @@ */ package com.facebook.presto.server.security.oauth2; +import com.facebook.airlift.units.Duration; import com.facebook.presto.server.security.oauth2.TokenPairSerializer.TokenPair; import com.nimbusds.jose.KeyLengthException; -import io.airlift.units.Duration; import io.jsonwebtoken.ExpiredJwtException; import io.jsonwebtoken.Jwts; import org.testng.annotations.Test; @@ -31,8 +31,8 @@ import java.util.Map; import java.util.Optional; +import static com.facebook.airlift.units.Duration.succinctDuration; import static com.facebook.presto.server.security.oauth2.TokenPairSerializer.TokenPair.accessAndRefreshTokens; -import static io.airlift.units.Duration.succinctDuration; import static java.time.temporal.ChronoUnit.MILLIS; import static java.util.concurrent.TimeUnit.MINUTES; import static java.util.concurrent.TimeUnit.SECONDS; diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2AuthenticationFilterWithJwt.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2AuthenticationFilterWithJwt.java index a4eb2f6069038..71806b590bc00 100644 --- a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2AuthenticationFilterWithJwt.java +++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2AuthenticationFilterWithJwt.java @@ -16,11 +16,16 @@ import com.google.common.collect.ImmutableMap; import com.google.common.io.Resources; +import java.io.IOException; + +import static java.io.File.createTempFile; + public class TestOAuth2AuthenticationFilterWithJwt extends BaseOAuth2AuthenticationFilterTest { @Override protected ImmutableMap getOAuth2Config(String idpUrl) + throws IOException { return ImmutableMap.builder() .put("http-server.authentication.allow-forwarded-https", "true") @@ -35,6 +40,7 @@ protected ImmutableMap getOAuth2Config(String idpUrl) .put("http-server.authentication.oauth2.max-clock-skew", "0s") .put("http-server.authentication.oauth2.user-mapping.pattern", "(.*)(@.*)?") .put("http-server.authentication.oauth2.oidc.discovery", "false") + .put("configuration-based-authorizer.role-regex-map.file-path", createTempFile("regex-map", null).getAbsolutePath().toString()) .put("oauth2-jwk.http-client.trust-store-path", Resources.getResource("cert/localhost.pem").getPath()) .build(); } diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2AuthenticationFilterWithOpaque.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2AuthenticationFilterWithOpaque.java index b36453cccfef5..114b9b0ed20d1 100644 --- a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2AuthenticationFilterWithOpaque.java +++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2AuthenticationFilterWithOpaque.java @@ -16,11 +16,16 @@ import com.google.common.collect.ImmutableMap; import com.google.common.io.Resources; +import java.io.IOException; + +import static java.io.File.createTempFile; + public class TestOAuth2AuthenticationFilterWithOpaque extends BaseOAuth2AuthenticationFilterTest { @Override protected ImmutableMap getOAuth2Config(String idpUrl) + throws IOException { return ImmutableMap.builder() .put("http-server.authentication.allow-forwarded-https", "true") @@ -38,6 +43,7 @@ protected ImmutableMap getOAuth2Config(String idpUrl) .put("http-server.authentication.oauth2.max-clock-skew", "0s") .put("http-server.authentication.oauth2.user-mapping.pattern", "(.*)(@.*)?") .put("http-server.authentication.oauth2.oidc.discovery", "false") + .put("configuration-based-authorizer.role-regex-map.file-path", createTempFile("regex-map", null).getAbsolutePath().toString()) .put("oauth2-jwk.http-client.trust-store-path", Resources.getResource("cert/localhost.pem").getPath()) .build(); } diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2Config.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2Config.java index e8dca428ab98b..bb2735b95e7f6 100644 --- a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2Config.java +++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2Config.java @@ -13,8 +13,8 @@ */ package com.facebook.presto.server.security.oauth2; +import com.facebook.airlift.units.Duration; import com.google.common.collect.ImmutableMap; -import io.airlift.units.Duration; import org.testng.annotations.Test; import java.io.IOException; diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2Utils.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2Utils.java index 0d21db65f4ec7..119fd7213bd96 100644 --- a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2Utils.java +++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOAuth2Utils.java @@ -15,11 +15,10 @@ import com.facebook.presto.server.MockHttpServletRequest; import com.google.common.collect.ImmutableListMultimap; +import jakarta.servlet.http.HttpServletRequest; +import jakarta.ws.rs.core.UriBuilder; import org.testng.annotations.Test; -import javax.servlet.http.HttpServletRequest; -import javax.ws.rs.core.UriBuilder; - import static com.facebook.presto.server.security.oauth2.OAuth2Utils.getSchemeUriBuilder; import static com.google.common.net.HttpHeaders.X_FORWARDED_PROTO; import static org.testng.Assert.assertEquals; diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOidcDiscovery.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOidcDiscovery.java index d6e65f9d73c4a..679095b9c7cd4 100644 --- a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOidcDiscovery.java +++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOidcDiscovery.java @@ -24,13 +24,12 @@ import com.google.common.io.Resources; import com.google.inject.Key; import com.google.inject.TypeLiteral; +import jakarta.servlet.http.HttpServlet; +import jakarta.servlet.http.HttpServletRequest; +import jakarta.servlet.http.HttpServletResponse; import org.testng.annotations.DataProvider; import org.testng.annotations.Test; -import javax.servlet.http.HttpServlet; -import javax.servlet.http.HttpServletRequest; -import javax.servlet.http.HttpServletResponse; - import java.io.IOException; import java.net.URI; import java.time.Instant; @@ -41,10 +40,11 @@ import static com.facebook.airlift.http.client.HttpStatus.TOO_MANY_REQUESTS; import static com.facebook.presto.server.security.oauth2.BaseOAuth2AuthenticationFilterTest.PRESTO_CLIENT_ID; import static com.facebook.presto.server.security.oauth2.BaseOAuth2AuthenticationFilterTest.PRESTO_CLIENT_SECRET; +import static jakarta.ws.rs.core.HttpHeaders.CONTENT_TYPE; +import static jakarta.ws.rs.core.MediaType.APPLICATION_JSON; +import static java.io.File.createTempFile; import static java.nio.charset.StandardCharsets.UTF_8; import static java.util.Objects.requireNonNull; -import static javax.ws.rs.core.HttpHeaders.CONTENT_TYPE; -import static javax.ws.rs.core.MediaType.APPLICATION_JSON; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; @@ -63,7 +63,8 @@ public void testStaticConfiguration(Optional accessTokenPath, Optional properties.put("http-server.authentication.oauth2.access-token-issuer", uri)); userinfoUrl.map(URI::toString).ifPresent(uri -> properties.put("http-server.authentication.oauth2.userinfo-url", uri)); try (TestingPrestoServer server = createServer(properties.build())) { @@ -94,6 +95,7 @@ public void testOidcDiscovery(String configuration, Optional accessToken ImmutableMap.builder() .put("http-server.authentication.oauth2.issuer", metadataServer.getBaseUrl().toString()) .put("http-server.authentication.oauth2.oidc.discovery", "true") + .put("configuration-based-authorizer.role-regex-map.file-path", createTempFile("regex-map", null).getAbsolutePath().toString()) .build())) { URI issuer = metadataServer.getBaseUrl(); assertConfiguration(server, issuer, accessTokenIssuer.map(issuer::resolve), userinfoUrl.map(issuer::resolve)); @@ -123,6 +125,7 @@ public void testIssuerCheck() ImmutableMap.builder() .put("http-server.authentication.oauth2.issuer", metadataServer.getBaseUrl().toString()) .put("http-server.authentication.oauth2.oidc.discovery", "true") + .put("configuration-based-authorizer.role-regex-map.file-path", createTempFile("regex-map", null).getAbsolutePath().toString()) .build())) { // should throw an exception server.getInstance(Key.get(OAuth2ServerConfigProvider.class)).get(); @@ -141,6 +144,7 @@ public void testStopOnClientError() ImmutableMap.builder() .put("http-server.authentication.oauth2.issuer", metadataServer.getBaseUrl().toString()) .put("http-server.authentication.oauth2.oidc.discovery", "true") + .put("configuration-based-authorizer.role-regex-map.file-path", createTempFile("regex-map", null).getAbsolutePath().toString()) .build())) { // should throw an exception server.getInstance(Key.get(OAuth2ServerConfigProvider.class)).get(); @@ -162,6 +166,7 @@ public void testOidcDiscoveryRetrying() .put("http-server.authentication.oauth2.issuer", metadataServer.getBaseUrl().toString()) .put("http-server.authentication.oauth2.oidc.discovery", "true") .put("http-server.authentication.oauth2.oidc.discovery.timeout", "10s") + .put("configuration-based-authorizer.role-regex-map.file-path", createTempFile("regex-map", null).getAbsolutePath().toString()) .build())) { URI issuer = metadataServer.getBaseUrl(); assertConfiguration(server, issuer, Optional.empty(), Optional.of(issuer.resolve("/connect/userinfo"))); @@ -182,6 +187,7 @@ public void testOidcDiscoveryTimesOut() .put("http-server.authentication.oauth2.issuer", metadataServer.getBaseUrl().toString()) .put("http-server.authentication.oauth2.oidc.discovery", "true") .put("http-server.authentication.oauth2.oidc.discovery.timeout", "5s") + .put("configuration-based-authorizer.role-regex-map.file-path", createTempFile("regex-map", null).getAbsolutePath().toString()) .build())) { // should throw an exception server.getInstance(Key.get(OAuth2ServerConfigProvider.class)).get(); @@ -203,6 +209,7 @@ public void testIgnoringUserinfoUrl() .put("http-server.authentication.oauth2.issuer", metadataServer.getBaseUrl().toString()) .put("http-server.authentication.oauth2.oidc.discovery", "true") .put("http-server.authentication.oauth2.oidc.use-userinfo-endpoint", "false") + .put("configuration-based-authorizer.role-regex-map.file-path", createTempFile("regex-map", null).getAbsolutePath().toString()) .build())) { URI issuer = metadataServer.getBaseUrl(); assertConfiguration(server, issuer, Optional.empty(), Optional.empty()); @@ -233,6 +240,7 @@ public void testBackwardCompatibility() .put("http-server.authentication.oauth2.jwks-url", jwksUrl.toString()) .put("http-server.authentication.oauth2.access-token-issuer", accessTokenIssuer) .put("http-server.authentication.oauth2.userinfo-url", userinfoUrl.toString()) + .put("configuration-based-authorizer.role-regex-map.file-path", createTempFile("regex-map", null).getAbsolutePath().toString()) .build())) { assertComponents(server); OAuth2ServerConfig config = server.getInstance(Key.get(OAuth2ServerConfigProvider.class)).get(); diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOidcDiscoveryConfig.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOidcDiscoveryConfig.java index 5144b3f934ea9..37c4dd317924b 100644 --- a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOidcDiscoveryConfig.java +++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestOidcDiscoveryConfig.java @@ -13,8 +13,8 @@ */ package com.facebook.presto.server.security.oauth2; +import com.facebook.airlift.units.Duration; import com.google.common.collect.ImmutableMap; -import io.airlift.units.Duration; import org.testng.annotations.Test; import java.util.Map; diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestRefreshTokensConfig.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestRefreshTokensConfig.java index d365a513573de..29c7d5b11c603 100644 --- a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestRefreshTokensConfig.java +++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestRefreshTokensConfig.java @@ -25,7 +25,7 @@ import static com.facebook.airlift.configuration.testing.ConfigAssertions.assertFullMapping; import static com.facebook.airlift.configuration.testing.ConfigAssertions.assertRecordedDefaults; import static com.facebook.airlift.configuration.testing.ConfigAssertions.recordDefaults; -import static io.airlift.units.Duration.succinctDuration; +import static com.facebook.airlift.units.Duration.succinctDuration; import static io.jsonwebtoken.io.Encoders.BASE64; import static java.util.concurrent.TimeUnit.HOURS; diff --git a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestingHydraIdentityProvider.java b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestingHydraIdentityProvider.java index ee2cb16482900..0daa570a29842 100644 --- a/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestingHydraIdentityProvider.java +++ b/presto-main/src/test/java/com/facebook/presto/server/security/oauth2/TestingHydraIdentityProvider.java @@ -29,6 +29,10 @@ import com.google.common.io.Resources; import com.google.inject.Key; import com.nimbusds.oauth2.sdk.GrantType; +import jakarta.servlet.http.HttpServlet; +import jakarta.servlet.http.HttpServletRequest; +import jakarta.servlet.http.HttpServletResponse; +import jakarta.ws.rs.core.HttpHeaders; import okhttp3.Credentials; import okhttp3.FormBody; import okhttp3.HttpUrl; @@ -46,11 +50,6 @@ import org.testcontainers.containers.wait.strategy.WaitAllStrategy; import org.testcontainers.utility.MountableFile; -import javax.servlet.http.HttpServlet; -import javax.servlet.http.HttpServletRequest; -import javax.servlet.http.HttpServletResponse; -import javax.ws.rs.core.HttpHeaders; - import java.io.Closeable; import java.io.IOException; import java.net.URI; @@ -62,10 +61,10 @@ import static com.facebook.presto.server.security.oauth2.TokenEndpointAuthMethod.CLIENT_SECRET_BASIC; import static com.google.common.base.Preconditions.checkState; import static com.google.common.base.Throwables.throwIfUnchecked; +import static jakarta.servlet.http.HttpServletResponse.SC_NOT_FOUND; +import static jakarta.servlet.http.HttpServletResponse.SC_OK; +import static jakarta.ws.rs.core.MediaType.APPLICATION_JSON; import static java.util.Objects.requireNonNull; -import static javax.servlet.http.HttpServletResponse.SC_NOT_FOUND; -import static javax.servlet.http.HttpServletResponse.SC_OK; -import static javax.ws.rs.core.MediaType.APPLICATION_JSON; public class TestingHydraIdentityProvider implements Closeable @@ -229,7 +228,14 @@ private TestingHttpServer createTestingLoginAndConsentServer() NodeInfo nodeInfo = new NodeInfo("test"); HttpServerConfig config = new HttpServerConfig().setHttpPort(0); HttpServerInfo httpServerInfo = new HttpServerInfo(config, nodeInfo); - return new TestingHttpServer(httpServerInfo, nodeInfo, config, new AcceptAllLoginsAndConsentsServlet(), ImmutableMap.of(), ImmutableMap.of(), Optional.empty()); + return new TestingHttpServer( + httpServerInfo, + nodeInfo, + config, + new AcceptAllLoginsAndConsentsServlet(), + ImmutableMap.of(), + ImmutableMap.of(), + Optional.empty()); } private class AcceptAllLoginsAndConsentsServlet diff --git a/presto-native-execution/pom.xml b/presto-native-execution/pom.xml index 3dab4aa3b4e23..f9150f63a4daf 100644 --- a/presto-native-execution/pom.xml +++ b/presto-native-execution/pom.xml @@ -100,12 +100,7 @@ com.facebook.presto presto-main - - - org.apache.commons - commons-lang3 - - + test diff --git a/presto-router/src/main/java/com/facebook/presto/router/RouterModule.java b/presto-router/src/main/java/com/facebook/presto/router/RouterModule.java index efc7cc783ddb7..168668fe6ed3f 100755 --- a/presto-router/src/main/java/com/facebook/presto/router/RouterModule.java +++ b/presto-router/src/main/java/com/facebook/presto/router/RouterModule.java @@ -15,6 +15,7 @@ import com.facebook.airlift.configuration.AbstractConfigurationAwareModule; import com.facebook.airlift.units.Duration; +import com.facebook.presto.ClientRequestFilterManager; import com.facebook.presto.client.NodeVersion; import com.facebook.presto.router.cluster.ClusterManager; import com.facebook.presto.router.cluster.ClusterManager.ClusterStatusTracker; @@ -70,6 +71,7 @@ protected void setup(Binder binder) { ServerConfig serverConfig = buildConfigObject(ServerConfig.class); + binder.bind(ClientRequestFilterManager.class).in(Scopes.SINGLETON); binder.bind(RouterPluginManager.class).in(Scopes.SINGLETON); configBinder(binder).bindConfig(RouterConfig.class); diff --git a/presto-main/src/main/resources/oauth2/failure.html b/presto-ui/src/static/oauth2/failure.html similarity index 100% rename from presto-main/src/main/resources/oauth2/failure.html rename to presto-ui/src/static/oauth2/failure.html diff --git a/presto-ui/src/static/ouath2/logout.html b/presto-ui/src/static/oauth2/logout.html similarity index 100% rename from presto-ui/src/static/ouath2/logout.html rename to presto-ui/src/static/oauth2/logout.html diff --git a/presto-main/src/main/resources/oauth2/success.html b/presto-ui/src/static/oauth2/success.html similarity index 100% rename from presto-main/src/main/resources/oauth2/success.html rename to presto-ui/src/static/oauth2/success.html From 20ed52bb11b202795ab8a8911782ed9a9301fe04 Mon Sep 17 00:00:00 2001 From: auden-woolfson Date: Tue, 1 Jul 2025 11:46:13 -0700 Subject: [PATCH 109/113] Update security docs to include oauth2 --- presto-docs/src/main/sphinx/security.rst | 1 + .../main/sphinx/security/authorization.rst | 6 +- .../src/main/sphinx/security/oauth2.rst | 69 +++++++++++++++++-- 3 files changed, 68 insertions(+), 8 deletions(-) diff --git a/presto-docs/src/main/sphinx/security.rst b/presto-docs/src/main/sphinx/security.rst index f8c5432555e97..6d793c4c785e6 100644 --- a/presto-docs/src/main/sphinx/security.rst +++ b/presto-docs/src/main/sphinx/security.rst @@ -13,3 +13,4 @@ Security security/built-in-system-access-control security/internal-communication security/authorization + security/oauth2 diff --git a/presto-docs/src/main/sphinx/security/authorization.rst b/presto-docs/src/main/sphinx/security/authorization.rst index a354c72ed0401..82649077ba3e3 100644 --- a/presto-docs/src/main/sphinx/security/authorization.rst +++ b/presto-docs/src/main/sphinx/security/authorization.rst @@ -45,9 +45,9 @@ so make sure you have authentication enabled. http-server.authentication.type=CERTIFICATE - - It is also possible to specify other authentication types such as - ``KERBEROS``, ``PASSWORD``, ``JWT``, and ``OAUTH2``. Additional configuration may be - needed. +- It is also possible to specify other authentication types such as + ``KERBEROS``, ``PASSWORD``, ``JWT``, and ``OAUTH2``. Additional configuration may be + needed. .. code-block:: none diff --git a/presto-docs/src/main/sphinx/security/oauth2.rst b/presto-docs/src/main/sphinx/security/oauth2.rst index 0858ebc23ed7b..65b676c42d394 100644 --- a/presto-docs/src/main/sphinx/security/oauth2.rst +++ b/presto-docs/src/main/sphinx/security/oauth2.rst @@ -1,9 +1,68 @@ -============= +======================== Oauth 2.0 Authentication -============= +======================== -Presto server configuration ---------------------------- +Presto can be configured to enable frontend OAuth2 authentication over HTTPS for clients such as the CLI, JDBC, and ODBC drivers. OAuth2 provides a secure and flexible way to authenticate users by using an external identity provider (IdP), such as Okta, Auth0, Azure AD, or Google. + +OAuth2 authentication in Presto uses the Authorization Code Flow with PKCE and OpenID Connect (OIDC). The Presto coordinator initiates an OAuth2 challenge, and the client completes the flow by obtaining an access token from the identity provider. -Refresh Tokens +Presto Server Configuration --------------------------- + +To enable OAuth2 authentication, configuration changes are made **only on the Presto coordinator**. No changes are required on the workers. + +Secure Communication +-------------------- + +Access to the Presto coordinator must be secured with HTTPS. You must configure a valid TLS certificate and keystore on the coordinator. See the `TLS setup guide `_ for details. + +OAuth2 Configuration +-------------------- + +Below are the key configuration properties for enabling OAuth2 authentication in ``config.properties``: + +.. code-block:: properties + + http-server.authentication.type=OAUTH2 + + http-server.authentication.oauth2.issuer=https://your-idp.com/oauth2/default + http-server.authentication.oauth2.client-id=your-client-id + http-server.authentication.oauth2.client-secret=your-client-secret + http-server.authentication.oauth2.scopes=openid,email,profile + http-server.authentication.oauth2.principal-field=sub + http-server.authentication.oauth2.groups-field=groups + http-server.authentication.oauth2.challenge-timeout=15m + http-server.authentication.oauth2.max-clock-skew=1m + http-server.authentication.oauth2.refresh-tokens=true + http-server.authentication.oauth2.oidc.discovery=true + http-server.authentication.oauth2.state-key=your-hmac-secret + http-server.authentication.oauth2.additional-audiences=your-client-id,another-audience + http-server.authentication.oauth2.user-mapping.pattern=(.*) + +It is worth noting that ``configuration-based-authorizer.role-regex-map.file-path`` must be configured if +authentication type is set to ``OAUTH2``. + +TLS Truststore for IdP +---------------------- + +If your IdP uses a custom or self-signed certificate, import it into the Java truststore on the Presto coordinator: + +.. code-block:: bash + + keytool -import \ + -keystore $JAVA_HOME/lib/security/cacerts \ + -trustcacerts \ + -alias idp_cert \ + -file idp_cert.crt + +Notes +----- + +- **Issuer**: The base URL of your IdP’s OIDC discovery endpoint. +- **Client ID/Secret**: Registered credentials for Presto in your IdP. +- **Scopes**: Must include ``openid``; others like ``email``, ``profile``, or ``groups`` are optional. +- **Principal Field**: The claim in the ID token used as the Presto username. +- **Groups Field**: Optional claim used for role-based access control. +- **State Key**: A secret used to sign the OAuth2 state parameter (HMAC). +- **Refresh Tokens**: Enable if your IdP supports issuing refresh tokens. +- **Callback**: When configuring your IdP the callback URI must be set to ``[presto]/oauth2/callback`` From 4a1d83f7e90f977a7f9e3f6e29b5063ce78e44a0 Mon Sep 17 00:00:00 2001 From: Feilong Liu Date: Sun, 17 Aug 2025 23:38:37 -0700 Subject: [PATCH 110/113] Add plan support for using uniqueness of row_id --- .../facebook/presto/hive/HiveMetadata.java | 4 +- .../hive/TestHiveBucketedTablesWithRowId.java | 240 +++++++++++++ .../presto/hive/TestHiveLogicalPlanner.java | 337 ++++++++++++++++++ .../presto/SystemSessionProperties.java | 10 + .../facebook/presto/metadata/TableLayout.java | 5 + .../presto/sql/analyzer/FeaturesConfig.java | 14 + .../optimizations/ActualProperties.java | 91 ++++- .../planner/optimizations/AddExchanges.java | 13 +- .../optimizations/AddLocalExchanges.java | 3 +- .../optimizations/LocalProperties.java | 6 + .../optimizations/PropertyDerivations.java | 67 +++- .../StreamPreferredProperties.java | 6 +- .../StreamPropertyDerivations.java | 139 ++++++-- .../sanity/ValidateStreamingAggregations.java | 4 +- .../sql/analyzer/TestFeaturesConfig.java | 3 + .../planner/assertions/PlanMatchPattern.java | 5 + .../optimizations/TestLocalProperties.java | 53 +++ .../presto/spi/ConnectorTableLayout.java | 21 ++ .../facebook/presto/spi/UniqueProperty.java | 98 +++++ 19 files changed, 1081 insertions(+), 38 deletions(-) create mode 100644 presto-hive/src/test/java/com/facebook/presto/hive/TestHiveBucketedTablesWithRowId.java create mode 100644 presto-spi/src/main/java/com/facebook/presto/spi/UniqueProperty.java diff --git a/presto-hive/src/main/java/com/facebook/presto/hive/HiveMetadata.java b/presto-hive/src/main/java/com/facebook/presto/hive/HiveMetadata.java index 1efa0d67f3aaa..d8940cb5e0b9c 100644 --- a/presto-hive/src/main/java/com/facebook/presto/hive/HiveMetadata.java +++ b/presto-hive/src/main/java/com/facebook/presto/hive/HiveMetadata.java @@ -184,6 +184,7 @@ import static com.facebook.presto.hive.HiveColumnHandle.FILE_SIZE_COLUMN_NAME; import static com.facebook.presto.hive.HiveColumnHandle.PATH_COLUMN_NAME; import static com.facebook.presto.hive.HiveColumnHandle.ROW_ID_COLUMN_NAME; +import static com.facebook.presto.hive.HiveColumnHandle.rowIdColumnHandle; import static com.facebook.presto.hive.HiveColumnHandle.updateRowIdHandle; import static com.facebook.presto.hive.HiveErrorCode.HIVE_COLUMN_ORDER_MISMATCH; import static com.facebook.presto.hive.HiveErrorCode.HIVE_CONCURRENT_MODIFICATION_DETECTED; @@ -2927,7 +2928,8 @@ && isOrderBasedExecutionEnabled(session)) { streamPartitionColumns, discretePredicates, localPropertyBuilder.build(), - Optional.of(combinedRemainingPredicate)); + Optional.of(combinedRemainingPredicate), + Optional.of(rowIdColumnHandle())); } @Override diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveBucketedTablesWithRowId.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveBucketedTablesWithRowId.java new file mode 100644 index 0000000000000..c32fa22ac1116 --- /dev/null +++ b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveBucketedTablesWithRowId.java @@ -0,0 +1,240 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.hive; + +import com.facebook.presto.Session; +import com.facebook.presto.testing.QueryRunner; +import com.facebook.presto.tests.AbstractTestQueryFramework; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + +import java.util.Optional; + +import static com.facebook.presto.SystemSessionProperties.UTILIZE_UNIQUE_PROPERTY_IN_QUERY_PLANNING; +import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.aggregation; +import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.anyTree; +import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.exchange; +import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.filter; +import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.join; +import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.project; +import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.tableScan; +import static io.airlift.tpch.TpchTable.CUSTOMER; +import static io.airlift.tpch.TpchTable.ORDERS; + +@Test(singleThreaded = true) +public class TestHiveBucketedTablesWithRowId + extends AbstractTestQueryFramework +{ + @Override + protected QueryRunner createQueryRunner() + throws Exception + { + return HiveQueryRunner.createQueryRunner( + ImmutableList.of(ORDERS, CUSTOMER), + ImmutableMap.of(), + Optional.empty()); + } + + @BeforeClass + public void setUp() + { + // Create bucketed customer table + assertUpdate("CREATE TABLE customer_bucketed WITH " + + "(bucketed_by = ARRAY['custkey'], bucket_count = 13) " + + "AS SELECT * FROM customer", 1500); + + // Create bucketed orders table + assertUpdate("CREATE TABLE orders_bucketed WITH " + + "(bucketed_by = ARRAY['orderkey'], bucket_count = 11) " + + "AS SELECT * FROM orders", 15000); + + // Verify tables are created + assertQuery("SELECT count(*) FROM customer_bucketed", "SELECT count(*) FROM customer"); + assertQuery("SELECT count(*) FROM orders_bucketed", "SELECT count(*) FROM orders"); + } + + @AfterClass(alwaysRun = true) + public void tearDown() + { + try { + assertUpdate("DROP TABLE IF EXISTS customer_bucketed"); + assertUpdate("DROP TABLE IF EXISTS orders_bucketed"); + } + catch (Exception e) { + // Ignore cleanup errors + } + } + + @Test + public void testRowIdWithBucketColumn() + { + Session session = Session.builder(getQueryRunner().getDefaultSession()) + .setSystemProperty(UTILIZE_UNIQUE_PROPERTY_IN_QUERY_PLANNING, "true") + .build(); + + // Test basic query with both $row_id and $bucket + String sql = "SELECT \"$row_id\", \"$bucket\", custkey, name " + + "FROM customer_bucketed " + + "WHERE \"$bucket\" = 5"; + + assertPlan(session, sql, anyTree( + project(filter(tableScan("customer_bucketed"))))); + + // Test aggregation grouping by both $row_id and $bucket + sql = "SELECT \"$row_id\", \"$bucket\", COUNT(*) " + + "FROM customer_bucketed " + + "GROUP BY \"$row_id\", \"$bucket\""; + + assertPlan(session, sql, anyTree( + aggregation(ImmutableMap.of(), + project(tableScan("customer_bucketed"))))); + + // Test join between bucketed tables using both $row_id and $bucket + sql = "SELECT c.\"$row_id\" AS customer_row_id, " + + "c.\"$bucket\" AS customer_bucket, " + + "o.\"$row_id\" AS order_row_id, " + + "o.\"$bucket\" AS order_bucket, " + + "c.name, o.orderkey " + + "FROM customer_bucketed c " + + "JOIN orders_bucketed o " + + "ON c.custkey = o.custkey " + + "WHERE c.\"$bucket\" IN (1, 3, 5) " + + "AND o.\"$bucket\" IN (2, 4, 6)"; + + assertPlan(session, sql, anyTree( + join( + project(filter(tableScan("customer_bucketed"))), + exchange(anyTree(tableScan("orders_bucketed")))))); + } + + @Test + public void testRowIdUniquePropertyWithBucketing() + { + Session session = Session.builder(getQueryRunner().getDefaultSession()) + .setSystemProperty(UTILIZE_UNIQUE_PROPERTY_IN_QUERY_PLANNING, "true") + .build(); + + // Test unique grouping by $row_id with bucket filtering + String sql = "SELECT " + + "customer_row_id, " + + "ARBITRARY(name) AS customer_name, " + + "ARBITRARY(bucket_num) AS customer_bucket, " + + "ARRAY_AGG(orderkey) AS orders_info " + + "FROM (" + + " SELECT " + + " c.\"$row_id\" AS customer_row_id, " + + " c.\"$bucket\" AS bucket_num, " + + " c.name, " + + " o.orderkey " + + " FROM customer_bucketed c " + + " LEFT JOIN orders_bucketed o " + + " ON c.custkey = o.custkey " + + " AND o.orderstatus IN ('O', 'F') " + + " WHERE c.\"$bucket\" < 5 " + + " AND c.nationkey IN (1, 2, 3) " + + ") " + + "GROUP BY customer_row_id"; + + assertPlan(session, sql, anyTree( + aggregation(ImmutableMap.of(), + join( + anyTree(tableScan("customer_bucketed")), + anyTree(tableScan("orders_bucketed")))))); + } + + @Test + public void testRowIdAndBucketInComplexQuery() + { + Session session = Session.builder(getQueryRunner().getDefaultSession()) + .setSystemProperty(UTILIZE_UNIQUE_PROPERTY_IN_QUERY_PLANNING, "true") + .build(); + + // Complex query with both row_id and bucket columns + String sql = "SELECT " + + " unique_id, " + + " bucket_group, " + + " COUNT(*) AS order_count, " + + " AVG(totalprice) AS avg_price " + + "FROM (" + + " SELECT " + + " c.\"$row_id\" AS unique_id, " + + " CASE " + + " WHEN c.\"$bucket\" < 5 THEN 'low' " + + " WHEN c.\"$bucket\" < 10 THEN 'medium' " + + " ELSE 'high' " + + " END AS bucket_group, " + + " o.totalprice " + + " FROM customer_bucketed c " + + " JOIN orders_bucketed o " + + " ON c.custkey = o.custkey " + + " WHERE o.\"$bucket\" % 2 = 0 " + + ") t " + + "GROUP BY unique_id, bucket_group"; + + assertPlan(session, sql, anyTree( + aggregation(ImmutableMap.of(), + project(project(join( + project(tableScan("customer_bucketed")), + anyTree(tableScan("orders_bucketed")))))))); + } + + @Test + public void testDistinctRowIdWithBucketFilter() + { + Session session = Session.builder(getQueryRunner().getDefaultSession()) + .setSystemProperty(UTILIZE_UNIQUE_PROPERTY_IN_QUERY_PLANNING, "true") + .build(); + + // Test DISTINCT with row_id and bucket filtering + String sql = "SELECT " + + " DISTINCT c.\"$row_id\" AS unique_id, " + + " c.\"$bucket\" AS bucket_num, " + + " c.name " + + "FROM customer_bucketed c " + + "WHERE c.\"$bucket\" BETWEEN 3 AND 8 " + + " AND c.nationkey = 1"; + + assertPlan(session, sql, anyTree( + aggregation(ImmutableMap.of(), + project(filter(tableScan("customer_bucketed")))))); + } + + @Test + public void testRowIdJoinOnBucketColumn() + { + Session session = Session.builder(getQueryRunner().getDefaultSession()) + .setSystemProperty(UTILIZE_UNIQUE_PROPERTY_IN_QUERY_PLANNING, "true") + .build(); + + // Test joining on bucket column while selecting row_id + String sql = "SELECT " + + " c.\"$row_id\" AS customer_row_id, " + + " o.\"$row_id\" AS order_row_id, " + + " c.\"$bucket\" AS shared_bucket, " + + " c.name, " + + " o.orderkey " + + "FROM customer_bucketed c " + + "JOIN orders_bucketed o " + + " ON c.\"$bucket\" = o.\"$bucket\" " + + "WHERE c.\"$bucket\" < 5"; + + assertPlan(session, sql, anyTree( + join( + exchange(anyTree(tableScan("customer_bucketed"))), + exchange(anyTree(tableScan("orders_bucketed")))))); + } +} diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveLogicalPlanner.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveLogicalPlanner.java index 07d0e3a7c6385..534beafe5a0b6 100644 --- a/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveLogicalPlanner.java +++ b/presto-hive/src/test/java/com/facebook/presto/hive/TestHiveLogicalPlanner.java @@ -85,6 +85,7 @@ import static com.facebook.presto.SystemSessionProperties.PUSHDOWN_DEREFERENCE_ENABLED; import static com.facebook.presto.SystemSessionProperties.PUSHDOWN_SUBFIELDS_ENABLED; import static com.facebook.presto.SystemSessionProperties.PUSHDOWN_SUBFIELDS_FOR_MAP_FUNCTIONS; +import static com.facebook.presto.SystemSessionProperties.UTILIZE_UNIQUE_PROPERTY_IN_QUERY_PLANNING; import static com.facebook.presto.common.function.OperatorType.EQUAL; import static com.facebook.presto.common.predicate.Domain.create; import static com.facebook.presto.common.predicate.Domain.multipleValues; @@ -130,6 +131,7 @@ import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.node; import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.output; import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.project; +import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.semiJoin; import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.strictTableScan; import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.values; import static com.facebook.presto.sql.planner.optimizations.PlanNodeSearcher.searchFrom; @@ -2153,6 +2155,336 @@ public void testPartialAggregatePushdown() } } + @Test + public void testRowId() + { + Session session = Session.builder(getQueryRunner().getDefaultSession()) + .setSystemProperty(UTILIZE_UNIQUE_PROPERTY_IN_QUERY_PLANNING, "true") + .build(); + String sql; + sql = "SELECT\n" + + " unique_id AS unique_id,\n" + + " ARBITRARY(name) AS customer_name,\n" + + " ARRAY_AGG(orderkey) AS orders_info\n" + + "FROM (\n" + + " SELECT\n" + + " customer.name,\n" + + " orders.orderkey,\n" + + " customer.\"$row_id\" AS unique_id\n" + + " FROM customer\n" + + " LEFT JOIN orders\n" + + " ON customer.custkey = orders.custkey\n" + + " AND orders.orderstatus IN ('O', 'F')\n" + + " AND orders.orderdate BETWEEN DATE '1995-01-01' AND DATE '1995-12-31'\n" + + " WHERE\n" + + " customer.nationkey IN (1, 2, 3, 4, 5)\n" + + ")\n" + + "GROUP BY\n" + + " unique_id"; + + assertPlan(session, + sql, + anyTree( + aggregation( + ImmutableMap.of(), + join( + anyTree(tableScan("customer")), + anyTree(tableScan("orders")))))); + + sql = "select \"$row_id\", count(*) from orders group by 1"; + assertPlan(sql, anyTree(aggregation(ImmutableMap.of(), tableScan("orders")))); + sql = "SELECT\n" + + " customer.\"$row_id\" AS unique_id,\n" + + " COUNT(orders.orderkey) AS order_count,\n" + + " ARRAY_AGG(orders.orderkey) AS order_keys\n" + + "FROM customer\n" + + "LEFT JOIN orders\n" + + " ON customer.custkey = orders.custkey\n" + + "WHERE customer.acctbal > 1000\n" + + "GROUP BY customer.\"$row_id\""; + assertPlan(sql, anyTree( + aggregation(ImmutableMap.of(), + join( + anyTree(tableScan("customer")), + anyTree(tableScan("orders")))))); + + sql = "SELECT\n" + + " unique_id,\n" + + " MAX(orderdate) AS latest_order_date\n" + + "FROM (\n" + + " SELECT\n" + + " customer.\"$row_id\" AS unique_id,\n" + + " orders.orderdate\n" + + " FROM customer\n" + + " JOIN orders\n" + + " ON customer.custkey = orders.custkey\n" + + " WHERE orders.orderstatus = 'O'\n" + + ") t\n" + + "GROUP BY unique_id"; + assertPlan(sql, anyTree( + aggregation(ImmutableMap.of(), + join( + anyTree(tableScan("customer")), + anyTree(tableScan("orders")))))); + + sql = "SELECT\n" + + " DISTINCT customer.\"$row_id\" AS unique_id,\n" + + " customer.name\n" + + "FROM customer\n" + + "WHERE customer.nationkey = 1"; + + assertPlan(sql, anyTree( + aggregation(ImmutableMap.of(), + project(filter(tableScan("customer")))))); + + sql = "SELECT\n" + + " c.name,\n" + + " o.orderkey,\n" + + " c.\"$row_id\" AS customer_row_id,\n" + + " o.\"$row_id\" AS order_row_id\n" + + "FROM customer c\n" + + "JOIN orders o\n" + + " ON c.\"$row_id\" = o.\"$row_id\"\n" + + "WHERE o.totalprice > 10000"; + assertPlan(sql, anyTree( + join( + exchange(anyTree(tableScan("customer"))), + exchange(anyTree(tableScan("orders")))))); + + sql = "SELECT\n" + + " \"$row_id\" AS unique_id,\n" + + " orderstatus,\n" + + " COUNT(*) AS cnt\n" + + "FROM orders\n" + + "GROUP BY \"$row_id\", orderstatus"; + assertPlan(sql, anyTree( + aggregation(ImmutableMap.of(), + project(tableScan("orders"))))); + + sql = "SELECT\n" + + " o1.orderkey,\n" + + " o2.totalprice\n" + + "FROM orders o1\n" + + "JOIN orders o2\n" + + " ON o1.\"$row_id\" = o2.\"$row_id\"\n" + + "WHERE o1.orderstatus = 'O'"; + assertPlan(sql, anyTree( + join( + exchange(anyTree(tableScan("orders"))), + exchange(anyTree(tableScan("orders")))))); + + sql = "SELECT\n" + + " orderkey,\n" + + " totalprice\n" + + "FROM orders o1\n" + + "WHERE EXISTS (\n" + + " SELECT 1\n" + + " FROM orders o2\n" + + " WHERE o1.\"$row_id\" = o2.\"$row_id\"\n" + + " AND o2.orderstatus = 'F'\n" + + ")"; + assertPlan(sql, anyTree( + join( + exchange(anyTree(tableScan("orders"))), + exchange(anyTree(aggregation(ImmutableMap.of(), project(filter(tableScan("orders"))))))))); + + sql = "SELECT\n" + + " custkey,\n" + + " name\n" + + "FROM customer\n" + + "WHERE \"$row_id\" IN (\n" + + " SELECT \"$row_id\"\n" + + " FROM customer\n" + + " WHERE acctbal > 5000\n" + + ")"; + assertPlan(sql, anyTree( + semiJoin( + anyTree(tableScan("customer")), + anyTree(tableScan("customer"))))); + + sql = "SELECT \"$row_id\", orderkey FROM orders WHERE orderstatus = 'O'\n" + + "UNION\n" + + "SELECT \"$row_id\", orderkey FROM orders WHERE orderstatus = 'F'"; + assertPlan(sql, anyTree( + aggregation(ImmutableMap.of(), + project(filter(tableScan("orders")))))); + + sql = "SELECT\n" + + " c.\"$row_id\",\n" + + " COUNT(o.orderkey) AS order_count,\n" + + " SUM(o.totalprice) AS total_spent,\n" + + " MAX(o.orderdate) AS latest_order,\n" + + " MIN(o.orderdate) AS first_order\n" + + "FROM customer c\n" + + "LEFT JOIN orders o ON c.custkey = o.custkey\n" + + "GROUP BY c.\"$row_id\""; + assertPlan(sql, anyTree( + aggregation(ImmutableMap.of(), + join( + anyTree(tableScan("customer")), + anyTree(tableScan("orders")))))); + + sql = "SELECT\n" + + " \"$row_id\",\n" + + " COUNT(*) AS cnt\n" + + "FROM orders\n" + + "GROUP BY \"$row_id\"\n" + + "HAVING COUNT(*) = 1"; + assertPlan(sql, anyTree( + project(filter( + aggregation(ImmutableMap.of(), + tableScan("orders")))))); + + sql = "SELECT\n" + + " c.custkey,\n" + + " c.name,\n" + + " (\n" + + " SELECT COUNT(*)\n" + + " FROM orders o\n" + + " WHERE o.custkey = c.custkey\n" + + " AND o.\"$row_id\" IS NOT NULL\n" + + " ) AS order_count\n" + + "FROM customer c"; + assertPlan(sql, anyTree( + project( + join( + anyTree(tableScan("customer")), + anyTree(tableScan("orders")))))); + + sql = "SELECT\n" + + " outer_query.customer_id,\n" + + " outer_query.order_count\n" + + "FROM (\n" + + " SELECT\n" + + " c.\"$row_id\" AS customer_id,\n" + + " COUNT(DISTINCT o.\"$row_id\") AS order_count\n" + + " FROM customer c\n" + + " LEFT JOIN orders o ON c.custkey = o.custkey\n" + + " WHERE c.nationkey IN (1, 2, 3)\n" + + " GROUP BY c.\"$row_id\"\n" + + ") outer_query\n" + + "WHERE outer_query.order_count > 2"; + assertPlan(sql, anyTree( + aggregation(ImmutableMap.of(), + project( + join( + anyTree(tableScan("customer")), + anyTree(tableScan("orders"))))))); + + sql = "SELECT\n" + + " c.custkey,\n" + + " c.name\n" + + "FROM customer c\n" + + "WHERE NOT EXISTS (\n" + + " SELECT 1\n" + + " FROM orders o\n" + + " WHERE c.\"$row_id\" = o.\"$row_id\"\n" + + ")"; + assertPlan(sql, anyTree( + join( + exchange(anyTree(tableScan("customer"))), + exchange(anyTree(aggregation(ImmutableMap.of(), tableScan("orders"))))))); + + sql = "SELECT\n" + + " c.name,\n" + + " o.orderkey,\n" + + " l.linenumber\n" + + "FROM customer c\n" + + "JOIN orders o ON c.custkey = o.custkey\n" + + "JOIN lineitem l ON o.orderkey = l.orderkey\n" + + "WHERE c.\"$row_id\" IS NOT NULL\n" + + " AND o.\"$row_id\" IS NOT NULL\n" + + " AND l.\"$row_id\" IS NOT NULL"; + assertPlan(sql, anyTree( + join( + anyTree( + join( + anyTree(tableScan("customer")), + anyTree(tableScan("orders")))), + anyTree(tableScan("lineitem"))))); + + sql = "SELECT\n" + + " c.\"$row_id\" AS customer_row_id,\n" + + " o.\"$row_id\" AS order_row_id,\n" + + " c.name,\n" + + " o.orderkey\n" + + "FROM customer c\n" + + "CROSS JOIN orders o\n" + + "WHERE c.nationkey = 1 AND o.orderstatus = 'O'"; + assertPlan(sql, anyTree( + join( + anyTree(tableScan("customer")), + anyTree(tableScan("orders"))))); + + sql = "SELECT\n" + + " orderstatus,\n" + + " COUNT(\"$row_id\") AS row_count,\n" + + " MIN(\"$row_id\") AS min_row_id\n" + + "FROM orders\n" + + "GROUP BY orderstatus"; + assertPlan(sql, anyTree( + aggregation(ImmutableMap.of(), + exchange(anyTree(tableScan("orders")))))); + + sql = "SELECT\n" + + " xxhash64(\"$row_id\") AS bucket,\n" + + " COUNT(*) AS cnt\n" + + "FROM orders\n" + + "GROUP BY xxhash64(\"$row_id\")"; + assertPlan(sql, anyTree( + aggregation(ImmutableMap.of(), + exchange(anyTree(tableScan("orders")))))); + + sql = "SELECT \"$row_id\", 'customer' AS source FROM customer\n" + + "UNION ALL\n" + + "SELECT \"$row_id\", 'orders' AS source FROM orders"; + assertPlan(sql, anyTree( + exchange( + anyTree(tableScan("customer")), + anyTree(tableScan("orders"))))); + + sql = "SELECT\n" + + " o.\"$row_id\" AS order_row_id,\n" + + " o.orderkey,\n" + + " l.linenumber\n" + + "FROM orders o\n" + + "JOIN lineitem l ON o.orderkey = l.orderkey\n" + + "WHERE o.orderstatus = 'O'"; + assertPlan(sql, anyTree( + join( + exchange(anyTree(tableScan("lineitem"))), + exchange(anyTree(tableScan("orders")))))); + + sql = "SELECT\n" + + " orderkey,\n" + + " totalprice\n" + + "FROM orders o1\n" + + "WHERE \"$row_id\" > (\n" + + " SELECT MIN(\"$row_id\")\n" + + " FROM orders o2\n" + + " WHERE o2.orderstatus = 'F'\n" + + ")"; + assertPlan(sql, anyTree( + join( + tableScan("orders"), + exchange(anyTree(tableScan("orders")))))); + + sql = "SELECT\n" + + " CASE \n" + + " WHEN nationkey = 1 THEN \"$row_id\"\n" + + " ELSE cast('default' as varbinary)\n" + + " END AS conditional_row_id,\n" + + " COUNT(*) AS cnt\n" + + "FROM customer\n" + + "GROUP BY CASE \n" + + " WHEN nationkey = 1 THEN \"$row_id\"\n" + + " ELSE cast('default' as varbinary)\n" + + " END"; + assertPlan(sql, anyTree( + aggregation(ImmutableMap.of(), + exchange(anyTree(tableScan("customer")))))); + } + private static Set toSubfields(String... subfieldPaths) { return Arrays.stream(subfieldPaths) @@ -2170,6 +2502,11 @@ private void assertPushdownSubfields(Session session, String query, String table assertPlan(session, query, anyTree(tableScan(tableName, requiredSubfields))); } + private static PlanMatchPattern tableScan(String expectedTableName) + { + return PlanMatchPattern.tableScan(expectedTableName); + } + private static PlanMatchPattern tableScan(String expectedTableName, Map> expectedRequiredSubfields) { return PlanMatchPattern.tableScan(expectedTableName).with(new HiveTableScanMatcher(expectedRequiredSubfields)); diff --git a/presto-main-base/src/main/java/com/facebook/presto/SystemSessionProperties.java b/presto-main-base/src/main/java/com/facebook/presto/SystemSessionProperties.java index 1a30a76798d37..f335f1f718c57 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/SystemSessionProperties.java +++ b/presto-main-base/src/main/java/com/facebook/presto/SystemSessionProperties.java @@ -337,6 +337,7 @@ public final class SystemSessionProperties public static final String QUERY_CLIENT_TIMEOUT = "query_client_timeout"; public static final String REWRITE_MIN_MAX_BY_TO_TOP_N = "rewrite_min_max_by_to_top_n"; public static final String ADD_DISTINCT_BELOW_SEMI_JOIN_BUILD = "add_distinct_below_semi_join_build"; + public static final String UTILIZE_UNIQUE_PROPERTY_IN_QUERY_PLANNING = "utilize_unique_property_in_query_planning"; public static final String PUSHDOWN_SUBFIELDS_FOR_MAP_FUNCTIONS = "pushdown_subfields_for_map_functions"; public static final String MAX_SERIALIZABLE_OBJECT_SIZE = "max_serializable_object_size"; @@ -1950,6 +1951,10 @@ public SystemSessionProperties( false, value -> Duration.valueOf((String) value), Duration::toString), + booleanProperty(UTILIZE_UNIQUE_PROPERTY_IN_QUERY_PLANNING, + "Utilize the unique property of input columns in query planning", + featuresConfig.isUtilizeUniquePropertyInQueryPlanning(), + false), booleanProperty(ADD_DISTINCT_BELOW_SEMI_JOIN_BUILD, "Add distinct aggregation below semi join build", featuresConfig.isAddDistinctBelowSemiJoinBuild(), @@ -3310,6 +3315,11 @@ public static boolean isPushSubfieldsForMapFunctionsEnabled(Session session) return session.getSystemProperty(PUSHDOWN_SUBFIELDS_FOR_MAP_FUNCTIONS, Boolean.class); } + public static boolean isUtilizeUniquePropertyInQueryPlanningEnabled(Session session) + { + return session.getSystemProperty(UTILIZE_UNIQUE_PROPERTY_IN_QUERY_PLANNING, Boolean.class); + } + public static boolean isAddDistinctBelowSemiJoinBuildEnabled(Session session) { return session.getSystemProperty(ADD_DISTINCT_BELOW_SEMI_JOIN_BUILD, Boolean.class); diff --git a/presto-main-base/src/main/java/com/facebook/presto/metadata/TableLayout.java b/presto-main-base/src/main/java/com/facebook/presto/metadata/TableLayout.java index 553f5b5f8da1c..19026f586c035 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/metadata/TableLayout.java +++ b/presto-main-base/src/main/java/com/facebook/presto/metadata/TableLayout.java @@ -78,6 +78,11 @@ public List> getLocalProperties() return layout.getLocalProperties(); } + public Optional getUniqueColumn() + { + return layout.getUniqueColumn(); + } + public ConnectorTableLayoutHandle getLayoutHandle() { return layout.getHandle(); diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/FeaturesConfig.java b/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/FeaturesConfig.java index 06b1fee1fdce7..cdecf697bda5e 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/FeaturesConfig.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/FeaturesConfig.java @@ -309,6 +309,7 @@ public class FeaturesConfig private boolean addDistinctBelowSemiJoinBuild; private boolean pushdownSubfieldForMapFunctions = true; private long maxSerializableObjectSize = 1000; + private boolean utilizeUniquePropertyInQueryPlanning = true; private boolean builtInSidecarFunctionsEnabled; @@ -3101,6 +3102,19 @@ public boolean isPushdownSubfieldForMapFunctions() return pushdownSubfieldForMapFunctions; } + @Config("optimizer.utilize-unique-property-in-query-planning") + @ConfigDescription("Utilize the unique property of input columns in query planning") + public FeaturesConfig setUtilizeUniquePropertyInQueryPlanning(boolean utilizeUniquePropertyInQueryPlanning) + { + this.utilizeUniquePropertyInQueryPlanning = utilizeUniquePropertyInQueryPlanning; + return this; + } + + public boolean isUtilizeUniquePropertyInQueryPlanning() + { + return utilizeUniquePropertyInQueryPlanning; + } + @Config("max_serializable_object_size") @ConfigDescription("Configure the maximum byte size of a serializable object in expression interpreters") public FeaturesConfig setMaxSerializableObjectSize(long maxSerializableObjectSize) diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/ActualProperties.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/ActualProperties.java index 32fc8913fa12c..f562d79b50c32 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/ActualProperties.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/ActualProperties.java @@ -49,6 +49,7 @@ import static com.google.common.base.MoreObjects.toStringHelper; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.collect.Iterables.transform; +import static java.lang.String.format; import static java.util.Objects.requireNonNull; public class ActualProperties @@ -56,11 +57,22 @@ public class ActualProperties private final Global global; private final List> localProperties; private final Map constants; + // Used to track the properties of the unique row_id + private final Optional propertiesFromUniqueColumn; private ActualProperties( Global global, List> localProperties, Map constants) + { + this(global, localProperties, constants, Optional.empty()); + } + + private ActualProperties( + Global global, + List> localProperties, + Map constants, + Optional propertiesFromUniqueColumn) { requireNonNull(global, "globalProperties is null"); requireNonNull(localProperties, "localProperties is null"); @@ -85,6 +97,8 @@ private ActualProperties( this.localProperties = ImmutableList.copyOf(updatedLocalProperties); this.constants = ImmutableMap.copyOf(constants); + propertiesFromUniqueColumn.ifPresent(actualProperties -> checkArgument(!actualProperties.getPropertiesFromUniqueColumn().isPresent())); + this.propertiesFromUniqueColumn = propertiesFromUniqueColumn; } public boolean isCoordinatorOnly() @@ -92,6 +106,11 @@ public boolean isCoordinatorOnly() return global.isCoordinatorOnly(); } + public Optional getPropertiesFromUniqueColumn() + { + return propertiesFromUniqueColumn; + } + /** * @return true if the plan will only execute on a single node */ @@ -120,6 +139,16 @@ public boolean isStreamPartitionedOn(Collection col } } + public boolean isStreamPartitionedOnAdditionalProperty(Collection columns, boolean exactly) + { + if (exactly) { + return propertiesFromUniqueColumn.isPresent() && propertiesFromUniqueColumn.get().global.isStreamPartitionedOnExactly(columns, ImmutableSet.of(), false); + } + else { + return propertiesFromUniqueColumn.isPresent() && propertiesFromUniqueColumn.get().global.isStreamPartitionedOn(columns, ImmutableSet.of(), false); + } + } + public boolean isNodePartitionedOn(Collection columns, boolean exactly) { return isNodePartitionedOn(columns, false, exactly); @@ -135,6 +164,16 @@ public boolean isNodePartitionedOn(Collection colum } } + public boolean isNodePartitionedOnAdditionalProperty(Collection columns, boolean exactly) + { + if (exactly) { + return propertiesFromUniqueColumn.isPresent() && propertiesFromUniqueColumn.get().global.isNodePartitionedOnExactly(columns, ImmutableSet.of(), false); + } + else { + return propertiesFromUniqueColumn.isPresent() && propertiesFromUniqueColumn.get().global.isNodePartitionedOn(columns, ImmutableSet.of(), false); + } + } + @Deprecated public boolean isCompatibleTablePartitioningWith(Partitioning partitioning, boolean nullsAndAnyReplicated, Metadata metadata, Session session) { @@ -194,6 +233,13 @@ public ActualProperties translateVariable(Function newAdditionalProperty = Optional.empty(); + if (propertiesFromUniqueColumn.isPresent()) { + ActualProperties translatedAdditionalProperty = propertiesFromUniqueColumn.get().translateVariable(translator); + if (!translatedAdditionalProperty.getLocalProperties().isEmpty()) { + newAdditionalProperty = Optional.of(translatedAdditionalProperty); + } + } return builder() .global(global.translateVariableToRowExpression(variable -> { Optional translated = translator.apply(variable).map(RowExpression.class::cast); @@ -204,6 +250,7 @@ public ActualProperties translateVariable(Function !inputToOutputVariables.containsKey(entry.getKey())) .forEach(inputToOutputMappings::put); + + Optional newAdditionalProperty = Optional.empty(); + if (propertiesFromUniqueColumn.isPresent()) { + ActualProperties translatedAdditionalProperty = propertiesFromUniqueColumn.get().translateRowExpression(assignments); + if (!translatedAdditionalProperty.getLocalProperties().isEmpty()) { + newAdditionalProperty = Optional.of(translatedAdditionalProperty); + } + } + return builder() .global(global.translateRowExpression(inputToOutputMappings.build(), assignments)) .local(LocalProperties.translate(localProperties, variable -> Optional.ofNullable(inputToOutputVariables.get(variable)))) .constants(translatedConstants) + .propertiesFromUniqueColumn(newAdditionalProperty) .build(); } @@ -274,6 +331,7 @@ public static class Builder private List> localProperties; private Map constants; private boolean unordered; + private Optional propertiesFromUniqueColumn; public Builder() { @@ -281,10 +339,16 @@ public Builder() } public Builder(Global global, List> localProperties, Map constants) + { + this(global, localProperties, constants, Optional.empty()); + } + + public Builder(Global global, List> localProperties, Map constants, Optional propertiesFromUniqueColumn) { this.global = requireNonNull(global, "global is null"); this.localProperties = ImmutableList.copyOf(localProperties); this.constants = ImmutableMap.copyOf(constants); + this.propertiesFromUniqueColumn = propertiesFromUniqueColumn; } public Builder global(Global global) @@ -317,6 +381,18 @@ public Builder unordered(boolean unordered) return this; } + public Builder propertiesFromUniqueColumn(Optional propertiesFromUniqueColumn) + { + if (propertiesFromUniqueColumn.isPresent() && !propertiesFromUniqueColumn.get().getLocalProperties().isEmpty()) { + checkArgument(propertiesFromUniqueColumn.get().getLocalProperties().size() == 1); + this.propertiesFromUniqueColumn = propertiesFromUniqueColumn; + } + else { + this.propertiesFromUniqueColumn = Optional.empty(); + } + return this; + } + public ActualProperties build() { List> localProperties = this.localProperties; @@ -332,14 +408,19 @@ public ActualProperties build() } localProperties = newLocalProperties.build(); } - return new ActualProperties(global, localProperties, constants); + if (propertiesFromUniqueColumn.isPresent() && unordered) { + propertiesFromUniqueColumn = Optional.of(ActualProperties.builderFrom(propertiesFromUniqueColumn.get()) + .unordered(unordered) + .build()); + } + return new ActualProperties(global, localProperties, constants, propertiesFromUniqueColumn); } } @Override public int hashCode() { - return Objects.hash(global, localProperties, constants.keySet()); + return Objects.hash(global, localProperties, constants.keySet(), propertiesFromUniqueColumn); } @Override @@ -354,7 +435,8 @@ public boolean equals(Object obj) final ActualProperties other = (ActualProperties) obj; return Objects.equals(this.global, other.global) && Objects.equals(this.localProperties, other.localProperties) - && Objects.equals(this.constants.keySet(), other.constants.keySet()); + && Objects.equals(this.constants.keySet(), other.constants.keySet()) + && Objects.equals(this.propertiesFromUniqueColumn, other.propertiesFromUniqueColumn); } @Override @@ -364,6 +446,7 @@ public String toString() .add("globalProperties", global) .add("localProperties", localProperties) .add("constants", constants) + .add("propertiesFromUniqueColumn", propertiesFromUniqueColumn) .toString(); } @@ -391,7 +474,7 @@ private Global(Optional nodePartitioning, Optional s || !streamPartitioning.isPresent() || nodePartitioning.get().getVariableReferences().containsAll(streamPartitioning.get().getVariableReferences()) || streamPartitioning.get().getVariableReferences().containsAll(nodePartitioning.get().getVariableReferences()), - "Global stream partitioning columns should match node partitioning columns"); + format("Global stream partitioning columns should match node partitioning columns, nodePartitioning: %s, streamPartitioning: %s", nodePartitioning, streamPartitioning)); this.nodePartitioning = requireNonNull(nodePartitioning, "nodePartitioning is null"); this.streamPartitioning = requireNonNull(streamPartitioning, "streamPartitioning is null"); this.nullsAndAnyReplicated = nullsAndAnyReplicated; diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/AddExchanges.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/AddExchanges.java index 84c151ac6e7a0..9aba136f0f1c9 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/AddExchanges.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/AddExchanges.java @@ -312,7 +312,8 @@ public PlanWithProperties visitAggregation(AggregationNode node, PreferredProper child.getProperties()); } else if (hasMixedGroupingSets - || !isStreamPartitionedOn(child.getProperties(), partitioningRequirement) && !isNodePartitionedOn(child.getProperties(), partitioningRequirement)) { + || !isStreamPartitionedOn(child.getProperties(), partitioningRequirement) && !isNodePartitionedOn(child.getProperties(), partitioningRequirement) + && !isNodePartitionedOnAdditionalProperty(child.getProperties(), partitioningRequirement) && !isStreamPartitionedOnAdditionalProperty(child.getProperties(), partitioningRequirement)) { child = withDerivedProperties( partitionedExchange( idAllocator.getNextId(), @@ -1631,11 +1632,21 @@ private boolean isNodePartitionedOn(ActualProperties properties, Collection columns) + { + return properties.isNodePartitionedOnAdditionalProperty(columns, isExactPartitioningPreferred(session)); + } + private boolean isStreamPartitionedOn(ActualProperties properties, Collection columns) { return properties.isStreamPartitionedOn(columns, isExactPartitioningPreferred(session)); } + private boolean isStreamPartitionedOnAdditionalProperty(ActualProperties properties, Collection columns) + { + return properties.isStreamPartitionedOnAdditionalProperty(columns, isExactPartitioningPreferred(session)); + } + private boolean shouldAggregationMergePartitionPreferences(AggregationPartitioningMergingStrategy aggregationPartitioningMergingStrategy) { if (isExactPartitioningPreferred(session)) { diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/AddLocalExchanges.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/AddLocalExchanges.java index 1618c69b543e6..65fa710226c48 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/AddLocalExchanges.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/AddLocalExchanges.java @@ -399,7 +399,8 @@ public PlanWithProperties visitAggregation(AggregationNode node, StreamPreferred // [A, B] [(A, C)] -> List.of(Optional.of(GroupingProperty(C))) // [A, B] [(D, A, C)] -> List.of(Optional.of(GroupingProperty(D, C))) List>> matchResult = LocalProperties.match(child.getProperties().getLocalProperties(), LocalProperties.grouped(groupingKeys)); - if (!matchResult.get(0).isPresent()) { + List>> matchResultForAdditional = LocalProperties.match(child.getProperties().getAdditionalLocalProperties(), LocalProperties.grouped(groupingKeys)); + if (!matchResult.get(0).isPresent() || !matchResultForAdditional.get(0).isPresent()) { // !isPresent() indicates the property was satisfied completely preGroupedSymbols = groupingKeys; } diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/LocalProperties.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/LocalProperties.java index e8c02f02216e4..00fa4935ee21a 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/LocalProperties.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/LocalProperties.java @@ -18,6 +18,7 @@ import com.facebook.presto.spi.GroupingProperty; import com.facebook.presto.spi.LocalProperty; import com.facebook.presto.spi.SortingProperty; +import com.facebook.presto.spi.UniqueProperty; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; import com.google.common.collect.PeekingIterator; @@ -55,6 +56,11 @@ public static List> sorted(Collection columns, SortOrder return columns.stream().map(column -> new SortingProperty<>(column, order)).collect(toImmutableList()); } + public static List> unique(T column) + { + return ImmutableList.of(new UniqueProperty<>(column)); + } + public static List> stripLeadingConstants(List> properties) { PeekingIterator> iterator = peekingIterator(properties.iterator()); diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/PropertyDerivations.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/PropertyDerivations.java index f086283a140da..bb4f27b09bb0e 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/PropertyDerivations.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/PropertyDerivations.java @@ -23,6 +23,7 @@ import com.facebook.presto.spi.GroupingProperty; import com.facebook.presto.spi.LocalProperty; import com.facebook.presto.spi.SortingProperty; +import com.facebook.presto.spi.UniqueProperty; import com.facebook.presto.spi.plan.AggregationNode; import com.facebook.presto.spi.plan.DeleteNode; import com.facebook.presto.spi.plan.DistinctLimitNode; @@ -90,6 +91,7 @@ import static com.facebook.presto.SystemSessionProperties.isJoinSpillingEnabled; import static com.facebook.presto.SystemSessionProperties.isSpillEnabled; +import static com.facebook.presto.SystemSessionProperties.isUtilizeUniquePropertyInQueryPlanningEnabled; import static com.facebook.presto.SystemSessionProperties.planWithTableNodePartitioning; import static com.facebook.presto.common.predicate.TupleDomain.toLinkedMap; import static com.facebook.presto.spi.relation.DomainTranslator.BASIC_COLUMN_EXTRACTOR; @@ -142,6 +144,22 @@ public static ActualProperties streamBackdoorDeriveProperties(PlanNode node, Lis return node.accept(new Visitor(metadata, session), inputProperties); } + public static Optional uniqueToGroupProperties(ActualProperties properties) + { + // We only call uniqueToGroupProperties on derived properties from propertiesFromUniqueColumn, which can have one local property if the column is preserved in a node + // output, or no local property if the column is not preserved in a node output + checkArgument(properties.getLocalProperties().size() <= 1); + if (properties.getLocalProperties().isEmpty()) { + return Optional.empty(); + } + LocalProperty localProperty = Iterables.getOnlyElement(properties.getLocalProperties()); + if (localProperty instanceof UniqueProperty) { + return Optional.of(ActualProperties.builderFrom(properties).local(ImmutableList.of(new GroupingProperty<>(ImmutableList.of(((UniqueProperty) localProperty).getColumn())))).build()); + } + checkState(localProperty instanceof GroupingProperty, "returned actual properties should have grouping property"); + return Optional.of(properties); + } + private static class Visitor extends InternalPlanVisitor> { @@ -196,12 +214,14 @@ public ActualProperties visitAssignUniqueId(AssignUniqueId node, List in inputToOutputMappings.putIfAbsent(argument, argument); } - return Iterables.getOnlyElement(inputProperties).translateVariable(column -> Optional.ofNullable(inputToOutputMappings.get(column))); + ActualProperties properties = Iterables.getOnlyElement(inputProperties); + return ActualProperties.builderFrom(properties.translateVariable(column -> Optional.ofNullable(inputToOutputMappings.get(column)))) + .propertiesFromUniqueColumn(properties.getPropertiesFromUniqueColumn().flatMap(x -> uniqueToGroupProperties(x.translateVariable(column -> Optional.ofNullable(inputToOutputMappings.get(column)))))) + .build(); } @Override @@ -291,10 +314,9 @@ public ActualProperties visitAggregation(AggregationNode node, List node.getGroupingKeys().contains(variable) ? Optional.of(variable) : Optional.empty()); - return ActualProperties.builderFrom(translated) .local(LocalProperties.grouped(node.getGroupingKeys())) - .build(); + .propertiesFromUniqueColumn(uniqueProperties(translated.getPropertiesFromUniqueColumn())).build(); } @Override @@ -303,6 +325,14 @@ public ActualProperties visitRowNumber(RowNumberNode node, List uniqueProperties(Optional properties) + { + if (properties.isPresent() && properties.get().getLocalProperties().size() == 1 && properties.get().getLocalProperties().get(0) instanceof UniqueProperty) { + return properties; + } + return Optional.empty(); + } + @Override public ActualProperties visitTopNRowNumber(TopNRowNumberNode node, List inputProperties) { @@ -316,6 +346,7 @@ public ActualProperties visitTopNRowNumber(TopNRowNumberNode node, List inputPro return ActualProperties.builderFrom(properties) .local(localProperties) + .propertiesFromUniqueColumn(uniqueProperties(properties.getPropertiesFromUniqueColumn())) .build(); } @@ -344,6 +376,7 @@ public ActualProperties visitSort(SortNode node, List inputPro return ActualProperties.builderFrom(properties) .local(localProperties) + .propertiesFromUniqueColumn(uniqueProperties(properties.getPropertiesFromUniqueColumn())) .build(); } @@ -360,6 +393,7 @@ public ActualProperties visitDistinctLimit(DistinctLimitNode node, List inputPro return ActualProperties.builderFrom(probeProperties) .constants(constants) + .propertiesFromUniqueColumn(probeProperties.getPropertiesFromUniqueColumn().flatMap(x -> uniqueToGroupProperties(x.translateVariable(column -> filterOrRewrite(outputVariableReferences, node.getCriteria(), column))))) .unordered(unordered) .build(); case LEFT: return ActualProperties.builderFrom(probeProperties.translateVariable(column -> filterIfMissing(outputVariableReferences, column))) + .propertiesFromUniqueColumn(probeProperties.getPropertiesFromUniqueColumn().flatMap(x -> uniqueToGroupProperties(x.translateVariable(column -> filterIfMissing(outputVariableReferences, column))))) .unordered(unordered) .build(); case RIGHT: @@ -518,10 +554,12 @@ public ActualProperties visitIndexJoin(IndexJoinNode node, List uniqueToGroupProperties(x))) .build(); case SOURCE_OUTER: return ActualProperties.builderFrom(probeProperties) .constants(probeProperties.getConstants()) + .propertiesFromUniqueColumn(probeProperties.getPropertiesFromUniqueColumn().flatMap(x -> uniqueToGroupProperties(x))) .build(); default: throw new UnsupportedOperationException("Unsupported join type: " + node.getType()); @@ -553,16 +591,19 @@ public ActualProperties visitMergeJoin(MergeJoinNode node, List uniqueToGroupProperties(x.translateVariable(column -> filterOrRewrite(outputVariableReferences, node.getCriteria(), column))))) .constants(constants) .build(); case LEFT: return ActualProperties.builderFrom(leftProperties.translateVariable(column -> filterIfMissing(outputVariableReferences, column))) + .propertiesFromUniqueColumn(leftProperties.getPropertiesFromUniqueColumn().flatMap(x -> uniqueToGroupProperties(x.translateVariable(column -> filterIfMissing(outputVariableReferences, column))))) .build(); case RIGHT: rightProperties = rightProperties.translateVariable(column -> filterIfMissing(node.getOutputVariables(), column)); return ActualProperties.builderFrom(rightProperties.translateVariable(column -> filterIfMissing(outputVariableReferences, column))) .local(ImmutableList.of()) + .propertiesFromUniqueColumn(rightProperties.getPropertiesFromUniqueColumn().flatMap(x -> uniqueToGroupProperties(x.translateVariable(column -> filterIfMissing(outputVariableReferences, column))))) .unordered(true) .build(); case FULL: @@ -598,6 +639,7 @@ public ActualProperties visitExchange(ExchangeNode node, List checkArgument(!node.getScope().isRemote() || inputProperties.stream().noneMatch(ActualProperties::isNullsAndAnyReplicated), "Null-and-any replicated inputs should not be remotely exchanged"); Set> entries = null; + ActualProperties translated = null; for (int sourceIndex = 0; sourceIndex < node.getSources().size(); sourceIndex++) { List inputVariables = node.getInputs().get(sourceIndex); Map inputToOutput = new HashMap<>(); @@ -605,7 +647,7 @@ public ActualProperties visitExchange(ExchangeNode node, List inputToOutput.put(inputVariables.get(i), node.getOutputVariables().get(i)); } - ActualProperties translated = inputProperties.get(sourceIndex).translateVariable(variable -> Optional.ofNullable(inputToOutput.get(variable))); + translated = inputProperties.get(sourceIndex).translateVariable(variable -> Optional.ofNullable(inputToOutput.get(variable))); entries = (entries == null) ? translated.getConstants().entrySet() : Sets.intersection(entries, translated.getConstants().entrySet()); } @@ -621,6 +663,8 @@ public ActualProperties visitExchange(ExchangeNode node, List .forEach(localProperties::add); } + boolean additionalPropertyIsUnique = inputProperties.size() == 1 && uniqueProperties(translated.getPropertiesFromUniqueColumn()).isPresent() && !node.getType().equals(ExchangeNode.Type.REPLICATE); + // Local exchanges are only created in AddLocalExchanges, at the end of optimization, and // local exchanges do not produce all global properties as represented by ActualProperties. // This is acceptable because AddLocalExchanges does not use global properties and is only @@ -640,6 +684,10 @@ else if (inputProperties.stream().anyMatch(ActualProperties::isSingleNode)) { builder.global(coordinatorSingleStreamPartition()); } + if (additionalPropertyIsUnique) { + builder.propertiesFromUniqueColumn(translated.getPropertiesFromUniqueColumn()); + } + return builder.build(); } @@ -650,6 +698,7 @@ else if (inputProperties.stream().anyMatch(ActualProperties::isSingleNode)) { .global(coordinatorOnly ? coordinatorSingleStreamPartition() : singleStreamPartition()) .local(localProperties.build()) .constants(constants) + .propertiesFromUniqueColumn(additionalPropertyIsUnique ? translated.getPropertiesFromUniqueColumn() : Optional.empty()) .build(); case REPARTITION: { Global globalPartitioning; @@ -666,6 +715,7 @@ else if (inputProperties.stream().anyMatch(ActualProperties::isSingleNode)) { return ActualProperties.builder() .global(globalPartitioning) .constants(constants) + .propertiesFromUniqueColumn(additionalPropertyIsUnique ? translated.getPropertiesFromUniqueColumn() : Optional.empty()) .build(); } case REPLICATE: @@ -691,6 +741,7 @@ public ActualProperties visitFilter(FilterNode node, List inpu return ActualProperties.builderFrom(properties) .constants(constants) + .propertiesFromUniqueColumn(properties.getPropertiesFromUniqueColumn()) .build(); } @@ -728,6 +779,7 @@ else if (!(value instanceof RowExpression)) { return ActualProperties.builderFrom(translatedProperties) .constants(constants) + .propertiesFromUniqueColumn(properties.getPropertiesFromUniqueColumn().map(x -> x.translateRowExpression(node.getAssignments().getMap()))) .build(); } @@ -814,6 +866,13 @@ public ActualProperties visitTableScan(TableScanNode node, List Optional.ofNullable(assignments.get(column)))); + if (isUtilizeUniquePropertyInQueryPlanningEnabled(session) && layout.getUniqueColumn().isPresent() && assignments.containsKey(layout.getUniqueColumn().get())) { + VariableReferenceExpression uniqueVariable = assignments.get(layout.getUniqueColumn().get()); + ActualProperties.Builder propertiesFromUniqueColumn = ActualProperties.builder(); + propertiesFromUniqueColumn.global(partitionedOn(ARBITRARY_DISTRIBUTION, ImmutableList.of(uniqueVariable), Optional.empty())); + propertiesFromUniqueColumn.local(LocalProperties.unique(uniqueVariable)); + properties.propertiesFromUniqueColumn(Optional.of(propertiesFromUniqueColumn.build())); + } return properties.build(); } diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/StreamPreferredProperties.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/StreamPreferredProperties.java index 221d224ea6d40..96288e6af729c 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/StreamPreferredProperties.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/StreamPreferredProperties.java @@ -188,9 +188,11 @@ else if (actualProperties.getDistribution() == SINGLE) { // is there a preference for a specific partitioning scheme? if (partitioningColumns.isPresent()) { if (exactColumnOrder) { - return actualProperties.isExactlyPartitionedOn(partitioningColumns.get()); + return actualProperties.isExactlyPartitionedOn(partitioningColumns.get()) + || actualProperties.getStreamPropertiesFromUniqueColumn().isPresent() && actualProperties.getStreamPropertiesFromUniqueColumn().get().isExactlyPartitionedOn(partitioningColumns.get()); } - return actualProperties.isPartitionedOn(partitioningColumns.get()); + return actualProperties.isPartitionedOn(partitioningColumns.get()) + || actualProperties.getStreamPropertiesFromUniqueColumn().isPresent() && actualProperties.getStreamPropertiesFromUniqueColumn().get().isPartitionedOn(partitioningColumns.get()); } return true; diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/StreamPropertyDerivations.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/StreamPropertyDerivations.java index 58f8e32977de6..51e6d3de35ab2 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/StreamPropertyDerivations.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/StreamPropertyDerivations.java @@ -18,6 +18,7 @@ import com.facebook.presto.metadata.TableLayout; import com.facebook.presto.spi.ColumnHandle; import com.facebook.presto.spi.LocalProperty; +import com.facebook.presto.spi.UniqueProperty; import com.facebook.presto.spi.plan.AggregationNode; import com.facebook.presto.spi.plan.DeleteNode; import com.facebook.presto.spi.plan.DistinctLimitNode; @@ -78,6 +79,7 @@ import java.util.function.Function; import java.util.stream.Collectors; +import static com.facebook.presto.SystemSessionProperties.isUtilizeUniquePropertyInQueryPlanningEnabled; import static com.facebook.presto.sql.planner.SystemPartitioningHandle.FIXED_ARBITRARY_DISTRIBUTION; import static com.facebook.presto.sql.planner.optimizations.PropertyDerivations.extractFixedValuesToConstantExpressions; import static com.facebook.presto.sql.planner.optimizations.StreamPropertyDerivations.StreamProperties.StreamDistribution.FIXED; @@ -123,7 +125,10 @@ public static StreamProperties deriveProperties(PlanNode node, List properties.otherActualProperties) + .map(properties -> { + checkState(properties.otherActualProperties.isPresent(), "otherActualProperties should always exist"); + return properties.otherActualProperties.get(); + }) .collect(toImmutableList()), metadata, session); @@ -225,11 +230,8 @@ public StreamProperties visitIndexJoin(IndexJoinNode node, List streamPropertiesFromUniqueColumn = Optional.empty(); + if (isUtilizeUniquePropertyInQueryPlanningEnabled(session) && layout.getUniqueColumn().isPresent() && assignments.containsKey(layout.getUniqueColumn().get())) { + streamPropertiesFromUniqueColumn = Optional.of(new StreamProperties(streamDistribution, Optional.of(ImmutableList.of(assignments.get(layout.getUniqueColumn().get()))), false)); + } + // if we are partitioned on empty set, we must say multiple of unknown partitioning, because // the connector does not guarantee a single split in this case (since it might not understand // that the value is a constant). if (streamPartitionSymbols.isPresent() && streamPartitionSymbols.get().isEmpty()) { - return new StreamProperties(streamDistribution, Optional.empty(), false); + return new StreamProperties(streamDistribution, Optional.empty(), false, Optional.empty(), streamPropertiesFromUniqueColumn); } - return new StreamProperties(streamDistribution, streamPartitionSymbols, false); + return new StreamProperties(streamDistribution, streamPartitionSymbols, false, Optional.empty(), streamPropertiesFromUniqueColumn); } private Optional> getNonConstantVariables(Set columnHandles, Map assignments, Set globalConstants) @@ -337,20 +344,32 @@ public StreamProperties visitExchange(ExchangeNode node, List return new StreamProperties(MULTIPLE, Optional.empty(), false); } + Optional additionalUniqueProperty = Optional.empty(); + if (inputProperties.size() == 1 && inputProperties.get(0).hasUniqueProperties() && !node.getType().equals(ExchangeNode.Type.REPLICATE)) { + List inputVariables = node.getInputs().get(0); + Map inputToOutput = new HashMap<>(); + for (int i = 0; i < node.getOutputVariables().size(); i++) { + inputToOutput.put(inputVariables.get(i), node.getOutputVariables().get(i)); + } + checkState(inputProperties.get(0).getStreamPropertiesFromUniqueColumn().isPresent(), + "when unique columns exists, the stream is also partitioned by the unique column and should be represented in the streamPropertiesFromUniqueColumn field"); + additionalUniqueProperty = Optional.of(inputProperties.get(0).getStreamPropertiesFromUniqueColumn().get().translate(column -> Optional.ofNullable(inputToOutput.get(column)))); + } + if (node.getScope().isRemote()) { // TODO: correctly determine if stream is parallelised // based on session properties - return StreamProperties.fixedStreams(); + return StreamProperties.fixedStreams().withStreamPropertiesFromUniqueColumn(additionalUniqueProperty); } switch (node.getType()) { case GATHER: - return StreamProperties.singleStream(); + return StreamProperties.singleStream().withStreamPropertiesFromUniqueColumn(additionalUniqueProperty); case REPARTITION: if (node.getPartitioningScheme().getPartitioning().getHandle().equals(FIXED_ARBITRARY_DISTRIBUTION) || // no strict partitioning guarantees when multiple writers per partitions are allows (scaled writers) node.getPartitioningScheme().isScaleWriters()) { - return new StreamProperties(FIXED, Optional.empty(), false); + return new StreamProperties(FIXED, Optional.empty(), false).withStreamPropertiesFromUniqueColumn(additionalUniqueProperty); } checkArgument( node.getPartitioningScheme().getPartitioning().getArguments().stream().allMatch(VariableReferenceExpression.class::isInstance), @@ -359,7 +378,7 @@ public StreamProperties visitExchange(ExchangeNode node, List FIXED, Optional.of(node.getPartitioningScheme().getPartitioning().getArguments().stream() .map(VariableReferenceExpression.class::cast) - .collect(toImmutableList())), false); + .collect(toImmutableList())), false).withStreamPropertiesFromUniqueColumn(additionalUniqueProperty); case REPLICATE: return new StreamProperties(MULTIPLE, Optional.empty(), false); } @@ -668,21 +687,33 @@ public enum StreamDistribution // We are only interested in the local properties, but PropertyDerivations requires input // ActualProperties, so we hold on to the whole object - private final ActualProperties otherActualProperties; + private final Optional otherActualProperties; // NOTE: Partitioning on zero columns (or effectively zero columns if the columns are constant) indicates that all // the rows will be partitioned into a single stream. + private final Optional streamPropertiesFromUniqueColumn; + private StreamProperties(StreamDistribution distribution, Optional> partitioningColumns, boolean ordered) { - this(distribution, partitioningColumns, ordered, null); + this(distribution, partitioningColumns, ordered, Optional.empty()); + } + + private StreamProperties( + StreamDistribution distribution, + Optional> partitioningColumns, + boolean ordered, + Optional otherActualProperties) + { + this(distribution, partitioningColumns, ordered, otherActualProperties, Optional.empty()); } private StreamProperties( StreamDistribution distribution, Optional> partitioningColumns, boolean ordered, - ActualProperties otherActualProperties) + Optional otherActualProperties, + Optional streamPropertiesFromUniqueColumn) { this.distribution = requireNonNull(distribution, "distribution is null"); @@ -697,13 +728,38 @@ private StreamProperties( this.ordered = ordered; checkArgument(!ordered || distribution == SINGLE, "Ordered must be a single stream"); - this.otherActualProperties = otherActualProperties; + this.otherActualProperties = requireNonNull(otherActualProperties); + requireNonNull(streamPropertiesFromUniqueColumn).ifPresent(properties -> checkArgument(!properties.streamPropertiesFromUniqueColumn.isPresent())); + // When unique properties exists, the stream is also partitioned on the unique column + if (otherActualProperties.isPresent() && otherActualProperties.get().getPropertiesFromUniqueColumn().isPresent()) { + ActualProperties propertiesFromUniqueColumn = otherActualProperties.get().getPropertiesFromUniqueColumn().get(); + if (Iterables.getOnlyElement(propertiesFromUniqueColumn.getLocalProperties()) instanceof UniqueProperty) { + VariableReferenceExpression uniqueVariable = (VariableReferenceExpression) ((UniqueProperty) Iterables.getOnlyElement(propertiesFromUniqueColumn.getLocalProperties())).getColumn(); + checkState(streamPropertiesFromUniqueColumn.isPresent() && streamPropertiesFromUniqueColumn.get().partitioningColumns.isPresent() + && Iterables.getOnlyElement(streamPropertiesFromUniqueColumn.get().partitioningColumns.get()).equals(uniqueVariable)); + } + } + this.streamPropertiesFromUniqueColumn = streamPropertiesFromUniqueColumn; } public List> getLocalProperties() { - checkState(otherActualProperties != null, "otherActualProperties not set"); - return otherActualProperties.getLocalProperties(); + checkState(otherActualProperties.isPresent(), "otherActualProperties not set"); + return otherActualProperties.get().getLocalProperties(); + } + + public List> getAdditionalLocalProperties() + { + checkState(otherActualProperties.isPresent(), "otherActualProperties not set"); + if (!otherActualProperties.get().getPropertiesFromUniqueColumn().isPresent()) { + return ImmutableList.of(); + } + return otherActualProperties.get().getPropertiesFromUniqueColumn().get().getLocalProperties(); + } + + public Optional getStreamPropertiesFromUniqueColumn() + { + return streamPropertiesFromUniqueColumn; } private static StreamProperties singleStream() @@ -725,8 +781,9 @@ private StreamProperties unordered(boolean unordered) { if (unordered) { ActualProperties updatedProperties = null; - if (otherActualProperties != null) { - updatedProperties = ActualProperties.builderFrom(otherActualProperties) + if (otherActualProperties.isPresent()) { + updatedProperties = ActualProperties.builderFrom(otherActualProperties.get()) + .propertiesFromUniqueColumn(otherActualProperties.get().getPropertiesFromUniqueColumn().map(x -> ActualProperties.builderFrom(x).unordered(true).build())) .unordered(true) .build(); } @@ -734,11 +791,33 @@ private StreamProperties unordered(boolean unordered) distribution, partitioningColumns, false, - updatedProperties); + Optional.ofNullable(updatedProperties), + streamPropertiesFromUniqueColumn.map(x -> x.unordered(true))); } return this; } + public StreamProperties uniqueToGroupProperties() + { + if (otherActualProperties.isPresent() && otherActualProperties.get().getPropertiesFromUniqueColumn().isPresent()) { + if (Iterables.getOnlyElement(otherActualProperties.get().getPropertiesFromUniqueColumn().get().getLocalProperties()) instanceof UniqueProperty) { + Optional groupedProperties = PropertyDerivations.uniqueToGroupProperties(otherActualProperties.get().getPropertiesFromUniqueColumn().get()); + return new StreamProperties(distribution, partitioningColumns, ordered, + otherActualProperties.map(x -> ActualProperties.builderFrom(x).propertiesFromUniqueColumn(groupedProperties).build()), + streamPropertiesFromUniqueColumn.map(StreamProperties::uniqueToGroupProperties)); + } + } + return this; + } + + public boolean hasUniqueProperties() + { + if (otherActualProperties.isPresent() && otherActualProperties.get().getPropertiesFromUniqueColumn().isPresent()) { + return Iterables.getOnlyElement(otherActualProperties.get().getPropertiesFromUniqueColumn().get().getLocalProperties()) instanceof UniqueProperty; + } + return false; + } + public boolean isSingleStream() { return distribution == SINGLE; @@ -783,7 +862,12 @@ private StreamProperties withUnspecifiedPartitioning() private StreamProperties withOtherActualProperties(ActualProperties actualProperties) { - return new StreamProperties(distribution, partitioningColumns, ordered, actualProperties); + return new StreamProperties(distribution, partitioningColumns, ordered, Optional.ofNullable(actualProperties), streamPropertiesFromUniqueColumn); + } + + private StreamProperties withStreamPropertiesFromUniqueColumn(Optional streamPropertiesFromUniqueColumn) + { + return new StreamProperties(distribution, partitioningColumns, ordered, otherActualProperties, streamPropertiesFromUniqueColumn); } public StreamProperties translate(Function> translator) @@ -801,7 +885,8 @@ public StreamProperties translate(Function x.translateVariable(translator)), + streamPropertiesFromUniqueColumn.map(x -> x.translate(translator))); } public Optional> getPartitioningColumns() @@ -812,7 +897,7 @@ public Optional> getPartitioningColumns() @Override public int hashCode() { - return Objects.hash(distribution, partitioningColumns); + return Objects.hash(distribution, partitioningColumns, ordered, otherActualProperties, streamPropertiesFromUniqueColumn); } @Override @@ -826,7 +911,10 @@ public boolean equals(Object obj) } StreamProperties other = (StreamProperties) obj; return Objects.equals(this.distribution, other.distribution) && - Objects.equals(this.partitioningColumns, other.partitioningColumns); + Objects.equals(this.partitioningColumns, other.partitioningColumns) && + this.ordered == other.ordered && + Objects.equals(this.otherActualProperties, other.otherActualProperties) && + Objects.equals(this.streamPropertiesFromUniqueColumn, other.streamPropertiesFromUniqueColumn); } @Override @@ -835,6 +923,9 @@ public String toString() return toStringHelper(this) .add("distribution", distribution) .add("partitioningColumns", partitioningColumns) + .add("ordered", ordered) + .add("otherActualProperties", otherActualProperties) + .add("streamPropertiesFromUniqueColumn", streamPropertiesFromUniqueColumn) .toString(); } } diff --git a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/sanity/ValidateStreamingAggregations.java b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/sanity/ValidateStreamingAggregations.java index cfe040826bdcf..f5ca4ab176918 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/sql/planner/sanity/ValidateStreamingAggregations.java +++ b/presto-main-base/src/main/java/com/facebook/presto/sql/planner/sanity/ValidateStreamingAggregations.java @@ -86,8 +86,10 @@ public Void visitAggregation(AggregationNode node, Void context) List> desiredProperties = ImmutableList.of(new GroupingProperty<>(node.getPreGroupedVariables())); Iterator>> matchIterator = LocalProperties.match(properties.getLocalProperties(), desiredProperties).iterator(); + Iterator>> additionalMatchIterator = LocalProperties.match(properties.getAdditionalLocalProperties(), desiredProperties).iterator(); Optional> unsatisfiedRequirement = Iterators.getOnlyElement(matchIterator); - checkArgument(!unsatisfiedRequirement.isPresent(), "Streaming aggregation with input not grouped on the grouping keys"); + Optional> additionalUnsatisfiedRequirement = Iterators.getOnlyElement(additionalMatchIterator); + checkArgument(!unsatisfiedRequirement.isPresent() || !additionalUnsatisfiedRequirement.isPresent(), "Streaming aggregation with input not grouped on the grouping keys"); return null; } } diff --git a/presto-main-base/src/test/java/com/facebook/presto/sql/analyzer/TestFeaturesConfig.java b/presto-main-base/src/test/java/com/facebook/presto/sql/analyzer/TestFeaturesConfig.java index cc29fe7160751..56be62a79306a 100644 --- a/presto-main-base/src/test/java/com/facebook/presto/sql/analyzer/TestFeaturesConfig.java +++ b/presto-main-base/src/test/java/com/facebook/presto/sql/analyzer/TestFeaturesConfig.java @@ -263,6 +263,7 @@ public void testDefaults() .setAddExchangeBelowPartialAggregationOverGroupId(false) .setAddDistinctBelowSemiJoinBuild(false) .setPushdownSubfieldForMapFunctions(true) + .setUtilizeUniquePropertyInQueryPlanning(true) .setInnerJoinPushdownEnabled(false) .setBroadcastSemiJoinForDelete(true) .setInEqualityJoinPushdownEnabled(false) @@ -482,6 +483,7 @@ public void testExplicitPropertyMappings() .put("exclude-invalid-worker-session-properties", "true") .put("optimizer.add-distinct-below-semi-join-build", "true") .put("optimizer.pushdown-subfield-for-map-functions", "false") + .put("optimizer.utilize-unique-property-in-query-planning", "false") .put("optimizer.add-exchange-below-partial-aggregation-over-group-id", "true") .put("max_serializable_object_size", "50") .build(); @@ -692,6 +694,7 @@ public void testExplicitPropertyMappings() .setAddExchangeBelowPartialAggregationOverGroupId(true) .setAddDistinctBelowSemiJoinBuild(true) .setPushdownSubfieldForMapFunctions(false) + .setUtilizeUniquePropertyInQueryPlanning(false) .setInEqualityJoinPushdownEnabled(true) .setBroadcastSemiJoinForDelete(false) .setRewriteMinMaxByToTopNEnabled(true) diff --git a/presto-main-base/src/test/java/com/facebook/presto/sql/planner/assertions/PlanMatchPattern.java b/presto-main-base/src/test/java/com/facebook/presto/sql/planner/assertions/PlanMatchPattern.java index 9c6eb9a528e56..c69c4ea815c92 100644 --- a/presto-main-base/src/test/java/com/facebook/presto/sql/planner/assertions/PlanMatchPattern.java +++ b/presto-main-base/src/test/java/com/facebook/presto/sql/planner/assertions/PlanMatchPattern.java @@ -412,6 +412,11 @@ public static PlanMatchPattern strictProject(Map assi .withExactAssignments(assignments.values()); } + public static PlanMatchPattern semiJoin(PlanMatchPattern source, PlanMatchPattern filtering) + { + return node(SemiJoinNode.class, source, filtering); + } + public static PlanMatchPattern semiJoin(String sourceSymbolAlias, String filteringSymbolAlias, String outputAlias, PlanMatchPattern source, PlanMatchPattern filtering) { return semiJoin(sourceSymbolAlias, filteringSymbolAlias, outputAlias, Optional.empty(), source, filtering); diff --git a/presto-main-base/src/test/java/com/facebook/presto/sql/planner/optimizations/TestLocalProperties.java b/presto-main-base/src/test/java/com/facebook/presto/sql/planner/optimizations/TestLocalProperties.java index db0600d2d6380..edbeeb3e0cdc2 100644 --- a/presto-main-base/src/test/java/com/facebook/presto/sql/planner/optimizations/TestLocalProperties.java +++ b/presto-main-base/src/test/java/com/facebook/presto/sql/planner/optimizations/TestLocalProperties.java @@ -20,6 +20,7 @@ import com.facebook.presto.spi.GroupingProperty; import com.facebook.presto.spi.LocalProperty; import com.facebook.presto.spi.SortingProperty; +import com.facebook.presto.spi.UniqueProperty; import com.facebook.presto.spi.relation.VariableReferenceExpression; import com.facebook.presto.testing.TestingMetadata.TestingColumnHandle; import com.fasterxml.jackson.core.JsonParser; @@ -78,6 +79,10 @@ public void testConstantProcessing() input = ImmutableList.of(constant("a"), constant("b")); assertEquals(stripLeadingConstants(input), ImmutableList.of()); assertEquals(extractLeadingConstants(input), ImmutableSet.of("a", "b")); + + input = ImmutableList.of(unique("a")); + assertEquals(stripLeadingConstants(input), ImmutableList.of(unique("a"))); + assertEquals(extractLeadingConstants(input), ImmutableSet.of()); } @Test @@ -138,6 +143,14 @@ public void testTranslate() map = ImmutableMap.of("a", "a1", "b", "b1", "c", "c1"); input = ImmutableList.of(grouped("a"), constant("b"), grouped("c")); assertEquals(LocalProperties.translate(input, translateWithMap(map)), ImmutableList.of(grouped("a1"), constant("b1"), grouped("c1"))); + + map = ImmutableMap.of(); + input = ImmutableList.of(unique("a")); + assertEquals(LocalProperties.translate(input, translateWithMap(map)), ImmutableList.of()); + + map = ImmutableMap.of("a", "a1"); + input = ImmutableList.of(unique("a")); + assertEquals(LocalProperties.translate(input, translateWithMap(map)), ImmutableList.of(unique("a1"))); } private static Function> translateWithMap(Map translateMap) @@ -177,6 +190,35 @@ public void testNormalizeOverlappingSymbol() assertNormalizeAndFlatten( localProperties, grouped("a")); + + localProperties = builder() + .unique("a") + .sorted("a", SortOrder.ASC_NULLS_FIRST) + .constant("a") + .build(); + assertNormalize( + localProperties, + Optional.of(unique("a")), + Optional.empty(), + Optional.empty()); + assertNormalizeAndFlatten( + localProperties, + unique("a")); + + localProperties = builder() + .grouped("a") + .unique("a") + .constant("a") + .build(); + assertNormalize( + localProperties, + Optional.of(grouped("a")), + Optional.of(unique("a")), + Optional.empty()); + assertNormalizeAndFlatten( + localProperties, + grouped("a"), + unique("a")); } @Test @@ -780,6 +822,11 @@ private static GroupingProperty grouped(String... columns) return new GroupingProperty<>(Arrays.asList(columns)); } + private static UniqueProperty unique(String column) + { + return new UniqueProperty<>(column); + } + private static SortingProperty sorted(String column, SortOrder order) { return new SortingProperty<>(column, order); @@ -812,6 +859,12 @@ public Builder constant(String column) return this; } + public Builder unique(String column) + { + properties.add(new UniqueProperty<>(column)); + return this; + } + public List> build() { return new ArrayList<>(properties); diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/ConnectorTableLayout.java b/presto-spi/src/main/java/com/facebook/presto/spi/ConnectorTableLayout.java index f355b52f7e993..d794088f14f79 100644 --- a/presto-spi/src/main/java/com/facebook/presto/spi/ConnectorTableLayout.java +++ b/presto-spi/src/main/java/com/facebook/presto/spi/ConnectorTableLayout.java @@ -34,6 +34,7 @@ public class ConnectorTableLayout private final Optional discretePredicates; private final List> localProperties; private final Optional remainingPredicate; + private final Optional uniqueColumn; public ConnectorTableLayout(ConnectorTableLayoutHandle handle) { @@ -68,6 +69,20 @@ public ConnectorTableLayout( Optional discretePredicates, List> localProperties, Optional remainingPredicate) + { + this(handle, columns, predicate, tablePartitioning, streamPartitioningColumns, discretePredicates, localProperties, remainingPredicate, Optional.empty()); + } + + public ConnectorTableLayout( + ConnectorTableLayoutHandle handle, + Optional> columns, + TupleDomain predicate, + Optional tablePartitioning, + Optional> streamPartitioningColumns, + Optional discretePredicates, + List> localProperties, + Optional remainingPredicate, + Optional uniqueColumn) { requireNonNull(handle, "handle is null"); requireNonNull(columns, "columns is null"); @@ -86,6 +101,7 @@ public ConnectorTableLayout( this.discretePredicates = discretePredicates; this.localProperties = localProperties; this.remainingPredicate = remainingPredicate; + this.uniqueColumn = uniqueColumn; } public ConnectorTableLayoutHandle getHandle() @@ -123,6 +139,11 @@ public Optional getRemainingPredicate() return remainingPredicate; } + public Optional getUniqueColumn() + { + return uniqueColumn; + } + /** * The partitioning of the table across the worker nodes. *

diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/UniqueProperty.java b/presto-spi/src/main/java/com/facebook/presto/spi/UniqueProperty.java new file mode 100644 index 0000000000000..946cdb6dc1dab --- /dev/null +++ b/presto-spi/src/main/java/com/facebook/presto/spi/UniqueProperty.java @@ -0,0 +1,98 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.spi; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; + +import java.util.Collections; +import java.util.Objects; +import java.util.Optional; +import java.util.Set; +import java.util.function.Function; + +import static java.util.Objects.requireNonNull; + +public final class UniqueProperty + implements LocalProperty +{ + private final E column; + + @JsonCreator + public UniqueProperty(@JsonProperty("column") E column) + { + this.column = requireNonNull(column, "column is null"); + } + + @Override + public boolean isOrderSensitive() + { + return false; + } + + @JsonProperty + public E getColumn() + { + return column; + } + + public Set getColumns() + { + return Collections.singleton(column); + } + + @Override + public Optional> translate(Function> translator) + { + return translator.apply(column) + .map(UniqueProperty::new); + } + + @Override + public boolean isSimplifiedBy(LocalProperty known) + { + return known instanceof UniqueProperty && known.equals(this); + } + + @Override + public Optional> withConstants(Set constants) + { + return Optional.of(this); + } + + @Override + public String toString() + { + return "U(" + column + ")"; + } + + @Override + public boolean equals(Object o) + { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + UniqueProperty that = (UniqueProperty) o; + return Objects.equals(column, that.column); + } + + @Override + public int hashCode() + { + return Objects.hash(column); + } +} From 0788f73f6c95378261641274d6cd3969d4df5375 Mon Sep 17 00:00:00 2001 From: Ke Wang Date: Tue, 16 Sep 2025 02:09:12 -0700 Subject: [PATCH 111/113] [native] Introduce firstTimeReceiveTaskUpdateMs in PrestoTask PrestoTask can be created in different endpoints: - getTaskStatus - getTaskInfo - receive task update etc PrestoTask can be created in getTaskStatus, but it won't be able to create velox plan and start. It has to wait until receiving taskUpdate Make taskCreationTime represent the time between receiving first taskUpdate and task creation time --- presto-native-execution/presto_cpp/main/PrestoTask.cpp | 3 ++- presto-native-execution/presto_cpp/main/PrestoTask.h | 2 ++ presto-native-execution/presto_cpp/main/TaskManager.cpp | 6 ++++++ 3 files changed, 10 insertions(+), 1 deletion(-) diff --git a/presto-native-execution/presto_cpp/main/PrestoTask.cpp b/presto-native-execution/presto_cpp/main/PrestoTask.cpp index c48ace86bad8e..3621a763bdbdb 100644 --- a/presto-native-execution/presto_cpp/main/PrestoTask.cpp +++ b/presto-native-execution/presto_cpp/main/PrestoTask.cpp @@ -769,7 +769,8 @@ void PrestoTask::updateTimeInfoLocked( taskRuntimeStats["endTime"].addValue(veloxTaskStats.endTimeMs); } taskRuntimeStats.insert({"nativeProcessCpuTime", fromNanos(processCpuTime_)}); - taskRuntimeStats.insert({"taskCreationTime", fromNanos((createFinishTimeMs - createTimeMs) * 1'000'000)}); + // Represents the time between receiving first taskUpdate and task creation time + taskRuntimeStats.insert({"taskCreationTime", fromNanos((createFinishTimeMs - firstTimeReceiveTaskUpdateMs) * 1'000'000)}); } void PrestoTask::updateMemoryInfoLocked( diff --git a/presto-native-execution/presto_cpp/main/PrestoTask.h b/presto-native-execution/presto_cpp/main/PrestoTask.h index 3732ccffedde2..b41666fa90a93 100644 --- a/presto-native-execution/presto_cpp/main/PrestoTask.h +++ b/presto-native-execution/presto_cpp/main/PrestoTask.h @@ -123,6 +123,8 @@ struct PrestoTask { uint64_t lastMemoryReservation{0}; /// Time point (in ms) when the time we start task creating. uint64_t createTimeMs{0}; + /// Time point (in ms) when the first time we receive task update. + uint64_t firstTimeReceiveTaskUpdateMs{0}; /// Time point (in ms) when the time we finish task creating. uint64_t createFinishTimeMs{0}; uint64_t startTimeMs{0}; diff --git a/presto-native-execution/presto_cpp/main/TaskManager.cpp b/presto-native-execution/presto_cpp/main/TaskManager.cpp index da39aefd316f1..0b7fb7cb87b3c 100644 --- a/presto-native-execution/presto_cpp/main/TaskManager.cpp +++ b/presto-native-execution/presto_cpp/main/TaskManager.cpp @@ -536,9 +536,13 @@ std::unique_ptr TaskManager::createOrUpdateTaskImpl( bool summarize, std::shared_ptr queryCtx, long startProcessCpuTime) { + auto receiveTaskUpdateMs = getCurrentTimeMs(); std::shared_ptr execTask; bool startTask = false; auto prestoTask = findOrCreateTask(taskId, startProcessCpuTime); + if (prestoTask->firstTimeReceiveTaskUpdateMs == 0) { + prestoTask->firstTimeReceiveTaskUpdateMs = receiveTaskUpdateMs; + } { std::lock_guard l(prestoTask->mutex); prestoTask->updateCoordinatorHeartbeatLocked(); @@ -770,6 +774,8 @@ void TaskManager::startTaskLocked(std::shared_ptr& prestoTask) { // Record the time we spent between task creation and start, which is the // planned (queued) time. + // Note task could be created at getTaskStatus/getTaskInfo endpoint and later + // receive taskUpdate to create and start task. const auto queuedTimeInMs = velox::getCurrentTimeMs() - prestoTask->createTimeMs; prestoTask->info.stats.queuedTimeInNanos = queuedTimeInMs * 1'000'000; From 41016fee70319839e0b74e894c04bbb2cc75186f Mon Sep 17 00:00:00 2001 From: Jialiang Tan Date: Tue, 16 Sep 2025 16:49:16 -0700 Subject: [PATCH 112/113] [pos] Fix TestPrestoSparkNativeArrayFunctionQueries --- ...PrestoSparkNativeArrayFunctionQueries.java | 19 ++++++++------- .../TestPrestoSparkNativeGeneralQueries.java | 23 +++++++++---------- 2 files changed, 22 insertions(+), 20 deletions(-) diff --git a/presto-native-execution/src/test/java/com/facebook/presto/spark/TestPrestoSparkNativeArrayFunctionQueries.java b/presto-native-execution/src/test/java/com/facebook/presto/spark/TestPrestoSparkNativeArrayFunctionQueries.java index 89871668882e5..f513715b25f79 100644 --- a/presto-native-execution/src/test/java/com/facebook/presto/spark/TestPrestoSparkNativeArrayFunctionQueries.java +++ b/presto-native-execution/src/test/java/com/facebook/presto/spark/TestPrestoSparkNativeArrayFunctionQueries.java @@ -14,9 +14,9 @@ package com.facebook.presto.spark; import com.facebook.presto.nativeworker.AbstractTestNativeArrayFunctionQueries; +import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin; import com.facebook.presto.testing.ExpectedQueryRunner; import com.facebook.presto.testing.QueryRunner; -import org.testng.annotations.Ignore; public class TestPrestoSparkNativeArrayFunctionQueries extends AbstractTestNativeArrayFunctionQueries @@ -24,18 +24,21 @@ public class TestPrestoSparkNativeArrayFunctionQueries @Override protected QueryRunner createQueryRunner() { - return PrestoSparkNativeQueryRunnerUtils.createHiveRunner(); + QueryRunner queryRunner = PrestoSparkNativeQueryRunnerUtils.createHiveRunner(); + + // Install plugins needed for extra array functions. + queryRunner.installPlugin(new SqlInvokedFunctionsPlugin()); + return queryRunner; } @Override protected ExpectedQueryRunner createExpectedQueryRunner() throws Exception { - return PrestoSparkNativeQueryRunnerUtils.createJavaQueryRunner(); - } + QueryRunner queryRunner = PrestoSparkNativeQueryRunnerUtils.createJavaQueryRunner(); - // Caused by: com.facebook.presto.sql.analyzer.SemanticException: line 1:32: Function array_sort_desc not registered - @Override - @Ignore - public void testArraySort() {} + // Install plugins needed for extra array functions. + queryRunner.installPlugin(new SqlInvokedFunctionsPlugin()); + return queryRunner; + } } diff --git a/presto-native-execution/src/test/java/com/facebook/presto/spark/TestPrestoSparkNativeGeneralQueries.java b/presto-native-execution/src/test/java/com/facebook/presto/spark/TestPrestoSparkNativeGeneralQueries.java index 448ddde08bd1f..d67dba20dfa93 100644 --- a/presto-native-execution/src/test/java/com/facebook/presto/spark/TestPrestoSparkNativeGeneralQueries.java +++ b/presto-native-execution/src/test/java/com/facebook/presto/spark/TestPrestoSparkNativeGeneralQueries.java @@ -16,6 +16,7 @@ import com.facebook.airlift.log.Level; import com.facebook.airlift.log.Logging; import com.facebook.presto.nativeworker.AbstractTestNativeGeneralQueries; +import com.facebook.presto.scalar.sql.SqlInvokedFunctionsPlugin; import com.facebook.presto.testing.ExpectedQueryRunner; import com.facebook.presto.testing.QueryRunner; import org.testng.annotations.Ignore; @@ -31,14 +32,22 @@ public class TestPrestoSparkNativeGeneralQueries @Override protected QueryRunner createQueryRunner() { - return PrestoSparkNativeQueryRunnerUtils.createHiveRunner(); + QueryRunner queryRunner = PrestoSparkNativeQueryRunnerUtils.createHiveRunner(); + + // Install plugins needed for extra array functions. + queryRunner.installPlugin(new SqlInvokedFunctionsPlugin()); + return queryRunner; } @Override protected ExpectedQueryRunner createExpectedQueryRunner() throws Exception { - return PrestoSparkNativeQueryRunnerUtils.createJavaQueryRunner(); + QueryRunner queryRunner = PrestoSparkNativeQueryRunnerUtils.createJavaQueryRunner(); + + // Install plugins needed for extra array functions. + queryRunner.installPlugin(new SqlInvokedFunctionsPlugin()); + return queryRunner; } @Override @@ -107,18 +116,8 @@ public void testRowWiseExchange() {} @Ignore public void testAnalyzeStatsOnDecimals() {} - //Caused by: com.facebook.presto.sql.analyzer.SemanticException: line 1:31: Function array_duplicates not registered - @Override - @Ignore - public void testArrayAndMapFunctions() {} - // VeloxRuntimeError: it != connectors().end() Connector with ID 'hivecached' not registered @Override @Ignore public void testCatalogWithCacheEnabled() {} - - // Caused by: com.facebook.presto.spi.PrestoException: Sampling function: key_sampling_percent not cannot be resolved - @Override - @Ignore - public void testKeyBasedSamplingInlined() {} } From 6f1fd8a2851664e00821d8e44f5435847c37cd6a Mon Sep 17 00:00:00 2001 From: Vamsi Karnika Date: Wed, 17 Sep 2025 21:46:59 +0530 Subject: [PATCH 113/113] Use batch call to get partition information from the metastore --- .../facebook/presto/hudi/HudiMetadata.java | 3 +- .../presto/hudi/HudiPartitionManager.java | 44 +++++++-- .../presto/hudi/HudiSplitManager.java | 45 +-------- .../facebook/presto/hudi/HudiSplitSource.java | 8 +- .../facebook/presto/hudi/HudiTableHandle.java | 23 +++++ .../hudi/split/HudiBackgroundSplitLoader.java | 16 ++-- .../split/HudiPartitionSplitGenerator.java | 60 +++++++++--- .../presto/hudi/TestHudiPartitionManager.java | 96 +++++++++++++++++-- .../hudi/TestingExtendedHiveMetastore.java | 15 ++- 9 files changed, 221 insertions(+), 89 deletions(-) diff --git a/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiMetadata.java b/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiMetadata.java index 24a04a82c44e1..ac6018c3f142c 100644 --- a/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiMetadata.java +++ b/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiMetadata.java @@ -101,6 +101,7 @@ public ConnectorTableHandle getTableHandle(ConnectorSession session, SchemaTable } return new HudiTableHandle( + Optional.of(table), table.getDatabaseName(), table.getTableName(), table.getStorage().getLocation(), @@ -245,7 +246,7 @@ static List fromPartitionColumns(List partitionColumns return builder.build(); } - static List fromDataColumns(List dataColumns) + public static List fromDataColumns(List dataColumns) { ImmutableList.Builder builder = ImmutableList.builder(); int id = 0; diff --git a/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiPartitionManager.java b/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiPartitionManager.java index ba80fd2f516c7..0c1451d89479b 100644 --- a/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiPartitionManager.java +++ b/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiPartitionManager.java @@ -14,6 +14,7 @@ package com.facebook.presto.hudi; +import com.facebook.airlift.log.Logger; import com.facebook.presto.common.predicate.Domain; import com.facebook.presto.common.predicate.NullableValue; import com.facebook.presto.common.predicate.TupleDomain; @@ -23,6 +24,8 @@ import com.facebook.presto.hive.metastore.Column; import com.facebook.presto.hive.metastore.ExtendedHiveMetastore; import com.facebook.presto.hive.metastore.MetastoreContext; +import com.facebook.presto.hive.metastore.Partition; +import com.facebook.presto.hive.metastore.StorageFormat; import com.facebook.presto.hive.metastore.Table; import com.facebook.presto.spi.ColumnHandle; import com.facebook.presto.spi.ConnectorSession; @@ -45,10 +48,10 @@ import static com.facebook.presto.hudi.HudiMetadata.fromPartitionColumns; import static com.facebook.presto.hudi.HudiMetadata.toMetastoreContext; import static java.util.Objects.requireNonNull; -import static java.util.stream.Collectors.toList; public class HudiPartitionManager { + private static final Logger log = Logger.get(HudiPartitionManager.class); private final TypeManager typeManager; @Inject @@ -57,10 +60,11 @@ public HudiPartitionManager(TypeManager typeManager) this.typeManager = requireNonNull(typeManager, "typeManager is null"); } - public List getEffectivePartitions( + public Map getEffectivePartitions( ConnectorSession connectorSession, ExtendedHiveMetastore metastore, SchemaTableName schemaTableName, + String tableBasePath, TupleDomain constraintSummary) { MetastoreContext metastoreContext = toMetastoreContext(connectorSession); @@ -68,7 +72,17 @@ public List getEffectivePartitions( Verify.verify(table.isPresent()); List partitionColumns = table.get().getPartitionColumns(); if (partitionColumns.isEmpty()) { - return ImmutableList.of(""); + return ImmutableMap.of( + "", Partition.builder() + .setCatalogName(Optional.empty()) + .setDatabaseName(schemaTableName.getSchemaName()) + .setTableName(schemaTableName.getTableName()) + .withStorage(storageBuilder -> + storageBuilder.setLocation(tableBasePath) + .setStorageFormat(StorageFormat.VIEW_STORAGE_FORMAT)) + .setColumns(ImmutableList.of()) + .setValues(ImmutableList.of()) + .build()); } Map partitionPredicate = new HashMap<>(); @@ -87,17 +101,29 @@ public List getEffectivePartitions( List partitionNames = metastore.getPartitionNamesByFilter(metastoreContext, schemaTableName.getSchemaName(), schemaTableName.getTableName(), partitionPredicate); List partitionTypes = partitionColumns.stream() .map(column -> typeManager.getType(column.getType().getTypeSignature())) - .collect(toList()); + .toList(); - return partitionNames.stream() - .map(PartitionNameWithVersion::getPartitionName) + List filteredPartitionNames = partitionNames.stream() // Apply extra filters which could not be done by getPartitionNamesByFilter, similar to filtering in HivePartitionManager#getPartitionsIterator - .filter(partitionName -> parseValuesAndFilterPartition( - partitionName, + .filter(partitionNameWithVersion -> parseValuesAndFilterPartition( + partitionNameWithVersion.getPartitionName(), hudiColumnHandles, partitionTypes, constraintSummary)) - .collect(toList()); + .toList(); + Map> partitionsByNames = metastore.getPartitionsByNames(metastoreContext, schemaTableName.getSchemaName(), schemaTableName.getTableName(), filteredPartitionNames); + List partitionsNotFound = partitionsByNames.entrySet().stream().filter(e -> e.getValue().isEmpty()).map(Map.Entry::getKey).toList(); + if (!partitionsNotFound.isEmpty()) { + log.warn("The following partitions were not found in the metastore for table {}.{}: {}", + schemaTableName.getSchemaName(), + schemaTableName.getTableName(), + partitionsNotFound); + } + + return partitionsByNames + .entrySet().stream() + .filter(e -> e.getValue().isPresent()) + .collect(Collectors.toMap(Map.Entry::getKey, e -> e.getValue().get())); } private boolean parseValuesAndFilterPartition( diff --git a/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiSplitManager.java b/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiSplitManager.java index c6e3cae63915f..24d14fbaa74dc 100644 --- a/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiSplitManager.java +++ b/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiSplitManager.java @@ -19,9 +19,7 @@ import com.facebook.presto.hive.HdfsEnvironment; import com.facebook.presto.hive.filesystem.ExtendedFileSystem; import com.facebook.presto.hive.metastore.ExtendedHiveMetastore; -import com.facebook.presto.hive.metastore.MetastoreContext; import com.facebook.presto.hive.metastore.Partition; -import com.facebook.presto.hive.metastore.Table; import com.facebook.presto.hudi.split.ForHudiBackgroundSplitLoader; import com.facebook.presto.hudi.split.ForHudiSplitAsyncQueue; import com.facebook.presto.hudi.split.ForHudiSplitSource; @@ -33,8 +31,6 @@ import com.facebook.presto.spi.connector.ConnectorSplitManager; import com.facebook.presto.spi.connector.ConnectorTransactionHandle; import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableMap; -import com.google.common.collect.Streams; import jakarta.inject.Inject; import org.apache.hadoop.conf.Configuration; import org.apache.hadoop.fs.Path; @@ -48,19 +44,13 @@ import org.apache.hudi.storage.StorageConfiguration; import java.io.IOException; -import java.util.List; import java.util.Map; import java.util.concurrent.ExecutorService; import java.util.concurrent.ScheduledExecutorService; -import static com.facebook.presto.hive.metastore.MetastoreUtil.extractPartitionValues; import static com.facebook.presto.hudi.HudiErrorCode.HUDI_FILESYSTEM_ERROR; -import static com.facebook.presto.hudi.HudiErrorCode.HUDI_INVALID_METADATA; -import static com.facebook.presto.hudi.HudiMetadata.fromDataColumns; import static com.facebook.presto.hudi.HudiSessionProperties.getMaxOutstandingSplits; import static com.facebook.presto.hudi.HudiSessionProperties.isHudiMetadataTableEnabled; -import static com.google.common.base.Preconditions.checkArgument; -import static java.lang.String.format; import static java.util.Objects.requireNonNull; import static org.apache.hudi.common.table.view.FileSystemViewManager.createInMemoryFileSystemViewWithTimeline; import static org.apache.hudi.hadoop.fs.HadoopFSUtils.getStorageConfWithCopy; @@ -107,7 +97,7 @@ public ConnectorSplitSource getSplits( // Retrieve and prune partitions HoodieTimer timer = HoodieTimer.start(); - List partitions = hudiPartitionManager.getEffectivePartitions(session, metastore, table.getSchemaTableName(), layout.getTupleDomain()); + Map partitions = hudiPartitionManager.getEffectivePartitions(session, metastore, table.getSchemaTableName(), table.getPath(), layout.getTupleDomain()); log.debug("Took %d ms to get %d partitions", timer.endTimer(), partitions.size()); if (partitions.isEmpty()) { return new FixedSplitSource(ImmutableList.of()); @@ -129,7 +119,6 @@ public ConnectorSplitSource getSplits( return new HudiSplitSource( session, - metastore, layout, fsView, partitions, @@ -155,36 +144,4 @@ private ExtendedFileSystem getFileSystem(ConnectorSession session, HudiTableHand throw new PrestoException(HUDI_FILESYSTEM_ERROR, "Could not open file system for " + table, e); } } - - public static HudiPartition getHudiPartition(ExtendedHiveMetastore metastore, MetastoreContext context, HudiTableLayoutHandle tableLayout, String partitionName) - { - String databaseName = tableLayout.getTable().getSchemaName(); - String tableName = tableLayout.getTable().getTableName(); - List partitionColumns = tableLayout.getPartitionColumns(); - - if (partitionColumns.isEmpty()) { - // non-partitioned tableLayout - Table table = metastore.getTable(context, databaseName, tableName) - .orElseThrow(() -> new PrestoException(HUDI_INVALID_METADATA, format("Table %s.%s expected but not found", databaseName, tableName))); - return new HudiPartition(partitionName, ImmutableList.of(), ImmutableMap.of(), table.getStorage(), tableLayout.getDataColumns()); - } - else { - // partitioned tableLayout - List partitionValues = extractPartitionValues(partitionName); - checkArgument(partitionColumns.size() == partitionValues.size(), - format("Invalid partition name %s for partition columns %s", partitionName, partitionColumns)); - Partition partition = metastore.getPartition(context, databaseName, tableName, partitionValues) - .orElseThrow(() -> new PrestoException(HUDI_INVALID_METADATA, format("Partition %s expected but not found", partitionName))); - Map keyValues = zipPartitionKeyValues(partitionColumns, partitionValues); - return new HudiPartition(partitionName, partitionValues, keyValues, partition.getStorage(), fromDataColumns(partition.getColumns())); - } - } - - private static Map zipPartitionKeyValues(List partitionColumns, List partitionValues) - { - ImmutableMap.Builder builder = ImmutableMap.builder(); - Streams.forEachPair(partitionColumns.stream(), partitionValues.stream(), - (column, value) -> builder.put(column.getName(), value)); - return builder.build(); - } } diff --git a/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiSplitSource.java b/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiSplitSource.java index 989870b2bee7b..c69cfa3f3903c 100644 --- a/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiSplitSource.java +++ b/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiSplitSource.java @@ -14,7 +14,7 @@ package com.facebook.presto.hudi; -import com.facebook.presto.hive.metastore.ExtendedHiveMetastore; +import com.facebook.presto.hive.metastore.Partition; import com.facebook.presto.hive.util.AsyncQueue; import com.facebook.presto.hudi.split.HudiBackgroundSplitLoader; import com.facebook.presto.spi.ConnectorSession; @@ -24,7 +24,7 @@ import com.google.common.util.concurrent.Futures; import org.apache.hudi.common.table.view.HoodieTableFileSystemView; -import java.util.List; +import java.util.Map; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutorService; import java.util.concurrent.ScheduledExecutorService; @@ -43,10 +43,9 @@ public class HudiSplitSource public HudiSplitSource( ConnectorSession session, - ExtendedHiveMetastore metastore, HudiTableLayoutHandle layout, HoodieTableFileSystemView fsView, - List partitions, + Map partitions, String latestInstant, ExecutorService asyncQueueExecutor, ScheduledExecutorService splitLoaderExecutorService, @@ -56,7 +55,6 @@ public HudiSplitSource( this.queue = new AsyncQueue<>(maxOutstandingSplits, asyncQueueExecutor); this.splitLoader = new HudiBackgroundSplitLoader( session, - metastore, splitGeneratorExecutorService, layout, fsView, diff --git a/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiTableHandle.java b/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiTableHandle.java index 686267d2d1c8a..d8a21b27f37e3 100644 --- a/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiTableHandle.java +++ b/presto-hudi/src/main/java/com/facebook/presto/hudi/HudiTableHandle.java @@ -14,13 +14,16 @@ package com.facebook.presto.hudi; +import com.facebook.presto.hive.metastore.Table; import com.facebook.presto.spi.ConnectorTableHandle; import com.facebook.presto.spi.SchemaTableName; import com.fasterxml.jackson.annotation.JsonCreator; import com.fasterxml.jackson.annotation.JsonProperty; import java.util.Objects; +import java.util.Optional; +import static com.google.common.base.Preconditions.checkArgument; import static java.util.Objects.requireNonNull; public class HudiTableHandle @@ -30,6 +33,7 @@ public class HudiTableHandle private final String tableName; private final String path; private final HudiTableType hudiTableType; + private final transient Optional

table; @JsonCreator public HudiTableHandle( @@ -38,12 +42,31 @@ public HudiTableHandle( @JsonProperty("path") String path, @JsonProperty("tableType") HudiTableType hudiTableType) { + this(Optional.empty(), schemaName, tableName, path, hudiTableType); + } + + public HudiTableHandle( + Optional
table, + String schemaName, + String tableName, + String path, + HudiTableType hudiTableType) + { + this.table = requireNonNull(table, "table is null"); this.schemaName = requireNonNull(schemaName, "schemaName is null"); this.tableName = requireNonNull(tableName, "tableName is null"); this.path = requireNonNull(path, "path is null"); this.hudiTableType = requireNonNull(hudiTableType, "tableType is null"); } + public Table getTable() + { + checkArgument(table.isPresent(), + "getTable() called on a table handle that has no metastore table object; " + + "this is likely because it is called on the worker."); + return table.get(); + } + @JsonProperty public String getSchemaName() { diff --git a/presto-hudi/src/main/java/com/facebook/presto/hudi/split/HudiBackgroundSplitLoader.java b/presto-hudi/src/main/java/com/facebook/presto/hudi/split/HudiBackgroundSplitLoader.java index 14894adf26318..afbff033d5e6c 100644 --- a/presto-hudi/src/main/java/com/facebook/presto/hudi/split/HudiBackgroundSplitLoader.java +++ b/presto-hudi/src/main/java/com/facebook/presto/hudi/split/HudiBackgroundSplitLoader.java @@ -15,7 +15,7 @@ package com.facebook.presto.hudi.split; import com.facebook.airlift.log.Logger; -import com.facebook.presto.hive.metastore.ExtendedHiveMetastore; +import com.facebook.presto.hive.metastore.Partition; import com.facebook.presto.hive.util.AsyncQueue; import com.facebook.presto.hudi.HudiTableLayoutHandle; import com.facebook.presto.spi.ConnectorSession; @@ -26,6 +26,7 @@ import java.util.ArrayList; import java.util.List; +import java.util.Map; import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; @@ -44,27 +45,24 @@ public class HudiBackgroundSplitLoader private static final Logger log = Logger.get(HudiBackgroundSplitLoader.class); private final ConnectorSession session; - private final ExtendedHiveMetastore metastore; private final HudiTableLayoutHandle layout; private final HoodieTableFileSystemView fsView; private final AsyncQueue asyncQueue; - private final List partitions; + private final Map partitions; private final String latestInstant; private final int splitGeneratorNumThreads; private final ExecutorService splitGeneratorExecutorService; public HudiBackgroundSplitLoader( ConnectorSession session, - ExtendedHiveMetastore metastore, ExecutorService splitGeneratorExecutorService, HudiTableLayoutHandle layout, HoodieTableFileSystemView fsView, AsyncQueue asyncQueue, - List partitions, + Map partitions, String latestInstant) { this.session = requireNonNull(session, "session is null"); - this.metastore = requireNonNull(metastore, "metastore is null"); this.layout = requireNonNull(layout, "layout is null"); this.fsView = requireNonNull(fsView, "fsView is null"); this.asyncQueue = requireNonNull(asyncQueue, "asyncQueue is null"); @@ -78,15 +76,15 @@ public HudiBackgroundSplitLoader( @Override public void run() { - HoodieTimer timer = new HoodieTimer().startTimer(); + HoodieTimer timer = HoodieTimer.start(); List splitGeneratorList = new ArrayList<>(); List splitGeneratorFutures = new ArrayList<>(); - ConcurrentLinkedQueue concurrentPartitionQueue = new ConcurrentLinkedQueue<>(partitions); + ConcurrentLinkedQueue concurrentPartitionQueue = new ConcurrentLinkedQueue<>(partitions.keySet()); // Start a number of partition split generators to generate the splits in parallel for (int i = 0; i < splitGeneratorNumThreads; i++) { HudiPartitionSplitGenerator generator = new HudiPartitionSplitGenerator( - session, metastore, layout, fsView, asyncQueue, concurrentPartitionQueue, latestInstant); + session, layout, fsView, partitions, asyncQueue, concurrentPartitionQueue, latestInstant); splitGeneratorList.add(generator); splitGeneratorFutures.add(splitGeneratorExecutorService.submit(generator)); } diff --git a/presto-hudi/src/main/java/com/facebook/presto/hudi/split/HudiPartitionSplitGenerator.java b/presto-hudi/src/main/java/com/facebook/presto/hudi/split/HudiPartitionSplitGenerator.java index 21ca4f7c36f16..d42eae11c84ef 100644 --- a/presto-hudi/src/main/java/com/facebook/presto/hudi/split/HudiPartitionSplitGenerator.java +++ b/presto-hudi/src/main/java/com/facebook/presto/hudi/split/HudiPartitionSplitGenerator.java @@ -16,9 +16,10 @@ import com.facebook.airlift.log.Logger; import com.facebook.airlift.units.DataSize; -import com.facebook.presto.hive.metastore.ExtendedHiveMetastore; -import com.facebook.presto.hive.metastore.MetastoreContext; +import com.facebook.presto.hive.metastore.Partition; +import com.facebook.presto.hive.metastore.Table; import com.facebook.presto.hive.util.AsyncQueue; +import com.facebook.presto.hudi.HudiColumnHandle; import com.facebook.presto.hudi.HudiFile; import com.facebook.presto.hudi.HudiPartition; import com.facebook.presto.hudi.HudiSplit; @@ -27,24 +28,31 @@ import com.facebook.presto.hudi.HudiTableType; import com.facebook.presto.spi.ConnectorSession; import com.facebook.presto.spi.ConnectorSplit; +import com.facebook.presto.spi.PrestoException; import com.facebook.presto.spi.schedule.NodeSelectionStrategy; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.Streams; import org.apache.hadoop.fs.Path; import org.apache.hudi.common.model.FileSlice; import org.apache.hudi.common.table.view.HoodieTableFileSystemView; import org.apache.hudi.common.util.HoodieTimer; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.Queue; import java.util.stream.Stream; -import static com.facebook.presto.hudi.HudiMetadata.toMetastoreContext; +import static com.facebook.presto.hive.metastore.MetastoreUtil.extractPartitionValues; +import static com.facebook.presto.hudi.HudiErrorCode.HUDI_INVALID_METADATA; +import static com.facebook.presto.hudi.HudiMetadata.fromDataColumns; import static com.facebook.presto.hudi.HudiSessionProperties.getMinimumAssignedSplitWeight; import static com.facebook.presto.hudi.HudiSessionProperties.getStandardSplitWeightSize; import static com.facebook.presto.hudi.HudiSessionProperties.isSizeBasedSplitWeightsEnabled; -import static com.facebook.presto.hudi.HudiSplitManager.getHudiPartition; +import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.collect.ImmutableList.toImmutableList; +import static java.lang.String.format; import static java.util.Objects.requireNonNull; import static org.apache.hudi.common.fs.FSUtils.getRelativePartitionPath; import static org.apache.hudi.hadoop.fs.HadoopFSUtils.convertToStoragePath; @@ -59,8 +67,6 @@ public class HudiPartitionSplitGenerator { private static final Logger log = Logger.get(HudiPartitionSplitGenerator.class); - private final ExtendedHiveMetastore metastore; - private final MetastoreContext metastoreContext; private final HudiTableLayoutHandle layout; private final HudiTableHandle table; private final Path tablePath; @@ -69,22 +75,22 @@ public class HudiPartitionSplitGenerator private final Queue concurrentPartitionQueue; private final String latestInstant; private final HudiSplitWeightProvider splitWeightProvider; + private final Map partitions; public HudiPartitionSplitGenerator( ConnectorSession session, - ExtendedHiveMetastore metastore, HudiTableLayoutHandle layout, HoodieTableFileSystemView fsView, + Map partitions, AsyncQueue asyncQueue, Queue concurrentPartitionQueue, String latestInstant) { - this.metastore = requireNonNull(metastore, "metastore is null"); - this.metastoreContext = toMetastoreContext(requireNonNull(session, "session is null")); this.layout = requireNonNull(layout, "layout is null"); this.table = layout.getTable(); this.tablePath = new Path(table.getPath()); this.fsView = requireNonNull(fsView, "fsView is null"); + this.partitions = requireNonNull(partitions, "partitionMap is null"); this.asyncQueue = requireNonNull(asyncQueue, "asyncQueue is null"); this.concurrentPartitionQueue = requireNonNull(concurrentPartitionQueue, "concurrentPartitionQueue is null"); this.latestInstant = requireNonNull(latestInstant, "latestInstant is null"); @@ -106,7 +112,7 @@ public void run() private void generateSplitsFromPartition(String partitionName) { - HudiPartition hudiPartition = getHudiPartition(metastore, metastoreContext, layout, partitionName); + HudiPartition hudiPartition = getHudiPartition(layout, partitionName); Path partitionPath = new Path(hudiPartition.getStorage().getLocation()); String relativePartitionPath = getRelativePartitionPath(convertToStoragePath(tablePath), convertToStoragePath(partitionPath)); Stream fileSlices = HudiTableType.MOR.equals(table.getTableType()) ? @@ -118,6 +124,38 @@ private void generateSplitsFromPartition(String partitionName) .forEach(asyncQueue::offer); } + private HudiPartition getHudiPartition(HudiTableLayoutHandle tableLayout, String partitionName) + { + String databaseName = tableLayout.getTable().getSchemaName(); + String tableName = tableLayout.getTable().getTableName(); + List partitionColumns = tableLayout.getPartitionColumns(); + + if (partitionColumns.isEmpty()) { + // non-partitioned tableLayout + Table metastoreTable = Optional.ofNullable(table.getTable()) + .orElseThrow(() -> new PrestoException(HUDI_INVALID_METADATA, format("Table %s.%s expected but not found", databaseName, tableName))); + return new HudiPartition(partitionName, ImmutableList.of(), ImmutableMap.of(), metastoreTable.getStorage(), tableLayout.getDataColumns()); + } + else { + // partitioned tableLayout + List partitionValues = extractPartitionValues(partitionName); + checkArgument(partitionColumns.size() == partitionValues.size(), + format("Invalid partition name %s for partition columns %s", partitionName, partitionColumns)); + Partition partition = Optional.ofNullable(partitions.get(partitionName)) + .orElseThrow(() -> new PrestoException(HUDI_INVALID_METADATA, format("Partition %s expected but not found", partitionName))); + Map keyValues = zipPartitionKeyValues(partitionColumns, partitionValues); + return new HudiPartition(partitionName, partitionValues, keyValues, partition.getStorage(), fromDataColumns(partition.getColumns())); + } + } + + private Map zipPartitionKeyValues(List partitionColumns, List partitionValues) + { + ImmutableMap.Builder builder = ImmutableMap.builder(); + Streams.forEachPair(partitionColumns.stream(), partitionValues.stream(), + (column, value) -> builder.put(column.getName(), value)); + return builder.build(); + } + private Optional createHudiSplit( HudiTableHandle table, FileSlice slice, @@ -132,7 +170,7 @@ private Optional createHudiSplit( List logFiles = slice.getLogFiles() .map(logFile -> new HudiFile(logFile.getPath().toString(), 0, logFile.getFileSize())) .collect(toImmutableList()); - long logFilesSize = logFiles.size() > 0 ? logFiles.stream().map(HudiFile::getLength).reduce(0L, Long::sum) : 0L; + long logFilesSize = logFiles.isEmpty() ? 0L : logFiles.stream().mapToLong(HudiFile::getLength).sum(); long sizeInBytes = baseFile != null ? baseFile.getLength() + logFilesSize : logFilesSize; return Optional.of(new HudiSplit( diff --git a/presto-hudi/src/test/java/com/facebook/presto/hudi/TestHudiPartitionManager.java b/presto-hudi/src/test/java/com/facebook/presto/hudi/TestHudiPartitionManager.java index ba445bc1d2054..abe1e3a10dd5a 100644 --- a/presto-hudi/src/test/java/com/facebook/presto/hudi/TestHudiPartitionManager.java +++ b/presto-hudi/src/test/java/com/facebook/presto/hudi/TestHudiPartitionManager.java @@ -23,18 +23,21 @@ import com.facebook.presto.hive.OrcFileWriterConfig; import com.facebook.presto.hive.ParquetFileWriterConfig; import com.facebook.presto.hive.metastore.Column; +import com.facebook.presto.hive.metastore.Partition; import com.facebook.presto.hive.metastore.PrestoTableType; import com.facebook.presto.hive.metastore.Storage; +import com.facebook.presto.hive.metastore.StorageFormat; import com.facebook.presto.hive.metastore.Table; import com.facebook.presto.spi.ColumnHandle; import com.facebook.presto.spi.ConnectorSession; -import com.facebook.presto.spi.SchemaTableName; import com.facebook.presto.testing.TestingConnectorSession; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; import org.testng.annotations.Test; -import java.util.List; +import java.util.Collections; +import java.util.Map; import java.util.Optional; import static com.facebook.presto.common.type.VarcharType.VARCHAR; @@ -51,6 +54,7 @@ public class TestHudiPartitionManager { private static final String SCHEMA_NAME = "schema"; private static final String TABLE_NAME = "table"; + private static final String NON_PARTITIONED_TABLE_NAME = "non_partitioned_table"; private static final String USER_NAME = "user"; private static final String LOCATION = "somewhere/over/the/rainbow"; private static final Column PARTITION_COLUMN = new Column("ds", HIVE_STRING, Optional.empty(), Optional.empty()); @@ -78,13 +82,65 @@ public class TestHudiPartitionManager Optional.empty(), Optional.empty()); - private static final List PARTITIONS = ImmutableList.of("ds=2019-07-23", "ds=2019-08-23"); + private static final Map> PARTITION_MAP = ImmutableMap.of( + "ds=2019-07-23", + Optional.of( + Partition.builder() + .setCatalogName(Optional.empty()) + .setDatabaseName(SCHEMA_NAME) + .setTableName(TABLE_NAME) + .withStorage(storageBuilder -> + storageBuilder.setLocation(LOCATION) + .setStorageFormat(StorageFormat.VIEW_STORAGE_FORMAT)) + .setColumns(ImmutableList.of(PARTITION_COLUMN)) + .setValues(Collections.singletonList("2019-07-23")) + .build()), + "ds=2019-08-23", + Optional.of( + Partition.builder() + .setCatalogName(Optional.empty()) + .setDatabaseName(SCHEMA_NAME) + .setTableName(TABLE_NAME) + .withStorage(storageBuilder -> + storageBuilder.setLocation(LOCATION) + .setStorageFormat(StorageFormat.VIEW_STORAGE_FORMAT)) + .setColumns(ImmutableList.of(PARTITION_COLUMN)) + .setValues(Collections.singletonList("2019-08-23")) + .build())); + + private static final Table NON_PARTITIONED_TABLE = new Table( + Optional.of("catalogName"), + SCHEMA_NAME, + NON_PARTITIONED_TABLE_NAME, + USER_NAME, + PrestoTableType.MANAGED_TABLE, + new Storage(fromHiveStorageFormat(PARQUET), + LOCATION, + Optional.of(new HiveBucketProperty( + ImmutableList.of(BUCKET_COLUMN.getName()), + 100, + ImmutableList.of(), + HIVE_COMPATIBLE, + Optional.empty())), + false, + ImmutableMap.of(), + ImmutableMap.of()), + ImmutableList.of(BUCKET_COLUMN), + ImmutableList.of(), + ImmutableMap.of(), + Optional.empty(), + Optional.empty()); + + private static final Map> NON_PARTITION_MAP = ImmutableMap.of( + "", + Optional.empty()); + private final HudiPartitionManager hudiPartitionManager = new HudiPartitionManager(new TestingTypeManager()); - private final TestingExtendedHiveMetastore metastore = new TestingExtendedHiveMetastore(TABLE, PARTITIONS); @Test public void testParseValuesAndFilterPartition() { + TestingExtendedHiveMetastore metastore = new TestingExtendedHiveMetastore(TABLE, PARTITION_MAP); ConnectorSession session = new TestingConnectorSession( new HiveSessionProperties( new HiveClientConfig().setMaxBucketsForGroupedExecution(100), @@ -100,11 +156,37 @@ public void testParseValuesAndFilterPartition() Optional.empty(), HudiColumnHandle.ColumnType.PARTITION_KEY), Domain.singleValue(VARCHAR, utf8Slice("2019-07-23")))); - List actualPartitions = hudiPartitionManager.getEffectivePartitions( + HudiTableHandle tableHandle = new HudiTableHandle(SCHEMA_NAME, TABLE_NAME, LOCATION, HudiTableType.COW); + Map actualPartitions = hudiPartitionManager.getEffectivePartitions( + session, + metastore, + tableHandle.getSchemaTableName(), + tableHandle.getPath(), + constraintSummary); + assertEquals(actualPartitions.keySet(), ImmutableSet.of("ds=2019-07-23")); + assertEquals(actualPartitions.get("ds=2019-07-23"), PARTITION_MAP.get("ds=2019-07-23").get()); + + // Non-partitioned table case + HudiTableHandle nonPartitionedTableHandle = new HudiTableHandle(SCHEMA_NAME, NON_PARTITIONED_TABLE_NAME, LOCATION, HudiTableType.COW); + metastore = new TestingExtendedHiveMetastore(NON_PARTITIONED_TABLE, NON_PARTITION_MAP); + actualPartitions = hudiPartitionManager.getEffectivePartitions( session, metastore, - new SchemaTableName(SCHEMA_NAME, TABLE_NAME), + nonPartitionedTableHandle.getSchemaTableName(), + nonPartitionedTableHandle.getPath(), constraintSummary); - assertEquals(actualPartitions, ImmutableList.of("ds=2019-07-23")); + assertEquals(actualPartitions.keySet(), ImmutableSet.of("")); + assertEquals( + actualPartitions.get(""), + Partition.builder() + .setCatalogName(Optional.empty()) + .setDatabaseName(nonPartitionedTableHandle.getSchemaName()) + .setTableName(nonPartitionedTableHandle.getTableName()) + .withStorage(storageBuilder -> + storageBuilder.setLocation(LOCATION) + .setStorageFormat(StorageFormat.VIEW_STORAGE_FORMAT)) + .setColumns(ImmutableList.of()) + .setValues(ImmutableList.of()) + .build()); } } diff --git a/presto-hudi/src/test/java/com/facebook/presto/hudi/TestingExtendedHiveMetastore.java b/presto-hudi/src/test/java/com/facebook/presto/hudi/TestingExtendedHiveMetastore.java index 40f7fe9425475..ac34d1d70d37f 100644 --- a/presto-hudi/src/test/java/com/facebook/presto/hudi/TestingExtendedHiveMetastore.java +++ b/presto-hudi/src/test/java/com/facebook/presto/hudi/TestingExtendedHiveMetastore.java @@ -18,6 +18,7 @@ import com.facebook.presto.hive.PartitionNameWithVersion; import com.facebook.presto.hive.metastore.Column; import com.facebook.presto.hive.metastore.MetastoreContext; +import com.facebook.presto.hive.metastore.Partition; import com.facebook.presto.hive.metastore.Table; import com.facebook.presto.hive.metastore.UnimplementedHiveMetastore; @@ -26,15 +27,16 @@ import java.util.Optional; import static com.facebook.presto.hive.metastore.MetastoreUtil.getPartitionNamesWithEmptyVersion; +import static com.google.common.collect.ImmutableMap.toImmutableMap; import static java.util.Objects.requireNonNull; public class TestingExtendedHiveMetastore extends UnimplementedHiveMetastore { private final Table table; - private final List partitions; + private final Map> partitions; - public TestingExtendedHiveMetastore(Table table, List partitions) + public TestingExtendedHiveMetastore(Table table, Map> partitions) { this.table = requireNonNull(table, "table is null"); this.partitions = requireNonNull(partitions, "partitions is null"); @@ -53,6 +55,13 @@ public List getPartitionNamesByFilter( String tableName, Map partitionPredicates) { - return getPartitionNamesWithEmptyVersion(partitions); + return getPartitionNamesWithEmptyVersion(partitions.keySet()); + } + + @Override + public Map> getPartitionsByNames(MetastoreContext metastoreContext, String databaseName, String tableName, List partitionNames) + { + return partitionNames.stream().filter(partitionName -> partitions.containsKey(partitionName.getPartitionName())) + .collect(toImmutableMap(PartitionNameWithVersion::getPartitionName, partitionName -> partitions.get(partitionName.getPartitionName()))); } }