Redact sensitive information in catalog queries

core/trino-main/src/main/java/io/trino/FeaturesConfig.java

@@ -122,6 +122,8 @@ public class FeaturesConfig
 
     private boolean faultTolerantExecutionExchangeEncryptionEnabled = true;
 
+    private boolean statementRedactingEnabled = true;
+
     public enum DataIntegrityVerification
     {
         NONE,
@@ -514,6 +516,18 @@ public FeaturesConfig setFaultTolerantExecutionExchangeEncryptionEnabled(boolean
         return this;
     }
 
+    public boolean isStatementRedactingEnabled()
+    {
+        return statementRedactingEnabled;
+    }
+
+    @Config("statement-redacting-enabled")
+    public FeaturesConfig setStatementRedactingEnabled(boolean statementRedactingEnabled)
+    {
+        this.statementRedactingEnabled = statementRedactingEnabled;
+        return this;
+    }
+
     public void applyFaultTolerantExecutionDefaults()
     {
         exchangeCompressionCodec = LZ4;

core/trino-main/src/main/java/io/trino/connector/CatalogFactory.java

@@ -20,6 +20,8 @@
 import io.trino.spi.connector.ConnectorFactory;
 import io.trino.spi.connector.ConnectorName;
 
+import java.util.Set;
+
 @ThreadSafe
 public interface CatalogFactory
 {
@@ -28,4 +30,6 @@ public interface CatalogFactory
     CatalogConnector createCatalog(CatalogProperties catalogProperties);
 
     CatalogConnector createCatalog(CatalogHandle catalogHandle, ConnectorName connectorName, Connector connector);
+
+    Set<String> getRedactablePropertyNames(ConnectorName connectorName, Set<String> propertyNames);
 }

core/trino-main/src/main/java/io/trino/connector/DefaultCatalogFactory.java

@@ -45,6 +45,7 @@
 
 import java.util.Map;
 import java.util.Optional;
+import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 
@@ -144,6 +145,18 @@ public CatalogConnector createCatalog(CatalogHandle catalogHandle, ConnectorName
         return createCatalog(catalogHandle, connectorName, connector, Optional.empty());
     }
 
+    @Override
+    public Set<String> getRedactablePropertyNames(ConnectorName connectorName, Set<String> propertyNames)
+    {
+        ConnectorFactory connectorFactory = connectorFactories.get(connectorName);
+        if (connectorFactory == null) {
+            // If someone tries to use a non-existent connector, we assume they
+            // misspelled the name and, for safety, we redact all the properties.
+            return propertyNames;
+        }
+        return connectorFactory.getRedactablePropertyNames(propertyNames);
+    }
+
     private CatalogConnector createCatalog(CatalogHandle catalogHandle, ConnectorName connectorName, Connector connector, Optional<CatalogProperties> catalogProperties)
     {
         Tracer tracer = createTracer(catalogHandle);

core/trino-main/src/main/java/io/trino/connector/LazyCatalogFactory.java

@@ -19,6 +19,7 @@
 import io.trino.spi.connector.ConnectorFactory;
 import io.trino.spi.connector.ConnectorName;
 
+import java.util.Set;
 import java.util.concurrent.atomic.AtomicReference;
 
 import static com.google.common.base.Preconditions.checkState;
@@ -51,6 +52,12 @@ public CatalogConnector createCatalog(CatalogHandle catalogHandle, ConnectorName
         return getDelegate().createCatalog(catalogHandle, connectorName, connector);
     }
 
+    @Override
+    public Set<String> getRedactablePropertyNames(ConnectorName connectorName, Set<String> propertyNames)
+    {
+        return getDelegate().getRedactablePropertyNames(connectorName, propertyNames);
+    }
+
     private CatalogFactory getDelegate()
     {
         CatalogFactory catalogFactory = delegate.get();

core/trino-main/src/main/java/io/trino/dispatcher/DispatchManager.java

@@ -44,6 +44,7 @@
 import io.trino.spi.TrinoException;
 import io.trino.spi.resourcegroups.SelectionContext;
 import io.trino.spi.resourcegroups.SelectionCriteria;
+import io.trino.sql.SensitiveStatementRedactor;
 import jakarta.annotation.PostConstruct;
 import jakarta.annotation.PreDestroy;
 import org.weakref.jmx.Flatten;
@@ -84,6 +85,7 @@ public class DispatchManager
     private final SessionPropertyDefaults sessionPropertyDefaults;
     private final SessionPropertyManager sessionPropertyManager;
     private final Tracer tracer;
+    private final SensitiveStatementRedactor sensitiveStatementRedactor;
 
     private final int maxQueryLength;
 
@@ -107,6 +109,7 @@ public DispatchManager(
             SessionPropertyDefaults sessionPropertyDefaults,
             SessionPropertyManager sessionPropertyManager,
             Tracer tracer,
+            SensitiveStatementRedactor sensitiveStatementRedactor,
             QueryManagerConfig queryManagerConfig,
             DispatchExecutor dispatchExecutor,
             QueryMonitor queryMonitor)
@@ -121,6 +124,7 @@ public DispatchManager(
         this.sessionPropertyDefaults = requireNonNull(sessionPropertyDefaults, "sessionPropertyDefaults is null");
         this.sessionPropertyManager = sessionPropertyManager;
         this.tracer = requireNonNull(tracer, "tracer is null");
+        this.sensitiveStatementRedactor = requireNonNull(sensitiveStatementRedactor, "sensitiveStatementRedactor is null");
 
         this.maxQueryLength = queryManagerConfig.getMaxQueryLength();
 
@@ -207,6 +211,7 @@ private <C> void createQueryInternal(QueryId queryId, Span querySpan, Slug slug,
     {
         Session session = null;
         PreparedQuery preparedQuery = null;
+        String redactedQuery = null;
         try {
             if (query.length() > maxQueryLength) {
                 int queryLength = query.length();
@@ -223,6 +228,9 @@ private <C> void createQueryInternal(QueryId queryId, Span querySpan, Slug slug,
             // prepare query
             preparedQuery = queryPreparer.prepareQuery(session, query);
 
+            // redact security-sensitive information that query may contain
+            redactedQuery = sensitiveStatementRedactor.redact(query, preparedQuery.getStatement());
+
             // select resource group
             Optional<String> queryType = getQueryType(preparedQuery.getStatement()).map(Enum::name);
             SelectionContext<C> selectionContext = resourceGroupManager.selectGroup(new SelectionCriteria(
@@ -240,7 +248,7 @@ private <C> void createQueryInternal(QueryId queryId, Span querySpan, Slug slug,
             DispatchQuery dispatchQuery = dispatchQueryFactory.createDispatchQuery(
                     session,
                     sessionContext.getTransactionId(),
-                    query,
+                    redactedQuery,
                     preparedQuery,
                     slug,
                     selectionContext.getResourceGroupId());
@@ -266,8 +274,16 @@ private <C> void createQueryInternal(QueryId queryId, Span querySpan, Slug slug,
                         .setSource(sessionContext.getSource().orElse(null))
                         .build();
             }
+            if (redactedQuery == null) {
+                redactedQuery = sensitiveStatementRedactor.redact(query);
+            }
             Optional<String> preparedSql = Optional.ofNullable(preparedQuery).flatMap(PreparedQuery::getPrepareSql);
-            DispatchQuery failedDispatchQuery = failedDispatchQueryFactory.createFailedDispatchQuery(session, query, preparedSql, Optional.empty(), throwable);
+            DispatchQuery failedDispatchQuery = failedDispatchQueryFactory.createFailedDispatchQuery(
+                    session,
+                    redactedQuery,
+                    preparedSql,
+                    Optional.empty(),
+                    throwable);
             queryCreated(failedDispatchQuery);
             // maintain proper order of calls such that EventListener has access to QueryInfo
             // - add query to tracker

core/trino-main/src/main/java/io/trino/server/CoordinatorModule.java

@@ -113,6 +113,7 @@
 import io.trino.server.ui.WorkerResource;
 import io.trino.spi.VersionEmbedder;
 import io.trino.sql.PlannerContext;
+import io.trino.sql.SensitiveStatementRedactor;
 import io.trino.sql.analyzer.AnalyzerFactory;
 import io.trino.sql.analyzer.QueryExplainerFactory;
 import io.trino.sql.planner.OptimizerStatsMBeanExporter;
@@ -304,6 +305,9 @@ List<OutputStatsEstimatorFactory> getCompositeOutputDataSizeEstimatorDelegateFac
         rewriteBinder.addBinding().to(ShowStatsRewrite.class).in(Scopes.SINGLETON);
         rewriteBinder.addBinding().to(ExplainRewrite.class).in(Scopes.SINGLETON);
 
+        // security-sensitive statement redactor
+        binder.bind(SensitiveStatementRedactor.class).in(Scopes.SINGLETON);
+
         // planner
         binder.bind(PlanFragmenter.class).in(Scopes.SINGLETON);
         binder.bind(PlanOptimizersFactory.class).to(PlanOptimizers.class).in(Scopes.SINGLETON);

core/trino-main/src/main/java/io/trino/server/ui/TrimmedBasicQueryInfo.java

@@ -13,6 +13,7 @@
  */
 package io.trino.server.ui;
 
+import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.google.errorprone.annotations.Immutable;
 import io.trino.execution.QueryState;
@@ -54,6 +55,45 @@ public class TrimmedBasicQueryInfo
     private final Optional<QueryType> queryType;
     private final RetryPolicy retryPolicy;
 
+    @JsonCreator
+    public TrimmedBasicQueryInfo(
+            @JsonProperty("queryId") QueryId queryId,
+            @JsonProperty("sessionUser") String sessionUser,
+            @JsonProperty("sessionPrincipal") Optional<String> sessionPrincipal,
+            @JsonProperty("sessionSource") Optional<String> sessionSource,
+            @JsonProperty("resourceGroupId") Optional<ResourceGroupId> resourceGroupId,
+            @JsonProperty("queryDataEncoding") Optional<String> queryDataEncoding,
+            @JsonProperty("state") QueryState state,
+            @JsonProperty("scheduled") boolean scheduled,
+            @JsonProperty("self") URI self,
+            @JsonProperty("queryTextPreview") String queryTextPreview,
+            @JsonProperty("updateType") Optional<String> updateType,
+            @JsonProperty("preparedQuery") Optional<String> preparedQuery,
+            @JsonProperty("queryStats") BasicQueryStats queryStats,
+            @JsonProperty("errorType") Optional<ErrorType> errorType,
+            @JsonProperty("errorCode") Optional<ErrorCode> errorCode,
+            @JsonProperty("queryType") Optional<QueryType> queryType,
+            @JsonProperty("retryPolicy") RetryPolicy retryPolicy)
+    {
+        this.queryId = requireNonNull(queryId, "queryId is null");
+        this.sessionUser = requireNonNull(sessionUser, "sessionUser is null");
+        this.sessionPrincipal = requireNonNull(sessionPrincipal, "sessionPrincipal is null");
+        this.sessionSource = requireNonNull(sessionSource, "sessionSource is null");
+        this.resourceGroupId = requireNonNull(resourceGroupId, "resourceGroupId is null");
+        this.queryDataEncoding = requireNonNull(queryDataEncoding, "queryDataEncoding is null");
+        this.state = requireNonNull(state, "state is null");
+        this.scheduled = scheduled;
+        this.self = requireNonNull(self, "self is null");
+        this.queryTextPreview = requireNonNull(queryTextPreview, "queryTextPreview is null");
+        this.updateType = requireNonNull(updateType, "updateType is null");
+        this.preparedQuery = requireNonNull(preparedQuery, "preparedQuery is null");
+        this.queryStats = requireNonNull(queryStats, "queryStats is null");
+        this.errorType = requireNonNull(errorType, "errorType is null");
+        this.errorCode = requireNonNull(errorCode, "errorCode is null");
+        this.queryType = requireNonNull(queryType, "queryType is null");
+        this.retryPolicy = requireNonNull(retryPolicy, "retryPolicy is null");
+    }
+
     public TrimmedBasicQueryInfo(BasicQueryInfo queryInfo)
     {
         this.queryId = requireNonNull(queryInfo.getQueryId(), "queryId is null");

core/trino-main/src/main/java/io/trino/sql/SensitiveStatementRedactor.java

@@ -0,0 +1,131 @@
+/*
+ * 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 io.trino.sql;
+
+import com.google.inject.Inject;
+import io.trino.FeaturesConfig;
+import io.trino.connector.CatalogFactory;
+import io.trino.spi.connector.ConnectorName;
+import io.trino.sql.tree.AstVisitor;
+import io.trino.sql.tree.CreateCatalog;
+import io.trino.sql.tree.Explain;
+import io.trino.sql.tree.ExplainAnalyze;
+import io.trino.sql.tree.Identifier;
+import io.trino.sql.tree.Node;
+import io.trino.sql.tree.Property;
+import io.trino.sql.tree.Statement;
+import io.trino.sql.tree.StringLiteral;
+
+import java.util.List;
+import java.util.Set;
+
+import static com.google.common.collect.ImmutableList.toImmutableList;
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
+
+public class SensitiveStatementRedactor
+{
+    public static final String REDACTED_VALUE = "***";
+
+    private final boolean enabled;
+    private final CatalogFactory catalogFactory;
+
+    @Inject
+    public SensitiveStatementRedactor(FeaturesConfig config, CatalogFactory catalogFactory)
+    {
+        this.enabled = config.isStatementRedactingEnabled();
+        this.catalogFactory = catalogFactory;
+    }
+
+    public String redact(String rawQuery, Statement statement)
+    {
+        if (enabled) {
+            RedactingVisitor visitor = new RedactingVisitor();
+            Node redactedStatement = visitor.process(statement);
+            if (visitor.isRedacted()) {
+                return SqlFormatter.formatSql(redactedStatement);
+            }
+        }
+        return rawQuery;
+    }
+
+    public String redact(String rawQuery)
+    {
+        if (enabled) {
+            return REDACTED_VALUE;
+        }
+        return rawQuery;
+    }
+
+    private class RedactingVisitor
+            extends AstVisitor<Node, Void>
+    {
+        private boolean redacted;
+
+        public boolean isRedacted()
+        {
+            return redacted;
+        }
+
+        @Override
+        protected Node visitNode(Node node, Void context)
+        {
+            return node;
+        }
+
+        @Override
+        protected Node visitExplain(Explain explain, Void context)
+        {
+            Statement statement = (Statement) process(explain.getStatement());
+            return new Explain(explain.getLocation().orElseThrow(), statement, explain.getOptions());
+        }
+
+        @Override
+        protected Node visitExplainAnalyze(ExplainAnalyze explainAnalyze, Void context)
+        {
+            Statement statement = (Statement) process(explainAnalyze.getStatement());
+            return new ExplainAnalyze(explainAnalyze.getLocation().orElseThrow(), statement, explainAnalyze.isVerbose());
+        }
+
+        @Override
+        protected Node visitCreateCatalog(CreateCatalog createCatalog, Void context)
+        {
+            ConnectorName connectorName = new ConnectorName(createCatalog.getConnectorName().getValue());
+            List<Property> redactedProperties = redact(connectorName, createCatalog.getProperties());
+            return createCatalog.withProperties(redactedProperties);
+        }
+
+        private List<Property> redact(ConnectorName connectorName, List<Property> properties)
+        {
+            redacted = true;
+            Set<String> propertyNames = properties.stream()
+                    .map(Property::getName)
+                    .map(Identifier::getValue)
+                    .collect(toImmutableSet());
+            Set<String> redactableProperties = catalogFactory.getRedactablePropertyNames(connectorName, propertyNames);
+            return redactProperties(properties, redactableProperties);
+        }
+
+        private List<Property> redactProperties(List<Property> properties, Set<String> redactableProperties)
+        {
+            return properties.stream()
+                    .map(property -> {
+                        if (redactableProperties.contains(property.getName().getValue())) {
+                            return new Property(property.getName(), new StringLiteral(REDACTED_VALUE));
+                        }
+                        return property;
+                    })
+                    .collect(toImmutableList());
+        }
+    }
+}