Ensure queries in system.runtime.queries are redacted

Author: piotrrzysko

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;
@@ -138,6 +139,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> securitySensitivePropertyNames;
 
     // access control
     private final ListRoleGrants roleGrants;
@@ -193,6 +195,7 @@ private MockConnectorFactory(
             Supplier<List<PropertyMetadata<?>>> tableProperties,
             Supplier<List<PropertyMetadata<?>>> columnProperties,
             Optional<ConnectorNodePartitioningProvider> partitioningProvider,
+            Set<String> securitySensitivePropertyNames,
             ListRoleGrants roleGrants,
             Optional<ConnectorAccessControl> accessControl,
             boolean allowMissingColumnsOnInsert,
@@ -240,6 +243,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.securitySensitivePropertyNames = requireNonNull(securitySensitivePropertyNames, "securitySensitivePropertyNames is null");
         this.roleGrants = requireNonNull(roleGrants, "roleGrants is null");
         this.accessControl = requireNonNull(accessControl, "accessControl is null");
         this.data = requireNonNull(data, "data is null");
@@ -320,6 +324,12 @@ public Connector create(String catalogName, Map<String, String> config, Connecto
                 allowSplittingReadIntoMultipleSubQueries);
     }
 
+    @Override
+    public Set<String> getSecuritySensitivePropertyNames(String catalogName, Map<String, String> config, ConnectorContext context)
+    {
+        return Sets.intersection(securitySensitivePropertyNames, config.keySet());
+    }
+
     public static MockConnectorFactory create()
     {
         return builder().build();
@@ -459,6 +469,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> securitySensitivePropertyNames = ImmutableSet.of();
 
         // access control
         private boolean provideAccessControl;
@@ -822,6 +833,12 @@ public Builder withAllowSplittingReadIntoMultipleSubQueries(boolean allowSplitti
             return this;
         }
 
+        public Builder withSecuritySensitivePropertyNames(Set<String> securitySensitivePropertyNames)
+        {
+            this.securitySensitivePropertyNames = securitySensitivePropertyNames;
+            return this;
+        }
+
         public MockConnectorFactory build()
         {
             Optional<ConnectorAccessControl> accessControl = Optional.empty();
@@ -872,6 +889,7 @@ public MockConnectorFactory build()
                     tableProperties,
                     columnProperties,
                     partitioningProvider,
+                    securitySensitivePropertyNames,
                     roleGrants,
                     accessControl,
                     allowMissingColumnsOnInsert,

testing/trino-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))
+                        .withSecuritySensitivePropertyNames(ImmutableSet.of("password"))
                         .build();
                 return ImmutableList.of(connectorFactory);
             }
@@ -299,6 +302,75 @@ public void testTasksTable()
         getQueryRunner().execute("SELECT * FROM system.runtime.tasks");
     }
 
+    @Test
+    public void testRedactedCreateCatalog()
+    {
+        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));
+    }
+
+    @Test
+    public void testRedactedPrepareCreateCatalog()
+    {
+        String catalog = "catalog_" + randomNameSuffix();
+        getQueryRunner().execute(
+                """
+                PREPARE create_catalog FROM
+                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 'PREPARE create_catalog FROM
+                CREATE CATALOG %s USING mock
+                WITH (
+                   "user" = ''bob'',
+                   "password" = ''***''
+                )'
+                """.formatted(catalog));
+
+        Session executeSession = testSessionBuilder()
+                .setCatalog("mock")
+                .setSchema("default")
+                .addPreparedStatement(
+                        "create_catalog",
+                        """
+                        CREATE CATALOG %s USING mock
+                        WITH (
+                           "user" = 'bob',
+                           "password" = '1234'
+                        )
+                        """.formatted(catalog))
+                .build();
+        getQueryRunner().execute(executeSession, "EXECUTE create_catalog");
+
+        assertQuery(
+                "SELECT query FROM system.runtime.queries WHERE query LIKE '%%EXECUTE create_catalog%%' AND query NOT LIKE '%%system.runtime.queries%%'",
+                "VALUES 'EXECUTE create_catalog'");
+    }
+
     private static void run(int repetitions, double successRate, Runnable test)
     {
         AssertionError lastError = null;