From c83dd8a5eb9a146ef8072932333a0bca98f9577a Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Fri, 25 Jan 2019 07:50:20 -0800
Subject: [PATCH] sqlite-utils csv data.db "select ..." command

---
 docs/cli.rst        | 11 +++++++++++
 sqlite_utils/cli.py | 23 +++++++++++++++++++++++
 tests/test_cli.py   | 21 +++++++++++++++++++++
 3 files changed, 55 insertions(+)

diff --git a/docs/cli.rst b/docs/cli.rst
index 7b4ace9..3b982bc 100644
--- a/docs/cli.rst
+++ b/docs/cli.rst
@@ -6,6 +6,17 @@
 
 The ``sqlite-utils`` command-line tool can be used to manipulate SQLite databases in a number of different ways.
 
+Running queries and returning CSV
+=================================
+
+You can execute a SQL query against a database and get the results back as CSV like this::
+
+    $ sqlite-utils csv docs.db "select id, title, author from documents"
+
+This will default to including the column names as a header row. To exclude the headers, use ``--no-headers``:
+
+    $ sqlite-utils csv docs.db "select id, title, author from documents" --no-headers
+
 Listing tables
 ==============
 
diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py
index 74f8a21..ed4089a 100644
--- a/sqlite_utils/cli.py
+++ b/sqlite_utils/cli.py
@@ -1,6 +1,8 @@
 import click
 import sqlite_utils
 import json
+import sys
+import csv as csv_std
 
 
 @click.group()
@@ -92,3 +94,24 @@ def upsert(path, table, json_file, pk):
     if isinstance(docs, dict):
         docs = [docs]
     db[table].upsert_all(docs, pk=pk)
+
+
+@cli.command()
+@click.argument(
+    "path",
+    type=click.Path(file_okay=True, dir_okay=False, allow_dash=False),
+    required=True,
+)
+@click.argument("sql")
+@click.option(
+    "--no-headers", help="Exclude headers from CSV output", is_flag=True, default=False
+)
+def csv(path, sql, no_headers):
+    "Execute SQL query and return the results as CSV"
+    db = sqlite_utils.Database(path)
+    cursor = db.conn.execute(sql)
+    writer = csv_std.writer(sys.stdout)
+    if not no_headers:
+        writer.writerow([c[0] for c in cursor.description])
+    for row in cursor:
+        writer.writerow(row)
diff --git a/tests/test_cli.py b/tests/test_cli.py
index c89bfc1..17d6e89 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -130,3 +130,24 @@ def test_upsert(db_path, tmpdir):
     assert upsert_dogs == db.execute_returning_dicts(
         "select * from dogs where id in (1, 2, 21) order by id"
     )
+
+
+def test_csv(db_path):
+    db = Database(db_path)
+    with db.conn:
+        db["dogs"].insert_all(
+            [
+                {"id": 1, "age": 4, "name": "Cleo"},
+                {"id": 2, "age": 2, "name": "Pancakes"},
+            ]
+        )
+    result = CliRunner().invoke(
+        cli.cli, ["csv", db_path, "select id, name, age from dogs"]
+    )
+    assert 0 == result.exit_code
+    assert "id,name,age\n1,Cleo,4\n2,Pancakes,2\n" == result.output
+    # Test the no-headers option:
+    result = CliRunner().invoke(
+        cli.cli, ["csv", db_path, "select id, name, age from dogs", "--no-headers"]
+    )
+    assert "1,Cleo,4\n2,Pancakes,2\n" == result.output