From c659b7c9996b6a7beb088821a6d1c876eeb16e32 Mon Sep 17 00:00:00 2001
From: Shawn Davis <shawn@ptltd.co>
Date: Thu, 11 Apr 2024 11:31:14 -0500
Subject: [PATCH] Added pgsql.sql command for scripting raw SQL.

---
 Makefile                          |  4 ++--
 scripttease/lib/commands/pgsql.py | 16 ++++++++++++++--
 tests/test_lib_commands_pgsql.py  |  8 ++++++++
 tests/test_lib_factories.py       |  2 +-
 4 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/Makefile b/Makefile
index 0780532..17ad68d 100644
--- a/Makefile
+++ b/Makefile
@@ -51,7 +51,7 @@ docs: lines
 	cd docs && make html;
 	cd docs && make coverage;
 	open docs/build/coverage/python.txt;
-	open docs/build/html/index.html;
+	firefox docs/build/html/index.html;
 
 #> clean - Remove pyc files and documentation builds.
 clean:
@@ -91,5 +91,5 @@ support:
 tests:
 	coverage run --source=. -m pytest;
 	coverage html --directory=$(COVERAGE_PATH);
-	open $(COVERAGE_PATH)/index.html;
+	firefox $(COVERAGE_PATH)/index.html;
 
diff --git a/scripttease/lib/commands/pgsql.py b/scripttease/lib/commands/pgsql.py
index d63f9fe..61b5f88 100644
--- a/scripttease/lib/commands/pgsql.py
+++ b/scripttease/lib/commands/pgsql.py
@@ -14,6 +14,7 @@ __all__ = (
     "pgsql_exists",
     "pgsql_grant",
     "pgsql_load",
+    "pgsql_sql",
     "pgsql_user",
 )
 
@@ -242,6 +243,17 @@ def pgsql_load(database, path, **kwargs):
     return pgsql("psql", **kwargs)
 
 
+def pgsql_sql(statement, database="template1", **kwargs):
+    kwargs.setdefault("comment", "run SQL statement")
+
+    kwargs['dbname'] = database
+
+    command = pgsql("psql", **kwargs)
+    command.statement += ' -c "%s"' % statement
+
+    return command
+
+
 def pgsql_user(name, admin_pass=None, admin_user="postgres", op="create", password=None, **kwargs):
     """Work with a PostgreSQL user.
 
@@ -278,7 +290,7 @@ def pgsql_user(name, admin_pass=None, admin_user="postgres", op="create", passwo
         return pgsql("dropuser", name, password=admin_pass, user=admin_user, **kwargs)
     elif op == "exists":
         kwargs.setdefault("comment", "determine if %s postgres user exists" % name)
-        kwargs.setdefault("register", "pgsql_use_exists")
+        kwargs.setdefault("register", "pgsql_user_exists")
 
         command = pgsql("psql", password=admin_pass, user=admin_user, **kwargs)
 
@@ -297,6 +309,6 @@ PGSQL_MAPPINGS = {
     'pgsql.dump': pgsql_dump,
     'pgsql.exists': pgsql_exists,
     'pgsql.grant': pgsql_grant,
-    # 'pgsql.sql': pgsql_exec,
+    'pgsql.sql': pgsql_sql,
     'pgsql.user': pgsql_user,
 }
diff --git a/tests/test_lib_commands_pgsql.py b/tests/test_lib_commands_pgsql.py
index 3d33902..5fd47db 100644
--- a/tests/test_lib_commands_pgsql.py
+++ b/tests/test_lib_commands_pgsql.py
@@ -71,6 +71,14 @@ def test_pgsql_load():
     assert '--bogus=1' in s  # to test passing any key
 
 
+def test_pgsql_sql():
+    c = pgsql_sql("UPDATE etl_testing SET is_processed = 't'", database="example_app")
+    s = c.get_statement()
+    # psql -U postgres --host=localhost --port=5432 --dbname="example_app" -c "UPDATE etl_testing SET is_processed = 't'"
+    assert '--dbname="example_app"' in s
+    assert '-c "UPDATE etl_testing SET is_processed = \'t\'"' in s
+
+
 def test_pgsql_user():
     c = pgsql_user("testing", password="secret")
     s = c.get_statement()
diff --git a/tests/test_lib_factories.py b/tests/test_lib_factories.py
index b495677..6816e00 100644
--- a/tests/test_lib_factories.py
+++ b/tests/test_lib_factories.py
@@ -21,7 +21,7 @@ def test_command_factory():
     ini = INILoader("tests/examples/kitchen_sink.ini")
     ini.load()
     commands = command_factory(ini)
-    assert len(commands) == 48
+    assert len(commands) == 49
 
     ini = INILoader("tests/examples/bad_command.ini")
     ini.load()