Refactored for code readability

datasette/write_sql.py

@@ -0,0 +1,255 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from .permissions import Resource
+from .resources import DatabaseResource, TableResource
+from .utils import named_parameters, sqlite3
+from .utils.asgi import Forbidden
+from .utils.sql_analysis import Operation, SQLAnalysis
+
+if TYPE_CHECKING:
+    from .app import Datasette
+
+
+class QueryWriteRejected(Exception):
+    def __init__(self, message: str):
+        self.message = message
+        super().__init__(message)
+
+
+@dataclass(frozen=True)
+class PermissionRequirement:
+    action: str
+    resource: Resource
+
+
+PermissionRequirements = tuple[PermissionRequirement, ...]
+
+
+class WriteSqlOperationDecision:
+    """What Datasette should do with one operation in user-supplied write SQL."""
+
+
+@dataclass(frozen=True)
+class IgnoreWriteSqlOperation(WriteSqlOperationDecision):
+    reason: str
+
+
+@dataclass(frozen=True)
+class RequireWriteSqlPermissions(WriteSqlOperationDecision):
+    permissions: PermissionRequirements
+
+
+@dataclass(frozen=True)
+class RejectWriteSqlOperation(WriteSqlOperationDecision):
+    message: str
+
+
+@dataclass(frozen=True)
+class UnsupportedWriteSqlOperation(WriteSqlOperationDecision):
+    message: str
+
+
+def row_mutation_requirements(database: str, table: str) -> PermissionRequirements:
+    resource = TableResource(database=database, table=table)
+    return tuple(
+        PermissionRequirement(action=action, resource=resource)
+        for action in ("insert-row", "update-row", "delete-row")
+    )
+
+
+def decision_for_write_sql_operation(
+    operation: Operation,
+) -> WriteSqlOperationDecision:
+    unsupported_message = (
+        f"Unsupported SQL operation: {operation.operation} {operation.target_type}"
+    )
+    if operation.internal:
+        return IgnoreWriteSqlOperation("internal SQLite operation")
+    if operation.operation == "select":
+        return IgnoreWriteSqlOperation("select statement")
+    if operation.operation == "vacuum":
+        return RejectWriteSqlOperation("VACUUM is not allowed in user-supplied SQL")
+    if operation.operation in {"insert", "update", "delete"}:
+        if operation.table_kind == "virtual":
+            return RejectWriteSqlOperation(
+                "Writes to virtual tables are not allowed in user-supplied SQL"
+            )
+        if operation.table_kind == "shadow":
+            return RejectWriteSqlOperation(
+                "Writes to shadow tables are not allowed in user-supplied SQL"
+            )
+    if operation.operation == "function":
+        # SQL functions currently have no Datasette permission mapping. They are
+        # rejected by the user-supplied write SQL allow-list as unsupported.
+        return UnsupportedWriteSqlOperation(unsupported_message)
+    if (
+        operation.operation == "read"
+        and operation.target_type == "table"
+        and operation.database is not None
+        and operation.table is not None
+    ):
+        return RequireWriteSqlPermissions(
+            (
+                PermissionRequirement(
+                    action="view-table",
+                    resource=TableResource(
+                        database=operation.database, table=operation.table
+                    ),
+                ),
+            )
+        )
+    if (
+        operation.operation in {"insert", "update"}
+        and operation.target_type == "table"
+        and operation.database is not None
+        and operation.table is not None
+    ):
+        return RequireWriteSqlPermissions(
+            row_mutation_requirements(
+                database=operation.database,
+                table=operation.table,
+            )
+        )
+    if (
+        operation.operation == "delete"
+        and operation.target_type == "table"
+        and operation.database is not None
+        and operation.table is not None
+    ):
+        return RequireWriteSqlPermissions(
+            (
+                PermissionRequirement(
+                    action="delete-row",
+                    resource=TableResource(
+                        database=operation.database, table=operation.table
+                    ),
+                ),
+            )
+        )
+    if operation.operation == "create" and operation.target_type == "table":
+        if operation.database is None:
+            return UnsupportedWriteSqlOperation(unsupported_message)
+        return RequireWriteSqlPermissions(
+            (
+                PermissionRequirement(
+                    action="create-table",
+                    resource=DatabaseResource(database=operation.database),
+                ),
+            )
+        )
+    if (
+        operation.operation == "alter"
+        and operation.target_type == "table"
+        and operation.database is not None
+        and operation.table is not None
+    ):
+        return RequireWriteSqlPermissions(
+            (
+                PermissionRequirement(
+                    action="alter-table",
+                    resource=TableResource(
+                        database=operation.database, table=operation.table
+                    ),
+                ),
+            )
+        )
+    if (
+        operation.operation == "drop"
+        and operation.target_type == "table"
+        and operation.database is not None
+        and operation.table is not None
+    ):
+        return RequireWriteSqlPermissions(
+            (
+                PermissionRequirement(
+                    action="drop-table",
+                    resource=TableResource(
+                        database=operation.database, table=operation.table
+                    ),
+                ),
+            )
+        )
+    if (
+        operation.operation in {"create", "drop"}
+        and operation.target_type == "index"
+        and operation.database is not None
+        and operation.table is not None
+    ):
+        return RequireWriteSqlPermissions(
+            (
+                PermissionRequirement(
+                    action="alter-table",
+                    resource=TableResource(
+                        database=operation.database, table=operation.table
+                    ),
+                ),
+            )
+        )
+    return UnsupportedWriteSqlOperation(unsupported_message)
+
+
+def operation_is_write(operation: Operation) -> bool:
+    return operation.operation in {
+        "insert",
+        "update",
+        "delete",
+        "create",
+        "alter",
+        "drop",
+        "begin",
+        "commit",
+        "rollback",
+        "savepoint",
+        "attach",
+        "detach",
+        "pragma",
+        "analyze",
+        "reindex",
+        "vacuum",
+        "unknown",
+    }
+
+
+async def ensure_query_write_permissions(
+    datasette: Datasette,
+    database: str,
+    sql: str,
+    *,
+    actor: dict[str, object] | None = None,
+    params: dict[str, object] | None = None,
+    analysis: SQLAnalysis | None = None,
+) -> SQLAnalysis:
+    db = datasette.get_database(database)
+    if analysis is None:
+        if params is None:
+            params = {name: "" for name in named_parameters(sql)}
+        try:
+            analysis = await db.analyze_sql(sql, params)
+        except sqlite3.DatabaseError as ex:
+            raise Forbidden(f"Could not analyze query: {ex}") from ex
+
+    for operation in analysis.operations:
+        decision = decision_for_write_sql_operation(operation)
+        if isinstance(decision, IgnoreWriteSqlOperation):
+            continue
+        if isinstance(decision, RejectWriteSqlOperation):
+            raise QueryWriteRejected(decision.message)
+        if isinstance(decision, UnsupportedWriteSqlOperation):
+            raise Forbidden(decision.message)
+        permissions = decision.permissions
+        if operation.database != database:
+            raise Forbidden("Writable queries may not access attached databases")
+        for permission in permissions:
+            if not await datasette.allowed(
+                action=permission.action,
+                resource=permission.resource,
+                actor=actor,
+            ):
+                raise Forbidden(
+                    f"Permission denied: need {permission.action} "
+                    f"on {permission.resource}"
+                )
+    return analysis