diff --git a/presto-pinot-toolkit/src/main/java/com/facebook/presto/pinot/PinotMetadata.java b/presto-pinot-toolkit/src/main/java/com/facebook/presto/pinot/PinotMetadata.java index 91e2cced93e64..ad32fa502b8da 100644 --- a/presto-pinot-toolkit/src/main/java/com/facebook/presto/pinot/PinotMetadata.java +++ b/presto-pinot-toolkit/src/main/java/com/facebook/presto/pinot/PinotMetadata.java @@ -23,6 +23,7 @@ import com.facebook.presto.spi.ConnectorTableLayoutResult; import com.facebook.presto.spi.ConnectorTableMetadata; import com.facebook.presto.spi.Constraint; +import com.facebook.presto.spi.PrestoException; import com.facebook.presto.spi.SchemaTableName; import com.facebook.presto.spi.SchemaTablePrefix; import com.facebook.presto.spi.TableNotFoundException; @@ -38,13 +39,16 @@ import static com.facebook.presto.pinot.PinotColumnHandle.PinotColumnType.REGULAR; import static com.facebook.presto.pinot.PinotErrorCode.PINOT_UNCLASSIFIED_ERROR; +import static com.facebook.presto.spi.StandardErrorCode.NOT_FOUND; import static com.google.common.base.Preconditions.checkArgument; +import static java.lang.String.format; import static java.util.Locale.ROOT; import static java.util.Objects.requireNonNull; public class PinotMetadata implements ConnectorMetadata { + private static final String SCHEMA_NAME = "default"; private final String connectorId; private final PinotConnection pinotPrestoConnection; private final PinotConfig pinotConfig; @@ -60,7 +64,7 @@ public PinotMetadata(ConnectorId connectorId, PinotConnection pinotPrestoConnect @Override public List listSchemaNames(ConnectorSession session) { - return ImmutableList.of("default"); + return ImmutableList.of(SCHEMA_NAME); } private String getPinotTableNameFromPrestoTableName(String prestoTableName) @@ -77,6 +81,9 @@ private String getPinotTableNameFromPrestoTableName(String prestoTableName) @Override public PinotTableHandle getTableHandle(ConnectorSession session, SchemaTableName tableName) { + if (!SCHEMA_NAME.equals(normalizeIdentifier(session, tableName.getSchemaName()))) { + throw new PrestoException(NOT_FOUND, format("Schema %s does not exist", tableName.getSchemaName())); + } String pinotTableName = getPinotTableNameFromPrestoTableName(tableName.getTableName()); return new PinotTableHandle(connectorId, tableName.getSchemaName(), pinotTableName); } @@ -114,8 +121,11 @@ public ConnectorTableMetadata getTableMetadata(ConnectorSession session, Connect public List listTables(ConnectorSession session, String schemaNameOrNull) { ImmutableList.Builder builder = ImmutableList.builder(); + if (schemaNameOrNull != null && !SCHEMA_NAME.equals(normalizeIdentifier(session, schemaNameOrNull))) { + return builder.build(); + } for (String table : pinotPrestoConnection.getTableNames()) { - builder.add(new SchemaTableName("default", table)); + builder.add(new SchemaTableName(SCHEMA_NAME, table)); } return builder.build(); } diff --git a/presto-pinot-toolkit/src/test/java/com/facebook/presto/pinot/TestPinotMetadata.java b/presto-pinot-toolkit/src/test/java/com/facebook/presto/pinot/TestPinotMetadata.java index d2337fd64e6b9..225faeb30736f 100644 --- a/presto-pinot-toolkit/src/test/java/com/facebook/presto/pinot/TestPinotMetadata.java +++ b/presto-pinot-toolkit/src/test/java/com/facebook/presto/pinot/TestPinotMetadata.java @@ -14,6 +14,7 @@ package com.facebook.presto.pinot; import com.facebook.presto.spi.ConnectorSession; +import com.facebook.presto.spi.PrestoException; import com.facebook.presto.spi.SchemaTableName; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; @@ -22,7 +23,10 @@ import java.util.List; import java.util.concurrent.Executors; +import static com.facebook.presto.pinot.TestPinotQueryBase.realtimeOnlyTable; import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertThrows; +import static org.testng.Assert.assertTrue; public class TestPinotMetadata { @@ -35,12 +39,17 @@ public void testTables() { ConnectorSession session = TestPinotSplitManager.createSessionWithNumSplits(1, false, pinotConfig); List schemaTableNames = metadata.listTables(session, (String) null); - assertEquals(ImmutableSet.copyOf(schemaTableNames), ImmutableSet.of(new SchemaTableName("default", TestPinotSplitManager.realtimeOnlyTable.getTableName()), new SchemaTableName("default", TestPinotSplitManager.hybridTable.getTableName()))); + assertEquals(ImmutableSet.copyOf(schemaTableNames), ImmutableSet.of(new SchemaTableName("default", realtimeOnlyTable.getTableName()), new SchemaTableName("default", TestPinotSplitManager.hybridTable.getTableName()))); + // Validate schemas List schemas = metadata.listSchemaNames(session); assertEquals(ImmutableList.copyOf(schemas), ImmutableList.of("default")); - PinotTableHandle withWeirdSchema = metadata.getTableHandle(session, new SchemaTableName("foo", TestPinotSplitManager.realtimeOnlyTable.getTableName())); - assertEquals(withWeirdSchema.getTableName(), TestPinotSplitManager.realtimeOnlyTable.getTableName()); - PinotTableHandle withAnotherSchema = metadata.getTableHandle(session, new SchemaTableName(TestPinotSplitManager.realtimeOnlyTable.getTableName(), TestPinotSplitManager.realtimeOnlyTable.getTableName())); - assertEquals(withAnotherSchema.getTableName(), TestPinotSplitManager.realtimeOnlyTable.getTableName()); + // Invalid schema should now fail + assertThrows(PrestoException.class, () -> metadata.getTableHandle(session, + new SchemaTableName("foo", realtimeOnlyTable.getTableName()))); + // Also invalid because schema != "default" + assertThrows(PrestoException.class, () -> metadata.getTableHandle(session, + new SchemaTableName(realtimeOnlyTable.getTableName(), realtimeOnlyTable.getTableName()))); + List otherSchemaTables = metadata.listTables(session, "other_schema"); + assertTrue(otherSchemaTables.isEmpty(), "Expected no tables for non-existent schema"); } }