1 files changed, 34 insertions, 6 deletions
diff --git a/src/web/models/tag.py b/src/web/models/tag.py
index 8d7fe4d4..d43ce300 100644
--- a/src/web/models/tag.py
+++ b/src/web/models/tag.py
@@ -8,15 +8,43 @@ from bootstrap import db
 
 
 class Tag(db.Model):
-    text = Column(String, primary_key=True, unique=False)
+    id = db.Column(db.Integer(), primary_key=True)
+    text = Column(String, unique=False)
+    obj = db.Column(db.String, default='article')
+    user_id = db.Column(db.Integer(), db.ForeignKey('user.id'), nullable=False)
 
-    # foreign keys
-    article_id = Column(Integer, ForeignKey('article.id', ondelete='CASCADE'),
-                        primary_key=True)
+    __mapper_args__ = {
+            'polymorphic_on': obj
+        }
+
+
+    def __init__(self, text):
+        self.text = text
+
+
+class ArticleTag(Tag):
+    __mapper_args__ = {
+        'polymorphic_identity': 'article'
+    }
+    id = db.Column(db.Integer, db.ForeignKey('tag.id', ondelete='CASCADE'),
+                    primary_key=True)
 
     # relationships
+    article_id = Column(Integer, ForeignKey('article.id', ondelete='CASCADE'),
+                        nullable=False)
     article = relationship('Article', back_populates='tag_objs',
                            foreign_keys=[article_id])
 
-    def __init__(self, text):
-        self.text = text
+
+class BookmarkTag(Tag):
+    __mapper_args__ = {
+        'polymorphic_identity': 'bookmark'
+    }
+    id = db.Column(db.Integer, db.ForeignKey('tag.id', ondelete='CASCADE'),
+                    primary_key=True)
+
+    # relationships
+    bookmark_id = Column(Integer, ForeignKey('bookmark.id', ondelete='CASCADE'),
+                        nullable=False)
+    bookmark = relationship('Bookmark', back_populates='tag_objs',
+                           foreign_keys=[bookmark_id])