Skip to content

Commit 114c7bd

Browse files
authored
Update model primary key for dynamic retrieval (#23)
* 删除、更新、查询支持主键名称非ID,从sqlalchemy 模型定义中获取主键 * custom error and update Document * Update docs for primary key * fix lint
1 parent 3a390c0 commit 114c7bd

File tree

6 files changed

+49
-10
lines changed

6 files changed

+49
-10
lines changed

docs/advanced/primary_key.md

+15-1
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,22 @@
11
!!! note 主键参数命名
22

33
由于在 python 内部 id 的特殊性,我们设定 pk (参考 Django) 作为模型主键命名,所以在 crud 方法中,任何涉及到主键的地方,入参都为 `pk`
4-
4+
55
```py title="e.g." hl_lines="2"
66
async def delete(self, db: AsyncSession, primary_key: int) -> int:
77
return self.delete_model(db, pk=primary_key)
88
```
9+
10+
## 主键定义
11+
12+
!!! warning 自动主键
13+
14+
我们在 SQLAlchemy CRUD Plus 内部通过 [inspect()](https://docs.sqlalchemy.org/en/20/core/inspection.html) 自动搜索表主键,
15+
而非强制绑定主键列必须命名为 id,感谢 [@DavidSche](https://github.com/DavidSche) 提供帮助
16+
17+
```py title="e.g." hl_lines="4"
18+
class ModelIns(Base):
19+
# your sqlalchemy model
20+
# define your primary_key
21+
custom_id: Mapped[int] = mapped_column(primary_key=True, index=True, autoincrement=True)
22+
```

docs/usage/delete_model.md

+3-1
Original file line numberDiff line numberDiff line change
@@ -18,13 +18,15 @@ from pydantic import BaseModel
1818

1919
from sqlalchemy_crud_plus import CRUDPlus
2020

21+
from sqlalchemy import Mapped, mapped_column
2122
from sqlalchemy import DeclarativeBase as Base
2223
from sqlalchemy.ext.asyncio import AsyncSession
2324

2425

2526
class ModelIns(Base):
2627
# your sqlalchemy model
27-
pass
28+
# define your primary_key
29+
custom_id: Mapped[int] = mapped_column(primary_key=True, index=True, autoincrement=True)
2830

2931

3032
class CreateIns(BaseModel):

docs/usage/select_model.md

+4-2
Original file line numberDiff line numberDiff line change
@@ -15,13 +15,15 @@ from pydantic import BaseModel
1515

1616
from sqlalchemy_crud_plus import CRUDPlus
1717

18+
from sqlalchemy import Mapped, mapped_column
1819
from sqlalchemy import DeclarativeBase as Base
1920
from sqlalchemy.ext.asyncio import AsyncSession
2021

2122

2223
class ModelIns(Base):
2324
# your sqlalchemy model
24-
pass
25+
# define your primary_key
26+
custom_id: Mapped[int] = mapped_column(primary_key=True, index=True, autoincrement=True)
2527

2628

2729
class CreateIns(BaseModel):
@@ -30,6 +32,6 @@ class CreateIns(BaseModel):
3032

3133

3234
class CRUDIns(CRUDPlus[ModelIns]):
33-
async def create(self, db: AsyncSession, pk: int) -> ModelIns:
35+
async def select(self, db: AsyncSession, pk: int) -> ModelIns:
3436
return await self.select_model(db, pk)
3537
```

docs/usage/update_model.md

+3-1
Original file line numberDiff line numberDiff line change
@@ -19,13 +19,15 @@ from pydantic import BaseModel
1919

2020
from sqlalchemy_crud_plus import CRUDPlus
2121

22+
from sqlalchemy import Mapped, mapped_column
2223
from sqlalchemy import DeclarativeBase as Base
2324
from sqlalchemy.ext.asyncio import AsyncSession
2425

2526

2627
class ModelIns(Base):
2728
# your sqlalchemy model
28-
pass
29+
# define your primary_key
30+
custom_id: Mapped[int] = mapped_column(primary_key=True, index=True, autoincrement=True)
2931

3032

3133
class UpdateIns(BaseModel):

sqlalchemy_crud_plus/crud.py

+17-5
Original file line numberDiff line numberDiff line change
@@ -2,17 +2,29 @@
22
# -*- coding: utf-8 -*-
33
from typing import Any, Generic, Iterable, Sequence, Type
44

5-
from sqlalchemy import Row, RowMapping, Select, delete, select, update
5+
from sqlalchemy import Row, RowMapping, Select, delete, inspect, select, update
66
from sqlalchemy.ext.asyncio import AsyncSession
77

8-
from sqlalchemy_crud_plus.errors import MultipleResultsError
8+
from sqlalchemy_crud_plus.errors import CompositePrimaryKeysError, MultipleResultsError
99
from sqlalchemy_crud_plus.types import CreateSchema, Model, UpdateSchema
1010
from sqlalchemy_crud_plus.utils import apply_sorting, count, parse_filters
1111

1212

1313
class CRUDPlus(Generic[Model]):
1414
def __init__(self, model: Type[Model]):
1515
self.model = model
16+
self.primary_key = self._get_primary_key()
17+
18+
def _get_primary_key(self):
19+
"""
20+
Dynamically retrieve the primary key column(s) for the model.
21+
"""
22+
mapper = inspect(self.model)
23+
primary_key = mapper.primary_key
24+
if len(primary_key) == 1:
25+
return primary_key[0]
26+
else:
27+
raise CompositePrimaryKeysError('Composite primary keys are not supported')
1628

1729
async def create_model(
1830
self,
@@ -69,7 +81,7 @@ async def select_model(self, session: AsyncSession, pk: int) -> Model | None:
6981
:param pk: The database primary key value.
7082
:return:
7183
"""
72-
stmt = select(self.model).where(self.model.id == pk)
84+
stmt = select(self.model).where(self.primary_key == pk)
7385
query = await session.execute(stmt)
7486
return query.scalars().first()
7587

@@ -166,7 +178,7 @@ async def update_model(
166178
instance_data = obj
167179
else:
168180
instance_data = obj.model_dump(exclude_unset=True)
169-
stmt = update(self.model).where(self.model.id == pk).values(**instance_data)
181+
stmt = update(self.model).where(self.primary_key == pk).values(**instance_data)
170182
result = await session.execute(stmt)
171183
if commit:
172184
await session.commit()
@@ -218,7 +230,7 @@ async def delete_model(
218230
:param commit: If `True`, commits the transaction immediately. Default is `False`.
219231
:return:
220232
"""
221-
stmt = delete(self.model).where(self.model.id == pk)
233+
stmt = delete(self.model).where(self.primary_key == pk)
222234
result = await session.execute(stmt)
223235
if commit:
224236
await session.commit()

sqlalchemy_crud_plus/errors.py

+7
Original file line numberDiff line numberDiff line change
@@ -36,3 +36,10 @@ class MultipleResultsError(SQLAlchemyCRUDPlusException):
3636

3737
def __init__(self, msg: str) -> None:
3838
super().__init__(msg)
39+
40+
41+
class CompositePrimaryKeysError(SQLAlchemyCRUDPlusException):
42+
"""Error raised when a table have Composite primary keys."""
43+
44+
def __init__(self, msg: str) -> None:
45+
super().__init__(msg)

0 commit comments

Comments
 (0)