Added raw mode and other enhancements

viragtripathi · viragtripathi · commit ea8416b7843c · 2025-05-17T15:34:53.000-04:00
diff --git a/README.md b/README.md
@@ -26,13 +26,13 @@ A powerful CLI tool to extract, deduplicate, and analyze SQL logs for **Cockroac
 
 ## 🖼 Sample Output
 
-| Report Type | Preview                                           |
-|-------------|---------------------------------------------------|
-| HTML        | ![HTML Report Screenshot](docs/sample_report.png) |
-| Chart       | ![Bar Chart](docs/sample_chart.png)               |
-| CSV         | ![CSV Snippet](docs/sample_csv.png)               |
-| SQL         | ![SQL Snippet](docs/sample_sql.png)               |
-| Markdown    | ![Markdown Snippet](docs/sample_md.png)           |
+| Report Type | Preview                                                                                                               |
+|-------------|-----------------------------------------------------------------------------------------------------------------------|
+| HTML        | ![HTML Report Screenshot](https://raw.githubusercontent.com/viragtripathi/crdb-sql-audit/main/docs/sample_report.png) |
+| Chart       | ![Bar Chart](https://raw.githubusercontent.com/viragtripathi/crdb-sql-audit/main/docs/sample_chart.png)               |
+| CSV         | ![CSV Snippet](https://raw.githubusercontent.com/viragtripathi/crdb-sql-audit/main/docs/sample_csv.png)               |
+| SQL         | ![SQL Snippet](https://raw.githubusercontent.com/viragtripathi/crdb-sql-audit/main/docs/sample_sql.png)               |
+| Markdown    | ![Markdown Snippet](https://raw.githubusercontent.com/viragtripathi/crdb-sql-audit/main/docs/sample_md.png)           |
 
 
 ## 📦 Installation
@@ -53,21 +53,49 @@ pip install dist/crdb_sql_audit-0.2.0-py3-none-any.whl
 ```
 
 ## 🧪 Usage
+
 ```bash
 crdb-sql-audit \
   --dir /path/to/logs \
   --terms execute,pg_ \
   --out output/report
 ```
 
-### Additional Options
+You can also analyze a single file:
+
+```bash
+crdb-sql-audit \
+  --file /path/to/logfile.log \
+  --terms SELECT,INSERT \
+  --raw \
+  --out output/single_file_report
+```
+
+> ⚠️ You must provide either `--dir` or `--file`, but not both.
+
+### 🔧 Additional Options
+
+```bash
+--dir       Directory containing SQL log files (mutually exclusive with --file)
+--file      Single SQL log file (mutually exclusive with --dir)
+--terms     Comma-separated search keywords to extract SQL (default: 'execute,pg_')
+--raw       Treat each matching line as a raw SQL statement (default: False)
+--rules     Path to YAML rules file (optional, default: built-in PostgreSQL rules)
+--out       Output file prefix (default: crdb_audit_output/report)
+--help      Show usage help
+--version   Show current version
+```
+
+### 📘 CLI Help Example
+
 ```bash
---rules    Path to YAML rules file (optional, default: built-in PostgreSQL rules)
---help     Show usage help
---version  Show current version
+crdb-sql-audit --help
 ```
 
+![CLI help screenshot](https://raw.githubusercontent.com/viragtripathi/crdb-sql-audit/main/docs/cli_help.png)
+
 ### Custom Rules Example
+
 ```bash
 crdb-sql-audit \
   --dir ./logs \
diff --git a/crdb_sql_audit/audit.py b/crdb_sql_audit/audit.py
@@ -8,22 +8,34 @@
 
 from .rules_engine import load_rules, apply_rules
 
-def extract_sql(logs_dir, search_terms):
+def extract_sql(logs_path, search_terms, raw_mode=False):
     seen_sql = set()
-    for filename in os.listdir(logs_dir):
-        path = os.path.join(logs_dir, filename)
+
+    if os.path.isfile(logs_path):
+        paths = [logs_path]
+    elif os.path.isdir(logs_path):
+        paths = [os.path.join(logs_path, f) for f in os.listdir(logs_path)]
+    else:
+        raise ValueError(f"The path '{logs_path}' is not a valid file or directory.")
+
+    for path in paths:
         if os.path.isfile(path):
             with open(path, "r", encoding="utf-8", errors="ignore") as f:
                 for line in f:
                     if any(term in line for term in search_terms):
-                        match = re.search(r'execute [^:]+: (.+)', line)
-                        if match:
-                            sql = match.group(1).strip()
-                            seen_sql.add(sql)
+                        if raw_mode:
+                            sql = line.strip()
+                            if len(sql) > 5:
+                                seen_sql.add(sql)
                         else:
-                            func_match = re.findall(r'\b(pg_\w+\s*\(.*?\))', line)
-                            for func in func_match:
-                                seen_sql.add(func.strip())
+                            match = re.search(r'execute [^:]+: (.+)', line)
+                            if match:
+                                sql = match.group(1).strip()
+                                seen_sql.add(sql)
+                            else:
+                                func_match = re.findall(r'\b(pg_\w+\s*\(.*?\))', line)
+                                for func in func_match:
+                                    seen_sql.add(func.strip())
     return seen_sql
 
 
@@ -38,7 +50,7 @@ def analyze_compatibility(seen_sql, rules_path=None):
             logging.info(f"⚠️ MATCHED: {sql}")
             for m in matches:
                 logging.info(f"   🔸 Rule: {m['Rule_ID']} — {m['Issue']}")
-        all_issues.extend(matches)  # ✅ move this inside the loop
+        all_issues.extend(matches)
 
     return all_issues
 
diff --git a/crdb_sql_audit/main.py b/crdb_sql_audit/main.py
@@ -1,8 +1,10 @@
 import argparse
 import logging
+import os
+import sys
 from crdb_sql_audit.audit import extract_sql, analyze_compatibility, generate_reports
 
-__version__ = "0.2.5"
+__version__ = "0.2.6"
 
 logging.basicConfig(
     level=logging.INFO,
@@ -13,32 +15,47 @@
     ]
 )
 
-
 def main():
     parser = argparse.ArgumentParser(
         description="🔍 Analyze SQL logs for CockroachDB compatibility.",
         formatter_class=argparse.ArgumentDefaultsHelpFormatter
     )
 
     parser.add_argument("--version", action="version", version=f"%(prog)s {__version__}")
-    parser.add_argument("--dir", required=True, help="Directory with SQL log files")
+    parser.add_argument("--dir", help="Directory containing SQL log files")
+    parser.add_argument("--file", help="Single SQL log file")
     parser.add_argument("--terms", default="execute,pg_", help="Comma-separated search keywords to extract SQL (default: 'execute,pg_')")
+    parser.add_argument("--raw", action="store_true", help="Treat each matching line as raw SQL (skip extraction)")
     parser.add_argument("--out", default="crdb_audit_output/report", help="Output prefix for reports")
     parser.add_argument("--rules", required=False, help="Path to YAML rules file (defaults to Postgres rules)")
     args = parser.parse_args()
 
+    if not args.dir and not args.file:
+        parser.print_help()
+        sys.exit(1)
+    if args.dir and args.file:
+        parser.error("Please provide only one of --dir or --file, not both")
+    if args.dir and not os.path.isdir(args.dir):
+        parser.error(f"--dir expects a directory. Got: {args.dir}")
+    if args.file and not os.path.isfile(args.file):
+        parser.error(f"--file expects a file. Got: {args.file}")
+
     search_terms = [t.strip() for t in args.terms.split(",")]
-    seen_sql = extract_sql(args.dir, search_terms)
+    logs_path = args.dir if args.dir else args.file
+    seen_sql = extract_sql(logs_path, search_terms, raw_mode=args.raw)
     issues = analyze_compatibility(seen_sql, rules_path=args.rules)
+
     if args.rules:
         logging.info(f"🔍 Using rule file: {args.rules}")
         with open(args.rules, "r") as f:
             logging.info("🔍 Rule file contents:")
             logging.info(f.read())
     else:
         logging.warning("📦 No rule file provided. Falling back to built-in postgres_to_crdb.yaml.")
+
     logging.info(f"🧪 Total SQL statements analyzed: {len(seen_sql)}")
     logging.info(f"🚨 Total issues found: {len(issues)}")
+
     if not issues:
         logging.info("🕵️ Showing 5 sample SQL statements that didn't match:")
         for sql in list(seen_sql)[:5]:
@@ -47,4 +64,4 @@ def main():
     generate_reports(seen_sql, issues, args.out)
 
 if __name__ == "__main__":
-    main()
+    main()
diff --git a/docs/cli_help.png b/docs/cli_help.png
diff --git a/pyproject.toml b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "crdb-sql-audit"
-version = "0.2.5"
+version = "0.2.6"
 description = "Analyze PostgreSQL SQL logs for CockroachDB compatibility"
 authors = [
     { name = "Virag Tripathi", email = "virag.tripathi@gmail.com" }
