address review comments

awslabs · Mar 22, 2024 · e91af9f · e91af9f
1 parent 2bb77a0
commit e91af9f
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/...m/src/main/java/com/amazonaws/athena/connectors/timestream/TimestreamMetadataHandler.java b/...m/src/main/java/com/amazonaws/athena/connectors/timestream/TimestreamMetadataHandler.java
@@ -320,6 +320,7 @@ public GetTableResponse doGetQueryPassthroughSchema(BlockAllocator allocator, Ge
         queryPassthrough.verify(request.getQueryPassthroughArguments());
         String customerPassedQuery = request.getQueryPassthroughArguments().get(TimestreamQueryPassthrough.QUERY);
         QueryRequest queryRequest = new QueryRequest().withQueryString(customerPassedQuery).withMaxRows(1);
+        // Timestream Query does not provide a way to conduct a dry run or retrieve metadata results without execution. Therefore, we need to "seek" at least once before obtaining metadata.
         QueryResult queryResult = tsQuery.query(queryRequest);
         List<ColumnInfo> columnInfo = queryResult.getColumnInfo();
         SchemaBuilder schemaBuilder = SchemaBuilder.newBuilder();

diff --git a/.../main/java/com/amazonaws/athena/connectors/timestream/qpt/TimestreamQueryPassthrough.java b/.../main/java/com/amazonaws/athena/connectors/timestream/qpt/TimestreamQueryPassthrough.java
@@ -72,8 +72,8 @@ public Logger getLogger()
     @Override
     public void customConnectorVerifications(Map<String, String> engineQptArguments)
     {
-        String partiQLStatement = engineQptArguments.get(QUERY);
-        String upperCaseStatement = partiQLStatement.trim().toUpperCase(Locale.ENGLISH);
+        String customerPassedQuery = engineQptArguments.get(QUERY);
+        String upperCaseStatement = customerPassedQuery.trim().toUpperCase(Locale.ENGLISH);
 
         // Immediately check if the statement starts with "SELECT"
         if (!upperCaseStatement.startsWith("SELECT")) {
@@ -85,7 +85,7 @@ public void customConnectorVerifications(Map<String, String> engineQptArguments)
 
         // Check if the statement contains any disallowed keywords
         for (String keyword : disallowedKeywords) {
-            if (upperCaseStatement.contains(" " + keyword + " ") || upperCaseStatement.startsWith(keyword + " ")) {
+            if (upperCaseStatement.contains(keyword)) {
                 throw new UnsupportedOperationException("Unaccepted operation; only SELECT statements are allowed. Found: " + keyword);
             }
         }