2412 change to db models for data models (#591)

* feat: adjusting model_data dump script, to new model * feat(api): updating django model for app * fix: updating Cache in data_model, fix failing test
passportxyz · May 14, 2024 · cc486d9 · cc486d9
1 parent 1c3ed25
commit cc486d9
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 7 deletions.
diff --git a/api/data_model/management/commands/scorer_dump_data_model_score.py b/api/data_model/management/commands/scorer_dump_data_model_score.py
@@ -29,9 +29,8 @@ def __init__(self, file):
                     def write_batch(self, data):
                         for d in data:
                             try:
-                                key = d["key"]
                                 value = d["value"]
-                                address = key[1].lower()
+                                address = d["key_1"].lower()
                                 self.file.write(
                                     json.dumps(
                                         {
@@ -118,9 +117,9 @@ def handle(self, *args, **options):
         try:
             export_data_for_model(
                 Cache.objects.filter(
-                    key__0=data_model_name
+                    key_0=data_model_name
                 ),  # This will only filter the scores for eth_stamp_model (v1)
-                "key",
+                "id",
                 batch_size,
                 get_writer(filename),
                 jsonfields_as_str=False,

diff --git a/api/data_model/models.py b/api/data_model/models.py
@@ -10,10 +10,14 @@
 
 
 class Cache(models.Model):
-    key = models.JSONField(primary_key=True)
+    key = models.JSONField(blank=True, null=True)
     value = models.JSONField()
     updated_at = models.DateTimeField()
+    id = models.BigAutoField(primary_key=True)
+    key_0 = models.CharField(max_length=100, blank=True, null=True)
+    key_1 = models.CharField(max_length=100, blank=True, null=True)
 
     class Meta:
         managed = False
         db_table = "cache"
+        unique_together = (("key_0", "key_1"),)
diff --git a/api/data_model/test/test_cmd_scorer_dump_data_eth_model_score.py b/api/data_model/test/test_cmd_scorer_dump_data_eth_model_score.py
@@ -55,7 +55,8 @@ def test_dump_data_eth_model_score(self, mocker, create_cache_table):
             address = f"0x{i}"
             timestamp = datetime(2024, 4, 9, i, 0, 0, tzinfo=timezone.utc)
             Cache.objects.create(
-                key=["predict", address],
+                key_0="predict",
+                key_1=address,
                 value={
                     "data": {"human_probability": i},
                     "meta": {"version": "v1", "Training_date": "2023/12/27"},
@@ -77,7 +78,8 @@ def test_dump_data_eth_model_score(self, mocker, create_cache_table):
             # Also add data to cache for other models (differet value in first element in key)
             # The export command should filter correctly
             Cache.objects.create(
-                key=["predict_nft", address],
+                key_0="predict_nft",
+                key_1=address,
                 value={
                     "data": {"human_probability": i},
                     "meta": {"version": "v1", "Training_date": "2023/12/27"},