From 9ec12bb98c797bac465c5c00426f5c4f2a82282a Mon Sep 17 00:00:00 2001 From: zevaryx Date: Fri, 24 Mar 2023 13:43:25 -0600 Subject: [PATCH] Finish migrating to beanie, centralize partial fields --- jarvis_core/db/models/__init__.py | 30 +++-- jarvis_core/db/models/actions.py | 13 +-- jarvis_core/db/models/backups.py | 178 ++++++++++++++---------------- jarvis_core/db/models/mastodon.py | 32 +++--- jarvis_core/db/models/modlog.py | 25 +---- jarvis_core/db/models/reddit.py | 13 +-- jarvis_core/db/models/twitter.py | 13 +-- jarvis_core/db/utils.py | 8 +- 8 files changed, 126 insertions(+), 186 deletions(-) diff --git a/jarvis_core/db/models/__init__.py b/jarvis_core/db/models/__init__.py index 7d424c0..74c067f 100644 --- a/jarvis_core/db/models/__init__.py +++ b/jarvis_core/db/models/__init__.py @@ -1,6 +1,5 @@ """JARVIS database models.""" -from datetime import datetime, timezone -from functools import partial +from datetime import datetime from typing import Optional from beanie import Document, Link @@ -10,6 +9,7 @@ from jarvis_core.db.models.actions import Ban, Kick, Mute, Unban, Warning from jarvis_core.db.models.modlog import Action, Modlog, Note from jarvis_core.db.models.reddit import Subreddit, SubredditFollow from jarvis_core.db.models.twitter import TwitterAccount, TwitterFollow +from jarvis_core.db.utils import NowField __all__ = [ "Action", @@ -45,14 +45,6 @@ __all__ = [ ] -def get_now() -> datetime: - """Get proper timestamp.""" - return datetime.now(tz=timezone.utc) - - -NowField = partial(Field, default_factory=get_now) - - class Autopurge(Document): guild: int channel: int @@ -131,7 +123,7 @@ class Purge(Document): admin: int channel: int guild: int - count: int = 10 + count_: int = Field(10, alias="count") created_at: datetime = NowField() @@ -181,6 +173,15 @@ class Setting(Document): value: str | int | bool | list[int | str] +class Phishlist(Document): + """Phishlist database object.""" + + url: str + confirmed: bool = False + valid: bool = True + created_at: datetime = NowField() + + class Pinboard(Document): """Pinboard database object.""" @@ -233,15 +234,11 @@ class UserSetting(Document): user: int type: str - settings: str + setting: str value: str | int | bool - class Setting: - name = "usersetting" - all_models = [ - Action, Autopurge, Autoreact, Ban, @@ -252,7 +249,6 @@ all_models = [ Lockdown, Modlog, Mute, - Note, Pin, Pinboard, Purge, diff --git a/jarvis_core/db/models/actions.py b/jarvis_core/db/models/actions.py index 36f05c9..b2a66e0 100644 --- a/jarvis_core/db/models/actions.py +++ b/jarvis_core/db/models/actions.py @@ -1,17 +1,10 @@ """User action models.""" -from datetime import datetime, timezone -from functools import partial +from datetime import datetime from typing import Optional -from beanie import Document, Field +from beanie import Document - -def get_now() -> datetime: - """Get proper timestamp.""" - return datetime.now(tz=timezone.utc) - - -NowField = partial(Field, default_factory=get_now) +from jarvis_core.db.utils import NowField class Ban(Document): diff --git a/jarvis_core/db/models/backups.py b/jarvis_core/db/models/backups.py index 9a0b661..9926abb 100644 --- a/jarvis_core/db/models/backups.py +++ b/jarvis_core/db/models/backups.py @@ -1,122 +1,104 @@ from datetime import datetime -from typing import List, Optional +from typing import Optional -from umongo import Document, EmbeddedDocument, fields +from beanie import Document, Indexed, Link +from pydantic import BaseModel, Field from jarvis_core import __version__ -from jarvis_core.db import JARVIS_INST -from jarvis_core.db.fields import BinaryField -from jarvis_core.db.utils import get_id, get_now +from jarvis_core.db.utils import NanoField, NowField -@JARVIS_INST.register class Image(Document): - discord_id: int = fields.IntegerField(unique=True) - image_data: List[bytes] = BinaryField() - image_ext: str = fields.StringField() - created_at: datetime = fields.AwareDateTimeField(default=get_now) + discord_id: int = Indexed(int, unique=True) + image_data: list[bytes] + image_ext: str + created_at: datetime = NowField() -@JARVIS_INST.register -class PermissionOverwriteBackup(EmbeddedDocument): - id: int = fields.IntegerField() - type: int = fields.IntegerField() - allow: int = fields.IntegerField() - deny: int = fields.IntegerField() +class PermissionOverwriteBackup(BaseModel): + id: int + type: int + allow: int + deny: int -@JARVIS_INST.register -class WebhookBackup(EmbeddedDocument): - id: int = fields.IntegerField() - channel_id: int = fields.IntegerField() - type: int = fields.IntegerField() - avatar: Image = fields.ReferenceField(Image) - name: str = fields.StringField() +class WebhookBackup(BaseModel): + id: int + channel_id: int + type: int + avatar: Link[Image] + name: str -@JARVIS_INST.register -class ChannelBackup(EmbeddedDocument): - id: int = fields.IntegerField() - name: str = fields.StringField() - type: int = fields.IntegerField() - position: int = fields.IntegerField() - topic: Optional[str] = fields.StringField(default=None) - nsfw: bool = fields.BooleanField(default=False) - rate_limit_per_user: int = fields.IntegerField(default=None) - bitrate: Optional[int] = fields.IntegerField(default=None) - user_limit: Optional[int] = fields.IntegerField(default=None) - permission_overwrites: List[PermissionOverwriteBackup] = fields.ListField( - fields.EmbeddedField(PermissionOverwriteBackup), factory=list - ) - parent_id: Optional[int] = fields.IntegerField(default=None) - rtc_region: Optional[str] = fields.StringField(default=None) - video_quality_mode: Optional[int] = fields.IntegerField(default=None) - default_auto_archive_duration: Optional[int] = fields.IntegerField(default=None) - webhooks: List[WebhookBackup] = fields.ListField( - fields.EmbeddedField(WebhookBackup), factory=list - ) +class ChannelBackup(BaseModel): + id: int + name: str + type: int + position: int + topic: Optional[str] = None + nsfw: bool = False + rate_limit_per_user: Optional[int] = None + bitrate: Optional[int] = None + user_limit: Optional[int] = None + permission_overwrites: list[PermissionOverwriteBackup] = Field(default_factory=list) + parent_id: Optional[int] = None + rtc_region: Optional[str] = None + video_quality_mode: Optional[int] = None + default_auto_archive_duration: Optional[int] = None + webhooks: list[WebhookBackup] = Field(default_factory=list) -@JARVIS_INST.register -class RoleBackup(EmbeddedDocument): - id: int = fields.IntegerField() - name: str = fields.StringField() - permissions: int = fields.IntegerField() - color: str = fields.StringField() - hoist: bool = fields.BooleanField() - mentionable: bool = fields.BooleanField() +class RoleBackup(BaseModel): + id: int + name: str + permissions: int + color: str + hoist: bool + mentionable: bool -@JARVIS_INST.register -class EmojiBackup(EmbeddedDocument): - id: int = fields.IntegerField() - name: str = fields.StringField() - image: Image = fields.ReferenceField(Image) +class EmojiBackup(BaseModel): + id: int + name: str + image: Link[Image] -@JARVIS_INST.register -class StickerBackup(EmbeddedDocument): - id: int = fields.IntegerField() - name: str = fields.StringField() - format_type: int = fields.IntegerField() - tags: str = fields.StringField() - type: int = fields.IntegerField() - image: Image = fields.ReferenceField(Image) +class StickerBackup(BaseModel): + id: int + name: str + format_type: int + tags: str + type: int + image: Link[Image] -@JARVIS_INST.register -class GuildBackup(EmbeddedDocument): - name: str = fields.StringField(required=True) - description: str = fields.StringField(default=None) - default_message_notifications: Optional[int] = fields.IntegerField(default=None) - explicit_content_filter: Optional[int] = fields.IntegerField(default=None) - afk_channel: Optional[int] = fields.IntegerField(default=None) - afk_timeout: Optional[int] = fields.IntegerField(default=None) - icon: Optional[Image] = fields.ReferenceField(Image, default=None) - owner: int = fields.IntegerField(required=True) - splash: Optional[Image] = fields.ReferenceField(Image, default=None) - discovery_splash: Optional[Image] = fields.ReferenceField(Image, default=None) - banner: Optional[Image] = fields.ReferenceField(Image, default=None) - system_channel: Optional[int] = fields.IntegerField(default=None) - system_channel_flags: Optional[int] = fields.IntegerField(default=None) - rules_channel: Optional[int] = fields.IntegerField(default=None) - public_updates_channel: Optional[int] = fields.IntegerField(default=None) - preferred_locale: Optional[str] = fields.StringField(default=None) - features: List[str] = fields.ListField(fields.StringField, factory=list) - channels: List[ChannelBackup] = fields.ListField( - fields.EmbeddedField(ChannelBackup), factory=list - ) - roles: List[RoleBackup] = fields.ListField(fields.EmbeddedField(RoleBackup), factory=list) - emojis: List[EmojiBackup] = fields.ListField(fields.EmbeddedField(EmojiBackup), factory=list) - stickers: List[StickerBackup] = fields.ListField( - fields.EmbeddedField(StickerBackup), factory=list - ) +class GuildBackup(BaseModel): + name: str + description: Optional[str] = None + default_message_notifications: Optional[int] = None + explicit_content_filter: Optional[int] = None + afk_channel: Optional[int] = None + afk_timeout: Optional[int] = None + icon: Optional[Link[Image]] = None + owner: int + splash: Optional[Link[Image]] = None + discovery_splash: Optional[Link[Image]] = None + banner: Optional[Link[Image]] = None + system_channel: Optional[int] = None + system_channel_flags: Optional[int] = None + rules_channel: Optional[int] = None + public_updates_channel: Optional[int] = None + preferred_locale: Optional[str] = None + features: list[str] = Field(default_factory=list) + channels: list[ChannelBackup] = Field(default_factory=list) + roles: list[RoleBackup] = Field(default_factory=list) + emojis: list[EmojiBackup] = Field(default_factory=list) + stickers: list[StickerBackup] = Field(default_factory=list) -@JARVIS_INST.register class Backup(Document): - created_at: datetime = fields.AwareDateTimeField(default=get_now) - guild_id: int = fields.IntegerField() - bkid: str = fields.StringField(default=get_id) - guild: GuildBackup = fields.EmbeddedField(GuildBackup, required=True) - version: str = fields.StringField(default=__version__) + created_at: datetime = NowField() + guild_id: int + bkid: str = NanoField() + guild: GuildBackup + version: str = Field(default=__version__) diff --git a/jarvis_core/db/models/mastodon.py b/jarvis_core/db/models/mastodon.py index 1a79ed0..f7384b9 100644 --- a/jarvis_core/db/models/mastodon.py +++ b/jarvis_core/db/models/mastodon.py @@ -1,36 +1,30 @@ """Mastodon databaes models.""" from datetime import datetime, timezone -from umongo import Document, fields +from beanie import Document from jarvis_core.db import JARVIS_INST -from jarvis_core.db.utils import get_now +from jarvis_core.db.utils import NowField @JARVIS_INST.register class MastodonUser(Document): """User object.""" - user_id: int = fields.IntegerField(required=True) - acct: str = fields.StringField(required=True) - username: str = fields.StringField(required=True) - last_sync: datetime = fields.AwareDateTimeField(default=get_now) - - class Meta: - collection_name = "mastodonuser" + user_id: int + acct: str + username: str + last_sync: datetime = NowField() @JARVIS_INST.register class MastodonFollow(Document): """User Follow object.""" - active: bool = fields.BooleanField(default=True) - user_id: int = fields.IntegerField(required=True) - channel: int = fields.IntegerField(required=True) - guild: int = fields.IntegerField(required=True) - reblogged: bool = fields.BooleanField(default=True) - admin: int = fields.IntegerField(required=True) - created_at: datetime = fields.AwareDateTimeField(default=get_now) - - class Meta: - collection_name = "mastodonfollow" + active: bool = True + user_id: int + channel: int + guild: int + reblogged: bool = True + admin: int + created_at: datetime = NowField() diff --git a/jarvis_core/db/models/modlog.py b/jarvis_core/db/models/modlog.py index 6b53df3..4641162 100644 --- a/jarvis_core/db/models/modlog.py +++ b/jarvis_core/db/models/modlog.py @@ -1,34 +1,17 @@ """Modlog database models.""" -from datetime import datetime, timezone -from functools import partial +from datetime import datetime -import nanoid -from bson import ObjectId -from beanie import Document +from beanie import Document, PydanticObjectId from pydantic import BaseModel, Field -NANOID_ALPHA = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz" - - -def get_now() -> datetime: - """Get proper timestamp.""" - return datetime.now(tz=timezone.utc) - - -def get_id() -> str: - """Get nanoid.""" - return nanoid.generate(NANOID_ALPHA, 12) - - -NanoField = partial(Field, default_factory=get_id) -NowField = partial(Field, default_factory=get_now) +from jarvis_core.db.utils import NowField, NanoField class Action(BaseModel): """Modlog embedded action document.""" action_type: str - parent: ObjectId + parent: PydanticObjectId orphaned: bool = False diff --git a/jarvis_core/db/models/reddit.py b/jarvis_core/db/models/reddit.py index eeb8926..fc6be4e 100644 --- a/jarvis_core/db/models/reddit.py +++ b/jarvis_core/db/models/reddit.py @@ -1,16 +1,9 @@ """Reddit databaes models.""" -from datetime import datetime, timezone -from functools import partial +from datetime import datetime -from beanie import Document, Field +from beanie import Document - -def get_now() -> datetime: - """Get proper timestamp.""" - return datetime.now(tz=timezone.utc) - - -NowField = partial(Field, default_factory=get_now) +from jarvis_core.db.utils import NowField class Subreddit(Document): diff --git a/jarvis_core/db/models/twitter.py b/jarvis_core/db/models/twitter.py index e13dd8a..a494aae 100644 --- a/jarvis_core/db/models/twitter.py +++ b/jarvis_core/db/models/twitter.py @@ -1,16 +1,9 @@ """Twitter database models.""" -from datetime import datetime, timezone -from functools import partial +from datetime import datetime -from beanie import Document, Field +from beanie import Document - -def get_now() -> datetime: - """Get proper timestamp.""" - return datetime.now(tz=timezone.utc) - - -NowField = partial(Field, default_factory=get_now) +from jarvis_core.db.utils import NowField class TwitterAccount(Document): diff --git a/jarvis_core/db/utils.py b/jarvis_core/db/utils.py index c56a2e2..393889e 100644 --- a/jarvis_core/db/utils.py +++ b/jarvis_core/db/utils.py @@ -1,6 +1,8 @@ from datetime import datetime, timezone +from functools import partial import nanoid +from pydantic import Field NANOID_ALPHA = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz" @@ -12,4 +14,8 @@ def get_now() -> datetime: def get_id() -> str: """Get nanoid.""" - return nanoid.generate(NANOID_ALPHA, 12) + return nanoid.generate(NANOID_ALPHA, 12) + + +NowField = partial(Field, default_factory=get_now) +NanoField = partial(Field, default_factory=get_id)