From ffe0eddd9b1a3161573328cadf0a2e95dcb5f3bc Mon Sep 17 00:00:00 2001
From: Ari Brown <ari@airintech.com>
Date: Tue, 19 Sep 2023 16:12:58 -0400
Subject: [PATCH] output consolidation; added executable via poetry

---
 bin/minerva-console | 32 --------------------------------
 minerva/athena.py   |  7 ++++++-
 minerva/console.py  | 37 +++++++++++++++++++++++++++++++++++++
 pyproject.toml      |  4 +++-
 4 files changed, 46 insertions(+), 34 deletions(-)
 delete mode 100755 bin/minerva-console
 create mode 100755 minerva/console.py

diff --git a/bin/minerva-console b/bin/minerva-console
deleted file mode 100755
index d95ae1a..0000000
--- a/bin/minerva-console
+++ /dev/null
@@ -1,32 +0,0 @@
-#!/usr/bin/env python3
-import minerva
-import pprint
-import readline
-import argparse
-
-pp = pprint.PrettyPrinter(indent=4)
-
-parser = argparse.ArgumentParser(description="""
-REPL for the Athena SQL engine
-""")
-parser.add_argument("-p", "--profile", default="hay", help="The AWS profile to use")
-parser.add_argument("-o", "--output", default="s3://haystac-pmo-athena/output")
-args = parser.parse_args()
-
-m = minerva.Minerva(args.profile)
-athena = m.athena(args.output)
-
-text  = input("> ")
-while text != "\\q":
-  query = athena.query(text)
-
-  try:
-    data  = query.results()
-    pp.pprint(data.head(10))
-    print()
-    print(f"\t({'$%.2f' % query.cost}, {query.runtime})")
-  except Exception as e:
-    print(e)
-
-  text = input("> ")
-
diff --git a/minerva/athena.py b/minerva/athena.py
index 75e39e4..5073a92 100644
--- a/minerva/athena.py
+++ b/minerva/athena.py
@@ -50,7 +50,11 @@ class Execute:
 
     # Send the SQL to Athena for running
     def run(self):
-        config = {"OutputLocation": self.athena.output}
+        if self.__class__ == Query:
+            config = {"OutputLocation": os.path.join(self.athena.output, "results")}
+        else:
+            config = {"OutputLocation": self.athena.output}
+
         if self.params:
             resp  = self.client.start_query_execution(QueryString=self.query(),
                                                       ResultConfiguration=config,
@@ -94,6 +98,7 @@ class Query(Execute):
     # Automatically includes unloading the results to Parquet format
     def query(self):
         out    = os.path.join(self.athena.output,
+                              "results",
                               str(random.random()))
         query  = f"unload ({self.sql}) to {repr(out)} " + \
                  f"with (format = '{self.DATA_STYLE}')"
diff --git a/minerva/console.py b/minerva/console.py
new file mode 100755
index 0000000..d362902
--- /dev/null
+++ b/minerva/console.py
@@ -0,0 +1,37 @@
+import minerva
+import pprint
+import readline
+import argparse
+
+def main():
+  pp = pprint.PrettyPrinter(indent=4)
+
+  parser = argparse.ArgumentParser(description="""
+  REPL for the Athena SQL engine
+  """)
+  parser.add_argument("-p", "--profile", default="hay", help="The AWS profile to use")
+  parser.add_argument("-o", "--output", default="s3://haystac-pmo-athena/")
+  args = parser.parse_args()
+
+  m = minerva.Minerva(args.profile)
+  athena = m.athena(args.output)
+
+  text  = input("> ")
+  while text != "\\q":
+    try:
+      first_word = text.split()[0]
+      if first_word.lower() in ["drop", "alter", "create", "update"]:
+        query = athena.execute(text)
+        query.finish()
+      else:
+        query = athena.query(text)
+        data  = query.results()
+        pp.pprint(data.head(10))
+
+      print()
+      print(f"\t({'$%.2f' % query.cost}, {query.runtime})")
+    except Exception as e:
+      print(e)
+
+    text = input("> ")
+
diff --git a/pyproject.toml b/pyproject.toml
index ce4de04..6389a39 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -9,10 +9,12 @@ authors = [
 ]
 packages = [
   { include = "minerva/**/*.py" },
-  { include = "bin/*" }
 ]
 readme = "README.md"
 
+[tool.poetry.scripts]
+minerva-console = "minerva.console:main"
+
 [tool.poetry.dependencies]
 python = ">3.9, <3.11"
 boto3 = "^1.26.161"