From 777fb15a8017153e3b42807f565c0e56c78c1943 Mon Sep 17 00:00:00 2001
From: John Pruitt <jgpruitt@gmail.com>
Date: Thu, 25 Jul 2024 09:08:02 -0500
Subject: [PATCH] Add flag to ignore specific procedural languages (#185)

---
 .gitignore                |  1 +
 README.md                 |  1 +
 src/pgspot/cli.py         |  9 +++++++
 src/pgspot/visitors.py    |  2 ++
 tests/ignore_lang_test.py | 51 +++++++++++++++++++++++++++++++++++++++
 5 files changed, 64 insertions(+)
 create mode 100644 tests/ignore_lang_test.py

diff --git a/.gitignore b/.gitignore
index 6cbb3f9..d2ca7a9 100644
--- a/.gitignore
+++ b/.gitignore
@@ -4,3 +4,4 @@ __pycache__
 .pytest_cache
 .tox
 dist
+.idea
\ No newline at end of file
diff --git a/README.md b/README.md
index 3514abe..3c12709 100644
--- a/README.md
+++ b/README.md
@@ -61,6 +61,7 @@ options:
                         Analyze PLpgSQL code (default: True)
   --explain EXPLAIN     Describe an error/warning code
   --ignore IGNORE       Ignore error or warning code
+  --ignore-lang LANG    Ignore unknown procedural language
   --sql-accepting SQL_FN
                         Specify one or more sql-accepting functions
 ```
diff --git a/src/pgspot/cli.py b/src/pgspot/cli.py
index 70e3983..228435c 100755
--- a/src/pgspot/cli.py
+++ b/src/pgspot/cli.py
@@ -54,6 +54,14 @@ def run():
         type=str,
         help="Ignore error or warning code",
     )
+    parser.add_argument(
+        "--ignore-lang",
+        dest="ignore_lang",
+        action="append",
+        default=[],
+        type=str,
+        help="Ignore one or more unknown procedural languages",
+    )
     parser.add_argument(
         "--sql-accepting",
         dest="sql_fn",
@@ -70,6 +78,7 @@ def run():
     )
 
     args = parser.parse_args()
+    args.ignore_lang = [lang.lower() for lang in args.ignore_lang]
 
     counter = Counter(args)
     state = State(counter)
diff --git a/src/pgspot/visitors.py b/src/pgspot/visitors.py
index a10fbfe..892b834 100644
--- a/src/pgspot/visitors.py
+++ b/src/pgspot/visitors.py
@@ -267,6 +267,8 @@ def visit_CreateFunctionStmt(self, ancestors, node):
                 visit_plpgsql(state, node)
             case "c" | "internal":
                 pass
+            case language if language in self.state.args.ignore_lang:
+                pass
             case _:
                 self.state.unknown(f"Unknown function language: {language}")
 
diff --git a/tests/ignore_lang_test.py b/tests/ignore_lang_test.py
new file mode 100644
index 0000000..80a6392
--- /dev/null
+++ b/tests/ignore_lang_test.py
@@ -0,0 +1,51 @@
+from util import run
+
+
+sql = """
+CREATE FUNCTION python_max(a integer, b integer) RETURNS integer AS $$
+    return max(a, b)
+$$ LANGUAGE plpython3u SET search_path TO 'pg_catalog', 'pg_temp'
+;
+
+CREATE FUNCTION tcl_max(integer, integer) RETURNS integer AS $$
+    if {$1 > $2} {return $1}
+    return $2
+$$ LANGUAGE pltcl STRICT SET search_path TO 'pg_catalog', 'pg_temp'
+;
+"""
+
+
+def test_unknown_lang_plpython3u():
+    output = run(sql)
+    assert "Unknown function language: plpython3u" in output
+
+
+def test_unknown_lang_pltcl():
+    output = run(sql)
+    assert "Unknown function language: pltcl" in output
+
+
+def test_ignore_lang_plpython3u():
+    output = run(sql, list("--ignore-lang=plpython3u"))
+    assert "Unknown function language: plpython3u" not in output
+
+
+def test_ignore_lang_pltcl():
+    output = run(sql, list("--ignore-lang=pltcl"))
+    assert "Unknown function language: pltcl" not in output
+
+
+def test_ignore_lang_plpython3u_pltcl():
+    output = run(sql, ["--ignore-lang=plpython3u", "--ignore-lang=pltcl"])
+    assert (
+        "Unknown function language: pltcl" not in output
+        and "Unknown function language: plpython3u" not in output
+    )
+
+
+def test_ignore_lang_upper():
+    output = run(sql, ["--ignore-lang=PLPYTHON3U", "--ignore-lang=PLTCL"])
+    assert (
+        "Unknown function language: pltcl" not in output
+        and "Unknown function language: plpython3u" not in output
+    )