3 files changed, 81 insertions, 4 deletions

diff --git a/mediagoblin/db/sql/migrations.py b/mediagoblin/db/sql/migrations.py
index 49798a54..3db51ab9 100644
--- a/mediagoblin/db/sql/migrations.py
+++ b/mediagoblin/db/sql/migrations.py
@@ -14,7 +14,7 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
-from sqlalchemy import MetaData, Table, Column, Boolean, SmallInteger
+from sqlalchemy import MetaData, Table, Column, Boolean, SmallInteger, Integer
 
 from mediagoblin.db.sql.util import RegisterMigration
 
@@ -59,3 +59,15 @@ def add_transcoding_progress(db_conn):
     col = Column('transcoding_progress', SmallInteger)
     col.create(media_entry)
     db_conn.commit()
+
+
+@RegisterMigration(4, MIGRATIONS)
+def add_mediaentry_collections(db_conn):
+    metadata = MetaData(bind=db_conn.bind)
+
+    media_entry = Table('core__media_entries', metadata, autoload=True,
+            autoload_with=db_conn.bind)
+
+    col = Column('collections', Integer)
+    col.create(media_entry)
+    db_conn.commit()
diff --git a/mediagoblin/db/sql/models.py b/mediagoblin/db/sql/models.py
index 7c8c0f99..5862f722 100644
--- a/mediagoblin/db/sql/models.py
+++ b/mediagoblin/db/sql/models.py
@@ -33,7 +33,7 @@ from sqlalchemy.util import memoized_property
 
 from mediagoblin.db.sql.extratypes import PathTupleWithSlashes, JSONEncoded
 from mediagoblin.db.sql.base import Base, DictReadAttrProxy
-from mediagoblin.db.mixin import UserMixin, MediaEntryMixin, MediaCommentMixin
+from mediagoblin.db.mixin import UserMixin, MediaEntryMixin, MediaCommentMixin, CollectionMixin, CollectionItemMixin
 from mediagoblin.db.sql.base import Session
 
 # It's actually kind of annoying how sqlalchemy-migrate does this, if
@@ -103,6 +103,7 @@ class MediaEntry(Base, MediaEntryMixin):
     state = Column(Unicode, default=u'unprocessed', nullable=False)
         # or use sqlalchemy.types.Enum?
     license = Column(Unicode)
+    collected = Column(Integer, default=0)
 
     fail_error = Column(Unicode)
     fail_metadata = Column(JSONEncoded)
@@ -143,6 +144,11 @@ class MediaEntry(Base, MediaEntryMixin):
         creator=lambda v: MediaTag(name=v["name"], slug=v["slug"])
         )
 
+    collections_helper = relationship("CollectionItem",
+        cascade="all, delete-orphan"
+        )
+    collections = association_proxy("collections_helper", "in_collection")
+
     ## TODO
     # media_data
     # fail_error
@@ -348,8 +354,58 @@ class MediaComment(Base, MediaCommentMixin):
     _id = SimpleFieldAlias("id")
 
 
+class Collection(Base, CollectionMixin):
+    __tablename__ = "core__collections"
+
+    id = Column(Integer, primary_key=True)
+    title = Column(Unicode, nullable=False)
+    slug = Column(Unicode)
+    created = Column(DateTime, nullable=False, default=datetime.datetime.now,
+        index=True)
+    description = Column(UnicodeText) 
+    creator = Column(Integer, ForeignKey(User.id), nullable=False)
+    items = Column(Integer, default=0)
+
+    get_creator = relationship(User)
+    
+    def get_collection_items(self, ascending=False):
+        order_col = CollectionItem.position
+        if not ascending:
+            order_col = desc(order_col)
+        return CollectionItem.query.filter_by(
+            collection=self.id).order_by(order_col)
+
+    _id = SimpleFieldAlias("id")
+
+
+class CollectionItem(Base, CollectionItemMixin):
+    __tablename__ = "core__collection_items"
+
+    id = Column(Integer, primary_key=True)
+    media_entry = Column(
+        Integer, ForeignKey(MediaEntry.id), nullable=False, index=True)
+    collection = Column(Integer, ForeignKey(Collection.id), nullable=False)
+    note = Column(UnicodeText, nullable=True)
+    added = Column(DateTime, nullable=False, default=datetime.datetime.now)
+    position = Column(Integer)
+    in_collection = relationship("Collection")
+
+    get_media_entry = relationship(MediaEntry)
+
+    _id = SimpleFieldAlias("id")
+
+    __table_args__ = (
+        UniqueConstraint('collection', 'media_entry'),
+        {})
+
+    @property
+    def dict_view(self):
+        """A dict like view on this object"""
+        return DictReadAttrProxy(self)
+
+
 MODELS = [
-    User, MediaEntry, Tag, MediaTag, MediaComment, MediaFile, FileKeynames,
+    User, MediaEntry, Tag, MediaTag, MediaComment, Collection, CollectionItem, MediaFile, FileKeynames,
     MediaAttachmentFile]
 
 
diff --git a/mediagoblin/db/sql/util.py b/mediagoblin/db/sql/util.py
index 53260db2..74b5d73e 100644
--- a/mediagoblin/db/sql/util.py
+++ b/mediagoblin/db/sql/util.py
@@ -17,7 +17,7 @@
 
 import sys
 from mediagoblin.db.sql.base import Session
-from mediagoblin.db.sql.models import MediaEntry, Tag, MediaTag
+from mediagoblin.db.sql.models import MediaEntry, Tag, MediaTag, Collection
 
 from mediagoblin.tools.common import simple_printer
 
@@ -310,6 +310,15 @@ def clean_orphan_tags():
     Session.commit()
 
 
+def check_collection_slug_used(dummy_db, creator_id, slug, ignore_c_id):
+    filt = (Collection.creator == creator_id) \
+        & (Collection.slug == slug)
+    if ignore_c_id is not None:
+        filt = filt & (Collection.id != ignore_c_id)
+    does_exist = Session.query(Collection.id).filter(filt).first() is not None
+    return does_exist
+
+
 if __name__ == '__main__':
     from mediagoblin.db.sql.open import setup_connection_and_db_from_config