Repository: dbcli/athenacli

Issue:
Feature request: output cost and size of query
Something like:

        submission_date = stats['QueryExecution']['Status']['SubmissionDateTime']
        completion_date = stats['QueryExecution']['Status']['CompletionDateTime']
        execution_time = stats['QueryExecution']['Statistics']['EngineExecutionTimeInMillis']
        data_scanned = stats['QueryExecution']['Statistics']['DataScannedInBytes']
        query_cost = data_scanned / 1000000000000.0 * 5.0

        print('Time: {}, CPU Time: {}ms total, Data Scanned: {}, Cost: ${:,.2f}\n'.format(
            str(completion_date - submission_date).split('.')[0],
            execution_time,
            human_readable(data_scanned),
            query_cost
        ))

from

https://github.com/guardian/athena-cli/blob/0d5cc98a1cac21094077ed29cf168686ca61f10d/athena_cli.py#L217-L228

(Currently we only print `Time`.)

Comments:

Patch:
diff --git a/athenacli/packages/format_utils.py b/athenacli/packages/format_utils.py
new file mode 100644
index 0000000..8450913
--- /dev/null
+++ b/athenacli/packages/format_utils.py
@@ -0,0 +1,34 @@
+# -*- coding: utf-8 -*-
+
+
+def format_status(rows_length=None, cursor=None):
+    return rows_status(rows_length) + statistics(cursor)
+
+def rows_status(rows_length):
+    if rows_length:
+        return '%d row%s in set' % (rows_length, '' if rows_length == 1 else 's')
+    else:
+        return 'Query OK'
+
+def statistics(cursor):
+    if cursor:
+        # Most regions are $5 per TB: https://aws.amazon.com/athena/pricing/
+        approx_cost = cursor.data_scanned_in_bytes / (1024 ** 4) * 5
+
+        return '\nExecution time: %d ms, Data scanned: %s, Approximate cost: $%.2f' % (
+                cursor.execution_time_in_millis,
+                humanize_size(cursor.data_scanned_in_bytes),
+                approx_cost)
+    else:
+        return ''
+
+def humanize_size(num_bytes):
+    suffixes = ['B', 'KB', 'MB', 'GB', 'TB']
+
+    suffix_index = 0
+    while num_bytes >= 1024 and suffix_index < len(suffixes) - 1:
+        num_bytes /= 1024.0
+        suffix_index += 1
+
+    num = ('%.2f' % num_bytes).rstrip('0').rstrip('.')
+    return '%s %s' % (num, suffixes[suffix_index])
diff --git a/athenacli/sqlexecute.py b/athenacli/sqlexecute.py
index 3a0d8c5..d560ecd 100644
--- a/athenacli/sqlexecute.py
+++ b/athenacli/sqlexecute.py
@@ -5,6 +5,7 @@ import sqlparse
 import pyathena
 
 from athenacli.packages import special
+from athenacli.packages.format_utils import format_status
 
 logger = logging.getLogger(__name__)
 
@@ -92,11 +93,11 @@ class SQLExecute(object):
         if cursor.description is not None:
             headers = [x[0] for x in cursor.description]
             rows = cursor.fetchall()
-            status = '%d row%s in set' % (len(rows), '' if len(rows) == 1 else 's')
+            status = format_status(rows_length=len(rows), cursor=cursor)
         else:
             logger.debug('No rows in result.')
             rows = None
-            status = 'Query OK'
+            status = format_status(rows_length=None, cursor=cursor)
         return (title, rows, headers, status)
 
     def tables(self):
diff --git a/changelog.md b/changelog.md
index 22d8273..cafd121 100644
--- a/changelog.md
+++ b/changelog.md
@@ -1,6 +1,13 @@
 (Unreleased; add upcoming change notes here)
 ==============================================
 
+1.3.0
+========
+
+Features
+----------
+* Show query execution statistics, such as the amount of data scanned and the approximate cost. (Thanks: @pgr0ss)
+
 1.2.0
 ========
 


Test patch:
diff --git a/test/test_format_utils.py b/test/test_format_utils.py
new file mode 100644
index 0000000..abe8d75
--- /dev/null
+++ b/test/test_format_utils.py
@@ -0,0 +1,26 @@
+# -*- coding: utf-8 -*-
+
+
+from collections import namedtuple
+from athenacli.packages.format_utils import format_status, humanize_size
+
+
+def test_format_status_plural():
+    assert format_status(rows_length=1) == "1 row in set"
+    assert format_status(rows_length=2) == "2 rows in set"
+
+def test_format_status_no_results():
+    assert format_status(rows_length=None) == "Query OK"
+
+def test_format_status_with_stats():
+    FakeCursor = namedtuple("FakeCursor", ["execution_time_in_millis", "data_scanned_in_bytes"])
+
+    assert format_status(rows_length=1, cursor=FakeCursor(10, 12345678900)) == "1 row in set\nExecution time: 10 ms, Data scanned: 11.5 GB, Approximate cost: $0.06"
+    assert format_status(rows_length=2, cursor=FakeCursor(1000, 1234)) == "2 rows in set\nExecution time: 1000 ms, Data scanned: 1.21 KB, Approximate cost: $0.00"
+
+def test_humanize_size():
+    assert humanize_size(20) == "20 B"
+    assert humanize_size(2000) == "1.95 KB"
+    assert humanize_size(200000) == "195.31 KB"
+    assert humanize_size(20000000) == "19.07 MB"
+    assert humanize_size(200000000000) == "186.26 GB"