Container, bug fixes

src/oidc_test/models.py

@@ -1,19 +1,19 @@
 from functools import cached_property
 from typing import Self
 
-from pydantic import BaseModel, EmailStr, AnyHttpUrl, Field, computed_field
+from pydantic import computed_field, AnyHttpUrl, EmailStr
 from authlib.integrations.starlette_client.apps import StarletteOAuth2App
-
-# from app.models import User
+from sqlmodel import SQLModel, Field
 
 
-class Role(BaseModel, extra="ignore"):
+class Role(SQLModel, extra="ignore"):
     name: str
 
 
-class UserBase(BaseModel, extra="ignore"):
+class UserBase(SQLModel, extra="ignore"):
 
     id: str | None = None
+    sid: str | None = None
     name: str
     email: EmailStr | None = None
     picture: AnyHttpUrl | None = None
@@ -46,3 +46,33 @@ class User(UserBase):
     @cached_property
     def roles_as_set(self) -> set[str]:
         return set([role.name for role in self.roles])
+
+
+class OAuth2Token(SQLModel):
+    name: str = Field(primary_key=True)
+    token_type: str  # = Field(max_length=40)
+    access_token: str  # = Field(max_length=2000)
+    raw_id_token: str
+    refresh_token: str  # = Field(max_length=200)
+    expires_at: int  # = PositiveIntegerField()
+    user: User  # = ForeignKey(User)
+
+    def to_token(self):
+        return dict(
+            access_token=self.access_token,
+            token_type=self.token_type,
+            refresh_token=self.refresh_token,
+            expires_at=self.expires_at,
+        )
+
+    @classmethod
+    def from_dict(cls, token_dict: dict, user: User) -> Self:
+        return cls(
+            name=token_dict["access_token"],
+            access_token=token_dict["access_token"],
+            raw_id_token=token_dict["id_token"],
+            token_type=token_dict["token_type"],
+            refresh_token=token_dict["refresh_token"],
+            expires_at=token_dict["expires_at"],
+            user=user,
+        )