Roll our own version of create_or_get, replacing the one broken/removed one from Peewee

1 files changed, 25 insertions, 0 deletions

diff --git a/archivist/model.py b/archivist/model.py
index fa24304..b76783e 100644
--- a/archivist/model.py
+++ b/archivist/model.py
@@ -28,6 +28,29 @@ def drop_tables():
     db.drop_tables(__tables__, True)
 
 class BaseModel(Model):
+    __keys__ = []
+
+    @classmethod
+    def try_create(cls, **kwargs):
+        try:
+            return cls.create(**kwargs)
+        except IntegrityError:
+            return None
+
+    @classmethod
+    def create_or_get(cls, **kwargs):
+        if not __keys__:
+            raise ValueError("create_or_get not supported on %s" % cls.__name__)
+        for key in __keys__:
+            if key not in kwargs:
+                raise ValueError("required key '%s' missing" % key)
+
+        try:
+            return cls.create(**kwargs), True
+        except IntegrityError:
+            sel = {k:kwargs[k] for k in __keys__}
+            return cls.get(**sel), False
+
     class Meta:
         database = db
 
@@ -61,6 +84,7 @@ class Document(BaseModel):
 
 @table
 class Prefix(BaseModel):
+    __keys__ = ('name',)
     name = CharField(primary_key = True)
     virtual = BooleanField(default = False)
     description = CharField(null=True)
@@ -84,6 +108,7 @@ class Prefix(BaseModel):
 
 @table
 class Tag(BaseModel):
+    __keys__ = ('name', 'prefix')
     name = CharField()
     prefix = ForeignKeyField(Prefix, null=True, related_name = 'tag', db_column = 'prefix')
     description = CharField(null=True)