DRILL-8395: Add Support for Insert and Drop Table to GoogleSheets Plugin (#2748)
diff --git a/contrib/storage-googlesheets/README.md b/contrib/storage-googlesheets/README.md
index 5f26582..38aaae6 100644
--- a/contrib/storage-googlesheets/README.md
+++ b/contrib/storage-googlesheets/README.md
@@ -110,6 +110,12 @@
WHERE SCHEMA_NAME LIKE 'googlesheets%'
```
+### Implicit Metadata Fields
+GoogleSheets has two implicit metadata fields which are:
+
+* `_sheets`: This will return a list of sheet (tab) names in a given GS document
+* `_title`: You can also access the file name with the `_title` field. Note that the file name is NOT unique and should not be used for querying data.
+
Due to rate limits from Google, the tabs are not reported to the `INFORMATION_SCHEMA`. However, it is possible to obtain a list of all available tabs with the following query:
```sql
@@ -118,8 +124,6 @@
LIMIT 1
```
-You can also access the file name with the `_title` field. Note that the file name is NOT
-unique and should not be used for querying data.
### Using Aliases
Since the sheet IDs from Google are not human readable, one way to make your life easier is to use Drill's aliasing features to provide a better name for the actual sheet name.
@@ -150,19 +154,36 @@
If this is incorrect you can set the `extractHeaders` parameter to `false`and Drill will name each field `field_n` where `n` is the column index.
# Writing Data To Google Sheets
-When Drill is connected to Google Sheets, you can also write data to Google Sheets. The basic procedure is
+When Drill is connected to GoogleSheets, you can also write data to Google Sheets. The basic procedure is
the same as with any other data source. Simply write a `CREATE TABLE AS` (CTAS) query and your data will be
-written to Google Sheets.
+written to GoogleSheets.
-One challenge is that once you have created the new sheet, you will have to manually retrieve the spreadsheet ID
-from Google Sheets in order to query your new data.
-
-### Dropping Tables
-At the time of implementation, it is only possible to delete tables from within a Google Sheets document. You may encounter errors if you try to delete tables from documents
-that only have one table in them. The format for deleting a table is:
+If you use a GoogleSheets filetoken in your CTAS query, Drill will create a new tab in that GoogleSheets document. However, if you use a file name, Drill will create a new GoogleSheets document and then create a new tab within that document.
```sql
-DROP TABLE googlesheets.<sheet id>.<tab name>
+-- This will add a tab to an existing GoogleSheets Document
+CREATE TABLE googlesheets.`2384r7wuf2934iroeci2390ue2ur3r23948230948`.`tab_name` AS SELECT * FROM data
+
+-- This will create a new GoogleSheets Document with a single tab
+CREATE TABLE googlesheets.`new_doc`.`tab1` AS SELECT * FROM data
+
+```
+
+### Inserting (Appending) to Existing GoogleSheets Tabs
+GoogleSheets also supports inserting (appending) data to existing GoogleSheets tabs. Syntax is:
+
+```sql
+INSERT INTO googlesheets.`<file_token>`.`<tab name>` SELECT * FROM data
+```
+
+
+### Dropping Tables
+The `DROP TABLE` command will drop a tab from a GoogleSheet document. If the document only has one tab, the entire document will be deleted.
+
+The format for deleting a table is:
+
+```sql
+DROP TABLE googlesheets.<file_token>.<tab name>
```
# Possible Future Work
@@ -173,10 +194,6 @@
much more usable would be to automatically create an alias (either public) automatically mapping the unreadable sheetID to the document title.
This could be accomplished after the first query or after a CTAS query.
-### Google Drive Integration
-Integrating with Google Drive may allow additional functionality such as getting the actual document name, deleting documents and a few other basic functions. However, the
-Google Drive permissions require additional validation from Google.
-
### Additional Pushdowns
The current implementation supports pushdowns for projection and limit.
The Google Sheets API is quite complex and incredibly poorly documented. In this author's opinion, it is quite possibly one of the worst APIs he has ever seen.
diff --git a/contrib/storage-googlesheets/src/main/java/org/apache/drill/exec/store/googlesheets/GoogleSheetsBatchInsertWriter.java b/contrib/storage-googlesheets/src/main/java/org/apache/drill/exec/store/googlesheets/GoogleSheetsBatchInsertWriter.java
new file mode 100644
index 0000000..952557a
--- /dev/null
+++ b/contrib/storage-googlesheets/src/main/java/org/apache/drill/exec/store/googlesheets/GoogleSheetsBatchInsertWriter.java
@@ -0,0 +1,52 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.drill.exec.store.googlesheets;
+
+import org.apache.drill.common.exceptions.UserException;
+import org.apache.drill.exec.ops.OperatorContext;
+import org.apache.drill.exec.record.VectorAccessible;
+import org.apache.drill.exec.store.googlesheets.utils.GoogleSheetsUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+
+public class GoogleSheetsBatchInsertWriter extends GoogleSheetsBatchWriter {
+ private static final Logger logger = LoggerFactory.getLogger(GoogleSheetsBatchInsertWriter.class);
+
+ public GoogleSheetsBatchInsertWriter(OperatorContext context, String name, GoogleSheetsWriter config) {
+ super(context, name, config);
+ }
+
+ @Override
+ public void updateSchema(VectorAccessible batch) {
+ // no-op
+ }
+
+ @Override
+ public void cleanup() {
+ try {
+ GoogleSheetsUtils.appendDataToGoogleSheet(service, sheetName, tabName, values);
+ } catch (IOException e) {
+ throw UserException.dataWriteError(e)
+ .message("Error writing to GoogleSheets " + e.getMessage())
+ .build(logger);
+ }
+ }
+}
diff --git a/contrib/storage-googlesheets/src/main/java/org/apache/drill/exec/store/googlesheets/GoogleSheetsBatchWriter.java b/contrib/storage-googlesheets/src/main/java/org/apache/drill/exec/store/googlesheets/GoogleSheetsBatchWriter.java
index 6741a31..78e7375 100644
--- a/contrib/storage-googlesheets/src/main/java/org/apache/drill/exec/store/googlesheets/GoogleSheetsBatchWriter.java
+++ b/contrib/storage-googlesheets/src/main/java/org/apache/drill/exec/store/googlesheets/GoogleSheetsBatchWriter.java
@@ -69,12 +69,12 @@
public class GoogleSheetsBatchWriter extends AbstractRecordWriter {
private static final Logger logger = LoggerFactory.getLogger(GoogleSheetsBatchWriter.class);
- private final Sheets service;
- private final String tabName;
- private final String sheetName;
- private final List<List<Object>> values;
+ protected final Sheets service;
+ protected final String tabName;
+ protected final String sheetName;
+ protected final List<List<Object>> values;
private List<Object> rowList;
- private String spreadsheetID;
+ protected String spreadsheetID;
public GoogleSheetsBatchWriter(OperatorContext context, String name, GoogleSheetsWriter config) {
GoogleSheetsStoragePlugin plugin = config.getPlugin();
@@ -100,17 +100,24 @@
@Override
public void updateSchema(VectorAccessible batch) throws IOException {
- // Create the new GoogleSheet doc
- Spreadsheet spreadsheet = new Spreadsheet()
- .setProperties(new SpreadsheetProperties().setTitle(sheetName));
- spreadsheet = service.spreadsheets().create(spreadsheet)
- .setFields("spreadsheetId")
- .execute();
+ // If the incoming sheetName is actually a file token then simply add a new tab to the existing document.
+ if (GoogleSheetsUtils.isProbableFileToken(sheetName)) {
+ GoogleSheetsUtils.addTabToGoogleSheet(service, sheetName, tabName);
+ spreadsheetID = sheetName;
+ } else {
+ // Otherwise, create the new GoogleSheet document and add a tab.
+ Spreadsheet spreadsheet = new Spreadsheet()
+ .setProperties(new SpreadsheetProperties().setTitle(sheetName));
- this.spreadsheetID = spreadsheet.getSpreadsheetId();
- // Now add the tab
- GoogleSheetsUtils.addTabToGoogleSheet(service, spreadsheetID, tabName);
+ spreadsheet = service.spreadsheets().create(spreadsheet)
+ .setFields("spreadsheetId")
+ .execute();
+
+ this.spreadsheetID = spreadsheet.getSpreadsheetId();
+ // Now add the tab
+ GoogleSheetsUtils.addTabToGoogleSheet(service, spreadsheetID, tabName);
+ }
// Add the column names to the values list. GoogleSheets does not have any concept
// of column names, so we just insert the column names as the first row of data.
diff --git a/contrib/storage-googlesheets/src/main/java/org/apache/drill/exec/store/googlesheets/GoogleSheetsInsertWriter.java b/contrib/storage-googlesheets/src/main/java/org/apache/drill/exec/store/googlesheets/GoogleSheetsInsertWriter.java
new file mode 100644
index 0000000..6c5eedf
--- /dev/null
+++ b/contrib/storage-googlesheets/src/main/java/org/apache/drill/exec/store/googlesheets/GoogleSheetsInsertWriter.java
@@ -0,0 +1,57 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.drill.exec.store.googlesheets;
+
+import com.fasterxml.jackson.annotation.JacksonInject;
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import org.apache.drill.common.logical.StoragePluginConfig;
+import org.apache.drill.exec.physical.base.PhysicalOperator;
+import org.apache.drill.exec.store.StoragePluginRegistry;
+
+
+public class GoogleSheetsInsertWriter extends GoogleSheetsWriter {
+
+ public static final String OPERATOR_TYPE = "GOOGLESHEETS_INSERT_WRITER";
+
+ @JsonCreator
+ public GoogleSheetsInsertWriter(
+ @JsonProperty("child") PhysicalOperator child,
+ @JsonProperty("sheetName") String sheetName,
+ @JsonProperty("name") String name,
+ @JsonProperty("storage") StoragePluginConfig storageConfig,
+ @JsonProperty("queryUser") String queryUser,
+ @JacksonInject StoragePluginRegistry engineRegistry) {
+ super(child, sheetName, name, storageConfig, queryUser, engineRegistry);
+ }
+
+ public GoogleSheetsInsertWriter(PhysicalOperator child, String sheetName, String name, String queryUser, GoogleSheetsStoragePlugin plugin) {
+ super(child, sheetName, name, queryUser, plugin);
+ }
+
+ @Override
+ public String getOperatorType() {
+ return OPERATOR_TYPE;
+ }
+
+ @Override
+ protected PhysicalOperator getNewWithChild(PhysicalOperator child) {
+ return new GoogleSheetsInsertWriter(child, getSheetName(), getTableName(), getQueryUser(), getPlugin());
+ }
+}
diff --git a/contrib/storage-googlesheets/src/main/java/org/apache/drill/exec/store/googlesheets/GoogleSheetsInsertWriterBatchCreator.java b/contrib/storage-googlesheets/src/main/java/org/apache/drill/exec/store/googlesheets/GoogleSheetsInsertWriterBatchCreator.java
new file mode 100644
index 0000000..e2314c8
--- /dev/null
+++ b/contrib/storage-googlesheets/src/main/java/org/apache/drill/exec/store/googlesheets/GoogleSheetsInsertWriterBatchCreator.java
@@ -0,0 +1,42 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.drill.exec.store.googlesheets;
+
+import org.apache.drill.exec.ops.ExecutorFragmentContext;
+import org.apache.drill.exec.physical.impl.BatchCreator;
+import org.apache.drill.exec.physical.impl.InsertWriterRecordBatch;
+import org.apache.drill.exec.proto.UserBitShared.UserCredentials;
+import org.apache.drill.exec.record.CloseableRecordBatch;
+import org.apache.drill.exec.record.RecordBatch;
+
+import java.util.List;
+
+@SuppressWarnings("unused")
+public class GoogleSheetsInsertWriterBatchCreator implements BatchCreator<GoogleSheetsInsertWriter> {
+ @Override
+ public CloseableRecordBatch getBatch(ExecutorFragmentContext context, GoogleSheetsInsertWriter config, List<RecordBatch> children) {
+ assert children != null && children.size() == 1;
+
+ UserCredentials userCreds = context.getContextInformation().getQueryUserCredentials();
+
+ return new InsertWriterRecordBatch(config, children.iterator().next(), context,
+ new GoogleSheetsBatchInsertWriter(null, config.getTableName(), config)
+ );
+ }
+}
diff --git a/contrib/storage-googlesheets/src/main/java/org/apache/drill/exec/store/googlesheets/GoogleSheetsStoragePlugin.java b/contrib/storage-googlesheets/src/main/java/org/apache/drill/exec/store/googlesheets/GoogleSheetsStoragePlugin.java
index f7e7cc5..eba0003 100644
--- a/contrib/storage-googlesheets/src/main/java/org/apache/drill/exec/store/googlesheets/GoogleSheetsStoragePlugin.java
+++ b/contrib/storage-googlesheets/src/main/java/org/apache/drill/exec/store/googlesheets/GoogleSheetsStoragePlugin.java
@@ -191,6 +191,11 @@
return true;
}
+ @Override
+ public boolean supportsInsert() {
+ return true;
+ }
+
/**
* This method gets (and caches) the Google Service needed for API calls.
* @return An authenticated {@link Sheets} Google Sheets service.
diff --git a/contrib/storage-googlesheets/src/main/java/org/apache/drill/exec/store/googlesheets/schema/GoogleSheetsDrillSchema.java b/contrib/storage-googlesheets/src/main/java/org/apache/drill/exec/store/googlesheets/schema/GoogleSheetsDrillSchema.java
index 911f74b..462f85b 100644
--- a/contrib/storage-googlesheets/src/main/java/org/apache/drill/exec/store/googlesheets/schema/GoogleSheetsDrillSchema.java
+++ b/contrib/storage-googlesheets/src/main/java/org/apache/drill/exec/store/googlesheets/schema/GoogleSheetsDrillSchema.java
@@ -18,6 +18,7 @@
package org.apache.drill.exec.store.googlesheets.schema;
+import com.google.api.services.drive.Drive;
import com.google.api.services.sheets.v4.Sheets;
import com.google.api.services.sheets.v4.model.Sheet;
import org.apache.calcite.schema.Table;
@@ -27,9 +28,11 @@
import org.apache.drill.exec.physical.base.Writer;
import org.apache.drill.exec.planner.logical.CreateTableEntry;
import org.apache.drill.exec.planner.logical.DynamicDrillTable;
+import org.apache.drill.exec.planner.logical.ModifyTableEntry;
import org.apache.drill.exec.store.AbstractSchema;
import org.apache.drill.exec.store.SchemaConfig;
import org.apache.drill.exec.store.StorageStrategy;
+import org.apache.drill.exec.store.googlesheets.GoogleSheetsInsertWriter;
import org.apache.drill.exec.store.googlesheets.GoogleSheetsScanSpec;
import org.apache.drill.exec.store.googlesheets.GoogleSheetsStoragePlugin;
import org.apache.drill.exec.store.googlesheets.GoogleSheetsStoragePluginConfig;
@@ -64,15 +67,17 @@
private final String fileToken;
private final String fileName;
+ private List<Sheet> tabList;
+
public GoogleSheetsDrillSchema(AbstractSchema parent, String fileToken,
GoogleSheetsStoragePlugin plugin,
SchemaConfig schemaConfig,
Sheets sheetsService, String fileName) {
- super(parent.getSchemaPath(), fileToken);
+ super(parent.getSchemaPath(), GoogleSheetsRootSchema.getFileTokenWithCorrectCase(((GoogleSheetsRootSchema) parent).getTokenMap(), fileToken));
this.plugin = plugin;
this.schemaConfig = schemaConfig;
- this.fileToken = fileToken;
this.parent = (GoogleSheetsRootSchema) parent;
+ this.fileToken = GoogleSheetsRootSchema.getFileTokenWithCorrectCase(((GoogleSheetsRootSchema) parent).getTokenMap(), fileToken);
this.sheetsService = sheetsService;
this.tableList = new ArrayList<>();
this.fileName = fileName;
@@ -145,12 +150,11 @@
}
private void populateActiveTables() {
- List<Sheet> tabList;
try {
tabList = GoogleSheetsUtils.getTabList(sheetsService, fileToken);
} catch (IOException e) {
throw UserException.connectionError(e)
- .message("Unable to obtain tab list for Google Sheet document " + fileToken)
+ .message("Unable to obtain tab list for Google Sheet document " + fileToken + ". " + e.getMessage())
.build(logger);
}
// Add sub schemas to list, then create tables
@@ -178,7 +182,7 @@
.message(plugin.getName() + " is not writable.")
.build(logger);
}
- String documentName = this.name;
+ String documentName = this.fileToken;
return new CreateTableEntry() {
@Override
public Writer getWriter(PhysicalOperator child) {
@@ -192,6 +196,39 @@
};
}
+ @Override
+ public ModifyTableEntry modifyTable(String tableName) {
+ return child -> new GoogleSheetsInsertWriter(child, this.fileToken, tableName, schemaConfig.getUserName(), plugin);
+ }
+
+ @Override
+ public void dropTable(String indexName) {
+ logger.debug("Index name: {}", indexName);
+
+ // The GoogleSheets API will not allow you to delete a tab if the file only has one tab. In that case,
+ // we delete the entire file.
+ if (tabList.size() == 1) {
+ Drive driveService = plugin.getDriveService(schemaConfig.getUserName());
+ try {
+ driveService.files().delete(fileToken);
+ } catch (IOException e) {
+ throw UserException.internalError(e)
+ .message("Error deleting GoogleSheets file. " + e.getMessage())
+ .build(logger);
+ }
+ }
+
+ Sheet sheetToDrop = GoogleSheetsUtils.getSheetFromTabList(indexName, tabList);
+ try {
+ GoogleSheetsUtils.removeTabFromGoogleSheet(sheetsService, fileToken, sheetToDrop);
+ } catch (IOException e) {
+ throw UserException.internalError(e)
+ .message(e.getMessage())
+ .build(logger);
+ }
+ }
+
+
private void registerTable(String name, DynamicDrillTable table) {
activeTables.put(name, table);
tableList.add(table);
diff --git a/contrib/storage-googlesheets/src/main/java/org/apache/drill/exec/store/googlesheets/schema/GoogleSheetsRootSchema.java b/contrib/storage-googlesheets/src/main/java/org/apache/drill/exec/store/googlesheets/schema/GoogleSheetsRootSchema.java
index 0c00a7e..9ceeb2f 100644
--- a/contrib/storage-googlesheets/src/main/java/org/apache/drill/exec/store/googlesheets/schema/GoogleSheetsRootSchema.java
+++ b/contrib/storage-googlesheets/src/main/java/org/apache/drill/exec/store/googlesheets/schema/GoogleSheetsRootSchema.java
@@ -21,6 +21,7 @@
import com.google.api.services.drive.Drive;
import com.google.api.services.sheets.v4.Sheets;
import org.apache.calcite.schema.Table;
+import org.apache.commons.lang3.StringUtils;
import org.apache.drill.common.exceptions.UserException;
import org.apache.drill.exec.planner.logical.DynamicDrillTable;
import org.apache.drill.exec.store.AbstractSchema;
@@ -67,6 +68,11 @@
@Override
public AbstractSchema getSubSchema(String name) {
+ // If the name is a file token but not mixed case, get the correct file token
+ if (GoogleSheetsUtils.isProbableFileToken(name) && !StringUtils.isMixedCase(name)) {
+ name = getFileTokenWithCorrectCase(tokenMap, name);
+ }
+
GoogleSheetsDrillSchema schema = schemas.get(name);
// This level here represents the actual Google document. Attempt to validate that it exists, and
// if so, add it to the schema list. If not, throw an exception.
@@ -102,4 +108,30 @@
public String getTypeName() {
return GoogleSheetsStoragePluginConfig.NAME;
}
+
+ public Map<String, String> getTokenMap() {
+ return this.tokenMap;
+ }
+
+ /**
+ * Drill automatically converts the file token to lower case during DDL queries. Since Google is case-sensitive,
+ * this method insures that we are using the correct file token.
+ * @return A file token in the correct case.
+ */
+ public static String getFileTokenWithCorrectCase(Map<String, String> tokenMap, String lowercaseToken) {
+ if (StringUtils.isMixedCase(lowercaseToken)) {
+ return lowercaseToken;
+ }
+ Set<String> tokens = tokenMap.keySet();
+
+ for (String token : tokens) {
+ if (token.toLowerCase().contentEquals(lowercaseToken)) {
+ return token;
+ }
+ }
+
+ throw UserException.internalError()
+ .message("Could not find token: " + lowercaseToken)
+ .build(logger);
+ }
}
diff --git a/contrib/storage-googlesheets/src/main/java/org/apache/drill/exec/store/googlesheets/utils/GoogleSheetsUtils.java b/contrib/storage-googlesheets/src/main/java/org/apache/drill/exec/store/googlesheets/utils/GoogleSheetsUtils.java
index 971c149..e3cc597 100644
--- a/contrib/storage-googlesheets/src/main/java/org/apache/drill/exec/store/googlesheets/utils/GoogleSheetsUtils.java
+++ b/contrib/storage-googlesheets/src/main/java/org/apache/drill/exec/store/googlesheets/utils/GoogleSheetsUtils.java
@@ -35,7 +35,9 @@
import com.google.api.services.sheets.v4.Sheets.Spreadsheets.Values.BatchGet;
import com.google.api.services.sheets.v4.SheetsScopes;
import com.google.api.services.sheets.v4.model.AddSheetRequest;
+import com.google.api.services.sheets.v4.model.AppendValuesResponse;
import com.google.api.services.sheets.v4.model.BatchUpdateSpreadsheetRequest;
+import com.google.api.services.sheets.v4.model.DeleteSheetRequest;
import com.google.api.services.sheets.v4.model.Request;
import com.google.api.services.sheets.v4.model.Sheet;
import com.google.api.services.sheets.v4.model.SheetProperties;
@@ -266,7 +268,7 @@
/**
* Google Sheets tokens are strings of length 44 that contain upper and lower case letters, numbers and underscores.
* This function will attempt to identify file tokens.
- *
+ * <p>
* Given that Google's spec for file IDs is not officially published, and can change at any time, we will keep the
* validation as light as possible to prevent future issues, in the event Google changes their file Id structure.
* @param id A {@link String} containing an unknown identifier
@@ -371,6 +373,25 @@
}
/**
+ * Finds a {@link Sheet} from a list of tabs with a given title. If the sheet is not present,
+ * the function will throw a User Exception.
+ * @param tabName The name of the desired sheet.
+ * @param tabList A {@link List} of {@link Sheet} objects
+ * @return The desired Sheet.
+ */
+ public static Sheet getSheetFromTabList(String tabName, List<Sheet> tabList) {
+ for (Sheet sheet : tabList) {
+ if (sheet.getProperties().getTitle().contentEquals(tabName)) {
+ return sheet;
+ }
+ }
+
+ throw UserException.dataReadError()
+ .message("Could not find sheet " + tabName)
+ .build(logger);
+ }
+
+ /**
* This function is used to get data when projection is pushed down to Google Sheets.
* @param service The Authenticated GoogleSheets service
* @param sheetID The GoogleSheet ID. This can be found in the Sheet URL
@@ -592,6 +613,24 @@
}
/**
+ * Removes a sheet from an existing GoogleSheets document. This method should only be used if the GoogleSheets
+ * document has more than one tab.
+ * @param service An authenticated GoogleSheet {@link Sheets}
+ * @param fileToken The File token of the GoogleSheet containing the sheet to be deleted
+ * @param deletedTab A {@link Sheet} which will be removed
+ * @throws IOException If anything goes wrong.
+ */
+ public static void removeTabFromGoogleSheet(Sheets service, String fileToken, Sheet deletedTab) throws IOException {
+ List<Request> requests = new ArrayList<>();
+ requests.add(new Request()
+ .setDeleteSheet(new DeleteSheetRequest().setSheetId(deletedTab.getProperties().getSheetId()))
+ );
+
+ BatchUpdateSpreadsheetRequest body = new BatchUpdateSpreadsheetRequest().setRequests(requests);
+ service.spreadsheets().batchUpdate(fileToken, body).execute();
+ }
+
+ /**
* Accepts a list of data and writes this data to a GoogleSheet document.
* @param service An authenticated GoogleSheet service
* @param sheetID The SheetID. This can be obtained from the URL of the GoogleSheet Document
@@ -610,4 +649,24 @@
.setValueInputOption("RAW")
.execute();
}
+
+ /**
+ * Accepts a list of data and writes this data to a GoogleSheet document.
+ * @param service An authenticated GoogleSheet service
+ * @param sheetID The SheetID. This can be obtained from the URL of the GoogleSheet Document
+ * @param tabName The tab name within the aforementioned GoogleSheet
+ * @param data A list of rows of the data to be inserted.
+ * @throws IOException If anything goes wrong, throw an IO exception
+ */
+ public static void appendDataToGoogleSheet(Sheets service, String sheetID, String tabName, List<List<Object>> data)
+ throws IOException {
+ String range = tabName + "!A1";
+ ValueRange body = new ValueRange()
+ .setValues(data)
+ .setMajorDimension("ROWS");
+
+ AppendValuesResponse result = service.spreadsheets().values().append(sheetID, range, body)
+ .setValueInputOption("RAW")
+ .execute();
+ }
}
diff --git a/contrib/storage-googlesheets/src/test/java/org/apache/drill/exec/store/googlesheets/TestGoogleSheetsWriter.java b/contrib/storage-googlesheets/src/test/java/org/apache/drill/exec/store/googlesheets/TestGoogleSheetsWriter.java
index daaba71..2b06de1 100644
--- a/contrib/storage-googlesheets/src/test/java/org/apache/drill/exec/store/googlesheets/TestGoogleSheetsWriter.java
+++ b/contrib/storage-googlesheets/src/test/java/org/apache/drill/exec/store/googlesheets/TestGoogleSheetsWriter.java
@@ -20,8 +20,12 @@
import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.drill.categories.RowSetTest;
+import org.apache.drill.common.types.TypeProtos.MinorType;
import org.apache.drill.common.util.DrillFileUtils;
import org.apache.drill.exec.oauth.PersistentTokenTable;
+import org.apache.drill.exec.physical.rowSet.RowSet;
+import org.apache.drill.exec.record.metadata.SchemaBuilder;
+import org.apache.drill.exec.record.metadata.TupleMetadata;
import org.apache.drill.exec.store.StoragePluginRegistry;
import org.apache.drill.exec.store.StoragePluginRegistry.PluginException;
import org.apache.drill.shaded.guava.com.google.common.base.Charsets;
@@ -29,6 +33,7 @@
import org.apache.drill.test.ClusterFixture;
import org.apache.drill.test.ClusterTest;
import org.apache.drill.test.QueryBuilder.QuerySummary;
+import org.apache.drill.test.rowSet.RowSetComparison;
import org.junit.BeforeClass;
import org.junit.Ignore;
import org.junit.Test;
@@ -40,6 +45,7 @@
import java.util.List;
import java.util.Map;
+import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;
@@ -54,7 +60,9 @@
private static String accessToken;
private static String refreshToken;
- // Note on testing: Testing the writing capabilites of this plugin is challenging.
+ private static String sheetID;
+
+ // Note on testing: Testing the writing capabilities of this plugin is challenging.
// The primary issue is that when you execute a CTAS query, you do so using the file name.
// However, it does not seem possible to retrieve the created file's ID which is what you
// need to actually verify that the query successfully wrote the results. Therefore, at this
@@ -74,6 +82,7 @@
String clientSecret = tokenMap.get("client_secret");
accessToken = tokenMap.get("access_token");
refreshToken = tokenMap.get("refresh_token");
+ sheetID = tokenMap.get("sheet_id");
pluginRegistry = cluster.drillbit().getContext().getStorage();
GoogleSheetsStoragePluginConfig config = GoogleSheetsStoragePluginConfig.builder()
@@ -82,6 +91,8 @@
.redirectUris(REDIRECT_URI)
.authUri(AUTH_URI)
.tokenUri(TOKEN_URI)
+ .extractHeaders(true)
+ .allTextMode(false)
.build();
config.setEnabled(true);
@@ -103,6 +114,65 @@
assertTrue(insertResults.succeeded());
}
+ @Test
+ public void testCTASLifecycle() throws Exception {
+ // This test goes through the entire CTAS, INSERT, DROP lifecycle.
+ try {
+ initializeTokens();
+ } catch (PluginException e) {
+ fail(e.getMessage());
+ }
+
+ // We are creating a new tab in an existing GS document
+ String sql = String.format("CREATE TABLE googlesheets.`%s`.`Sheet3` AS SELECT * FROM cp.`data/insert_data.csvh`", sheetID);
+ QuerySummary results = queryBuilder().sql(sql).run();
+ assertTrue(results.succeeded());
+
+ // Verify the sheet was created
+ sql = String.format("SELECT * FROM googlesheets.`%s`.`Sheet3`", sheetID);
+ results = queryBuilder().sql(sql).run();
+ assertTrue(results.succeeded());
+ assertEquals(2, results.recordCount());
+
+ // Now Insert additional records into the sheet
+ sql = String.format("INSERT INTO googlesheets.`%s`.`Sheet3` SELECT * FROM cp.`data/insert_data2.csvh`", sheetID);
+ results = queryBuilder().sql(sql).run();
+ assertTrue(results.succeeded());
+
+ // Verify that the records were inserted
+ sql = String.format("SELECT * FROM googlesheets.`%s`.`Sheet3`", sheetID);
+ RowSet rowSet = queryBuilder().sql(sql).rowSet();
+
+ TupleMetadata expectedSchema = new SchemaBuilder()
+ .addNullable("col1", MinorType.FLOAT8)
+ .addNullable("col2", MinorType.FLOAT8)
+ .addNullable("col3", MinorType.FLOAT8)
+ .buildSchema();
+
+ RowSet expected = client.rowSetBuilder(expectedSchema)
+ .addRow(1,2,3)
+ .addRow(4,5,6)
+ .addRow(7,8,9)
+ .addRow(10,11,12)
+ .build();
+ new RowSetComparison(expected).verifyAndClearAll(rowSet);
+
+ // Drop the table
+ sql = String.format("DROP TABLE googlesheets.`%s`.`Sheet3`", sheetID);
+ results = queryBuilder().sql(sql).run();
+ assertTrue(results.succeeded());
+
+ // Verify that it's gone
+ sql = String.format("SELECT * FROM googlesheets.`%s`.`Sheet3`", sheetID);
+ try {
+ results = queryBuilder().sql(sql).run();
+ fail();
+ } catch (Exception e) {
+ assertTrue(e.getMessage().contains("'Sheet3' not found"));
+ }
+ }
+
+
/**
* This function is used for testing only. It initializes a {@link PersistentTokenTable} and populates it
* with a valid access and refresh token.
diff --git a/contrib/storage-googlesheets/src/test/resources/data/insert_data.csvh b/contrib/storage-googlesheets/src/test/resources/data/insert_data.csvh
new file mode 100644
index 0000000..e39967f
--- /dev/null
+++ b/contrib/storage-googlesheets/src/test/resources/data/insert_data.csvh
@@ -0,0 +1,3 @@
+col1, col2, col3
+1,2,3
+4,5,6
\ No newline at end of file
diff --git a/contrib/storage-googlesheets/src/test/resources/data/insert_data2.csvh b/contrib/storage-googlesheets/src/test/resources/data/insert_data2.csvh
new file mode 100644
index 0000000..7128ece
--- /dev/null
+++ b/contrib/storage-googlesheets/src/test/resources/data/insert_data2.csvh
@@ -0,0 +1,3 @@
+col1, col2, col3
+7,8,9
+10,11,12
\ No newline at end of file
diff --git a/exec/java-exec/src/main/java/org/apache/drill/exec/planner/sql/SchemaUtilities.java b/exec/java-exec/src/main/java/org/apache/drill/exec/planner/sql/SchemaUtilities.java
index 780270e..5afbf06 100644
--- a/exec/java-exec/src/main/java/org/apache/drill/exec/planner/sql/SchemaUtilities.java
+++ b/exec/java-exec/src/main/java/org/apache/drill/exec/planner/sql/SchemaUtilities.java
@@ -96,6 +96,7 @@
/** Utility method to search for schema path starting from the given <i>schema</i> reference */
public static SchemaPlus searchSchemaTree(SchemaPlus schema, final List<String> schemaPath) {
for (String schemaName : schemaPath) {
+
// schemas in Drill are case insensitive and stored in lower case
schema = schema.getSubSchema(schemaName.toLowerCase());
if (schema == null) {
