diff options
| -rw-r--r-- | archivist/model.py | 25 |
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) |