Merge pull request #271 from freedomofpress/file-counter-column

Added file counter column

alembic/versions/2f363b3d680e_init.py

@@ -22,15 +22,17 @@ def upgrade():
         sa.Column('id', sa.Integer(), nullable=False),
         sa.Column('uuid', sa.String(length=36), nullable=False),
         sa.Column('journalist_designation', sa.String(length=255), nullable=False),
-        sa.Column('document_count', sa.Integer(), server_default='0', nullable=False),
-        sa.Column('is_flagged', sa.Boolean(name='is_flagged'), server_default='0', nullable=True),
+        sa.Column('document_count', sa.Integer(), server_default=sa.text("0"), nullable=False),
+        sa.Column('is_flagged', sa.Boolean(name='is_flagged'), server_default=sa.text("0"),
+                  nullable=True),
         sa.Column('public_key', sa.Text(), nullable=True),
         sa.Column('fingerprint', sa.String(length=64), nullable=True),
-        sa.Column('interaction_count', sa.Integer(), server_default='0', nullable=False),
-        sa.Column('is_starred', sa.Boolean(name='is_starred'), server_default='0', nullable=True),
+        sa.Column('interaction_count', sa.Integer(), server_default=sa.text("0"), nullable=False),
+        sa.Column('is_starred', sa.Boolean(name='is_starred'), server_default=sa.text("0"),
+                  nullable=True),
         sa.Column('last_updated', sa.DateTime(), nullable=True),
         sa.PrimaryKeyConstraint('id', name=op.f('pk_sources')),
-        sa.UniqueConstraint('uuid', name=op.f('uq_sources_uuid'))
+        sa.UniqueConstraint('uuid', name=op.f('uq_sources_uuid')),
     )
 
     op.create_table(
@@ -39,24 +41,27 @@ def upgrade():
         sa.Column('uuid', sa.String(length=36), nullable=False),
         sa.Column('username', sa.String(length=255), nullable=False),
         sa.PrimaryKeyConstraint('id', name=op.f('pk_users')),
-        sa.UniqueConstraint('uuid', name=op.f('uq_users_uuid'))
+        sa.UniqueConstraint('uuid', name=op.f('uq_users_uuid')),
     )
 
     op.create_table(
         'files',
         sa.Column('id', sa.Integer(), nullable=False),
         sa.Column('uuid', sa.String(length=36), nullable=False),
         sa.Column('filename', sa.String(length=255), nullable=False),
+        sa.Column('file_counter', sa.Integer(), nullable=False),
         sa.Column('size', sa.Integer(), nullable=False),
         sa.Column('download_url', sa.String(length=255), nullable=False),
-        sa.Column('is_downloaded', sa.Boolean(name='is_downloaded'), server_default='0',
+        sa.Column('is_downloaded', sa.Boolean(name='is_downloaded'), server_default=sa.text("0"),
+                  nullable=False),
+        sa.Column('is_read', sa.Boolean(name='is_read'), server_default=sa.text("0"),
                   nullable=False),
-        sa.Column('is_read', sa.Boolean(name='is_read'), server_default='0', nullable=False),
         sa.Column('is_decrypted', sa.Boolean(name='is_decrypted'), nullable=True),
-        sa.Column('source_id', sa.Integer(), nullable=True),
+        sa.Column('source_id', sa.Integer(), nullable=False),
         sa.ForeignKeyConstraint(['source_id'], ['sources.id'],
                                 name=op.f('fk_files_source_id_sources')),
         sa.PrimaryKeyConstraint('id', name=op.f('pk_files')),
+        sa.UniqueConstraint('source_id', 'file_counter', name='uq_messages_source_id_file_counter'),
         sa.UniqueConstraint('uuid', name=op.f('uq_files_uuid')),
         sa.CheckConstraint('CASE WHEN is_downloaded = 0 THEN is_decrypted IS NULL ELSE 1 END',
                            name='files_compare_is_downloaded_vs_is_decrypted'),
@@ -67,20 +72,23 @@ def upgrade():
         sa.Column('id', sa.Integer(), nullable=False),
         sa.Column('uuid', sa.String(length=36), nullable=False),
         sa.Column('filename', sa.String(length=255), nullable=False),
+        sa.Column('file_counter', sa.Integer(), nullable=False),
         sa.Column('size', sa.Integer(), nullable=False),
         sa.Column('download_url', sa.String(length=255), nullable=False),
-        sa.Column('is_downloaded', sa.Boolean(name='is_downloaded'), server_default='0',
+        sa.Column('is_downloaded', sa.Boolean(name='is_downloaded'), server_default=sa.text("0"),
+                  nullable=False),
+        sa.Column('is_read', sa.Boolean(name='is_read'), server_default=sa.text("0"),
                   nullable=False),
-        sa.Column('is_read', sa.Boolean(name='is_read'), server_default='0', nullable=False),
         sa.Column('is_decrypted', sa.Boolean(name='is_decrypted'), nullable=True),
         sa.Column('content', sa.Text(), nullable=True),
-        sa.Column('source_id', sa.Integer(), nullable=True),
+        sa.Column('source_id', sa.Integer(), nullable=False),
         sa.ForeignKeyConstraint(['source_id'], ['sources.id'],
                                 name=op.f('fk_messages_source_id_sources')),
         sa.PrimaryKeyConstraint('id', name=op.f('pk_messages')),
+        sa.UniqueConstraint('source_id', 'file_counter', name='uq_messages_source_id_file_counter'),
         sa.UniqueConstraint('uuid', name=op.f('uq_messages_uuid')),
         sa.CheckConstraint('CASE WHEN is_downloaded = 0 THEN content IS NULL ELSE 1 END',
-                           name='messages_compare_download_vs_content'),
+                           name=op.f('ck_message_compare_download_vs_content')),
         sa.CheckConstraint('CASE WHEN is_downloaded = 0 THEN is_decrypted IS NULL ELSE 1 END',
                            name='messages_compare_is_downloaded_vs_is_decrypted'),
         sa.CheckConstraint('CASE WHEN is_decrypted = 0 THEN content IS NULL ELSE 1 END',
@@ -91,9 +99,10 @@ def upgrade():
         'replies',
         sa.Column('id', sa.Integer(), nullable=False),
         sa.Column('uuid', sa.String(length=36), nullable=False),
-        sa.Column('source_id', sa.Integer(), nullable=True),
+        sa.Column('source_id', sa.Integer(), nullable=False),
         sa.Column('journalist_id', sa.Integer(), nullable=True),
         sa.Column('filename', sa.String(length=255), nullable=False),
+        sa.Column('file_counter', sa.Integer(), nullable=False),
         sa.Column('size', sa.Integer(), nullable=True),
         sa.Column('content', sa.Text(), nullable=True),
         sa.Column('is_downloaded', sa.Boolean(name='is_downloaded'), nullable=True),
@@ -103,6 +112,7 @@ def upgrade():
         sa.ForeignKeyConstraint(['source_id'], ['sources.id'],
                                 name=op.f('fk_replies_source_id_sources')),
         sa.PrimaryKeyConstraint('id', name=op.f('pk_replies')),
+        sa.UniqueConstraint('source_id', 'file_counter', name='uq_messages_source_id_file_counter'),
         sa.UniqueConstraint('uuid', name=op.f('uq_replies_uuid')),
         sa.CheckConstraint('CASE WHEN is_downloaded = 0 THEN content IS NULL ELSE 1 END',
                            name='replies_compare_download_vs_content'),

securedrop_client/db.py

@@ -1,7 +1,7 @@
 import os
 
 from sqlalchemy import Boolean, Column, create_engine, DateTime, ForeignKey, Integer, String, \
-    Text, MetaData, CheckConstraint
+    Text, MetaData, CheckConstraint, text, UniqueConstraint
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import relationship, backref
 
@@ -31,12 +31,12 @@ class Source(Base):
     id = Column(Integer, primary_key=True)
     uuid = Column(String(36), unique=True, nullable=False)
     journalist_designation = Column(String(255), nullable=False)
-    document_count = Column(Integer, server_default="0", nullable=False)
-    is_flagged = Column(Boolean(name='is_flagged'), server_default="0")
+    document_count = Column(Integer, server_default=text("0"), nullable=False)
+    is_flagged = Column(Boolean(name='is_flagged'), server_default=text("0"))
     public_key = Column(Text, nullable=True)
     fingerprint = Column(String(64))
-    interaction_count = Column(Integer, server_default="0", nullable=False)
-    is_starred = Column(Boolean(name='is_starred'), server_default="0")
+    interaction_count = Column(Integer, server_default=text("0"), nullable=False)
+    is_starred = Column(Boolean(name='is_starred'), server_default=text("0"))
     last_updated = Column(DateTime)
 
     def __init__(self, uuid, journalist_designation, is_flagged, public_key,
@@ -61,22 +61,26 @@ def collection(self):
         collection.extend(self.messages)
         collection.extend(self.files)
         collection.extend(self.replies)
-        collection.sort(key=lambda x: int(x.filename.split('-')[0]))
+        collection.sort(key=lambda x: x.file_counter)
         return collection
 
 
 class Message(Base):
 
     __tablename__ = 'messages'
+    __table_args__ = (
+        UniqueConstraint('source_id', 'file_counter', name='uq_messages_source_id_file_counter'),
+    )
 
     id = Column(Integer, primary_key=True)
     uuid = Column(String(36), unique=True, nullable=False)
     filename = Column(String(255), nullable=False)
+    file_counter = Column(Integer, nullable=False)
     size = Column(Integer, nullable=False)
     download_url = Column(String(255), nullable=False)
 
     # This is whether the submission has been downloaded in the local database.
-    is_downloaded = Column(Boolean(name='is_downloaded'), nullable=False, server_default="0")
+    is_downloaded = Column(Boolean(name='is_downloaded'), nullable=False, server_default=text("0"))
 
     # This tracks if the file had been successfully decrypted after download.
     is_decrypted = Column(
@@ -89,36 +93,47 @@ class Message(Base):
     )
 
     # This reflects read status stored on the server.
-    is_read = Column(Boolean(name='is_read'), nullable=False, server_default="0")
+    is_read = Column(Boolean(name='is_read'), nullable=False, server_default=text("0"))
 
     content = Column(
         Text,
         # this check contraint ensures the state of the DB is what one would expect
         CheckConstraint('CASE WHEN is_downloaded = 0 THEN content IS NULL ELSE 1 END',
-                        name='messages_compare_download_vs_content')
+                        name='ck_message_compare_download_vs_content')
     )
 
-    source_id = Column(Integer, ForeignKey('sources.id'))
+    source_id = Column(Integer, ForeignKey('sources.id'), nullable=False)
     source = relationship("Source",
                           backref=backref("messages", order_by=id,
                                           cascade="delete"))
 
+    def __init__(self, **kwargs) -> None:
+        if 'file_counter' in kwargs:
+            raise TypeError('Cannot manually set file_counter')
+        filename = kwargs['filename']
+        kwargs['file_counter'] = int(filename.split('-')[0])
+        super().__init__(**kwargs)
+
     def __repr__(self):
         return '<Message {}>'.format(self.filename)
 
 
 class File(Base):
 
     __tablename__ = 'files'
+    __table_args__ = (
+        UniqueConstraint('source_id', 'file_counter', name='uq_messages_source_id_file_counter'),
+    )
 
     id = Column(Integer, primary_key=True)
     uuid = Column(String(36), unique=True, nullable=False)
     filename = Column(String(255), nullable=False)
+    file_counter = Column(Integer, nullable=False)
     size = Column(Integer, nullable=False)
     download_url = Column(String(255), nullable=False)
 
     # This is whether the submission has been downloaded in the local database.
-    is_downloaded = Column(Boolean(name='is_downloaded'), nullable=False, server_default="0")
+    is_downloaded = Column(Boolean(name='is_downloaded'), nullable=False, server_default=text("0"))
 
     # This tracks if the file had been successfully decrypted after download.
     is_decrypted = Column(
@@ -129,24 +144,34 @@ class File(Base):
     )
 
     # This reflects read status stored on the server.
-    is_read = Column(Boolean(name='is_read'), nullable=False, server_default="0")
+    is_read = Column(Boolean(name='is_read'), nullable=False, server_default=text("0"))
 
-    source_id = Column(Integer, ForeignKey('sources.id'))
+    source_id = Column(Integer, ForeignKey('sources.id'), nullable=False)
     source = relationship("Source",
                           backref=backref("files", order_by=id,
                                           cascade="delete"))
 
+    def __init__(self, **kwargs) -> None:
+        if 'file_counter' in kwargs:
+            raise TypeError('Cannot manually set file_counter')
+        filename = kwargs['filename']
+        kwargs['file_counter'] = int(filename.split('-')[0])
+        super().__init__(**kwargs)
+
     def __repr__(self):
         return '<File {}>'.format(self.filename)
 
 
 class Reply(Base):
 
     __tablename__ = 'replies'
+    __table_args__ = (
+        UniqueConstraint('source_id', 'file_counter', name='uq_messages_source_id_file_counter'),
+    )
 
     id = Column(Integer, primary_key=True)
     uuid = Column(String(36), unique=True, nullable=False)
-    source_id = Column(Integer, ForeignKey('sources.id'))
+    source_id = Column(Integer, ForeignKey('sources.id'), nullable=False)
     source = relationship("Source",
                           backref=backref("replies", order_by=id,
                                           cascade="delete"))
@@ -156,6 +181,7 @@ class Reply(Base):
         "User", backref=backref('replies', order_by=id))
 
     filename = Column(String(255), nullable=False)
+    file_counter = Column(Integer, nullable=False)
     size = Column(Integer)
 
     # This is whether the reply has been downloaded in the local database.
@@ -178,6 +204,13 @@ class Reply(Base):
         nullable=True,
     )
 
+    def __init__(self, **kwargs) -> None:
+        if 'file_counter' in kwargs:
+            raise TypeError('Cannot manually set file_counter')
+        filename = kwargs['filename']
+        kwargs['file_counter'] = int(filename.split('-')[0])
+        super().__init__(**kwargs)
+
     def __repr__(self):
         return '<Reply {}>'.format(self.filename)
 

tests/gui/test_main.py

@@ -249,7 +249,7 @@ def test_conversation_pending_message(mocker):
     mock_source.journalistic_designation = 'Testy McTestface'
 
     msg_uuid = str(uuid4())
-    message = Message(source=mock_source, uuid=msg_uuid, size=123, filename="test.msg.gpg",
+    message = Message(source=mock_source, uuid=msg_uuid, size=123, filename="1-test.msg.gpg",
                       download_url='http://test/test', is_downloaded=False)
 
     mock_source.collection = [message]

tests/gui/test_widgets.py

@@ -339,9 +339,9 @@ def test_SourceWidget_delete_source_when_user_chooses_cancel(mocker):
     mock_source.submissions = []
 
     submission_files = (
-        "submission_1-msg.gpg",
-        "submission_2-msg.gpg",
-        "submission_3-doc.gpg",
+        "1-submission-msg.gpg",
+        "2-submission-msg.gpg",
+        "3-submission-doc.gpg",
     )
     for filename in submission_files:
         submission = mocker.MagicMock()
@@ -682,7 +682,7 @@ def test_FileWidget_init_left(mocker):
     """
     mock_controller = mocker.MagicMock()
     source = factory.Source()
-    message = db.Message(source=source, uuid='uuid', size=123, filename='mah-reply.gpg',
+    message = db.Message(source=source, uuid='uuid', size=123, filename='1-mah-reply.gpg',
                          download_url='http://mah-server/mah-reply-url', is_downloaded=True)
 
     fw = FileWidget(source, message, mock_controller, align='left')
@@ -700,7 +700,7 @@ def test_FileWidget_init_right(mocker):
     """
     mock_controller = mocker.MagicMock()
     source = factory.Source()
-    message = db.Message(source=source, uuid='uuid', size=123, filename='mah-reply.gpg',
+    message = db.Message(source=source, uuid='uuid', size=123, filename='1-mah-reply.gpg',
                          download_url='http://mah-server/mah-reply-url', is_downloaded=True)
 
     fw = FileWidget(source, message, mock_controller, align='right')
@@ -717,7 +717,7 @@ def test_FileWidget_mousePressEvent_download(mocker):
     """
     mock_controller = mocker.MagicMock()
     source = factory.Source()
-    file_ = db.File(source=source, uuid='uuid', size=123, filename='mah-reply.gpg',
+    file_ = db.File(source=source, uuid='uuid', size=123, filename='1-mah-reply.gpg',
                     download_url='http://mah-server/mah-reply-url', is_downloaded=False)
 
     fw = FileWidget(source, file_, mock_controller)
@@ -731,7 +731,7 @@ def test_FileWidget_mousePressEvent_open(mocker):
     """
     mock_controller = mocker.MagicMock()
     source = factory.Source()
-    file_ = db.File(source=source, uuid='uuid', size=123, filename='mah-reply.gpg',
+    file_ = db.File(source=source, uuid='uuid', size=123, filename='1-mah-reply.gpg',
                     download_url='http://mah-server/mah-reply-url', is_downloaded=True)
 
     fw = FileWidget(source, file_, mock_controller)