ENH SQL CLI (#319)

Provide options to run SQL commands through the CLI. Users can either run a command with short output from a file or from the command line, or download the results of a query.

Author: Stephen Hoover

CHANGELOG.md

@@ -4,6 +4,7 @@ This project adheres to [Semantic Versioning](http://semver.org/).
 
 ## Unreleased
 ### Added
+- Add CLI command "sql" for command line SQL query execution. (#319)
 - Add helper function (run_template) to run a template given its id and return
   either the JSON output or the associated file ids. (#318)
 - Add helper function to list CivisML models. (#314)

civis/cli/__main__.py

@@ -27,7 +27,7 @@
 from civis.cli._cli_commands import (
     civis_ascii_art, files_download_cmd, files_upload_cmd,
     notebooks_download_cmd, notebooks_new_cmd,
-    notebooks_up, notebooks_down, notebooks_open)
+    notebooks_up, notebooks_down, notebooks_open, sql_cmd)
 from civis.resources import get_api_spec, CACHED_SPEC_PATH
 from civis.resources._resources import parse_method_name
 from civis._utils import open_session
@@ -226,6 +226,8 @@ def add_extra_commands(cli):
     notebooks_cmd.add_command(notebooks_open)
     cli.add_command(civis_ascii_art)
 
+    cli.add_command(sql_cmd)
+
 
 def configure_log_level():
     if os.getenv('CIVIS_LOG_LEVEL'):

civis/cli/_cli_commands.py

@@ -3,8 +3,10 @@
 """
 Additional commands to add to the CLI beyond the OpenAPI spec.
 """
-
+from __future__ import print_function
+import functools
 import os
+import sys
 
 import click
 import requests
@@ -70,6 +72,90 @@ def files_download_cmd(file_id, path):
         civis_to_file(file_id, f)
 
 
+@click.command('sql')
+@click.option('--dbname', '-d', type=str, required=True,
+              help='Execute the query on this Civis Platform database')
+@click.option('--command', '-c', type=str, default=None,
+              help='Execute a single input command string')
+@click.option('--filename', '-f', type=click.Path(exists=True),
+              help='Execute a query read from the given file')
+@click.option('--output', '-o', type=click.Path(),
+              help='Download query results to this file')
+@click.option('--quiet', '-q', is_flag=True, help='Suppress screen output')
+@click.option('-n', type=int, default=100,
+              help="Display up to this many rows of the result. Max 100.")
+def sql_cmd(dbname, command, filename, output, quiet, n):
+    """\b Execute a SQL query in Civis Platform
+
+    If neither a command nor an input file is specified, read
+    the SQL command from stdin.
+    If writing to an output file, use a Civis SQL script and write the
+    entire query output to the specified file.
+    If not writing to an output file, use a Civis Query, and return a
+    preview of the results, up to a maximum of 100 rows.
+    """
+    if filename:
+        with open(filename, 'rt') as f:
+            sql = f.read()
+    elif not command:
+        # Read the SQL query from user input. This also allows use of a heredoc
+        lines = []
+        while True:
+            try:
+                _i = input()
+            except (KeyboardInterrupt, EOFError):
+                # The end of a heredoc produces an EOFError.
+                break
+            if not _i:
+                break
+            else:
+                lines.append(_i)
+        sql = '\n'.join(lines)
+    else:
+        sql = command
+
+    if not sql:
+        # If the user didn't enter a query, exit.
+        if not quiet:
+            print('Did not receive a SQL query.', file=sys.stderr)
+        return
+
+    if not quiet:
+        print('\nExecuting query...', file=sys.stderr)
+    if output:
+        fut = civis.io.civis_to_csv(output, sql, database=dbname)
+        fut.result()  # Block for completion and raise exceptions if any
+        if not quiet:
+            print("Downloaded the result of the query to %s." % output,
+                  file=sys.stderr)
+    else:
+        fut = civis.io.query_civis(sql, database=dbname,
+                                   preview_rows=n, polling_interval=3)
+        cols = fut.result()['result_columns']
+        rows = fut.result()['result_rows']
+        if not quiet:
+            print('...Query complete.\n', file=sys.stderr)
+            print(_str_table_result(cols, rows))
+
+
+def _str_table_result(cols, rows):
+    """Turn a Civis Query result into a readable table."""
+    # Determine the maximum width of each column.
+    # First find the width of each element in each row, then find the max
+    # width in each position.
+    max_len = functools.reduce(
+        lambda x, y: [max(z) for z in zip(x, y)],
+        [[len(_v) for _v in _r] for _r in [cols] + rows])
+
+    header_str = " | ".join("{0:<{width}}".format(_v, width=_l)
+                            for _l, _v in zip(max_len, cols))
+    tb_strs = [header_str, len(header_str) * '-']
+    for row in rows:
+        tb_strs.append(" | ".join("{0:>{width}}".format(_v, width=_l)
+                                  for _l, _v in zip(max_len, row)))
+    return '\n'.join(tb_strs)
+
+
 @click.command('download')
 @click.argument('notebook_id', type=int)
 @click.argument('path')

civis/tests/test_cli.py

@@ -5,6 +5,7 @@
 import pytest
 
 from civis.cli.__main__ import generate_cli, invoke, make_operation_name
+from civis.cli._cli_commands import _str_table_result
 from civis.compat import mock
 from civis.tests import TEST_SPEC
 
@@ -118,3 +119,11 @@ def test_parameter_case(mock_session):
 )
 def test_make_operation_name(path, method, resource_name, exp):
     assert make_operation_name(path, method, resource_name) == exp
+
+
+def test_str_table_result():
+    cols = ['a', 'snake!']
+    rows = [['2', '3'], ['1.1', '3.3']]
+
+    out = _str_table_result(cols, rows)
+    assert out == "a   | snake!\n------------\n  2 |      3\n1.1 |    3.3"

docs/source/cli.rst

@@ -59,3 +59,27 @@ backend for your Jupyter notebooks.
 - ``civis notebooks open $NOTEBOOK_ID``
 
   Open an existing notebook (which may or may not be running) in your default browser.
+
+SQL
+---
+
+The Civis CLI allows for easy running of SQL queries on Civis Platform
+through the following commands:
+
+- ``civis sql [-n $MAX_LINES] -d $DATABASE_NAME -f $FILE_NAME``
+
+  Read a SQL query from a text file and run it on the specified database.
+  The results of the query, if any, will be shown after it completes
+  (up to a maximum of $MAX_LINES rows, defaulting to 100).
+
+- ``civis sql [-n $MAX_LINES] -d $DATABASE_NAME -c [$SQL_QUERY]``
+
+  Instead of reading from a file, read query text from a command line
+  argument. If you do not provide a query on the command line,
+  the query text will be taken from stdin.
+
+- ``civis sql -d $DATABASE_NAME [-f $SQL_FILE_NAME] -o $OUTPUT_FILE_NAME``
+
+  With the `-o` or `--output` option specified, the complete results
+  of the query will be downloaded to a CSV file at the requested location
+  after the query completes.