sqlsaber 0.27.0__py3-none-any.whl → 0.28.0__py3-none-any.whl

sqlsaber/tools/sql_guard.py

@@ -0,0 +1,225 @@
+"""SQL query validation and security using sqlglot AST analysis."""
+
+from dataclasses import dataclass
+from typing import Optional
+
+import sqlglot
+from sqlglot import exp
+from sqlglot.errors import ParseError
+
+# Prohibited AST node types that indicate write/mutation operations
+# Only include expression types that exist in sqlglot
+PROHIBITED_NODES = {
+    # DML operations
+    exp.Insert,
+    exp.Update,
+    exp.Delete,
+    exp.Merge,
+    # DDL operations
+    exp.Create,
+    exp.Drop,
+    exp.Alter,
+    exp.TruncateTable,
+    exp.AlterRename,
+    # MySQL specific
+    exp.Replace,
+    # Transaction control
+    exp.Transaction,
+    # Analysis and maintenance
+    exp.Analyze,
+    # Data loading/copying
+    exp.Copy,
+    exp.LoadData,
+    # Session and configuration
+    exp.Set,
+    exp.Use,
+    exp.Pragma,
+    # Security
+    exp.Grant,
+    exp.Revoke,
+    # Database operations
+    exp.Attach,
+    exp.Detach,
+    # Locking and process control
+    exp.Lock,
+    exp.Kill,
+    # Commands
+    exp.Command,
+}
+
+try:
+    # Add optional types that may not exist in all sqlglot versions
+    PROHIBITED_NODES.add(exp.Vacuum)
+except AttributeError:
+    pass
+
+# Dangerous functions by dialect that can read files or execute commands
+DANGEROUS_FUNCTIONS_BY_DIALECT = {
+    "postgres": {
+        "pg_read_file",
+        "pg_read_binary_file",
+        "pg_ls_dir",
+        "pg_stat_file",
+        "pg_logdir_ls",
+        "dblink",
+        "dblink_exec",
+    },
+    "mysql": {
+        "load_file",
+        "sys_eval",
+        "sys_exec",
+    },
+    "sqlite": {
+        "readfile",
+        "writefile",
+    },
+    "tsql": {
+        "xp_cmdshell",
+    },
+}
+
+
+@dataclass
+class GuardResult:
+    """Result of SQL query validation."""
+
+    allowed: bool
+    reason: Optional[str] = None
+    is_select: bool = False
+
+
+def is_select_like(stmt: exp.Expression) -> bool:
+    """Check if statement is a SELECT-like query.
+
+    Handles CTEs (WITH) and set operations (UNION/INTERSECT/EXCEPT).
+    """
+    root = stmt
+    # WITH wraps another statement
+    if isinstance(root, exp.With):
+        root = root.this
+    return isinstance(root, (exp.Select, exp.Union, exp.Except, exp.Intersect))
+
+
+def has_prohibited_nodes(stmt: exp.Expression) -> Optional[str]:
+    """Walk AST to find any prohibited operations.
+
+    Checks for:
+    - Write operations (INSERT/UPDATE/DELETE/etc)
+    - DDL operations (CREATE/DROP/ALTER/etc)
+    - SELECT INTO
+    - Locking clauses (FOR UPDATE/FOR SHARE)
+    """
+    for node in stmt.walk():
+        # Check prohibited node types
+        if isinstance(node, tuple(PROHIBITED_NODES)):
+            return f"Prohibited operation: {type(node).__name__}"
+
+        # Block SELECT INTO (Postgres-style table creation)
+        if isinstance(node, exp.Select) and node.args.get("into"):
+            return "SELECT INTO is not allowed"
+
+        # Block locking clauses (FOR UPDATE/FOR SHARE)
+        if isinstance(node, exp.Select):
+            locks = node.args.get("locks")
+            if locks:
+                return "SELECT with locking clause (FOR UPDATE/SHARE) is not allowed"
+
+    return None
+
+
+def has_dangerous_functions(stmt: exp.Expression, dialect: str) -> Optional[str]:
+    """Check for dangerous functions that can read files or execute commands."""
+    deny_set = DANGEROUS_FUNCTIONS_BY_DIALECT.get(dialect, set())
+    if not deny_set:
+        return None
+
+    deny_lower = {f.lower() for f in deny_set}
+
+    for fn in stmt.find_all(exp.Func):
+        name = fn.name
+        if name and name.lower() in deny_lower:
+            return f"Use of dangerous function '{name}' is not allowed"
+
+    return None
+
+
+def validate_read_only(sql: str, dialect: str = "ansi") -> GuardResult:
+    """Validate that SQL query is read-only using AST analysis.
+
+    Args:
+        sql: SQL query to validate
+        dialect: SQL dialect (postgres, mysql, sqlite, tsql, etc.)
+
+    Returns:
+        GuardResult with validation outcome
+    """
+    try:
+        statements = sqlglot.parse(sql, read=dialect)
+    except ParseError as e:
+        return GuardResult(False, f"Unable to parse query safely: {e}")
+    except Exception as e:
+        return GuardResult(False, f"Error parsing query: {e}")
+
+    # Only allow single statements
+    if len(statements) != 1:
+        return GuardResult(
+            False,
+            f"Only single SELECT statements are allowed (got {len(statements)} statements)",
+        )
+
+    stmt = statements[0]
+
+    # Must be a SELECT-like statement
+    if not is_select_like(stmt):
+        return GuardResult(False, "Only SELECT-like statements are allowed")
+
+    # Check for prohibited operations in the AST
+    reason = has_prohibited_nodes(stmt)
+    if reason:
+        return GuardResult(False, reason)
+
+    # Check for dangerous functions
+    reason = has_dangerous_functions(stmt, dialect)
+    if reason:
+        return GuardResult(False, reason)
+
+    return GuardResult(True, None, is_select=True)
+
+
+def add_limit(sql: str, dialect: str = "ansi", limit: int = 100) -> str:
+    """Add LIMIT clause to query if not already present.
+
+    Args:
+        sql: SQL query
+        dialect: SQL dialect for proper rendering
+        limit: Maximum number of rows to return
+
+    Returns:
+        SQL with LIMIT clause added (or original if LIMIT already exists)
+    """
+    try:
+        statements = sqlglot.parse(sql, read=dialect)
+        if len(statements) != 1:
+            return sql
+
+        stmt = statements[0]
+
+        # Check if LIMIT/TOP/FETCH already exists
+        has_limit = any(
+            isinstance(n, (exp.Limit, exp.Top, exp.Fetch)) for n in stmt.walk()
+        )
+        if has_limit:
+            return stmt.sql(dialect=dialect)
+
+        # Add LIMIT - sqlglot will render appropriately for dialect
+        # (LIMIT for most, TOP for SQL Server, FETCH FIRST for Oracle)
+        stmt = stmt.limit(limit)
+        return stmt.sql(dialect=dialect)
+
+    except Exception:
+        # If parsing/transformation fails, fall back to simple string append
+        # This maintains backward compatibility
+        sql_upper = sql.strip().upper()
+        if "LIMIT" not in sql_upper:
+            return f"{sql.rstrip(';')} LIMIT {limit};"
+        return sql

sqlsaber/tools/sql_tools.py

@@ -9,6 +9,7 @@ from sqlsaber.database.schema import SchemaManager
 from .base import Tool
 from .enums import ToolCategory, WorkflowPosition
 from .registry import register_tool
+from .sql_guard import add_limit, validate_read_only
 
 
 class SQLTool(Tool):
@@ -207,13 +208,17 @@ class ExecuteSQLTool(SQLTool):
         limit = kwargs.get("limit", self.DEFAULT_LIMIT)
 
         try:
-            # Security check - only allow SELECT queries unless write is enabled
-            write_error = self._validate_write_operation(query)
-            if write_error:
-                return json.dumps({"error": write_error})
+            # Get the dialect for this database
+            dialect = self.db.sqlglot_dialect
+
+            # Security check using sqlglot AST analysis
+            validation_result = validate_read_only(query, dialect)
+            if not validation_result.allowed:
+                return json.dumps({"error": validation_result.reason})
 
             # Add LIMIT if not present and it's a SELECT query
-            query = self._add_limit_to_query(query, limit)
+            if validation_result.is_select and limit:
+                query = add_limit(query, dialect, limit)
 
             # Execute the query
             results = await self.db.execute_query(query)
@@ -249,33 +254,3 @@ class ExecuteSQLTool(SQLTool):
                 )
 
             return json.dumps({"error": error_msg, "suggestions": suggestions})
-
-    def _validate_write_operation(self, query: str) -> str | None:
-        """Validate if a write operation is allowed."""
-        query_upper = query.strip().upper()
-
-        # Check for write operations
-        write_keywords = [
-            "INSERT",
-            "UPDATE",
-            "DELETE",
-            "DROP",
-            "CREATE",
-            "ALTER",
-            "TRUNCATE",
-        ]
-        is_write_query = any(query_upper.startswith(kw) for kw in write_keywords)
-
-        if is_write_query:
-            return (
-                "Write operations are not allowed. Only SELECT queries are permitted."
-            )
-
-        return None
-
-    def _add_limit_to_query(self, query: str, limit: int = 100) -> str:
-        """Add LIMIT clause to SELECT queries if not present."""
-        query_upper = query.strip().upper()
-        if query_upper.startswith("SELECT") and "LIMIT" not in query_upper:
-            return f"{query.rstrip(';')} LIMIT {limit};"
-        return query

{sqlsaber-0.27.0.dist-info → sqlsaber-0.28.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sqlsaber
-Version: 0.27.0
+Version: 0.28.0
 Summary: SQLsaber - Open-source agentic SQL assistant
 License-File: LICENSE
 Requires-Python: >=3.12
@@ -17,6 +17,7 @@ Requires-Dist: prompt-toolkit>3.0.51
 Requires-Dist: pydantic-ai
 Requires-Dist: questionary>=2.1.0
 Requires-Dist: rich>=13.7.0
+Requires-Dist: sqlglot[rs]>=27.20.0
 Description-Content-Type: text/markdown
 
 # SQLsaber

{sqlsaber-0.27.0.dist-info → sqlsaber-0.28.0.dist-info}/RECORD

@@ -5,44 +5,46 @@ sqlsaber/agents/base.py,sha256=40-MKEoz5rGrqVIylV1U2DaAUSPFcC75ohRin4E3-kk,2668
 sqlsaber/agents/mcp.py,sha256=Pn8tdDRUEVLYQyEi5nHRp9MKNePwHVVoeNI-uqWcr0Y,757
 sqlsaber/agents/pydantic_ai_agent.py,sha256=wBxKz0pjOkL-HI-TXV6B67bczZNgu7k26Rr3w5usR3o,10064
 sqlsaber/application/__init__.py,sha256=KY_-d5nEdQyAwNOsK5r-f7Tb69c63XbuEkHPeLpJal8,84
-sqlsaber/application/auth_setup.py,sha256=aIZ4vyfdjO5zCcMy4k4FmqjeGgruypPIRacaOJDuVRY,5080
-sqlsaber/application/db_setup.py,sha256=VipoXACDVs8auYs7XeuNYrUWjVHpxqxWphx9yUl1_2I,6826
-sqlsaber/application/model_selection.py,sha256=eRWga3vRpoyyTmw0BsVDRvUtyUnxvP8plAQ2Nv7tN3o,3163
+sqlsaber/application/auth_setup.py,sha256=cbS2VTwL7mwm24UFRz84PhC13XMzC1F7JkN-ze7tApY,5104
+sqlsaber/application/db_setup.py,sha256=qtMxCd_KO7GsD4W_iRBpDRvLriiyvOXPvZdcvm6KVDM,6849
+sqlsaber/application/model_selection.py,sha256=xZI-nvUgYZikaTK38SCmEWvWSfRsDpFu2VthbVdI95g,3187
 sqlsaber/application/prompts.py,sha256=4rMGcWpYJbNWPMzqVWseUMx0nwvXOkWS6GaTAJ5mhfc,3473
 sqlsaber/cli/__init__.py,sha256=qVSLVJLLJYzoC6aj6y9MFrzZvAwc4_OgxU9DlkQnZ4M,86
-sqlsaber/cli/auth.py,sha256=Hz7nuHdX1iMr2UR1VVtu4EAvBYYIhsgthluqpiOXH2A,6064
-sqlsaber/cli/commands.py,sha256=R63HF4P0JZwbGlI3TimkEUg3OUaFXYkVvLw1VIy0UFI,8516
+sqlsaber/cli/auth.py,sha256=ysDBXEFR8Jz7wYbIP6X7yWA2ivd8SDnUp_jUg_qYNWk,6088
+sqlsaber/cli/commands.py,sha256=-rTxr-kW7j2rR8wAg0tATKoh284pMDPKVMpQKaJwtqk,8540
 sqlsaber/cli/completers.py,sha256=g-hLDq5fiBx7gg8Bte1Lq8GU-ZxCYVs4dcPsmHPIcK4,6574
-sqlsaber/cli/database.py,sha256=6SdxNuP8w-e2nbV937q4hxfAIwKu1MqoirogD7USaMU,10639
-sqlsaber/cli/display.py,sha256=JQBSdLC_a3nkmJ4QfLMhhMPCYY-otS2eXF1XrfEe19E,17045
-sqlsaber/cli/interactive.py,sha256=RG4DJN9lprfKi60eomxVt-w3gTsWoqbfEcfO36iibXM,13694
-sqlsaber/cli/memory.py,sha256=OufHFJFwV0_GGn7LvKRTJikkWhV1IwNIUDOxFPHXOaQ,7794
-sqlsaber/cli/models.py,sha256=v_wrJWV5NXiM3zGLSMyNJf4YhgozK6hnf76Ua8HvysE,8480
-sqlsaber/cli/onboarding.py,sha256=5p-rHG2eGM-7t9yg_kNDrXpY13jr__dlFII3K69VtKc,11414
-sqlsaber/cli/streaming.py,sha256=YViLCxUv-7WN5TCphLYtAR02HXvuHYuPttGGDZKDUKU,6921
-sqlsaber/cli/threads.py,sha256=zVlbOuD3GjjEVNebXwANKeKt4I_Lunf6itiBUL0TaKA,12877
+sqlsaber/cli/database.py,sha256=hh8PdWnhaD0fO2jwvSSQyxsjwk-JyvmcY7f5tuHfnAQ,10663
+sqlsaber/cli/display.py,sha256=WB5JCumhXadziDEX1EZHG3vN1Chol5FNAaTXHieqFK0,17892
+sqlsaber/cli/interactive.py,sha256=u1gvdMZJkFWelZQ2urOS6-EH1uShKF4u_edfr_BzDNk,13479
+sqlsaber/cli/memory.py,sha256=gKP-JJ0w1ya1YTM_Lk7Gw-7wL9ptyj6cZtg-uoW8K7A,7818
+sqlsaber/cli/models.py,sha256=NozZbnisSjbPKo7PW7CltJMIkGcPqTDpDQEY-C_eLhk,8504
+sqlsaber/cli/onboarding.py,sha256=l6FFWn8J1OVQUxr-xIAzKaFhAz8rFh6IEWwIyPWqR6U,11438
+sqlsaber/cli/streaming.py,sha256=1XoZGPPMoTmBQVgp_Bqk483MR93j9oXxSV6Tx_-TpOg,6923
+sqlsaber/cli/threads.py,sha256=5EV4ckRzKqhWeTKpTfQSNCBuqs3onsJURKT09g4E4XM,13634
 sqlsaber/config/__init__.py,sha256=olwC45k8Nc61yK0WmPUk7XHdbsZH9HuUAbwnmKe3IgA,100
-sqlsaber/config/api_keys.py,sha256=7X63B-ow66GZWzYYgc_7f1T26F7gRhfyYQM4RpdkCpI,3647
+sqlsaber/config/api_keys.py,sha256=bjogRmIuxNNGusyKXKi0ZpJWeS5Fyn53zrAD8hsoYx4,3671
 sqlsaber/config/auth.py,sha256=b5qB2h1doXyO9Bn8z0CcL8LAR2jF431gGXBGKLgTmtQ,2756
 sqlsaber/config/database.py,sha256=Yec6_0wdzq-ADblMNnbgvouYCimYOY_DWHT9oweaISc,11449
-sqlsaber/config/oauth_flow.py,sha256=A3bSXaBLzuAfXV2ZPA94m9NV33c2MyL6M4ii9oEkswQ,10291
-sqlsaber/config/oauth_tokens.py,sha256=C9z35hyx-PvSAYdC1LNf3rg9_wsEIY56hkEczelbad0,6015
+sqlsaber/config/oauth_flow.py,sha256=VNbq4TZPk0hVJIcOh7JUO5qSxJnNzqDj0vwjCn-59Ok,10316
+sqlsaber/config/oauth_tokens.py,sha256=SC-lXVcKCV7uiWtBiU2mxvx1z7ryW8tOSKHBApPsXtE,5931
 sqlsaber/config/providers.py,sha256=JFjeJv1K5Q93zWSlWq3hAvgch1TlgoF0qFa0KJROkKY,2957
 sqlsaber/config/settings.py,sha256=iB4CnGQ4hw8gxkaa9CVLB_JEy6Y9h9FQTAams5OCVyI,6421
 sqlsaber/database/__init__.py,sha256=Gi9N_NOkD459WRWXDg3hSuGoBs3xWbMDRBvsTVmnGAg,2025
-sqlsaber/database/base.py,sha256=yxYcfeNhRPbO5jFRVZH7eRUGj_up-y3p1ZX_obZXi0w,3552
-sqlsaber/database/csv.py,sha256=45eH9mAkBtwSu1Rc_vvG1Z40L4xvfHWSb8OMG15TbCA,4340
-sqlsaber/database/duckdb.py,sha256=v6gFUhih5NMbHHpUv7By2nXyl9aqdPtLt0zhqS4-OKE,11120
-sqlsaber/database/mysql.py,sha256=5qd9gnSCP3umtBJcQDTzzJfMzwqYCJhWlbOeJZ9_-6c,14349
-sqlsaber/database/postgresql.py,sha256=R8I3Y-w0P9qPe47-lmae0X17syIwI8saxEG3etx6Rqc,13097
+sqlsaber/database/base.py,sha256=oaipLxlvoylX6oJCITPAWWqRqv09hRELqqEBufsmFic,3703
+sqlsaber/database/csv.py,sha256=41wuP40FaGPfj28HMiD0I69uG0JbUxArpoTLC3MG2uc,4464
+sqlsaber/database/duckdb.py,sha256=8HNKdx208aFK_YtwGjLz6LTne0xEmNevD-f9dRWlrFg,11244
+sqlsaber/database/mysql.py,sha256=wMzDQqq4GFbfEdqXtv_sCb4Qbr9GSWqYAvOLeo5UryY,14472
+sqlsaber/database/postgresql.py,sha256=fuf2Wl29NKXvD3mqsR08PDleNQ1PG-fNvWSxT6HDh2M,13223
 sqlsaber/database/resolver.py,sha256=wSCcn__aCqwIfpt_LCjtW2Zgb8RpG5PlmwwZHli1q_U,3628
 sqlsaber/database/schema.py,sha256=68PrNcA-5eR9PZB3i-TUQw5_E7QatwiDU2wv9GgXgM4,6928
-sqlsaber/database/sqlite.py,sha256=zdNj5i4mLJK21sWgftAHDHVihRUWevn__tVF9_nnLfQ,9297
+sqlsaber/database/sqlite.py,sha256=iReEIiSpkhhS1VzITd79ZWqSL3fHMyfe3DRCDpM0DvE,9421
 sqlsaber/mcp/__init__.py,sha256=COdWq7wauPBp5Ew8tfZItFzbcLDSEkHBJSMhxzy8C9c,112
 sqlsaber/mcp/mcp.py,sha256=tpNPHpkaCre1Xjp7c4DHXbTKeuYpDQ8qhmJZvAyr7Vk,3939
 sqlsaber/memory/__init__.py,sha256=GiWkU6f6YYVV0EvvXDmFWe_CxarmDCql05t70MkTEWs,63
 sqlsaber/memory/manager.py,sha256=p3fybMVfH-E4ApT1ZRZUnQIWSk9dkfUPCyfkmA0HALs,2739
 sqlsaber/memory/storage.py,sha256=ne8szLlGj5NELheqLnI7zu21V8YS4rtpYGGC7tOmi-s,5745
+sqlsaber/theme/__init__.py,sha256=qCICX1Cg4B6yCbZ1UrerxglWxcqldRFVSRrSs73na_8,188
+sqlsaber/theme/manager.py,sha256=0DWuVXn7JoC8NvAl5FSqc61eagKFTx5YnoY8SoCTxGM,7236
 sqlsaber/threads/__init__.py,sha256=Hh3dIG1tuC8fXprREUpslCIgPYz8_6o7aRLx4yNeO48,139
 sqlsaber/threads/storage.py,sha256=rsUdxT4CR52D7xtGir9UlsFnBMk11jZeflzDrk2q4ME,11183
 sqlsaber/tools/__init__.py,sha256=x3YdmX_7P0Qq_HtZHAgfIVKTLxYqKk6oc4tGsujQWsc,586
@@ -50,9 +52,10 @@ sqlsaber/tools/base.py,sha256=mHhvAj27BHmckyvuDLCPlAQdzABJyYxd9SJnaYAwwuA,1777
 sqlsaber/tools/enums.py,sha256=CH32mL-0k9ZA18911xLpNtsgpV6tB85TktMj6uqGz54,411
 sqlsaber/tools/instructions.py,sha256=X-x8maVkkyi16b6Tl0hcAFgjiYceZaSwyWTfmrvx8U8,9024
 sqlsaber/tools/registry.py,sha256=HWOQMsNIdL4XZS6TeNUyrL-5KoSDH6PHsWd3X66o-18,3211
-sqlsaber/tools/sql_tools.py,sha256=2xLD_pkd0t8wKndQAKIr4c9UpWzVWeHbAFpkwo5j4kY,9954
-sqlsaber-0.27.0.dist-info/METADATA,sha256=-QcrEttuC3vwUmNdtjHzA_laIfTkcMGoghmSCvNn3nM,7138
-sqlsaber-0.27.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-sqlsaber-0.27.0.dist-info/entry_points.txt,sha256=qEbOB7OffXPFgyJc7qEIJlMEX5RN9xdzLmWZa91zCQQ,162
-sqlsaber-0.27.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-sqlsaber-0.27.0.dist-info/RECORD,,
+sqlsaber/tools/sql_guard.py,sha256=dTDwcZP-N4xPGzcr7MQtKUxKrlDzlc1irr9aH5a4wvk,6182
+sqlsaber/tools/sql_tools.py,sha256=ujmAcfLkNaBrb5LWEgWcINQEQSX0LRPX3VK5Dag1Sj4,9178
+sqlsaber-0.28.0.dist-info/METADATA,sha256=oY9Awl9jLkdeJLa1oeTPHPGH94KktJ_HmBVUVLQs-do,7174
+sqlsaber-0.28.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+sqlsaber-0.28.0.dist-info/entry_points.txt,sha256=qEbOB7OffXPFgyJc7qEIJlMEX5RN9xdzLmWZa91zCQQ,162
+sqlsaber-0.28.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+sqlsaber-0.28.0.dist-info/RECORD,,