sqlink/tests/test_comments.py at main · JSLEEKR/sqlink · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
"""Tests for SQL comments and query labeling."""

import pytest
from sqlink import Query, F


class TestComment:
    def test_basic_comment(self):
        sql, _ = Query("users").select("*").comment("Fetch all users").build()
        assert "/* Fetch all users */" in sql

    def test_comment_position(self):
        sql, _ = Query("users").select("*").comment("Load users").build()
        assert sql.startswith("/* Load users */")

    def test_comment_with_where(self):
        sql, params = (
            Query("users")
            .select("*")
            .where(F("active") == True)
            .comment("Active users only")
            .build()
        )
        assert "/* Active users only */" in sql
        assert "WHERE" in sql

    def test_comment_with_insert(self):
        sql, _ = (
            Query("users")
            .insert("name")
            .values({"name": "John"})
            .comment("Create user")
            .build()
        )
        assert "/* Create user */" in sql
        assert "INSERT" in sql

    def test_comment_with_update(self):
        sql, _ = (
            Query("users")
            .update(name="New")
            .comment("Update name")
            .build()
        )
        assert "/* Update name */" in sql

    def test_comment_with_delete(self):
        sql, _ = (
            Query("users")
            .delete()
            .where(F("id") == 1)
            .comment("Remove user")
            .build()
        )
        assert "/* Remove user */" in sql


class TestLabel:
    def test_basic_label(self):
        sql, _ = Query("users").select("*").label("api:get-users").build()
        assert "/* api:get-users */" in sql

    def test_label_overrides_comment(self):
        sql, _ = (
            Query("users")
            .select("*")
            .comment("This comment")
            .label("this-label")
            .build()
        )
        assert "/* this-label */" in sql
        assert "This comment" not in sql

    def test_label_for_tracing(self):
        sql, _ = (
            Query("orders")
            .select("*")
            .where(F("user_id") == 42)
            .label("order-service:list-orders")
            .build()
        )
        assert "order-service:list-orders" in sql

    def test_label_preserved_in_clone(self):
        q = Query("users").select("*").label("base-query")
        q2 = q.clone().where(F("active") == True)
        sql, _ = q2.build()
        assert "/* base-query */" in sql

    def test_no_comment_by_default(self):
        sql, _ = Query("users").select("*").build()
        assert "/*" not in sql

    def test_comment_injection_sanitized(self):
        """Prevent SQL injection via comment-closing sequences."""
        sql, _ = Query("users").select("*").comment("evil */ DROP TABLE users; /*").build()
        # The closing */ should only appear once (at the end of the comment)
        assert sql.count("*/") == 1
        # The opening /* should only appear once (at the start of the comment)
        assert sql.count("/*") == 1

    def test_label_injection_sanitized(self):
        """Prevent SQL injection via label with comment-closing sequences."""
        sql, _ = Query("users").select("*").label("evil */ DROP TABLE users; /*").build()
        assert sql.count("*/") == 1
        assert sql.count("/*") == 1