Finish migrating to beanie, centralize partial fields

jarvis_core/db/models/__init__.py

@@ -1,6 +1,5 @@
 """JARVIS database models."""
-from datetime import datetime, timezone
+from datetime import datetime
-from functools import partial
 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,

jarvis_core/db/models/actions.py

@@ -1,17 +1,10 @@
 """User action models."""
-from datetime import datetime, timezone
+from datetime import datetime
-from functools import partial
 from typing import Optional
 
-from beanie import Document, Field
+from beanie import Document
 
-
+from jarvis_core.db.utils import NowField
-def get_now() -> datetime:
-    """Get proper timestamp."""
-    return datetime.now(tz=timezone.utc)
-
-
-NowField = partial(Field, default_factory=get_now)
 
 
 class Ban(Document):

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.utils import NanoField, NowField
-from jarvis_core.db.fields import BinaryField
-from jarvis_core.db.utils import get_id, get_now
 
 
-@JARVIS_INST.register
 class Image(Document):
-    discord_id: int = fields.IntegerField(unique=True)
+    discord_id: int = Indexed(int, unique=True)
-    image_data: List[bytes] = BinaryField()
+    image_data: list[bytes]
-    image_ext: str = fields.StringField()
+    image_ext: str
-    created_at: datetime = fields.AwareDateTimeField(default=get_now)
+    created_at: datetime = NowField()
 
 
-@JARVIS_INST.register
+class PermissionOverwriteBackup(BaseModel):
-class PermissionOverwriteBackup(EmbeddedDocument):
+    id: int
-    id: int = fields.IntegerField()
+    type: int
-    type: int = fields.IntegerField()
+    allow: int
-    allow: int = fields.IntegerField()
+    deny: int
-    deny: int = fields.IntegerField()
 
 
-@JARVIS_INST.register
+class WebhookBackup(BaseModel):
-class WebhookBackup(EmbeddedDocument):
+    id: int
-    id: int = fields.IntegerField()
+    channel_id: int
-    channel_id: int = fields.IntegerField()
+    type: int
-    type: int = fields.IntegerField()
+    avatar: Link[Image]
-    avatar: Image = fields.ReferenceField(Image)
+    name: str
-    name: str = fields.StringField()
 
 
-@JARVIS_INST.register
+class ChannelBackup(BaseModel):
-class ChannelBackup(EmbeddedDocument):
+    id: int
-    id: int = fields.IntegerField()
+    name: str
-    name: str = fields.StringField()
+    type: int
-    type: int = fields.IntegerField()
+    position: int
-    position: int = fields.IntegerField()
+    topic: Optional[str] = None
-    topic: Optional[str] = fields.StringField(default=None)
+    nsfw: bool = False
-    nsfw: bool = fields.BooleanField(default=False)
+    rate_limit_per_user: Optional[int] = None
-    rate_limit_per_user: int = fields.IntegerField(default=None)
+    bitrate: Optional[int] = None
-    bitrate: Optional[int] = fields.IntegerField(default=None)
+    user_limit: Optional[int] = None
-    user_limit: Optional[int] = fields.IntegerField(default=None)
+    permission_overwrites: list[PermissionOverwriteBackup] = Field(default_factory=list)
-    permission_overwrites: List[PermissionOverwriteBackup] = fields.ListField(
+    parent_id: Optional[int] = None
-        fields.EmbeddedField(PermissionOverwriteBackup), factory=list
+    rtc_region: Optional[str] = None
-    )
+    video_quality_mode: Optional[int] = None
-    parent_id: Optional[int] = fields.IntegerField(default=None)
+    default_auto_archive_duration: Optional[int] = None
-    rtc_region: Optional[str] = fields.StringField(default=None)
+    webhooks: list[WebhookBackup] = Field(default_factory=list)
-    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
-    )
 
 
-@JARVIS_INST.register
+class RoleBackup(BaseModel):
-class RoleBackup(EmbeddedDocument):
+    id: int
-    id: int = fields.IntegerField()
+    name: str
-    name: str = fields.StringField()
+    permissions: int
-    permissions: int = fields.IntegerField()
+    color: str
-    color: str = fields.StringField()
+    hoist: bool
-    hoist: bool = fields.BooleanField()
+    mentionable: bool
-    mentionable: bool = fields.BooleanField()
 
 
-@JARVIS_INST.register
+class EmojiBackup(BaseModel):
-class EmojiBackup(EmbeddedDocument):
+    id: int
-    id: int = fields.IntegerField()
+    name: str
-    name: str = fields.StringField()
+    image: Link[Image]
-    image: Image = fields.ReferenceField(Image)
 
 
-@JARVIS_INST.register
+class StickerBackup(BaseModel):
-class StickerBackup(EmbeddedDocument):
+    id: int
-    id: int = fields.IntegerField()
+    name: str
-    name: str = fields.StringField()
+    format_type: int
-    format_type: int = fields.IntegerField()
+    tags: str
-    tags: str = fields.StringField()
+    type: int
-    type: int = fields.IntegerField()
+    image: Link[Image]
-    image: Image = fields.ReferenceField(Image)
 
 
-@JARVIS_INST.register
+class GuildBackup(BaseModel):
-class GuildBackup(EmbeddedDocument):
+    name: str
-    name: str = fields.StringField(required=True)
+    description: Optional[str] = None
-    description: str = fields.StringField(default=None)
+    default_message_notifications: Optional[int] = None
-    default_message_notifications: Optional[int] = fields.IntegerField(default=None)
+    explicit_content_filter: Optional[int] = None
-    explicit_content_filter: Optional[int] = fields.IntegerField(default=None)
+    afk_channel: Optional[int] = None
-    afk_channel: Optional[int] = fields.IntegerField(default=None)
+    afk_timeout: Optional[int] = None
-    afk_timeout: Optional[int] = fields.IntegerField(default=None)
+    icon: Optional[Link[Image]] = None
-    icon: Optional[Image] = fields.ReferenceField(Image, default=None)
+    owner: int
-    owner: int = fields.IntegerField(required=True)
+    splash: Optional[Link[Image]] = None
-    splash: Optional[Image] = fields.ReferenceField(Image, default=None)
+    discovery_splash: Optional[Link[Image]] = None
-    discovery_splash: Optional[Image] = fields.ReferenceField(Image, default=None)
+    banner: Optional[Link[Image]] = None
-    banner: Optional[Image] = fields.ReferenceField(Image, default=None)
+    system_channel: Optional[int] = None
-    system_channel: Optional[int] = fields.IntegerField(default=None)
+    system_channel_flags: Optional[int] = None
-    system_channel_flags: Optional[int] = fields.IntegerField(default=None)
+    rules_channel: Optional[int] = None
-    rules_channel: Optional[int] = fields.IntegerField(default=None)
+    public_updates_channel: Optional[int] = None
-    public_updates_channel: Optional[int] = fields.IntegerField(default=None)
+    preferred_locale: Optional[str] = None
-    preferred_locale: Optional[str] = fields.StringField(default=None)
+    features: list[str] = Field(default_factory=list)
-    features: List[str] = fields.ListField(fields.StringField, factory=list)
+    channels: list[ChannelBackup] = Field(default_factory=list)
-    channels: List[ChannelBackup] = fields.ListField(
+    roles: list[RoleBackup] = Field(default_factory=list)
-        fields.EmbeddedField(ChannelBackup), factory=list
+    emojis: list[EmojiBackup] = Field(default_factory=list)
-    )
+    stickers: list[StickerBackup] = Field(default_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
-    )
 
 
-@JARVIS_INST.register
 class Backup(Document):
-    created_at: datetime = fields.AwareDateTimeField(default=get_now)
+    created_at: datetime = NowField()
-    guild_id: int = fields.IntegerField()
+    guild_id: int
-    bkid: str = fields.StringField(default=get_id)
+    bkid: str = NanoField()
-    guild: GuildBackup = fields.EmbeddedField(GuildBackup, required=True)
+    guild: GuildBackup
-    version: str = fields.StringField(default=__version__)
+    version: str = Field(default=__version__)

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)
+    user_id: int
-    acct: str = fields.StringField(required=True)
+    acct: str
-    username: str = fields.StringField(required=True)
+    username: str
-    last_sync: datetime = fields.AwareDateTimeField(default=get_now)
+    last_sync: datetime = NowField()
-
-    class Meta:
-        collection_name = "mastodonuser"
 
 
 @JARVIS_INST.register
 class MastodonFollow(Document):
     """User Follow object."""
 
-    active: bool = fields.BooleanField(default=True)
+    active: bool = True
-    user_id: int = fields.IntegerField(required=True)
+    user_id: int
-    channel: int = fields.IntegerField(required=True)
+    channel: int
-    guild: int = fields.IntegerField(required=True)
+    guild: int
-    reblogged: bool = fields.BooleanField(default=True)
+    reblogged: bool = True
-    admin: int = fields.IntegerField(required=True)
+    admin: int
-    created_at: datetime = fields.AwareDateTimeField(default=get_now)
+    created_at: datetime = NowField()
-
-    class Meta:
-        collection_name = "mastodonfollow"

jarvis_core/db/models/modlog.py

@@ -1,34 +1,17 @@
 """Modlog database models."""
-from datetime import datetime, timezone
+from datetime import datetime
-from functools import partial
 
-import nanoid
+from beanie import Document, PydanticObjectId
-from bson import ObjectId
-from beanie import Document
 from pydantic import BaseModel, Field
 
-NANOID_ALPHA = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
+from jarvis_core.db.utils import NowField, NanoField
-
-
-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)
 
 
 class Action(BaseModel):
     """Modlog embedded action document."""
 
     action_type: str
-    parent: ObjectId
+    parent: PydanticObjectId
     orphaned: bool = False
 
 

jarvis_core/db/models/reddit.py

@@ -1,16 +1,9 @@
 """Reddit databaes models."""
-from datetime import datetime, timezone
+from datetime import datetime
-from functools import partial
 
-from beanie import Document, Field
+from beanie import Document
 
-
+from jarvis_core.db.utils import NowField
-def get_now() -> datetime:
-    """Get proper timestamp."""
-    return datetime.now(tz=timezone.utc)
-
-
-NowField = partial(Field, default_factory=get_now)
 
 
 class Subreddit(Document):

jarvis_core/db/models/twitter.py

@@ -1,16 +1,9 @@
 """Twitter database models."""
-from datetime import datetime, timezone
+from datetime import datetime
-from functools import partial
 
-from beanie import Document, Field
+from beanie import Document
 
-
+from jarvis_core.db.utils import NowField
-def get_now() -> datetime:
-    """Get proper timestamp."""
-    return datetime.now(tz=timezone.utc)
-
-
-NowField = partial(Field, default_factory=get_now)
 
 
 class TwitterAccount(Document):

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)