From ac2f505505665a275f631f3fd7dd9651fe755a97 Mon Sep 17 00:00:00 2001 From: Piotr Rzysko Date: Mon, 23 Dec 2024 11:11:03 +0100 Subject: [PATCH] Ensure queries in system.runtime.queries are redacted --- .../trino/connector/MockConnectorFactory.java | 18 ++++++++++++++ .../runtime/TestSystemRuntimeConnector.java | 24 +++++++++++++++++++ 2 files changed, 42 insertions(+) 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 index 7a0c8e090360..0a29cfc707bb 100644 --- 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>> columnProperties; private final Optional partitioningProvider; private final Function tableFunctionSplitsSources; + private final Set redactablePropertyNames; // access control private final ListRoleGrants roleGrants; @@ -196,6 +198,7 @@ private MockConnectorFactory( Supplier>> tableProperties, Supplier>> columnProperties, Optional partitioningProvider, + Set redactablePropertyNames, ListRoleGrants roleGrants, Optional 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 config, Connecto allowSplittingReadIntoMultipleSubQueries); } + @Override + public Set getRedactablePropertyNames(Set propertyNames) + { + return Sets.intersection(redactablePropertyNames, propertyNames); + } + public static MockConnectorFactory create() { return builder().build(); @@ -465,6 +475,7 @@ public static final class Builder private Supplier>> columnProperties = ImmutableList::of; private Optional partitioningProvider = Optional.empty(); private Function tableFunctionSplitsSources = handle -> null; + private Set redactablePropertyNames = ImmutableSet.of(); // access control private boolean provideAccessControl; @@ -844,6 +855,12 @@ public Builder withAllowSplittingReadIntoMultipleSubQueries(boolean allowSplitti return this; } + public Builder withRedactablePropertyNames(Set redactablePropertyNames) + { + this.redactablePropertyNames = redactablePropertyNames; + return this; + } + public MockConnectorFactory build() { Optional accessControl = Optional.empty(); @@ -895,6 +912,7 @@ public MockConnectorFactory build() tableProperties, columnProperties, partitioningProvider, + redactablePropertyNames, roleGrants, accessControl, allowMissingColumnsOnInsert, diff --git a/testing/trino-tests/src/test/java/io/trino/connector/system/runtime/TestSystemRuntimeConnector.java b/testing/trino-tests/src/test/java/io/trino/connector/system/runtime/TestSystemRuntimeConnector.java index e5a386dbacd3..26046e059a47 100644 --- a/testing/trino-tests/src/test/java/io/trino/connector/system/runtime/TestSystemRuntimeConnector.java +++ b/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 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;