sql-data-guard 0.1.6__tar.gz → 0.1.8__tar.gz

{sql_data_guard-0.1.6/src/sql_data_guard.egg-info → sql_data_guard-0.1.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sql-data-guard
-Version: 0.1.6
+Version: 0.1.8
 Summary: Safety Layer for LLM Database Interactions
 Author-email: Imperva - Threat Reseach Infra <ww.dis.imperva.threat-research-infra@thalesgroup.com>
 License: MIT License

{sql_data_guard-0.1.6 → sql_data_guard-0.1.8}/pyproject.toml

@@ -9,7 +9,7 @@ package-dir = {"" = "src"}
 extra-index-url = ["https://pypi.org/simple"]
 [project]
 name = "sql-data-guard"
-version = "0.1.6"
+version = "0.1.8"
 dependencies = [
     "sqlglot"
 ]

sql_data_guard-0.1.8/src/sql_data_guard/mcpwrapper/mcp_wrapper.py

@@ -0,0 +1,86 @@
+import json
+import os
+import sys
+import threading
+from typing import Optional
+
+import docker
+from sql_data_guard import verify_sql
+
+
+def load_config() -> dict:
+    return json.load(open("/conf/config.json"))
+
+
+def start_inner_container():
+    client = docker.from_env()
+    container = client.containers.run(
+        config["mcp-server"]["image"],
+        " ".join(config["mcp-server"]["args"]),
+        volumes=[
+            v.replace("$PWD", os.environ["PWD"])
+            for v in config["mcp-server"]["volumes"]
+        ],
+        stdin_open=True,
+        auto_remove=True,
+        detach=True,
+        stdout=True,
+    )
+
+    def stream_output():
+        for line in container.logs(stream=True):
+            sys.stdout.write(line.decode("utf-8"))
+            sys.stdout.flush()
+
+    threading.Thread(target=stream_output, daemon=True).start()
+    return container
+
+
+def main():
+    container = start_inner_container()
+
+    try:
+        socket = container.attach_socket(params={"stdin": True, "stream": True})
+        # noinspection PyProtectedMember
+        socket._sock.setblocking(True)
+        for line in sys.stdin:
+            line = input_line(line)
+            # noinspection PyProtectedMember
+            socket._sock.sendall(line.encode("utf-8"))
+    except (KeyboardInterrupt, EOFError):
+        pass
+    finally:
+        container.stop()
+
+
+def get_sql(json_line: dict) -> Optional[str]:
+    sys.stderr.write(f"json_line: {json_line}\n")
+    if json_line["method"] == "tools/call":
+        for tool in config["mcp-tools"]:
+            if tool["tool-name"] == json_line["params"]["name"]:
+                return json_line["params"]["arguments"][tool["arg-name"]]
+    return None
+
+
+def input_line(line: str) -> str:
+    json_line = json.loads(line.encode("utf-8"))
+    sql = get_sql(json_line)
+    if sql:
+        result = verify_sql(
+            sql,
+            config["sql-data-guard"],
+            config["sql-data-guard"]["dialect"],
+        )
+        if not result["allowed"]:
+            sys.stderr.write(f"Blocked SQL: {sql}\nErrors: {list(result['errors'])}\n")
+            updated_sql = "SELECT 'Blocked by SQL Data Guard' AS message"
+            for error in result["errors"]:
+                updated_sql += f"\nUNION ALL SELECT '{error}' AS message"
+            json_line["params"]["arguments"]["query"] = updated_sql
+            line = json.dumps(json_line) + "\n"
+    return line
+
+
+if __name__ == "__main__":
+    config = load_config()
+    main()

{sql_data_guard-0.1.6 → sql_data_guard-0.1.8}/src/sql_data_guard/restriction_verification.py

@@ -3,8 +3,8 @@ from typing import List
 import sqlglot
 import sqlglot.expressions as expr
 
-from sql_data_guard.verification_context import VerificationContext
-from sql_data_guard.verification_utils import split_to_expressions
+from .verification_context import VerificationContext
+from .verification_utils import split_to_expressions
 
 
 def verify_restrictions(

{sql_data_guard-0.1.6 → sql_data_guard-0.1.8}/src/sql_data_guard/sql_data_guard.py

@@ -63,7 +63,7 @@ def verify_sql(sql: str, config: dict, dialect: str = None) -> dict:
         else:
             result.add_error("Could not find a query statement", False, 0.7)
     if result.can_fix and len(result.errors) > 0:
-        result.fixed = parsed.sql()
+        result.fixed = parsed.sql(dialect=dialect)
     return {
         "allowed": len(result.errors) == 0,
         "errors": result.errors,

{sql_data_guard-0.1.6 → sql_data_guard-0.1.8/src/sql_data_guard.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sql-data-guard
-Version: 0.1.6
+Version: 0.1.8
 Summary: Safety Layer for LLM Database Interactions
 Author-email: Imperva - Threat Reseach Infra <ww.dis.imperva.threat-research-infra@thalesgroup.com>
 License: MIT License

{sql_data_guard-0.1.6 → sql_data_guard-0.1.8}/src/sql_data_guard.egg-info/SOURCES.txt

@@ -12,6 +12,7 @@ src/sql_data_guard.egg-info/SOURCES.txt
 src/sql_data_guard.egg-info/dependency_links.txt
 src/sql_data_guard.egg-info/requires.txt
 src/sql_data_guard.egg-info/top_level.txt
+src/sql_data_guard/mcpwrapper/mcp_wrapper.py
 src/sql_data_guard/rest/__init__.py
 src/sql_data_guard/rest/sql_data_guard_rest.py
 test/test_duckdb_unit.py

{sql_data_guard-0.1.6 → sql_data_guard-0.1.8}/test/test_sql_guard_unit.py

@@ -369,14 +369,13 @@ class TestTrinoWithRestrictions:
             ]
         }
 
-    @pytest.mark.skip(reason="Bug in sqlglot. Waiting for fix")
     def test_date_add(self, config):
         verify_sql_test(
-            "SELECT id FROM accounts WHERE DATE(day) >= DATE_ADD('day', -7, CURRENT_DATE)",
+            "SELECT id FROM accounts WHERE DATE(day) >= DATE_ADD('DAY', -7, CURRENT_DATE)",
             config,
             dialect="trino",
             errors={"Missing restriction for table: accounts column: id value: 123"},
-            fix="SELECT id FROM accounts WHERE DATE(day) >= DATE_ADD('day', -7, CURRENT_DATE) AND id = 123",
+            fix="SELECT id FROM accounts WHERE (DATE(day) >= DATE_ADD('DAY', -7, CURRENT_DATE)) AND id = 123",
         )
 
 

{sql_data_guard-0.1.6 → sql_data_guard-0.1.8}/test/test_sql_guard_validation_unit.py

@@ -1,4 +1,5 @@
 import pytest
+
 from sql_data_guard.restriction_validation import (
     validate_restrictions,
     UnsupportedRestrictionError,