async def create_or_update(oauth_name: str, token: Dict[str, Any],
profile: OAuth2Profile) -> "UserOAuthAccount":
access_token = token["access_token"]
refresh_token = token.get("refresh_token", None)
expires_at = token.get("expires_at", None)
async with db_session() as session:
statement = (select(UserOAuthAccount).where(
UserOAuthAccount.oauth_name == oauth_name).where(
UserOAuthAccount.account_id == profile.account_id))
results = await session.exec(statement)
oauth_account: Optional[UserOAuthAccount] = results.one_or_none()
if oauth_account:
oauth_account.access_token = access_token
oauth_account.refresh_token = refresh_token
oauth_account.expires_at = expires_at
oauth_account.account_name = profile.account_name
else:
oauth_account = UserOAuthAccount(
oauth_name=oauth_name,
access_token=access_token,
refresh_token=refresh_token,
expires_at=expires_at,
account_id=profile.account_id,
account_name=profile.account_name,
account_email=profile.account_email,
)
session.sync_session.add(oauth_account)
await session.commit()
await session.refresh(oauth_account)
return oauth_account
async def fetch_related(self, *fields: str) -> None:
def sync_func(_: Any) -> None:
for field in fields:
getattr(self, field)
async with db_session() as session:
await session.run_sync(sync_func)
async def create(
cls,
user_create: "UserCreate",
jwt_access_token: Optional["JWTAccessToken"],
register_ip: str,
) -> "User":
oauth_account: Optional[UserOAuthAccount]
if user_create.oauth_name:
if (jwt_access_token is None
or jwt_access_token.category != "oauth"
or jwt_access_token.oauth_name != user_create.oauth_name
or jwt_access_token.id != user_create.oauth_account_id):
raise BizError(ErrorCode.UserRegisterError,
"oauth account not matched")
user, oauth_account = await cls._create_user_by_oauth(
user_create, jwt_access_token, register_ip)
else:
user = cls._create_user(user_create, register_ip)
oauth_account = None
async with db_session() as session:
session.sync_session.add(user)
if oauth_account: # pragma: no cover
oauth_account.user_id = user.id
session.sync_session.add(oauth_account)
await session.commit()
await session.refresh(user)
return user
async def save_model(self,
commit: bool = True,
refresh: bool = True) -> None:
async with db_session() as session:
session.sync_session.add(self)
if commit:
await session.commit()
if refresh:
await session.refresh(self)
async def get_latest_problem_config(self) -> Optional["ProblemConfig"]:
from joj.horse import models
statement = (
models.ProblemConfig.sql_select()
.where(models.ProblemConfig.problem_id == self.id)
.order_by(models.ProblemConfig.created_at.desc()) # type: ignore
.limit(1)
)
async with db_session() as session:
results = await session.exec(statement)
return results.one_or_none()
async def find_by_url_or_id(
cls, url_or_id: str) -> Optional["BaseORMModelType"]:
if is_uuid(url_or_id):
statement = select(cls).where(cls.id == url_or_id)
else:
statement = select(cls).where(cls.url == url_or_id)
async with db_session() as session:
try:
result = await session.exec(statement)
except StatementError:
return None
return result.one_or_none()
async def execute_list_statement(
cls,
statement: Select,
ordering: Optional["OrderingQuery"] = None,
pagination: Optional["PaginationQuery"] = None,
) -> Tuple[Union[List["BaseORMModelType"], List[Row]], int]:
count_statement = cls.apply_count(statement)
statement = cls.apply_ordering(statement, ordering)
statement = cls.apply_pagination(statement, pagination)
async with db_session() as session:
try:
row_count = await session.exec(count_statement)
results = await session.exec(statement)
except StatementError:
return [], 0
row_count_value = row_count.one()
if not isinstance(row_count_value, int):
row_count_value = row_count_value[0]
return results.all(), row_count_value
async def find_by_domain_url_or_id(
cls,
domain: "Domain",
url_or_id: str,
options: Any = None,
) -> Optional["BaseORMModelType"]:
if is_uuid(url_or_id):
statement = (select(cls).where(cls.id == url_or_id).where(
cls.domain_id == domain.id))
else:
statement = (select(cls).where(cls.url == url_or_id).where(
cls.domain_id == domain.id))
if options:
if isinstance(options, list):
statement = statement.options(*options)
else:
statement = statement.options(options)
async with db_session() as session:
try:
result = await session.exec(statement)
except StatementError:
return None
return result.one_or_none()
async def session_exec(
cls: Type["BaseORMModelType"],
statement: Select) -> ScalarResult["BaseORMModelType"]:
async with db_session() as session:
return await session.exec(statement)
async def refresh_model(self) -> None:
async with db_session() as session:
await session.refresh(self)
async def delete_model(self, commit: bool = True) -> None:
async with db_session() as session:
session.sync_session.delete(self)
if commit:
await session.commit()
