From ecec5d6c0a4df7b43c67e1fe4b7503d5962e4426 Mon Sep 17 00:00:00 2001 From: Andrew Godwin Date: Sun, 8 Jan 2023 18:05:29 -0700 Subject: [PATCH] Implement snowflake IDs for some models. Still needs fixes to the client API paginator. --- activities/migrations/0001_initial.py | 11 ++-- activities/models/post.py | 5 +- activities/models/post_interaction.py | 12 +++- core/snowflake.py | 81 +++++++++++++++++++++++++++ users/migrations/0001_initial.py | 11 ++-- users/migrations/0005_report.py | 6 +- users/models/follow.py | 3 + users/models/identity.py | 9 ++- users/models/report.py | 3 + 9 files changed, 119 insertions(+), 22 deletions(-) create mode 100644 core/snowflake.py diff --git a/activities/migrations/0001_initial.py b/activities/migrations/0001_initial.py index 962f728..fa7ba2a 100644 --- a/activities/migrations/0001_initial.py +++ b/activities/migrations/0001_initial.py @@ -10,6 +10,7 @@ import activities.models.fan_out import activities.models.post import activities.models.post_attachment import activities.models.post_interaction +import core.snowflake import core.uploads import stator.models @@ -28,11 +29,10 @@ class Migration(migrations.Migration): fields=[ ( "id", - models.BigAutoField( - auto_created=True, + models.BigIntegerField( + default=core.snowflake.Snowflake.generate_post, primary_key=True, serialize=False, - verbose_name="ID", ), ), ("state_ready", models.BooleanField(default=True)), @@ -111,11 +111,10 @@ class Migration(migrations.Migration): fields=[ ( "id", - models.BigAutoField( - auto_created=True, + models.BigIntegerField( + default=core.snowflake.Snowflake.generate_post_interaction, primary_key=True, serialize=False, - verbose_name="ID", ), ), ("state_ready", models.BooleanField(default=True)), diff --git a/activities/models/post.py b/activities/models/post.py index 74ea0b4..738e6ef 100644 --- a/activities/models/post.py +++ b/activities/models/post.py @@ -30,6 +30,7 @@ from core.ld import ( get_value_or_map, parse_ld_date, ) +from core.snowflake import Snowflake from stator.exceptions import TryAgainLater from stator.models import State, StateField, StateGraph, StatorModel from users.models.follow import FollowStates @@ -215,6 +216,8 @@ class Post(StatorModel): question = "Question" video = "Video" + id = models.BigIntegerField(primary_key=True, default=Snowflake.generate_post) + # The author (attributedTo) of the post author = models.ForeignKey( "users.Identity", @@ -998,7 +1001,7 @@ class Post(StatorModel): "id": self.pk, "uri": self.object_uri, "created_at": format_ld_date(self.published), - "account": self.author.to_mastodon_json(), + "account": self.author.to_mastodon_json(include_counts=False), "content": self.safe_content_remote(), "visibility": visibility_mapping[self.visibility], "sensitive": self.sensitive, diff --git a/activities/models/post_interaction.py b/activities/models/post_interaction.py index bef4ad6..2004561 100644 --- a/activities/models/post_interaction.py +++ b/activities/models/post_interaction.py @@ -4,6 +4,7 @@ from django.utils import timezone from activities.models.fan_out import FanOut from activities.models.post import Post from core.ld import format_ld_date, get_str_or_id, parse_ld_date +from core.snowflake import Snowflake from stator.models import State, StateField, StateGraph, StatorModel from users.models.identity import Identity @@ -125,6 +126,11 @@ class PostInteraction(StatorModel): like = "like" boost = "boost" + id = models.BigIntegerField( + primary_key=True, + default=Snowflake.generate_post_interaction, + ) + # The state the boost is in state = StateField(PostInteractionStates) @@ -327,13 +333,13 @@ class PostInteraction(StatorModel): raise ValueError( f"Cannot make status JSON for interaction of type {self.type}" ) - # Grab our subject post JSON, and just return it if we're a post + # Make a fake post for this boost (because mastodon treats boosts as posts) post_json = self.post.to_mastodon_json(interactions=interactions) return { - "id": f"interaction-{self.pk}", + "id": f"{self.pk}", "uri": post_json["uri"], "created_at": format_ld_date(self.published), - "account": self.identity.to_mastodon_json(), + "account": self.identity.to_mastodon_json(include_counts=False), "content": "", "visibility": post_json["visibility"], "sensitive": post_json["sensitive"], diff --git a/core/snowflake.py b/core/snowflake.py new file mode 100644 index 0000000..87d7734 --- /dev/null +++ b/core/snowflake.py @@ -0,0 +1,81 @@ +import secrets +import time + + +class Snowflake: + """ + Snowflake ID generator and parser. + """ + + # Epoch is 2022/1/1 at midnight, as these are used for _created_ times in our + # own database, not original publish times (which would need an earlier one) + EPOCH = 1641020400 + + TYPE_POST = 0b000 + TYPE_POST_INTERACTION = 0b001 + TYPE_IDENTITY = 0b010 + TYPE_REPORT = 0b011 + TYPE_FOLLOW = 0b100 + + @classmethod + def generate(cls, type_id: int) -> int: + """ + Generates a snowflake-style ID for the given "type". They are designed + to fit inside 63 bits (a signed bigint) + + ID layout is: + * 41 bits of millisecond-level timestamp (enough for EPOCH + 69 years) + * 19 bits of random data (1% chance of clash at 10000 per millisecond) + * 3 bits of type information + + We use random data rather than a sequence ID to try and avoid pushing + this job onto the DB - we may do that in future. If a clash does + occur, the insert will fail and Stator will retry the work for anything + that's coming in remotely, leaving us to just handle that scenario for + our own posts, likes, etc. + """ + # Get the current time in milliseconds + now: int = int((time.time() - cls.EPOCH) * 1000) + # Generate random data + rand_seq: int = secrets.randbits(19) + # Compose them together + return (now << 22) | (rand_seq << 3) | type_id + + @classmethod + def get_type(cls, snowflake: int) -> int: + """ + Returns the type of a given snowflake ID + """ + if snowflake < (1 << 22): + raise ValueError("Not a valid Snowflake ID") + return snowflake & 0b111 + + @classmethod + def get_time(cls, snowflake: int) -> float: + """ + Returns the generation time (in UNIX timestamp seconds) of the ID + """ + if snowflake < (1 << 22): + raise ValueError("Not a valid Snowflake ID") + return ((snowflake >> 22) / 1000) + cls.EPOCH + + # Handy pre-baked methods for django model defaults + @classmethod + def generate_post(cls) -> int: + return cls.generate(cls.TYPE_POST) + + @classmethod + def generate_post_interaction(cls) -> int: + return cls.generate(cls.TYPE_POST_INTERACTION) + + @classmethod + def generate_identity(cls) -> int: + return cls.generate(cls.TYPE_IDENTITY) + + @classmethod + def generate_report(cls) -> int: + return cls.generate(cls.TYPE_REPORT) + + @classmethod + def generate_follow(cls) -> int: + return cls.generate(cls.TYPE_FOLLOW) diff --git a/users/migrations/0001_initial.py b/users/migrations/0001_initial.py index ef30c9b..3dc5249 100644 --- a/users/migrations/0001_initial.py +++ b/users/migrations/0001_initial.py @@ -6,6 +6,7 @@ import django.db.models.deletion from django.conf import settings from django.db import migrations, models +import core.snowflake import core.uploads import stator.models import users.models.follow @@ -216,11 +217,10 @@ class Migration(migrations.Migration): fields=[ ( "id", - models.BigAutoField( - auto_created=True, + models.BigIntegerField( + default=core.snowflake.Snowflake.generate_identity, primary_key=True, serialize=False, - verbose_name="ID", ), ), ("state_ready", models.BooleanField(default=True)), @@ -350,11 +350,10 @@ class Migration(migrations.Migration): fields=[ ( "id", - models.BigAutoField( - auto_created=True, + models.BigIntegerField( + default=core.snowflake.Snowflake.generate_follow, primary_key=True, serialize=False, - verbose_name="ID", ), ), ("state_ready", models.BooleanField(default=True)), diff --git a/users/migrations/0005_report.py b/users/migrations/0005_report.py index a4556b2..d319128 100644 --- a/users/migrations/0005_report.py +++ b/users/migrations/0005_report.py @@ -3,6 +3,7 @@ import django.db.models.deletion from django.db import migrations, models +import core.snowflake import stator.models import users.models.report @@ -20,11 +21,10 @@ class Migration(migrations.Migration): fields=[ ( "id", - models.BigAutoField( - auto_created=True, + models.BigIntegerField( + default=core.snowflake.Snowflake.generate_report, primary_key=True, serialize=False, - verbose_name="ID", ), ), ("state_ready", models.BooleanField(default=True)), diff --git a/users/models/follow.py b/users/models/follow.py index ac6948d..019530a 100644 --- a/users/models/follow.py +++ b/users/models/follow.py @@ -4,6 +4,7 @@ import httpx from django.db import models, transaction from core.ld import canonicalise, get_str_or_id +from core.snowflake import Snowflake from stator.models import State, StateField, StateGraph, StatorModel from users.models.identity import Identity @@ -101,6 +102,8 @@ class Follow(StatorModel): When one user (the source) follows other (the target) """ + id = models.BigIntegerField(primary_key=True, default=Snowflake.generate_follow) + source = models.ForeignKey( "users.Identity", on_delete=models.CASCADE, diff --git a/users/models/identity.py b/users/models/identity.py index 856643c..632ee4f 100644 --- a/users/models/identity.py +++ b/users/models/identity.py @@ -22,6 +22,7 @@ from core.ld import ( ) from core.models import Config from core.signatures import HttpSignature, RsaKeys +from core.snowflake import Snowflake from core.uploads import upload_namer from core.uris import ( AutoAbsoluteUrl, @@ -149,6 +150,8 @@ class Identity(StatorModel): ACTOR_TYPES = ["person", "service", "application", "group", "organization"] + id = models.BigIntegerField(primary_key=True, default=Snowflake.generate_identity) + # The Actor URI is essentially also a PK - we keep the default numeric # one around as well for making nice URLs etc. actor_uri = models.CharField(max_length=500, unique=True) @@ -862,9 +865,9 @@ class Identity(StatorModel): self.created.replace(hour=0, minute=0, second=0, microsecond=0) ), "last_status_at": None, # TODO: populate - "statuses_count": self.posts.count(), - "followers_count": self.inbound_follows.count(), - "following_count": self.outbound_follows.count(), + "statuses_count": self.posts.count() if include_counts else 0, + "followers_count": self.inbound_follows.count() if include_counts else 0, + "following_count": self.outbound_follows.count() if include_counts else 0, } ### Cryptography ### diff --git a/users/models/report.py b/users/models/report.py index 3b60280..5a5c27f 100644 --- a/users/models/report.py +++ b/users/models/report.py @@ -10,6 +10,7 @@ from django.template.loader import render_to_string from core.ld import canonicalise, get_list from core.models import Config +from core.snowflake import Snowflake from stator.models import State, StateField, StateGraph, StatorModel from users.models import Domain @@ -84,6 +85,8 @@ class Report(StatorModel): remote = "remote" other = "other" + id = models.BigIntegerField(primary_key=True, default=Snowflake.generate_report) + state = StateField(ReportStates) subject_identity = models.ForeignKey(