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
|
from __future__ import annotations
from pathlib import Path
import pytest
from msgspec import Struct
from pydantic import BaseModel
from sqlalchemy import create_engine, select
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
from sqlalchemy.orm import DeclarativeBase, Mapped, Session
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,
]
here = Path(__file__).parent
fixture_path = here.parent.parent / "examples"
state_registry = base.create_registry()
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]
name: Mapped[str]
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
def test_sync_fixture_and_query() -> None:
engine = create_engine("sqlite://")
state_registry.metadata.create_all(engine)
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
async def test_async_fixture_and_query() -> None:
engine = create_async_engine("sqlite+aiosqlite://")
async with engine.begin() as conn:
await conn.run_sync(state_registry.metadata.create_all)
async with AsyncSession(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
|
