From 270bf999e8cfed700c3639b564f17135e67238a2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ren=C3=A9=20=27Necoro=27=20Neumann?= <necoro@necoro.net>
Date: Sat, 11 Mar 2017 15:21:12 +0100
Subject: Roll our own version of create_or_get, replacing the one
 broken/removed one from Peewee

---
 archivist/model.py | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)

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)
-- 
cgit v1.2.3