diff --git a/src/gsheets_extension.cpp b/src/gsheets_extension.cpp index 5756623..85a8b0f 100644 --- a/src/gsheets_extension.cpp +++ b/src/gsheets_extension.cpp @@ -17,6 +17,7 @@ #include #include #include +#include using namespace std; // OpenSSL linked through vcpkg @@ -187,13 +188,35 @@ static void ReadSheetFunction(ClientContext &context, TableFunctionInput &data_p output.SetCardinality(row_count); } +static std::string extract_sheet_id(const std::string& input) { + // Check if the input is already a sheet ID (no slashes) + if (input.find('/') == std::string::npos) { + return input; + } + + // Regular expression to match the sheet ID in a Google Sheets URL + if(input.find("docs.google.com/spreadsheets/d/") != std::string::npos) { + std::regex sheet_id_regex("/d/([a-zA-Z0-9-_]+)"); + std::smatch match; + + if (std::regex_search(input, match, sheet_id_regex) && match.size() > 1) { + return match.str(1); + } + } + + throw duckdb::InvalidInputException("Invalid Google Sheets URL or ID"); +} + // Update the ReadSheetBind function static unique_ptr ReadSheetBind(ClientContext &context, TableFunctionBindInput &input, vector &return_types, vector &names) { - auto sheet_id = input.inputs[0].GetValue(); + auto sheet_input = input.inputs[0].GetValue(); auto token_file_path = input.inputs[1].GetValue(); bool header = input.inputs.size() > 2 ? input.inputs[2].GetValue() : true; + // Extract the sheet ID from the input (URL or ID) + std::string sheet_id = extract_sheet_id(sheet_input); + // Use the read_token_from_file function from gsheets_auth.hpp std::string token = read_token_from_file(token_file_path); diff --git a/test/sql/gsheets.test b/test/sql/gsheets.test index 626181f..262d7b0 100644 --- a/test/sql/gsheets.test +++ b/test/sql/gsheets.test @@ -20,4 +20,15 @@ Bob 25 New York Charlie 45 Chicago Drake NULL NULL NULL NULL NULL +Archie 99 NULL + +# Test the full URL +query III +FROM read_gsheet('https://docs.google.com/spreadsheets/d/11QdEasMWbETbFVxry-SsD8jVcdYIT1zBQszcF84MdE8/edit#gid=0', 'token.txt', header=true); +---- +Alice 30 Toronto +Bob 25 New York +Charlie 45 Chicago +Drake NULL NULL +NULL NULL NULL Archie 99 NULL \ No newline at end of file