Skip to content

Commit efdc469

Browse files
frankie567frankie567
committed
Implement SQLAlchemy 2.0 support
1 parent 1a01a90 commit efdc469

File tree

6 files changed

+87
-63
lines changed

6 files changed

+87
-63
lines changed

‎fastapi_users_db_sqlalchemy/__init__.py‎

Lines changed: 39 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,9 @@
44

55
from fastapi_users.db.base import BaseUserDatabase
66
from fastapi_users.models import ID, OAP, UP
7-
from sqlalchemy import Boolean, Column, ForeignKey, Integer, String, func, select
7+
from sqlalchemy import Boolean, ForeignKey, Integer, String, func, select
88
from sqlalchemy.ext.asyncio import AsyncSession
9-
from sqlalchemy.orm import declarative_mixin, declared_attr
9+
from sqlalchemy.orm import Mapped, declared_attr, mapped_column
1010
from sqlalchemy.sql import Select
1111

1212
from fastapi_users_db_sqlalchemy.generics import GUID
@@ -16,7 +16,6 @@
1616
UUID_ID = uuid.UUID
1717

1818

19-
@declarative_mixin
2019
class SQLAlchemyBaseUserTable(Generic[ID]):
2120
"""Base SQLAlchemy users table definition."""
2221

@@ -30,22 +29,28 @@ class SQLAlchemyBaseUserTable(Generic[ID]):
3029
is_superuser: bool
3130
is_verified: bool
3231
else:
33-
email: str = Column(String(length=320), unique=True, index=True, nullable=False)
34-
hashed_password: str = Column(String(length=1024), nullable=False)
35-
is_active: bool = Column(Boolean, default=True, nullable=False)
36-
is_superuser: bool = Column(Boolean, default=False, nullable=False)
37-
is_verified: bool = Column(Boolean, default=False, nullable=False)
32+
email: Mapped[str] = mapped_column(
33+
String(length=320), unique=True, index=True, nullable=False
34+
)
35+
hashed_password: Mapped[str] = mapped_column(
36+
String(length=1024), nullable=False
37+
)
38+
is_active: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)
39+
is_superuser: Mapped[bool] = mapped_column(
40+
Boolean, default=False, nullable=False
41+
)
42+
is_verified: Mapped[bool] = mapped_column(
43+
Boolean, default=False, nullable=False
44+
)
3845

3946

40-
@declarative_mixin
4147
class SQLAlchemyBaseUserTableUUID(SQLAlchemyBaseUserTable[UUID_ID]):
4248
if TYPE_CHECKING: # pragma: no cover
4349
id: UUID_ID
4450
else:
45-
id: UUID_ID = Column(GUID, primary_key=True, default=uuid.uuid4)
51+
id: Mapped[UUID_ID] = mapped_column(GUID, primary_key=True, default=uuid.uuid4)
4652

4753

48-
@declarative_mixin
4954
class SQLAlchemyBaseOAuthAccountTable(Generic[ID]):
5055
"""Base SQLAlchemy OAuth account table definition."""
5156

@@ -60,24 +65,32 @@ class SQLAlchemyBaseOAuthAccountTable(Generic[ID]):
6065
account_id: str
6166
account_email: str
6267
else:
63-
oauth_name: str = Column(String(length=100), index=True, nullable=False)
64-
access_token: str = Column(String(length=1024), nullable=False)
65-
expires_at: Optional[int] = Column(Integer, nullable=True)
66-
refresh_token: Optional[str] = Column(String(length=1024), nullable=True)
67-
account_id: str = Column(String(length=320), index=True, nullable=False)
68-
account_email: str = Column(String(length=320), nullable=False)
68+
oauth_name: Mapped[str] = mapped_column(
69+
String(length=100), index=True, nullable=False
70+
)
71+
access_token: Mapped[str] = mapped_column(String(length=1024), nullable=False)
72+
expires_at: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)
73+
refresh_token: Mapped[Optional[str]] = mapped_column(
74+
String(length=1024), nullable=True
75+
)
76+
account_id: Mapped[str] = mapped_column(
77+
String(length=320), index=True, nullable=False
78+
)
79+
account_email: Mapped[str] = mapped_column(String(length=320), nullable=False)
6980

7081

71-
@declarative_mixin
7282
class SQLAlchemyBaseOAuthAccountTableUUID(SQLAlchemyBaseOAuthAccountTable[UUID_ID]):
7383
if TYPE_CHECKING: # pragma: no cover
7484
id: UUID_ID
85+
user_id: UUID_ID
7586
else:
76-
id: UUID_ID = Column(GUID, primary_key=True, default=uuid.uuid4)
87+
id: Mapped[UUID_ID] = mapped_column(GUID, primary_key=True, default=uuid.uuid4)
7788

78-
@declared_attr
79-
def user_id(cls) -> Column[GUID]:
80-
return Column(GUID, ForeignKey("user.id", ondelete="cascade"), nullable=False)
89+
@declared_attr
90+
def user_id(cls) -> Mapped[GUID]:
91+
return mapped_column(
92+
GUID, ForeignKey("user.id", ondelete="cascade"), nullable=False
93+
)
8194

8295

8396
class SQLAlchemyUserDatabase(Generic[UP, ID], BaseUserDatabase[UP, ID]):
@@ -120,24 +133,22 @@ async def get_by_oauth_account(self, oauth: str, account_id: str) -> Optional[UP
120133
statement = (
121134
select(self.user_table)
122135
.join(self.oauth_account_table)
123-
.where(self.oauth_account_table.oauth_name == oauth)
124-
.where(self.oauth_account_table.account_id == account_id)
136+
.where(self.oauth_account_table.oauth_name == oauth) # type: ignore
137+
.where(self.oauth_account_table.account_id == account_id) # type: ignore
125138
)
126139
return await self._get_user(statement)
127140

128141
async def create(self, create_dict: Dict[str, Any]) -> UP:
129142
user = self.user_table(**create_dict)
130143
self.session.add(user)
131144
await self.session.commit()
132-
await self.session.refresh(user)
133145
return user
134146

135147
async def update(self, user: UP, update_dict: Dict[str, Any]) -> UP:
136148
for key, value in update_dict.items():
137149
setattr(user, key, value)
138150
self.session.add(user)
139151
await self.session.commit()
140-
await self.session.refresh(user)
141152
return user
142153

143154
async def delete(self, user: UP) -> None:
@@ -148,6 +159,7 @@ async def add_oauth_account(self, user: UP, create_dict: Dict[str, Any]) -> UP:
148159
if self.oauth_account_table is None:
149160
raise NotImplementedError()
150161

162+
await self.session.refresh(user)
151163
oauth_account = self.oauth_account_table(**create_dict)
152164
self.session.add(oauth_account)
153165
user.oauth_accounts.append(oauth_account) # type: ignore
@@ -172,8 +184,4 @@ async def update_oauth_account(
172184

173185
async def _get_user(self, statement: Select) -> Optional[UP]:
174186
results = await self.session.execute(statement)
175-
user = results.first()
176-
if user is None:
177-
return None
178-
179-
return user[0]
187+
return results.unique().scalar_one_or_none()

‎fastapi_users_db_sqlalchemy/access_token.py‎

Lines changed: 11 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -4,14 +4,13 @@
44

55
from fastapi_users.authentication.strategy.db import AP, AccessTokenDatabase
66
from fastapi_users.models import ID
7-
from sqlalchemy import Column, ForeignKey, String, select
7+
from sqlalchemy import ForeignKey, String, select
88
from sqlalchemy.ext.asyncio import AsyncSession
9-
from sqlalchemy.orm import declarative_mixin, declared_attr
9+
from sqlalchemy.orm import Mapped, declared_attr, mapped_column
1010

1111
from fastapi_users_db_sqlalchemy.generics import GUID, TIMESTAMPAware, now_utc
1212

1313

14-
@declarative_mixin
1514
class SQLAlchemyBaseAccessTokenTable(Generic[ID]):
1615
"""Base SQLAlchemy access token table definition."""
1716

@@ -22,21 +21,20 @@ class SQLAlchemyBaseAccessTokenTable(Generic[ID]):
2221
created_at: datetime
2322
user_id: ID
2423
else:
25-
token: str = Column(String(length=43), primary_key=True)
26-
created_at: datetime = Column(
24+
token: Mapped[str] = mapped_column(String(length=43), primary_key=True)
25+
created_at: Mapped[datetime] = mapped_column(
2726
TIMESTAMPAware(timezone=True), index=True, nullable=False, default=now_utc
2827
)
2928

3029

31-
@declarative_mixin
3230
class SQLAlchemyBaseAccessTokenTableUUID(SQLAlchemyBaseAccessTokenTable[uuid.UUID]):
3331
if TYPE_CHECKING: # pragma: no cover
3432
user_id: uuid.UUID
3533
else:
3634

3735
@declared_attr
38-
def user_id(cls) -> Column[GUID]:
39-
return Column(
36+
def user_id(cls) -> Mapped[GUID]:
37+
return mapped_column(
4038
GUID, ForeignKey("user.id", ondelete="cascade"), nullable=False
4139
)
4240

@@ -61,30 +59,27 @@ async def get_by_token(
6159
self, token: str, max_age: Optional[datetime] = None
6260
) -> Optional[AP]:
6361
statement = select(self.access_token_table).where(
64-
self.access_token_table.token == token
62+
self.access_token_table.token == token # type: ignore
6563
)
6664
if max_age is not None:
67-
statement = statement.where(self.access_token_table.created_at >= max_age)
65+
statement = statement.where(
66+
self.access_token_table.created_at >= max_age # type: ignore
67+
)
6868

6969
results = await self.session.execute(statement)
70-
access_token = results.first()
71-
if access_token is None:
72-
return None
73-
return access_token[0]
70+
return results.scalar_one_or_none()
7471

7572
async def create(self, create_dict: Dict[str, Any]) -> AP:
7673
access_token = self.access_token_table(**create_dict)
7774
self.session.add(access_token)
7875
await self.session.commit()
79-
await self.session.refresh(access_token)
8076
return access_token
8177

8278
async def update(self, access_token: AP, update_dict: Dict[str, Any]) -> AP:
8379
for key, value in update_dict.items():
8480
setattr(access_token, key, value)
8581
self.session.add(access_token)
8682
await self.session.commit()
87-
await self.session.refresh(access_token)
8883
return access_token
8984

9085
async def delete(self, access_token: AP) -> None:

‎fastapi_users_db_sqlalchemy/generics.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import uuid
22
from datetime import datetime, timezone
3+
from typing import Optional
34

45
from pydantic import UUID4
56
from sqlalchemy import CHAR, TIMESTAMP, TypeDecorator
@@ -61,7 +62,7 @@ class TIMESTAMPAware(TypeDecorator): # pragma: no cover
6162
impl = TIMESTAMP
6263
cache_ok = True
6364

64-
def process_result_value(self, value: datetime, dialect):
65-
if dialect.name != "postgresql":
65+
def process_result_value(self, value: Optional[datetime], dialect):
66+
if value is not None and dialect.name != "postgresql":
6667
return value.replace(tzinfo=timezone.utc)
6768
return value

‎pyproject.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,7 @@ classifiers = [
8282
requires-python = ">=3.7"
8383
dependencies = [
8484
"fastapi-users >= 10.0.0",
85-
"sqlalchemy[asyncio] >=1.4,<2.0.0",
85+
"sqlalchemy[asyncio] >=2.0.0,<2.1.0",
8686
]
8787

8888
[project.urls]

‎tests/test_access_token.py‎

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,13 @@
55
import pytest
66
from pydantic import UUID4
77
from sqlalchemy import exc
8-
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine
9-
from sqlalchemy.orm import declarative_base, sessionmaker
8+
from sqlalchemy.ext.asyncio import (
9+
AsyncEngine,
10+
AsyncSession,
11+
async_sessionmaker,
12+
create_async_engine,
13+
)
14+
from sqlalchemy.orm import DeclarativeBase
1015

1116
from fastapi_users_db_sqlalchemy import SQLAlchemyBaseUserTableUUID
1217
from fastapi_users_db_sqlalchemy.access_token import (
@@ -15,7 +20,9 @@
1520
)
1621
from tests.conftest import DATABASE_URL
1722

18-
Base = declarative_base()
23+
24+
class Base(DeclarativeBase):
25+
pass
1926

2027

2128
class AccessToken(SQLAlchemyBaseAccessTokenTableUUID, Base):
@@ -27,7 +34,7 @@ class User(SQLAlchemyBaseUserTableUUID, Base):
2734

2835

2936
def create_async_session_maker(engine: AsyncEngine):
30-
return sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
37+
return async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
3138

3239

3340
@pytest.fixture

‎tests/test_users.py‎

Lines changed: 22 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,18 @@
11
from typing import Any, AsyncGenerator, Dict, List
22

33
import pytest
4-
from sqlalchemy import Column, String, exc
5-
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine
6-
from sqlalchemy.orm import declarative_base, relationship, sessionmaker
4+
from sqlalchemy import String, exc
5+
from sqlalchemy.ext.asyncio import (
6+
AsyncEngine,
7+
async_sessionmaker,
8+
create_async_engine,
9+
)
10+
from sqlalchemy.orm import (
11+
DeclarativeBase,
12+
Mapped,
13+
mapped_column,
14+
relationship,
15+
)
716

817
from fastapi_users_db_sqlalchemy import (
918
UUID_ID,
@@ -15,26 +24,30 @@
1524

1625

1726
def create_async_session_maker(engine: AsyncEngine):
18-
return sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
27+
return async_sessionmaker(engine, expire_on_commit=False)
1928

2029

21-
Base = declarative_base()
30+
class Base(DeclarativeBase):
31+
pass
2232

2333

2434
class User(SQLAlchemyBaseUserTableUUID, Base):
25-
first_name = Column(String(255), nullable=True)
35+
first_name: Mapped[str] = mapped_column(String(255), nullable=True)
2636

2737

28-
OAuthBase = declarative_base()
38+
class OAuthBase(DeclarativeBase):
39+
pass
2940

3041

3142
class OAuthAccount(SQLAlchemyBaseOAuthAccountTableUUID, OAuthBase):
3243
pass
3344

3445

3546
class UserOAuth(SQLAlchemyBaseUserTableUUID, OAuthBase):
36-
first_name = Column(String(255), nullable=True)
37-
oauth_accounts: List[OAuthAccount] = relationship("OAuthAccount", lazy="joined")
47+
first_name: Mapped[str] = mapped_column(String(255), nullable=True)
48+
oauth_accounts: Mapped[List[OAuthAccount]] = relationship(
49+
"OAuthAccount", lazy="joined"
50+
)
3851

3952

4053
@pytest.fixture

0 commit comments

Comments
 (0)