Add validation for schema names in Pinot

Author: agrawalreetika

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<String> 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<SchemaTableName> listTables(ConnectorSession session, String schemaN
     {
         ImmutableList.Builder<SchemaTableName> builder = ImmutableList.builder();
         for (String table : pinotPrestoConnection.getTableNames()) {
-            builder.add(new SchemaTableName("default", table));
+            builder.add(new SchemaTableName(SCHEMA_NAME, table));
         }
         return builder.build();
     }

presto-pinot-toolkit/src/test/java/com/facebook/presto/pinot/TestPinotMetadata.java

@@ -22,6 +22,8 @@
 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;
 
 public class TestPinotMetadata
@@ -33,14 +35,14 @@ public class TestPinotMetadata
     @Test
     public void testTables()
     {
-        ConnectorSession session = TestPinotSplitManager.createSessionWithNumSplits(1, false, pinotConfig);
+        ConnectorSession session = createSessionWithNumSplits(1, false, pinotConfig);
         List<SchemaTableName> 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<String> 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());
+        PinotTableHandle withWeirdSchema = metadata.getTableHandle(session, new SchemaTableName("foo", realtimeOnlyTable.getTableName()));
+        assertEquals(withWeirdSchema.getTableName(), realtimeOnlyTable.getTableName());
+        PinotTableHandle withAnotherSchema = metadata.getTableHandle(session, new SchemaTableName(realtimeOnlyTable.getTableName(), realtimeOnlyTable.getTableName()));
+        assertEquals(withAnotherSchema.getTableName(), realtimeOnlyTable.getTableName());
     }
 }