5500946083

pending completion

Build # 5500946083

Build Type

push

github

Committed by

web-flow

Commit Message

Merge pull request #7 from amoffat/feature/generalize-sql

Feature/generalize sql

Run Details

299 of 335 branches covered (89.25%)

Branch coverage included in aggregate %.

444 of 444 new or added lines in 30 files covered. (100.0%)

1427 of 1494 relevant lines covered (95.52%)

0.96 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

99.15

/heimdallm/bifrosts/sql/sqlite/tests/select/test_select.py

from typing import Sequence

import pytest

from heimdallm.bifrosts.sql import exc
from heimdallm.bifrosts.sql.common import FqColumn, JoinCondition, RequiredConstraint
from heimdallm.bifrosts.sql.sqlite.select.bifrost import Bifrost

from .utils import PermissiveConstraints


def test_aliased_select_column():
    """tests that we can select a column from a table that was aliased via a
    join"""
    query = """
    select thing.col from t1
    join t2 thing on t1.jid = t2.jid
    """

    class MyConstraints(PermissiveConstraints):
        def select_column_allowed(self, column: FqColumn) -> bool:
            return column.name in {"t2.col"}

    bifrost = Bifrost.mocked(MyConstraints())
    bifrost.traverse(query)


def test_unqualified_columns():
    query = "select col from t1"

    bifrost = Bifrost.mocked(PermissiveConstraints())
    with pytest.raises(exc.UnqualifiedColumn) as e:
        bifrost.traverse(query)
    assert e.value.column == "col"


def test_unqualified_where_column():
    """a column alias in the where class is valid, as long as the alias points to an
    actual aliased column. if it doesn't, we consider it unqualified."""
    query = """
SELECT payment_id,customer_id,staff_id,rental_id,amount,payment_date
FROM payment
WHERE customer_id=:customer_id
ORDER BY payment_date DESC
LIMIT 5;
    """

    bifrost = Bifrost.mocked(PermissiveConstraints())
    with pytest.raises(exc.UnqualifiedColumn) as e:
        bifrost.traverse(query)
    assert e.value.column == "customer_id"


@pytest.mark.parametrize(
    "query",
    (
        'select "t1"."col" from t1',
        'select t1.col from "t1"',
        'select t1.col from t1 where t1."id"=1',
        'select t1.col from t1 join "t2" on t1.jid = "t2".jid',
    ),
)
def test_escapes(query):
    class MyConstraints(PermissiveConstraints):
        def select_column_allowed(self, column: FqColumn):
            return column.name in {"t1.col"}

        def condition_column_allowed(self, column: FqColumn) -> bool:
            return column.name in {"t1.id", "t2.jid", "t1.jid"}

    bifrost = Bifrost.mocked(MyConstraints())
    bifrost.traverse(query)


def test_disallowed_select_column():
    """tests that our constraint validator works when trying to select a
    column"""
    query = "select t1.col from t1"

    class AllowColumnConstraints(PermissiveConstraints):
        def select_column_allowed(self, fq_column: FqColumn) -> bool:
            return fq_column.name in {"t1.col"}

    class DenyColumnConstraints(PermissiveConstraints):
        def select_column_allowed(self, fq_column: FqColumn) -> bool:
            return fq_column.name in {"t2.col"}

    bifrost = Bifrost.mocked(AllowColumnConstraints())
    bifrost.traverse(query)

    bifrost = Bifrost.mocked(DenyColumnConstraints())
    with pytest.raises(exc.IllegalSelectedColumn) as excinfo:
        bifrost.traverse(query)

    e = excinfo.value
    assert e.column == "t1.col"


def test_required_constraint():
    """tests that our required constraint restricts the query correctly"""

    # missing constraint
    query = "select t1.col from t1"

    class RequiredConstraints(PermissiveConstraints):
        def required_constraints(self):
            return [
                RequiredConstraint(
                    column="t1.id",
                    placeholder="id",
                )
            ]

    bifrost = Bifrost.mocked(RequiredConstraints())
    with pytest.raises(exc.MissingRequiredConstraint) as excinfo:
        bifrost.traverse(query)

    e = excinfo.value
    assert e.column.name == "t1.id"
    assert e.placeholder == "id"

    # same query, but with constraint
    query = "select t1.col from t1 where t1.id=:id"
    bifrost.traverse(query)

    # same query, but with aliased table
    query = "select t1.col from t1 as aliased where aliased.id=:id"
    bifrost.traverse(query)

    # same query, without a fully qualified constraint, which would be
    # impossible for us to trace back to the original table
    query = "select t1.col from t1 as aliased where id=:id"
    with pytest.raises(exc.UnqualifiedColumn):
        bifrost.traverse(query)


@pytest.mark.parametrize(
    "query",
    (
        "select t1.col from",  # unexpected EOF
        "select `` fro t1",  # just broken
    ),
)
def test_broken_query(query):
    """tests that we raise an exception when we cannot parse the query"""

    bifrost = Bifrost.mocked(PermissiveConstraints())
    with pytest.raises(exc.InvalidQuery):
        bifrost.traverse(query)


def test_select_column_arith():
    query = "select t1.col + 1 as plus_one from t1"
    bifrost = Bifrost.mocked(PermissiveConstraints())
    bifrost.traverse(query)


@pytest.mark.parametrize(
    "query",
    (
        "select 1+1 as two from t1",
        "select t1.id, (1 + 1) as two from t1",
    ),
)
def test_select_expr(query):
    """select a non-column expression"""
    bifrost = Bifrost.mocked(PermissiveConstraints())
    bifrost.traverse(query)


def test_conflicting_validations():
    query = "select t1.col from t1"

    class MyConstraints(PermissiveConstraints):
        def select_column_allowed(self, column: FqColumn) -> bool:
            return column.name in {"t1.col"}

        def allowed_joins(self) -> Sequence[JoinCondition]:
            return []

    bifrost = Bifrost.mocked(MyConstraints())
    bifrost.traverse(query)


def test_count_disallowed_column():
    """we're allowed to count a disallowed column"""

    class GeneralConstraints(PermissiveConstraints):
        def select_column_allowed(self, fq_column: FqColumn) -> bool:
            return fq_column.name != "film_actor.film_id"

    bifrost = Bifrost.mocked(GeneralConstraints())

    # a query that uses the disallowed column in a count is allowed
    query = """
SELECT
    actor.actor_id,
    actor.first_name,
    actor.last_name,
    COUNT(film_actor.film_id)
FROM actor
JOIN film_actor ON actor.actor_id = film_actor.actor_id
GROUP BY actor.actor_id
    """

    bifrost.traverse(query, autofix=False)

    # prove that the column isn't stripped by the reconstructor
    fixed = bifrost.traverse(query, autofix=True)
    assert "film_actor.film_id" in fixed

    # but a query that uses the disallowed column without a count is not allowed
    query = """
SELECT
    actor.actor_id,
    actor.first_name,
    actor.last_name,
    film_actor.film_id
FROM actor
JOIN film_actor ON actor.actor_id = film_actor.actor_id
    """

    with pytest.raises(exc.IllegalSelectedColumn) as e:
        bifrost.traverse(query, autofix=False)
        assert e.value.column == "film_actor.film_id"

    fixed = bifrost.traverse(query, autofix=True)
    assert "film_actor.film_id" not in fixed

1	from typing import Sequence	1✔
2
3	import pytest	1✔
4
5	from heimdallm.bifrosts.sql import exc	1✔
6	from heimdallm.bifrosts.sql.common import FqColumn, JoinCondition, RequiredConstraint	1✔
7	from heimdallm.bifrosts.sql.sqlite.select.bifrost import Bifrost	1✔
8
9	from .utils import PermissiveConstraints	1✔
10
11
12	def test_aliased_select_column():	1✔
13	"""tests that we can select a column from a table that was aliased via a
14	join"""
15	query = """	1✔
16	select thing.col from t1
17	join t2 thing on t1.jid = t2.jid
18	"""
19
20	class MyConstraints(PermissiveConstraints):	1✔
21	def select_column_allowed(self, column: FqColumn) -> bool:	1✔
22	return column.name in {"t2.col"}	1✔
23
24	bifrost = Bifrost.mocked(MyConstraints())	1✔
25	bifrost.traverse(query)	1✔
26
27
28	def test_unqualified_columns():	1✔
29	query = "select col from t1"	1✔
30
31	bifrost = Bifrost.mocked(PermissiveConstraints())	1✔
32	with pytest.raises(exc.UnqualifiedColumn) as e:	1✔
33	bifrost.traverse(query)	1✔
34	assert e.value.column == "col"	1✔
35
36
37	def test_unqualified_where_column():	1✔
38	"""a column alias in the where class is valid, as long as the alias points to an
39	actual aliased column. if it doesn't, we consider it unqualified."""
40	query = """	1✔
41	SELECT payment_id,customer_id,staff_id,rental_id,amount,payment_date
42	FROM payment
43	WHERE customer_id=:customer_id
44	ORDER BY payment_date DESC
45	LIMIT 5;
46	"""
47
48	bifrost = Bifrost.mocked(PermissiveConstraints())	1✔
49	with pytest.raises(exc.UnqualifiedColumn) as e:	1✔
50	bifrost.traverse(query)	1✔
51	assert e.value.column == "customer_id"	1✔
52
53
54	@pytest.mark.parametrize(	1✔
55	"query",
56	(
57	'select "t1"."col" from t1',
58	'select t1.col from "t1"',
59	'select t1.col from t1 where t1."id"=1',
60	'select t1.col from t1 join "t2" on t1.jid = "t2".jid',
61	),
62	)
63	def test_escapes(query):	1✔
64	class MyConstraints(PermissiveConstraints):	1✔
65	def select_column_allowed(self, column: FqColumn):	1✔
66	return column.name in {"t1.col"}	1✔
67
68	def condition_column_allowed(self, column: FqColumn) -> bool:	1✔
69	return column.name in {"t1.id", "t2.jid", "t1.jid"}	1✔
70
71	bifrost = Bifrost.mocked(MyConstraints())	1✔
72	bifrost.traverse(query)	1✔
73
74
75	def test_disallowed_select_column():	1✔
76	"""tests that our constraint validator works when trying to select a
77	column"""
78	query = "select t1.col from t1"	1✔
79
80	class AllowColumnConstraints(PermissiveConstraints):	1✔
81	def select_column_allowed(self, fq_column: FqColumn) -> bool:	1✔
82	return fq_column.name in {"t1.col"}	1✔
83
84	class DenyColumnConstraints(PermissiveConstraints):	1✔
85	def select_column_allowed(self, fq_column: FqColumn) -> bool:	1✔
86	return fq_column.name in {"t2.col"}	1✔
87
88	bifrost = Bifrost.mocked(AllowColumnConstraints())	1✔
89	bifrost.traverse(query)	1✔
90
91	bifrost = Bifrost.mocked(DenyColumnConstraints())	1✔
92	with pytest.raises(exc.IllegalSelectedColumn) as excinfo:	1✔
93	bifrost.traverse(query)	1✔
94
95	e = excinfo.value	1✔
96	assert e.column == "t1.col"	1✔
97
98
99	def test_required_constraint():	1✔
100	"""tests that our required constraint restricts the query correctly"""
101
102	# missing constraint
103	query = "select t1.col from t1"	1✔
104
105	class RequiredConstraints(PermissiveConstraints):	1✔
106	def required_constraints(self):	1✔
107	return [	1✔
108	RequiredConstraint(
109	column="t1.id",
110	placeholder="id",
111	)
112	]
113
114	bifrost = Bifrost.mocked(RequiredConstraints())	1✔
115	with pytest.raises(exc.MissingRequiredConstraint) as excinfo:	1✔
116	bifrost.traverse(query)	1✔
117
118	e = excinfo.value	1✔
119	assert e.column.name == "t1.id"	1✔
120	assert e.placeholder == "id"	1✔
121
122	# same query, but with constraint
123	query = "select t1.col from t1 where t1.id=:id"	1✔
124	bifrost.traverse(query)	1✔
125
126	# same query, but with aliased table
127	query = "select t1.col from t1 as aliased where aliased.id=:id"	1✔
128	bifrost.traverse(query)	1✔
129
130	# same query, without a fully qualified constraint, which would be
131	# impossible for us to trace back to the original table
132	query = "select t1.col from t1 as aliased where id=:id"	1✔
133	with pytest.raises(exc.UnqualifiedColumn):	1✔
134	bifrost.traverse(query)	1✔
135
136
137	@pytest.mark.parametrize(	1✔
138	"query",
139	(
140	"select t1.col from", # unexpected EOF
141	"select `` fro t1", # just broken
142	),
143	)
144	def test_broken_query(query):	1✔
145	"""tests that we raise an exception when we cannot parse the query"""
146
147	bifrost = Bifrost.mocked(PermissiveConstraints())	1✔
148	with pytest.raises(exc.InvalidQuery):	1✔
149	bifrost.traverse(query)	1✔
150
151
152	def test_select_column_arith():	1✔
153	query = "select t1.col + 1 as plus_one from t1"	1✔
154	bifrost = Bifrost.mocked(PermissiveConstraints())	1✔
155	bifrost.traverse(query)	1✔
156
157
158	@pytest.mark.parametrize(	1✔
159	"query",
160	(
161	"select 1+1 as two from t1",
162	"select t1.id, (1 + 1) as two from t1",
163	),
164	)
165	def test_select_expr(query):	1✔
166	"""select a non-column expression"""
167	bifrost = Bifrost.mocked(PermissiveConstraints())	1✔
168	bifrost.traverse(query)	1✔
169
170
171	def test_conflicting_validations():	1✔
172	query = "select t1.col from t1"	1✔
173
174	class MyConstraints(PermissiveConstraints):	1✔
175	def select_column_allowed(self, column: FqColumn) -> bool:	1✔
176	return column.name in {"t1.col"}	1✔
177
178	def allowed_joins(self) -> Sequence[JoinCondition]:	1✔
179	return []	1✔
180
181	bifrost = Bifrost.mocked(MyConstraints())	1✔
182	bifrost.traverse(query)	1✔
183
184
185	def test_count_disallowed_column():	1✔
186	"""we're allowed to count a disallowed column"""
187
188	class GeneralConstraints(PermissiveConstraints):	1✔
189	def select_column_allowed(self, fq_column: FqColumn) -> bool:	1✔
190	return fq_column.name != "film_actor.film_id"	1✔
191
192	bifrost = Bifrost.mocked(GeneralConstraints())	1✔
193
194	# a query that uses the disallowed column in a count is allowed
195	query = """	1✔
196	SELECT
197	actor.actor_id,
198	actor.first_name,
199	actor.last_name,
200	COUNT(film_actor.film_id)
201	FROM actor
202	JOIN film_actor ON actor.actor_id = film_actor.actor_id
203	GROUP BY actor.actor_id
204	"""
205
206	bifrost.traverse(query, autofix=False)	1✔
207
208	# prove that the column isn't stripped by the reconstructor
209	fixed = bifrost.traverse(query, autofix=True)	1✔
210	assert "film_actor.film_id" in fixed	1✔
211
212	# but a query that uses the disallowed column without a count is not allowed
213	query = """	1✔
214	SELECT
215	actor.actor_id,
216	actor.first_name,
217	actor.last_name,
218	film_actor.film_id
219	FROM actor
220	JOIN film_actor ON actor.actor_id = film_actor.actor_id
221	"""
222
223	with pytest.raises(exc.IllegalSelectedColumn) as e:	1✔
224	bifrost.traverse(query, autofix=False)	1✔
225	assert e.value.column == "film_actor.film_id"	×
226
227	fixed = bifrost.traverse(query, autofix=True)	1✔
228	assert "film_actor.film_id" not in fixed	1✔

amoffat / HeimdaLLM / 5500946083

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous