Ensure queries in system.runtime.queries are redacted

trinodb · Dec 23, 2024 · ac2f505 · ac2f505
1 parent 11a7ed1
commit ac2f505
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 0 deletions.
diff --git a/core/trino-main/src/test/java/io/trino/connector/MockConnectorFactory.java b/core/trino-main/src/test/java/io/trino/connector/MockConnectorFactory.java
@@ -16,6 +16,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Sets;
 import io.trino.spi.connector.AggregateFunction;
 import io.trino.spi.connector.AggregationApplicationResult;
 import io.trino.spi.connector.CatalogSchemaTableName;
@@ -140,6 +141,7 @@ public class MockConnectorFactory
     private final Supplier<List<PropertyMetadata<?>>> columnProperties;
     private final Optional<ConnectorNodePartitioningProvider> partitioningProvider;
     private final Function<ConnectorTableFunctionHandle, ConnectorSplitSource> tableFunctionSplitsSources;
+    private final Set<String> redactablePropertyNames;
 
     // access control
     private final ListRoleGrants roleGrants;
@@ -196,6 +198,7 @@ private MockConnectorFactory(
             Supplier<List<PropertyMetadata<?>>> tableProperties,
             Supplier<List<PropertyMetadata<?>>> columnProperties,
             Optional<ConnectorNodePartitioningProvider> partitioningProvider,
+            Set<String> redactablePropertyNames,
             ListRoleGrants roleGrants,
             Optional<ConnectorAccessControl> accessControl,
             boolean allowMissingColumnsOnInsert,
@@ -244,6 +247,7 @@ private MockConnectorFactory(
         this.tableProperties = requireNonNull(tableProperties, "tableProperties is null");
         this.columnProperties = requireNonNull(columnProperties, "columnProperties is null");
         this.partitioningProvider = requireNonNull(partitioningProvider, "partitioningProvider is null");
+        this.redactablePropertyNames = requireNonNull(redactablePropertyNames, "redactablePropertyNames is null");
         this.roleGrants = requireNonNull(roleGrants, "roleGrants is null");
         this.accessControl = requireNonNull(accessControl, "accessControl is null");
         this.data = requireNonNull(data, "data is null");
@@ -325,6 +329,12 @@ public Connector create(String catalogName, Map<String, String> config, Connecto
                 allowSplittingReadIntoMultipleSubQueries);
     }
 
+    @Override
+    public Set<String> getRedactablePropertyNames(Set<String> propertyNames)
+    {
+        return Sets.intersection(redactablePropertyNames, propertyNames);
+    }
+
     public static MockConnectorFactory create()
     {
         return builder().build();
@@ -465,6 +475,7 @@ public static final class Builder
         private Supplier<List<PropertyMetadata<?>>> columnProperties = ImmutableList::of;
         private Optional<ConnectorNodePartitioningProvider> partitioningProvider = Optional.empty();
         private Function<ConnectorTableFunctionHandle, ConnectorSplitSource> tableFunctionSplitsSources = handle -> null;
+        private Set<String> redactablePropertyNames = ImmutableSet.of();
 
         // access control
         private boolean provideAccessControl;
@@ -844,6 +855,12 @@ public Builder withAllowSplittingReadIntoMultipleSubQueries(boolean allowSplitti
             return this;
         }
 
+        public Builder withRedactablePropertyNames(Set<String> redactablePropertyNames)
+        {
+            this.redactablePropertyNames = redactablePropertyNames;
+            return this;
+        }
+
         public MockConnectorFactory build()
         {
             Optional<ConnectorAccessControl> accessControl = Optional.empty();
@@ -895,6 +912,7 @@ public MockConnectorFactory build()
                     tableProperties,
                     columnProperties,
                     partitioningProvider,
+                    redactablePropertyNames,
                     roleGrants,
                     accessControl,
                     allowMissingColumnsOnInsert,

diff --git a/...ino-tests/src/test/java/io/trino/connector/system/runtime/TestSystemRuntimeConnector.java b/...ino-tests/src/test/java/io/trino/connector/system/runtime/TestSystemRuntimeConnector.java
@@ -15,6 +15,7 @@
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
 import com.google.common.util.concurrent.SettableFuture;
 import io.airlift.units.Duration;
@@ -48,6 +49,7 @@
 import static io.airlift.concurrent.Threads.threadsNamed;
 import static io.trino.spi.type.BigintType.BIGINT;
 import static io.trino.spi.type.VarcharType.VARCHAR;
+import static io.trino.testing.TestingNames.randomNameSuffix;
 import static io.trino.testing.TestingSession.testSessionBuilder;
 import static io.trino.testing.assertions.Assert.assertEventually;
 import static java.lang.String.format;
@@ -89,6 +91,7 @@ public Iterable<ConnectorFactory> getConnectorFactories()
                         .withGetViews((session, schemaTablePrefix) -> ImmutableMap.of())
                         .withListTables((session, s) -> ImmutableList.of("test_table"))
                         .withGetColumns(tableName -> getColumns.apply(tableName))
+                        .withRedactablePropertyNames(ImmutableSet.of("password"))
                         .build();
                 return ImmutableList.of(connectorFactory);
             }
@@ -299,6 +302,27 @@ public void testTasksTable()
         getQueryRunner().execute("SELECT * FROM system.runtime.tasks");
     }
 
+    @Test
+    public void testRedactedRuntimeQueries()
+    {
+        String catalog = "catalog_" + randomNameSuffix();
+        getQueryRunner().execute("""
+                CREATE CATALOG %s USING mock
+                WITH (
+                   "user" = 'bob',
+                   "password" = '1234'
+                )""".formatted(catalog));
+
+        assertQuery(
+                format("SELECT query FROM system.runtime.queries WHERE query LIKE '%%%s%%' AND query NOT LIKE '%%system.runtime.queries%%'", catalog),
+                """
+                        VALUES 'CREATE CATALOG %s USING mock
+                        WITH (
+                           "user" = ''bob'',
+                           "password" = ''***''
+                        )'""".formatted(catalog));
+    }
+
     private static void run(int repetitions, double successRate, Runnable test)
     {
         AssertionError lastError = null;