Merge pull request #28 from archiewood/copy-headers

Include headers with COPY TO
evidence-dev · Oct 25, 2024 · dd18663 · dd18663
2 parents f37b08a + 5a7281b
commit dd18663
Show file tree

Hide file tree

Showing 4 changed files with 73 additions and 19 deletions.
diff --git a/TODO.md b/TODO.md
@@ -1,7 +1,7 @@
 # TODO
 
 ## Copy to
-- [ ] header
+- [x] header
 - [ ] sheet
 - [ ] types
 - [ ] implicit copy to when it sees a gsheets url

diff --git a/src/gsheets_copy.cpp b/src/gsheets_copy.cpp
@@ -21,26 +21,12 @@ namespace duckdb
         copy_to_sink = GSheetWriteSink;
     }
 
-    struct GSheetCopyGlobalState : public GlobalFunctionData
-    {
-        explicit GSheetCopyGlobalState(ClientContext &context, const string &sheet_id, const string &token, const string &sheet_name)
-            : sheet_id(sheet_id), token(token), sheet_name(sheet_name)
-        {
-        }
-
-    public:
-        string sheet_id;
-        string token;
-        string sheet_name;
-    };
-
-    struct GSheetWriteBindData : public TableFunctionData
-    {
-    };
 
     unique_ptr<FunctionData> GSheetCopyFunction::GSheetWriteBind(ClientContext &context, CopyFunctionBindInput &input, const vector<string> &names, const vector<LogicalType> &sql_types)
     {
-        return make_uniq<GSheetWriteBindData>();
+        string file_path = input.info.file_path;
+
+        return make_uniq<GSheetWriteBindData>(file_path, sql_types, names);
     }
 
     unique_ptr<GlobalFunctionData> GSheetCopyFunction::GSheetWriteInitializeGlobal(ClientContext &context, FunctionData &bind_data, const string &file_path)
@@ -96,9 +82,11 @@ namespace duckdb
         sheet_data["range"] = "Sheet1";
         sheet_data["majorDimension"] = "ROWS";
 
-        // TODO: Add column headers
+        vector<string> headers = bind_data_p.Cast<GSheetWriteBindData>().options.name_list;        
 
         vector<vector<string>> values;
+        values.push_back(headers);
+
         for (idx_t r = 0; r < input.size(); r++)
         {
             vector<string> row;

diff --git a/src/include/gsheets_copy.hpp b/src/include/gsheets_copy.hpp
@@ -6,6 +6,37 @@
 
 namespace duckdb
 {
+    struct GSheetCopyGlobalState : public GlobalFunctionData
+    {
+        explicit GSheetCopyGlobalState(ClientContext &context, const string &sheet_id, const string &token, const string &sheet_name)
+            : sheet_id(sheet_id), token(token), sheet_name(sheet_name)
+        {
+        }
+
+    public:
+        string sheet_id;
+        string token;
+        string sheet_name;
+    };
+
+    struct GSheetWriteOptions
+    {
+        vector<string> name_list;
+    };
+
+    struct GSheetWriteBindData : public TableFunctionData
+    {
+        vector<string> files;
+        GSheetWriteOptions options;
+        vector<LogicalType> sql_types;
+
+        GSheetWriteBindData(string file_path, vector<LogicalType> sql_types, vector<string> names)
+            : sql_types(std::move(sql_types))
+        {
+            files.push_back(std::move(file_path));
+            options.name_list = std::move(names);
+        }
+    };
 
     class GSheetCopyFunction : public CopyFunction
     {

diff --git a/test/sql/copyto.test b/test/sql/copyto.test
@@ -0,0 +1,35 @@
+# name: test/sql/copyto.test
+# description: test COPY TO function
+# group: [gsheets]
+
+require-env TOKEN
+
+require gsheets
+
+# Create a secret NB must substitute a token, do not commit!
+statement ok
+create secret test_secret (type gsheet, token '${TOKEN}');
+
+# Create a table to copy to Google Sheet
+statement ok
+create table spreadsheets as 
+select 'Microsoft' as company, 'Excel' as product, 1985 as year_founded
+union all
+select 'Google', 'Google Sheets', 2006
+union all
+select 'Apple', 'Numbers', 1984
+union all
+select 'LibreOffice', 'Calc', 2000;
+
+# Copy the table to Google Sheet
+statement ok
+copy spreadsheets to 'https://docs.google.com/spreadsheets/d/11QdEasMWbETbFVxry-SsD8jVcdYIT1zBQszcF84MdE8/edit#gid=0' (format gsheet);
+
+# Read the table from Google Sheet
+query III
+from read_gsheet('https://docs.google.com/spreadsheets/d/11QdEasMWbETbFVxry-SsD8jVcdYIT1zBQszcF84MdE8/edit#gid=0');
+----
+Microsoft	Excel	1985
+Google	Google Sheets	2006
+Apple	Numbers	1984
+LibreOffice	Calc	2000