Skip to content

Commit fbbb3e4

Browse files
heartsuckerheartsucker
committed
added alembic migration and tests
1 parent dfde2be commit fbbb3e4

File tree

2 files changed

+283
-0
lines changed

2 files changed

+283
-0
lines changed

‎securedrop/alembic/versions/b58139cfdc8c_add_checksum_columns.py‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
"""add checksum columns
2+
3+
Revision ID: b58139cfdc8c
4+
Revises: f2833ac34bb6
5+
Create Date: 2019-04-02 10:45:05.178481
6+
7+
"""
8+
import os
9+
from alembic import op
10+
import sqlalchemy as sa
11+
12+
# raise the errors if we're not in production
13+
raise_errors = os.environ.get('SECUREDROP_ENV', 'prod') != 'prod'
14+
15+
try:
16+
from journalist_app import create_app
17+
from models import Submission, Reply
18+
from sdconfig import config
19+
from store import queued_add_checksum_for_file
20+
from worker import rq_worker_queue
21+
except:
22+
if raise_errors:
23+
raise
24+
25+
# revision identifiers, used by Alembic.
26+
revision = 'b58139cfdc8c'
27+
down_revision = 'f2833ac34bb6'
28+
branch_labels = None
29+
depends_on = None
30+
31+
32+
def upgrade():
33+
with op.batch_alter_table('replies', schema=None) as batch_op:
34+
batch_op.add_column(sa.Column('checksum', sa.String(length=255), nullable=True))
35+
36+
with op.batch_alter_table('submissions', schema=None) as batch_op:
37+
batch_op.add_column(sa.Column('checksum', sa.String(length=255), nullable=True))
38+
39+
try:
40+
app = create_app(config)
41+
42+
# we need an app context for the rq worker extension to work properly
43+
with app.app_context():
44+
conn = op.get_bind()
45+
query = sa.text('''SELECT submissions.id, sources.filesystem_id, submissions.filename
46+
FROM submissions
47+
INNER JOIN sources
48+
ON submissions.source_id = sources.id
49+
''')
50+
for (sub_id, filesystem_id, filename) in conn.execute(query):
51+
full_path = app.storage.path(filesystem_id, filename)
52+
rq_worker_queue.enqueue(
53+
queued_add_checksum_for_file,
54+
Submission,
55+
int(sub_id),
56+
full_path,
57+
app.config['SQLALCHEMY_DATABASE_URI'],
58+
)
59+
60+
query = sa.text('''SELECT replies.id, sources.filesystem_id, replies.filename
61+
FROM replies
62+
INNER JOIN sources
63+
ON replies.source_id = sources.id
64+
''')
65+
for (rep_id, filesystem_id, filename) in conn.execute(query):
66+
full_path = app.storage.path(filesystem_id, filename)
67+
rq_worker_queue.enqueue(
68+
queued_add_checksum_for_file,
69+
Reply,
70+
int(rep_id),
71+
full_path,
72+
app.config['SQLALCHEMY_DATABASE_URI'],
73+
)
74+
except:
75+
if raise_errors:
76+
raise
77+
78+
79+
def downgrade():
80+
with op.batch_alter_table('submissions', schema=None) as batch_op:
81+
batch_op.drop_column('checksum')
82+
83+
with op.batch_alter_table('replies', schema=None) as batch_op:
84+
batch_op.drop_column('checksum')

‎securedrop/tests/migrations/migration_b58139cfdc8c.py‎

Lines changed: 199 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,199 @@
1+
# -*- coding: utf-8 -*-
2+
import io
3+
import os
4+
import random
5+
import uuid
6+
7+
from os import path
8+
from sqlalchemy import text
9+
from sqlalchemy.exc import NoSuchColumnError
10+
11+
from db import db
12+
from journalist_app import create_app
13+
from .helpers import random_chars, random_datetime
14+
15+
random.seed('ᕕ( ᐛ )ᕗ')
16+
17+
DATA = b'wat'
18+
DATA_CHECKSUM = 'sha256:f00a787f7492a95e165b470702f4fe9373583fbdc025b2c8bdf0262cc48fcff4'
19+
20+
21+
class Helper:
22+
23+
def __init__(self):
24+
self.journalist_id = None
25+
self.source_id = None
26+
self._counter = 0
27+
28+
@property
29+
def counter(self):
30+
self._counter += 1
31+
return self._counter
32+
33+
def create_journalist(self):
34+
if self.journalist_id is not None:
35+
raise RuntimeError('Journalist already created')
36+
37+
params = {
38+
'uuid': str(uuid.uuid4()),
39+
'username': random_chars(50),
40+
}
41+
sql = '''INSERT INTO journalists (uuid, username)
42+
VALUES (:uuid, :username)
43+
'''
44+
self.journalist_id = db.engine.execute(text(sql), **params).lastrowid
45+
46+
def create_source(self):
47+
if self.source_id is not None:
48+
raise RuntimeError('Source already created')
49+
50+
self.source_filesystem_id = 'aliruhglaiurhgliaurg-{}'.format(self.counter)
51+
params = {
52+
'filesystem_id': self.source_filesystem_id,
53+
'uuid': str(uuid.uuid4()),
54+
'journalist_designation': random_chars(50),
55+
'flagged': False,
56+
'last_updated': random_datetime(nullable=True),
57+
'pending': False,
58+
'interaction_count': 0,
59+
}
60+
sql = '''INSERT INTO sources (filesystem_id, uuid, journalist_designation, flagged,
61+
last_updated, pending, interaction_count)
62+
VALUES (:filesystem_id, :uuid, :journalist_designation, :flagged, :last_updated,
63+
:pending, :interaction_count)
64+
'''
65+
self.source_id = db.engine.execute(text(sql), **params).lastrowid
66+
67+
def create_submission(self, checksum=False):
68+
filename = str(uuid.uuid4())
69+
params = {
70+
'uuid': str(uuid.uuid4()),
71+
'source_id': self.source_id,
72+
'filename': filename,
73+
'size': random.randint(10, 1000),
74+
'downloaded': False,
75+
76+
}
77+
78+
if checksum:
79+
params['checksum'] = \
80+
'sha256:f00a787f7492a95e165b470702f4fe9373583fbdc025b2c8bdf0262cc48fcff4'
81+
sql = '''INSERT INTO submissions (uuid, source_id, filename, size, downloaded, checksum)
82+
VALUES (:uuid, :source_id, :filename, :size, :downloaded, :checksum)
83+
'''
84+
else:
85+
sql = '''INSERT INTO submissions (uuid, source_id, filename, size, downloaded)
86+
VALUES (:uuid, :source_id, :filename, :size, :downloaded)
87+
'''
88+
89+
return (db.engine.execute(text(sql), **params).lastrowid, filename)
90+
91+
def create_reply(self, checksum=False):
92+
filename = str(uuid.uuid4())
93+
params = {
94+
'uuid': str(uuid.uuid4()),
95+
'source_id': self.source_id,
96+
'journalist_id': self.journalist_id,
97+
'filename': filename,
98+
'size': random.randint(10, 1000),
99+
'deleted_by_source': False,
100+
}
101+
102+
if checksum:
103+
params['checksum'] = \
104+
'sha256:f00a787f7492a95e165b470702f4fe9373583fbdc025b2c8bdf0262cc48fcff4'
105+
sql = '''INSERT INTO replies (uuid, source_id, journalist_id, filename, size,
106+
deleted_by_source, checksum)
107+
VALUES (:uuid, :source_id, :journalist_id, :filename, :size,
108+
:deleted_by_source, :checksum)
109+
'''
110+
else:
111+
sql = '''INSERT INTO replies (uuid, source_id, journalist_id, filename, size,
112+
deleted_by_source)
113+
VALUES (:uuid, :source_id, :journalist_id, :filename, :size,
114+
:deleted_by_source)
115+
'''
116+
return (db.engine.execute(text(sql), **params).lastrowid, filename)
117+
118+
119+
class UpgradeTester(Helper):
120+
121+
def __init__(self, config):
122+
Helper.__init__(self)
123+
self.config = config
124+
self.app = create_app(config)
125+
126+
def load_data(self):
127+
global DATA
128+
with self.app.app_context():
129+
self.create_journalist()
130+
self.create_source()
131+
132+
submission_id, submission_filename = self.create_submission()
133+
reply_id, reply_filename = self.create_reply()
134+
135+
# we need to actually create files and write data to them so the RQ worker can hash them
136+
for fn in [submission_filename, reply_filename]:
137+
full_path = self.app.storage.path(self.source_filesystem_id, fn)
138+
139+
dirname = path.dirname(full_path)
140+
if not path.exists(dirname):
141+
os.mkdir(dirname)
142+
143+
with io.open(full_path, 'wb') as f:
144+
f.write(DATA)
145+
146+
def check_upgrade(self):
147+
'''
148+
We cannot inject the `SDConfig` object provided by the fixture `config` into the alembic
149+
subprocess that actually performs the migration. This is needed to get both the value of the
150+
DB URL and access to the function `storage.path`. These values are passed to the `rqworker`,
151+
and without being able to inject this config, the checksum function won't succeed. The above
152+
`load_data` function provides data that can be manually verified by checking the `rqworker`
153+
log file in `/tmp/`.
154+
'''
155+
pass
156+
157+
158+
class DowngradeTester(Helper):
159+
160+
def __init__(self, config):
161+
Helper.__init__(self)
162+
self.config = config
163+
self.app = create_app(config)
164+
165+
def load_data(self):
166+
with self.app.app_context():
167+
self.create_journalist()
168+
self.create_source()
169+
170+
# create a submission and a reply that we don't add checksums to
171+
self.create_submission(checksum=False)
172+
self.create_reply(checksum=False)
173+
174+
# create a submission and a reply that have checksums added
175+
self.create_submission(checksum=True)
176+
self.create_reply(checksum=True)
177+
178+
def check_downgrade(self):
179+
'''
180+
Verify that the checksum column is now gone.
181+
'''
182+
with self.app.app_context():
183+
sql = "SELECT * FROM submissions"
184+
submissions = db.engine.execute(text(sql)).fetchall()
185+
for submission in submissions:
186+
try:
187+
# this should produce an exception since the column is gone
188+
submission['checksum']
189+
except NoSuchColumnError:
190+
pass
191+
192+
sql = "SELECT * FROM replies"
193+
replies = db.engine.execute(text(sql)).fetchall()
194+
for reply in replies:
195+
try:
196+
# this should produce an exception since the column is gone
197+
submission['checksum']
198+
except NoSuchColumnError:
199+
pass

0 commit comments

Comments
 (0)