Skip to content

Commit

Permalink
Test INSERT INTO with row filter with missing columns support
Browse files Browse the repository at this point in the history
  • Loading branch information
@crpeter
crpeter authored Jan 21, 2022
1 parent aeaf976 commit c5112d1
Show file tree
Hide file tree
Showing 4 changed files with 96 additions and 5 deletions.
11 changes: 10 additions & 1 deletion core/trino-main/src/test/java/io/trino/connector/MockConnector.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -123,6 +123,7 @@ public class MockConnector
private final Optional<ConnectorAccessControl> accessControl;
private final Function<SchemaTableName, List<List<?>>> data;
private final Set<Procedure> procedures;
private final boolean allowMissingColumnsOnInsert;

MockConnector(
Function<ConnectorSession, List<String>> listSchemaNames,
Expand All @@ -148,7 +149,8 @@ public class MockConnector
MockConnectorFactory.ListRoleGrants roleGrants,
Optional<ConnectorAccessControl> accessControl,
Function<SchemaTableName, List<List<?>>> data,
Set<Procedure> procedures)
Set<Procedure> procedures,
boolean allowMissingColumnsOnInsert)
{
this.listSchemaNames = requireNonNull(listSchemaNames, "listSchemaNames is null");
this.listTables = requireNonNull(listTables, "listTables is null");
Expand All @@ -174,6 +176,7 @@ public class MockConnector
this.accessControl = requireNonNull(accessControl, "accessControl is null");
this.data = requireNonNull(data, "data is null");
this.procedures = requireNonNull(procedures, "procedures is null");
this.allowMissingColumnsOnInsert = allowMissingColumnsOnInsert;
}

@Override
Expand Down Expand Up @@ -488,6 +491,12 @@ public ConnectorInsertTableHandle beginInsert(ConnectorSession session, Connecto
return new MockConnectorInsertTableHandle(((MockConnectorTableHandle) tableHandle).getTableName());
}

@Override
public boolean supportsMissingColumnsOnInsert()
{
return allowMissingColumnsOnInsert;
}

@Override
public Optional<ConnectorOutputMetadata> finishInsert(ConnectorSession session, ConnectorInsertTableHandle insertHandle, Collection<Slice> fragments, Collection<ComputedStatistics> computedStatistics)
{
Expand Down
5 changes: 5 additions & 0 deletions core/trino-main/src/test/java/io/trino/connector/MockConnectorEntities.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,11 @@ private MockConnectorEntities() {}
.add(ColumnMetadata.builder().setName("$hidden").setType(BigintType.BIGINT).setHidden(true).build())
.build();

public static final List<ColumnMetadata> TPCH_NATION_WITH_OPTIONAL_COLUMN = ImmutableList.<ColumnMetadata>builder()
.addAll(TPCH_NATION_SCHEMA)
.add(ColumnMetadata.builder().setName("optional").setType(createUnboundedVarcharType()).setNullable(true).build())
.build();

public static final List<List<?>> TPCH_NATION_DATA = ImmutableList.<List<?>>builder()
.add(ImmutableList.of(0, "ALGERIA", 0, " haggle. carefully final deposits detect slyly agai"))
.add(ImmutableList.of(1, "ARGENTINA", 1, "al foxes promise slyly according to the regular accounts. bold requests alon"))
Expand Down
31 changes: 27 additions & 4 deletions core/trino-main/src/test/java/io/trino/connector/MockConnectorFactory.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,7 @@
public class MockConnectorFactory
implements ConnectorFactory
{
private final String name;
private final Function<ConnectorSession, List<String>> listSchemaNames;
private final BiFunction<ConnectorSession, String, List<SchemaTableName>> listTables;
private final Optional<BiFunction<ConnectorSession, SchemaTablePrefix, Stream<TableColumnsMetadata>>> streamTableColumns;
Expand All @@ -92,12 +93,14 @@ public class MockConnectorFactory
private final Supplier<Iterable<EventListener>> eventListeners;
private final Function<SchemaTableName, List<List<?>>> data;
private final Set<Procedure> procedures;
private final boolean allowMissingColumnsOnInsert;

// access control
private final ListRoleGrants roleGrants;
private final Optional<ConnectorAccessControl> accessControl;

private MockConnectorFactory(
String name,
Function<ConnectorSession, List<String>> listSchemaNames,
BiFunction<ConnectorSession, String, List<SchemaTableName>> listTables,
Optional<BiFunction<ConnectorSession, SchemaTablePrefix, Stream<TableColumnsMetadata>>> streamTableColumns,
Expand All @@ -121,8 +124,10 @@ private MockConnectorFactory(
Function<SchemaTableName, List<List<?>>> data,
Set<Procedure> procedures,
ListRoleGrants roleGrants,
Optional<ConnectorAccessControl> accessControl)
Optional<ConnectorAccessControl> accessControl,
boolean allowMissingColumnsOnInsert)
{
this.name = (name != null) ? name : "mock";
this.listSchemaNames = requireNonNull(listSchemaNames, "listSchemaNames is null");
this.listTables = requireNonNull(listTables, "listTables is null");
this.streamTableColumns = requireNonNull(streamTableColumns, "streamTableColumns is null");
Expand All @@ -147,12 +152,13 @@ private MockConnectorFactory(
this.accessControl = requireNonNull(accessControl, "accessControl is null");
this.data = requireNonNull(data, "data is null");
this.procedures = requireNonNull(procedures, "procedures is null");
this.allowMissingColumnsOnInsert = allowMissingColumnsOnInsert;
}

@Override
public String getName()
{
return "mock";
return name;
}

@Override
Expand Down Expand Up @@ -188,7 +194,8 @@ public Connector create(String catalogName, Map<String, String> config, Connecto
roleGrants,
accessControl,
data,
procedures);
procedures,
allowMissingColumnsOnInsert);
}

public static Builder builder()
Expand Down Expand Up @@ -261,6 +268,7 @@ public interface ListRoleGrants

public static final class Builder
{
private String name;
private Function<ConnectorSession, List<String>> listSchemaNames = defaultListSchemaNames();
private BiFunction<ConnectorSession, String, List<SchemaTableName>> listTables = defaultListTables();
private Optional<BiFunction<ConnectorSession, SchemaTablePrefix, Stream<TableColumnsMetadata>>> streamTableColumns = Optional.empty();
Expand Down Expand Up @@ -291,6 +299,13 @@ public static final class Builder
private Grants<SchemaTableName> tableGrants = new AllowAllGrants<>();
private Function<SchemaTableName, ViewExpression> rowFilter = (tableName) -> null;
private BiFunction<SchemaTableName, String, ViewExpression> columnMask = (tableName, columnName) -> null;
private boolean allowMissingColumnsOnInsert;

public Builder withName(String name)
{
this.name = requireNonNull(name, "name is null");
return this;
}

public Builder withListSchemaNames(Function<ConnectorSession, List<String>> listSchemaNames)
{
Expand Down Expand Up @@ -470,13 +485,20 @@ public Builder withColumnMask(BiFunction<SchemaTableName, String, ViewExpression
return this;
}

public Builder withAllowMissingColumnsOnInsert(boolean allowMissingColumnsOnInsert)
{
this.allowMissingColumnsOnInsert = allowMissingColumnsOnInsert;
return this;
}

public MockConnectorFactory build()
{
Optional<ConnectorAccessControl> accessControl = Optional.empty();
if (provideAccessControl) {
accessControl = Optional.of(new MockConnectorAccessControl(schemaGrants, tableGrants, rowFilter, columnMask));
}
return new MockConnectorFactory(
name,
listSchemaNames,
listTables,
streamTableColumns,
Expand All @@ -500,7 +522,8 @@ public MockConnectorFactory build()
data,
procedures,
roleGrants,
accessControl);
accessControl,
allowMissingColumnsOnInsert);
}

public static Function<ConnectorSession, List<String>> defaultListSchemaNames()
Expand Down
54 changes: 54 additions & 0 deletions core/trino-main/src/test/java/io/trino/sql/query/TestRowFilter.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@
import static io.trino.connector.MockConnectorEntities.TPCH_NATION_DATA;
import static io.trino.connector.MockConnectorEntities.TPCH_NATION_SCHEMA;
import static io.trino.connector.MockConnectorEntities.TPCH_NATION_WITH_HIDDEN_COLUMN;
import static io.trino.connector.MockConnectorEntities.TPCH_NATION_WITH_OPTIONAL_COLUMN;
import static io.trino.connector.MockConnectorEntities.TPCH_WITH_HIDDEN_COLUMN_DATA;
import static io.trino.plugin.tpch.TpchMetadata.TINY_SCHEMA_NAME;
import static io.trino.testing.TestingSession.testSessionBuilder;
Expand All @@ -53,6 +54,7 @@ public class TestRowFilter
{
private static final String CATALOG = "local";
private static final String MOCK_CATALOG = "mock";
private static final String MOCK_CATALOG_MISSING_COLUMNS = "mockmissingcolumns";
private static final String USER = "user";
private static final String VIEW_OWNER = "view-owner";
private static final String RUN_AS_USER = "run-as-user";
Expand Down Expand Up @@ -94,6 +96,9 @@ public void init()
if (schemaTableName.equals(new SchemaTableName("tiny", "nation_with_hidden_column"))) {
return TPCH_NATION_WITH_HIDDEN_COLUMN;
}
if (schemaTableName.equals(new SchemaTableName("tiny", "nation_with_optional_column"))) {
return TPCH_NATION_WITH_OPTIONAL_COLUMN;
}
throw new UnsupportedOperationException();
})
.withData(schemaTableName -> {
Expand All @@ -103,12 +108,37 @@ public void init()
if (schemaTableName.equals(new SchemaTableName("tiny", "nation_with_hidden_column"))) {
return TPCH_WITH_HIDDEN_COLUMN_DATA;
}
if (schemaTableName.equals(new SchemaTableName("tiny", "nation_with_optional_column"))) {
return TPCH_NATION_DATA;
}
throw new UnsupportedOperationException();
})
.build();

runner.createCatalog(MOCK_CATALOG, mock, ImmutableMap.of());

MockConnectorFactory mockMissingColumns = MockConnectorFactory.builder()
.withName("mockmissingcolumns")
.withGetViews((s, prefix) -> ImmutableMap.<SchemaTableName, ConnectorViewDefinition>builder()
.put(new SchemaTableName("default", "nation_view"), view)
.build())
.withGetColumns(schemaTableName -> {
if (schemaTableName.equals(new SchemaTableName("tiny", "nation_with_optional_column"))) {
return TPCH_NATION_WITH_OPTIONAL_COLUMN;
}
throw new UnsupportedOperationException();
})
.withData(schemaTableName -> {
if (schemaTableName.equals(new SchemaTableName("tiny", "nation_with_optional_column"))) {
return TPCH_NATION_DATA;
}
throw new UnsupportedOperationException();
})
.withAllowMissingColumnsOnInsert(true)
.build();

runner.createCatalog(MOCK_CATALOG_MISSING_COLUMNS, mockMissingColumns, ImmutableMap.of());

assertions = new QueryAssertions(runner);
accessControl = assertions.getQueryRunner().getAccessControl();
}
Expand Down Expand Up @@ -559,4 +589,28 @@ public void testRowFilterOnHiddenColumn()
.skippingTypesCheck()
.matches("SELECT BIGINT '25'");
}

@Test
public void testRowFilterOnOptionalColumn()
{
accessControl.reset();

accessControl.rowFilter(
new QualifiedObjectName(MOCK_CATALOG_MISSING_COLUMNS, "tiny", "nation_with_optional_column"),
USER,
new ViewExpression(USER, Optional.empty(), Optional.empty(), "length(optional) > 2"));

assertions.query("INSERT INTO mockmissingcolumns.tiny.nation_with_optional_column(nationkey, name, regionkey, comment, optional) VALUES (0, 'POLAND', 0, 'No comment', 'some string')")
.assertThat()
.skippingTypesCheck()
.matches("VALUES BIGINT '1'");

assertThatThrownBy(() -> assertions.query("INSERT INTO mockmissingcolumns.tiny.nation_with_optional_column(nationkey, name, regionkey, comment, optional) VALUES (0, 'POLAND', 0, 'No comment', 'so')"))
.isInstanceOf(TrinoException.class)
.hasMessage("Access Denied: Cannot insert row that does not match to a row filter");

assertThatThrownBy(() -> assertions.query("INSERT INTO mockmissingcolumns.tiny.nation_with_optional_column(nationkey, name, regionkey, comment, optional) VALUES (0, 'POLAND', 0, 'No comment', null)"))
.isInstanceOf(TrinoException.class)
.hasMessage("Access Denied: Cannot insert row that does not match to a row filter");
}
}

0 comments on commit c5112d1

Please sign in to comment.