From 64f4f62ae943851af8d93b306edf301f7b0239b2 Mon Sep 17 00:00:00 2001
From: Ryan Barrett <git@ryanb.org>
Date: Fri, 28 Apr 2023 12:02:26 -0700
Subject: [PATCH] bluesky: add models.AtpNode for storing repo blocks

---
 models.py            | 39 ++++++++++++++++++++++++++++++++++++---
 tests/test_models.py | 16 +++++++++++++---
 2 files changed, 49 insertions(+), 6 deletions(-)

diff --git a/models.py b/models.py
index 6e1cdc3..2d81a56 100644
--- a/models.py
+++ b/models.py
@@ -3,16 +3,15 @@ import base64
 from datetime import timedelta, timezone
 import difflib
 import itertools
+import json
 import logging
 import random
 import urllib.parse
 
-import requests
-from werkzeug.exceptions import BadRequest, NotFound
-
 from Crypto import Random
 from Crypto.PublicKey import ECC, RSA
 from Crypto.Util import number
+import dag_json
 from flask import g, request
 from google.cloud import ndb
 from granary import as1, as2, bluesky, microformats2
@@ -20,7 +19,10 @@ from oauth_dropins.webutil.appengine_info import DEBUG
 from oauth_dropins.webutil.models import ComputedJsonProperty, JsonProperty, StringIdModel
 from oauth_dropins.webutil import util
 from oauth_dropins.webutil.util import json_dumps, json_loads
+import requests
+from werkzeug.exceptions import BadRequest, NotFound
 
+from atproto_mst import dag_cbor_cid
 import common
 
 # https://github.com/snarfed/bridgy-fed/issues/314
@@ -450,6 +452,37 @@ class Object(StringIdModel):
           {util.ellipsize(name, chars=40)}
         </a>"""
 
+class AtpNode(StringIdModel):
+    """An AT Protocol (Bluesky) node.
+
+    May be a data record, an MST node, or a commit.
+
+    Key name is the DAG-CBOR base32 CID of the data.
+
+    Properties:
+    * data: JSON-decoded DAG-JSON value of this node
+    * obj: optional, Key of the corresponding :class:`Object`, only populated
+      for records
+    """
+    data = JsonProperty(required=True)
+    obj = ndb.KeyProperty(Object)
+
+    @staticmethod
+    def create(data):
+        """Writes a new AtpNode to the datastore.
+
+        Args:
+          data: dict value
+
+        Returns:
+          :class:`AtpNode`
+        """
+        data = json.loads(dag_json.encode(data))
+        cid = dag_cbor_cid(data)
+        node = AtpNode(id=cid.encode('base32'), data=data)
+        node.put()
+        return node
+
 
 class Follower(StringIdModel):
     """A follower of a Bridgy Fed user.
diff --git a/tests/test_models.py b/tests/test_models.py
index fdd9cc6..d9acafd 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -2,16 +2,18 @@
 """Unit tests for models.py."""
 from unittest import mock
 
+from Crypto.PublicKey import ECC
 from flask import g, get_flashed_messages
 from granary import as2
+from granary.tests.test_bluesky import ACTOR_PROFILE_BSKY
+from multiformats import CID
 from oauth_dropins.webutil.testutil import NOW, requests_response
 
 from app import app
+from atproto_mst import dag_cbor_cid
 import common
-from Crypto.PublicKey import ECC
-from models import Follower, Object, OBJECT_EXPIRE_AGE, User
+from models import AtpNode, Follower, Object, OBJECT_EXPIRE_AGE, User
 import protocol
-from protocol import Protocol
 from . import testutil
 
 from .test_activitypub import ACTOR
@@ -387,3 +389,11 @@ class FollowerTest(testutil.TestCase):
         self.assertIsNone(Follower().to_as2())
         self.assertEqual(ACTOR, self.inbound.to_as2())
         self.assertEqual(ACTOR, self.outbound.to_as2())
+
+
+class AtpNodeTest(testutil.TestCase):
+
+    def test_create(self):
+        AtpNode.create(ACTOR_PROFILE_BSKY)
+        stored = AtpNode.get_by_id(dag_cbor_cid(ACTOR_PROFILE_BSKY).encode('base32'))
+        self.assertEqual(ACTOR_PROFILE_BSKY, stored.data)