From d7b1055292d1bcf2d34be426a16142d0ca2af236 Mon Sep 17 00:00:00 2001 From: Reetika Agrawal Date: Mon, 1 Dec 2025 21:21:14 +0530 Subject: [PATCH] Add validation for schema names in Pinot --- .../com/facebook/presto/pinot/PinotMetadata.java | 11 +++++++++-- .../facebook/presto/pinot/TestPinotMetadata.java | 16 ++++++++++------ 2 files changed, 19 insertions(+), 8 deletions(-) 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..e35a301068220 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,7 +39,9 @@ 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; @@ -48,6 +51,7 @@ public class PinotMetadata private final String connectorId; private final PinotConnection pinotPrestoConnection; private final PinotConfig pinotConfig; + private static final String SCHEMA_NAME = "default"; @Inject public PinotMetadata(ConnectorId connectorId, PinotConnection pinotPrestoConnection, 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); } @@ -115,7 +122,7 @@ public List listTables(ConnectorSession session, String schemaN { ImmutableList.Builder builder = ImmutableList.builder(); 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..bf7e51428793d 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 com.facebook.presto.pinot.TestPinotSplitManager.createSessionWithNumSplits; import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertThrows; public class TestPinotMetadata { @@ -33,14 +37,14 @@ public class TestPinotMetadata @Test public void testTables() { - ConnectorSession session = TestPinotSplitManager.createSessionWithNumSplits(1, false, pinotConfig); + ConnectorSession session = 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()))); 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()); + assertThrows(PrestoException.class, () -> metadata.getTableHandle(session, + new SchemaTableName("foo", realtimeOnlyTable.getTableName()))); + assertThrows(PrestoException.class, () -> metadata.getTableHandle(session, + new SchemaTableName(realtimeOnlyTable.getTableName(), realtimeOnlyTable.getTableName()))); } }