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
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
|
from __future__ import annotations
from pathlib import Path
import pytest
from msgspec import Struct
from pydantic import BaseModel
from sqlalchemy import Engine, String, select
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession
from sqlalchemy.orm import DeclarativeBase, Mapped, Session, mapped_column
from advanced_alchemy import base, mixins
from advanced_alchemy.repository import (
SQLAlchemyAsyncRepository,
SQLAlchemySyncRepository,
)
from advanced_alchemy.service import SQLAlchemyAsyncQueryService, SQLAlchemySyncQueryService
from advanced_alchemy.service._async import SQLAlchemyAsyncRepositoryService
from advanced_alchemy.service._sync import SQLAlchemySyncRepositoryService
from advanced_alchemy.service.typing import (
is_msgspec_struct,
is_msgspec_struct_with_field,
is_msgspec_struct_without_field,
is_pydantic_model,
is_pydantic_model_with_field,
is_pydantic_model_without_field,
)
from advanced_alchemy.utils.fixtures import open_fixture, open_fixture_async
pytestmark = [ # type: ignore
pytest.mark.integration,
pytest.mark.xdist_group("sqlquery_service"),
]
here = Path(__file__).parent
fixture_path = here.parent.parent / "examples"
state_registry = base.create_registry()
@pytest.fixture()
def sqlquery_test_tables(engine: Engine) -> None:
"""Create sqlquery test tables for sync engines."""
if getattr(engine.dialect, "name", "") != "mock":
state_registry.metadata.create_all(engine)
@pytest.fixture()
async def sqlquery_test_tables_async(async_engine: AsyncEngine) -> None:
"""Create sqlquery test tables for async engines."""
if getattr(async_engine.dialect, "name", "") != "mock":
async with async_engine.begin() as conn:
await conn.run_sync(state_registry.metadata.create_all)
class UUIDBase(mixins.UUIDPrimaryKey, base.CommonTableAttributes, DeclarativeBase):
"""Base for all SQLAlchemy declarative models with UUID primary keys."""
registry = state_registry
class USState(UUIDBase):
__tablename__ = "us_state_lookup" # type: ignore[assignment]
abbreviation: Mapped[str] = mapped_column(String(5))
name: Mapped[str] = mapped_column(String(50))
class USStateStruct(Struct):
abbreviation: str
name: str
class USStateBaseModel(BaseModel):
abbreviation: str
name: str
class USStateSyncRepository(SQLAlchemySyncRepository[USState]):
"""US State repository."""
model_type = USState
class USStateSyncService(SQLAlchemySyncRepositoryService[USState, USStateSyncRepository]):
"""US State repository."""
repository_type = USStateSyncRepository
class USStateAsyncRepository(SQLAlchemyAsyncRepository[USState]):
"""US State repository."""
model_type = USState
class USStateAsyncService(SQLAlchemyAsyncRepositoryService[USState, USStateAsyncRepository]):
"""US State repository."""
repository_type = USStateAsyncRepository
class StateQuery(base.SQLQuery):
"""Nonsensical query to test custom SQL queries."""
__table__ = select( # type: ignore
USState.abbreviation.label("state_abbreviation"),
USState.name.label("state_name"),
).alias("state_lookup")
__mapper_args__ = {
"primary_key": [USState.abbreviation],
}
state_abbreviation: str
state_name: str
class StateQueryStruct(Struct):
state_abbreviation: str
state_name: str
class StateQueryBaseModel(BaseModel):
state_abbreviation: str
state_name: str
@pytest.mark.xdist_group("sqlquery")
def test_sync_fixture_and_query(engine: Engine, sqlquery_test_tables: None) -> None:
# Skip mock engines as they don't support proper query operations
if getattr(engine.dialect, "name", "") == "mock":
pytest.skip("Mock engines don't support proper query operations")
with Session(engine) as session:
state_service = USStateSyncService(session=session)
query_service = SQLAlchemySyncQueryService(session=session)
fixture = open_fixture(fixture_path, USStateSyncRepository.model_type.__tablename__) # type: ignore[has-type]
_add_objs = state_service.create_many(
data=[USStateStruct(**raw_obj) for raw_obj in fixture],
)
_ordered_objs = state_service.list(order_by=(USState.name, True))
assert _ordered_objs[0].name == "Wyoming"
_ordered_objs_2 = state_service.list_and_count(order_by=[(USState.name, True)])
assert _ordered_objs_2[0][0].name == "Wyoming"
query_count = query_service.repository.count(statement=select(StateQuery))
assert query_count > 0
list_query_objs, list_query_count = query_service.repository.list_and_count(
statement=select(StateQuery),
)
assert list_query_count >= 50
_paginated_objs = query_service.to_schema(
data=list_query_objs,
total=list_query_count,
)
_pydantic_paginated_objs = query_service.to_schema(
data=list_query_objs,
total=list_query_count,
schema_type=StateQueryBaseModel,
)
assert isinstance(_pydantic_paginated_objs.items[0], StateQueryBaseModel)
_msgspec_paginated_objs = query_service.to_schema(
data=list_query_objs,
total=list_query_count,
schema_type=StateQueryStruct,
)
assert isinstance(_msgspec_paginated_objs.items[0], StateQueryStruct)
_list_service_objs = query_service.repository.list(statement=select(StateQuery))
assert len(_list_service_objs) >= 50
_get_ones = query_service.repository.list(statement=select(StateQuery), state_name="Alabama")
assert len(_get_ones) == 1
_get_one = query_service.repository.get_one(statement=select(StateQuery), state_name="Alabama")
assert _get_one.state_name == "Alabama"
_get_one_or_none_1 = query_service.repository.get_one_or_none(
statement=select(StateQuery).where(StateQuery.state_name == "Texas"), # type: ignore
)
assert _get_one_or_none_1 is not None
assert _get_one_or_none_1.state_name == "Texas"
_obj = query_service.to_schema(
data=_get_one_or_none_1,
)
_pydantic_obj = query_service.to_schema(
data=_get_one_or_none_1,
schema_type=StateQueryBaseModel,
)
assert isinstance(_pydantic_obj, StateQueryBaseModel)
assert is_pydantic_model(_pydantic_obj)
assert is_pydantic_model_with_field(_pydantic_obj, "state_abbreviation")
assert not is_pydantic_model_without_field(_pydantic_obj, "state_abbreviation")
_msgspec_obj = query_service.to_schema(
data=_get_one_or_none_1,
schema_type=StateQueryStruct,
)
assert isinstance(_msgspec_obj, StateQueryStruct)
assert is_msgspec_struct(_msgspec_obj)
assert is_msgspec_struct_with_field(_msgspec_obj, "state_abbreviation")
assert not is_msgspec_struct_without_field(_msgspec_obj, "state_abbreviation")
_get_one_or_none = query_service.repository.get_one_or_none(
statement=select(StateQuery).filter_by(state_name="Nope"),
)
assert _get_one_or_none is None
@pytest.mark.xdist_group("sqlquery")
async def test_async_fixture_and_query(async_engine: AsyncEngine, sqlquery_test_tables_async: None) -> None:
# Skip mock engines as they don't support proper query operations
if getattr(async_engine.dialect, "name", "") == "mock":
pytest.skip("Mock engines don't support proper query operations")
async with AsyncSession(async_engine) as session:
state_service = USStateAsyncService(session=session)
query_service = SQLAlchemyAsyncQueryService(session=session)
fixture = await open_fixture_async(fixture_path, USStateSyncRepository.model_type.__tablename__)
_add_objs = await state_service.create_many(
data=[USStateBaseModel(**raw_obj) for raw_obj in fixture],
)
_ordered_objs = await state_service.list(order_by=(USState.name, True))
assert _ordered_objs[0].name == "Wyoming"
_ordered_objs_2 = await state_service.list_and_count(order_by=(USState.name, True))
assert _ordered_objs_2[0][0].name == "Wyoming"
query_count = await query_service.repository.count(statement=select(StateQuery))
assert query_count > 0
list_query_objs, list_query_count = await query_service.repository.list_and_count(
statement=select(StateQuery),
)
assert list_query_count >= 50
_paginated_objs = query_service.to_schema(
list_query_objs,
total=list_query_count,
)
_pydantic_paginated_objs = query_service.to_schema(
data=list_query_objs,
total=list_query_count,
schema_type=StateQueryBaseModel,
)
assert isinstance(_pydantic_paginated_objs.items[0], StateQueryBaseModel)
_msgspec_paginated_objs = query_service.to_schema(
data=list_query_objs,
total=list_query_count,
schema_type=StateQueryStruct,
)
assert isinstance(_msgspec_paginated_objs.items[0], StateQueryStruct)
_list_service_objs = await query_service.repository.list(statement=select(StateQuery))
assert len(_list_service_objs) >= 50
_get_ones = await query_service.repository.list(statement=select(StateQuery), state_name="Alabama")
assert len(_get_ones) == 1
_get_one = await query_service.repository.get_one(statement=select(StateQuery), state_name="Alabama")
assert _get_one.state_name == "Alabama"
_get_one_or_none_1 = await query_service.repository.get_one_or_none(
statement=select(StateQuery).where(StateQuery.state_name == "Texas"), # type: ignore
)
assert _get_one_or_none_1 is not None
assert _get_one_or_none_1.state_name == "Texas"
_obj = query_service.to_schema(
data=_get_one_or_none_1,
)
_pydantic_obj = query_service.to_schema(
data=_get_one_or_none_1,
schema_type=StateQueryBaseModel,
)
assert isinstance(_pydantic_obj, StateQueryBaseModel)
assert is_pydantic_model(_pydantic_obj)
assert is_pydantic_model_with_field(_pydantic_obj, "state_abbreviation")
assert not is_pydantic_model_without_field(_pydantic_obj, "state_abbreviation")
_msgspec_obj = query_service.to_schema(
data=_get_one_or_none_1,
schema_type=StateQueryStruct,
)
assert isinstance(_msgspec_obj, StateQueryStruct)
assert is_msgspec_struct(_msgspec_obj)
assert is_msgspec_struct_with_field(_msgspec_obj, "state_abbreviation")
_get_one_or_none = await query_service.repository.get_one_or_none(
select(StateQuery).filter_by(state_name="Nope")
)
assert not is_msgspec_struct_without_field(_msgspec_obj, "state_abbreviation")
assert _get_one_or_none is None
@pytest.mark.xdist_group("sqlquery")
async def test_async_query_repository_instantiation(async_engine: AsyncEngine) -> None:
"""Test that SQLAlchemyAsyncQueryRepository can be instantiated without super().__init__() error."""
from advanced_alchemy.repository import SQLAlchemyAsyncQueryRepository
async with AsyncSession(async_engine) as session:
# Test direct instantiation - this should not raise TypeError
repository = SQLAlchemyAsyncQueryRepository(session=session)
assert repository is not None
assert repository.session == session
assert repository.error_messages is None
assert repository.wrap_exceptions is True
# Test with optional parameters
repository_with_params = SQLAlchemyAsyncQueryRepository(
session=session, error_messages={"not_found": "Custom not found"}, wrap_exceptions=False
)
assert repository_with_params.session == session
assert repository_with_params.error_messages == {"not_found": "Custom not found"}
assert repository_with_params.wrap_exceptions is False
@pytest.mark.xdist_group("sqlquery")
def test_sync_query_repository_instantiation(engine: Engine) -> None:
"""Test that SQLAlchemySyncQueryRepository can be instantiated without super().__init__() error."""
from advanced_alchemy.repository import SQLAlchemySyncQueryRepository
with Session(engine) as session:
# Test direct instantiation - this should not raise TypeError
repository = SQLAlchemySyncQueryRepository(session=session)
assert repository is not None
assert repository.session == session
assert repository.error_messages is None
assert repository.wrap_exceptions is True
# Test with optional parameters
repository_with_params = SQLAlchemySyncQueryRepository(
session=session, error_messages={"not_found": "Custom not found"}, wrap_exceptions=False
)
assert repository_with_params.session == session
assert repository_with_params.error_messages == {"not_found": "Custom not found"}
assert repository_with_params.wrap_exceptions is False
@pytest.mark.xdist_group("sqlquery")
async def test_async_query_service_with_repository_instantiation(async_engine: AsyncEngine) -> None:
"""Test that SQLAlchemyAsyncQueryService using the repository works correctly."""
from advanced_alchemy.service import SQLAlchemyAsyncQueryService
async with AsyncSession(async_engine) as session:
# This should not raise TypeError when creating the repository internally
query_service = SQLAlchemyAsyncQueryService(session=session)
assert query_service is not None
assert query_service.repository is not None
assert query_service.repository.session == session
@pytest.mark.xdist_group("sqlquery")
def test_sync_query_service_with_repository_instantiation(engine: Engine) -> None:
"""Test that SQLAlchemySyncQueryService using the repository works correctly."""
from advanced_alchemy.service import SQLAlchemySyncQueryService
with Session(engine) as session:
# This should not raise TypeError when creating the repository internally
query_service = SQLAlchemySyncQueryService(session=session)
assert query_service is not None
assert query_service.repository is not None
assert query_service.repository.session == session
|
