Add support for Parquet export format

splitgraph · Oct 29, 2022 · 6d0326a · 6d0326a
1 parent c5c99d7
commit 6d0326a
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 6 deletions.
diff --git a/splitgraph/cloud/__init__.py b/splitgraph/cloud/__init__.py
@@ -902,17 +902,21 @@ def get_export_job_status(self, task_id: str) -> Optional[ExportJobStatus]:
             output=data["output"],
         )
 
-    def start_export(self, query: str) -> str:
+    def start_export(self, query: str, export_format: Optional[str] = None) -> str:
         query = query.strip()
         if query.endswith(";"):
             logging.warning("The query ends with ';', automatically removing")
             query = query[:-1]
 
+        variables = {"query": query}
+        if export_format:
+            variables["format"] = export_format
+
         response = self._gql(
             {
                 "query": START_EXPORT,
                 "operationName": "StartExport",
-                "variables": {"query": query},
+                "variables": variables,
             },
             handle_errors=True,
         )

diff --git a/splitgraph/cloud/queries.py b/splitgraph/cloud/queries.py
@@ -276,8 +276,8 @@
 }"""
 
 
-START_EXPORT = """mutation StartExport($query: String!) {
-  exportQuery(query: $query, exportFormat: "csv") {
+START_EXPORT = """mutation StartExport($query: String!, $format: String! = "csv") {
+  exportQuery(query: $query, exportFormat: $format) {
     id
   }
 }

diff --git a/splitgraph/commandline/cloud.py b/splitgraph/commandline/cloud.py
@@ -1010,7 +1010,12 @@ def wait_for_download(client: "GQLAPIClient", task_id: str) -> str:
 
 @click.command("download")
 @click.option("--remote", default="data.splitgraph.com", help="Name of the remote registry to use.")
-@click.option("--file-format", default="csv", type=click.Choice(["csv"]))
+@click.option(
+    "--file-format",
+    default="csv",
+    type=click.Choice(["csv", "parquet"]),
+    help="File format (csv or Parquet with ZSTD)",
+)
 @click.argument("query", type=str)
 @click.argument("output_path", type=str, default=None, required=False)
 def download_c(remote, file_format, query, output_path):
@@ -1023,7 +1028,7 @@ def download_c(remote, file_format, query, output_path):
 
     client = GQLAPIClient(remote)
 
-    task_id = client.start_export(query=query)
+    task_id = client.start_export(query=query, export_format=file_format)
     download_url = wait_for_download(client, task_id)
     download_file(download_url, output_path)
 

diff --git a/test/splitgraph/commandline/http_fixtures.py b/test/splitgraph/commandline/http_fixtures.py
@@ -1022,6 +1022,7 @@ def _gql_callback(request, uri, response_headers):
         if body["query"] == START_EXPORT:
             assert body["variables"] == {
                 "query": "SELECT * FROM some_table",
+                "format": "csv",
             }
             return [
                 200,