def mock_data():
"""Create mock tables temporarily for testing purposes."""
src_dirname = os.path.join(mkdtemp(), '') # ensure path ends with /
dest_dirname = mkdtemp()
os.mkdir(os.path.join(src_dirname, 'test_subfolder'))
with psycopg2.connect(port='5432', host='localhost', **db_args) as conn:
table_id = 'log_sensor_file'
column_names = ['log_sensor_file_id', 'sensor_file_path', 'another_column']
dtypes = ['VARCHAR (255)', 'text[]', 'VARCHAR (255)']
drop_table(table_id, conn)
create_table(table_id, conn, column_names, dtypes,
primary_key='log_sensor_file_id')
table_id = 'log_file'
column_names = ['operation_id', 'log_sensor_file_id', 'src_dirname',
'dest_dirname', 'fname', 'time_verified',
'rsync_operation', 'is_deleted', 'is_finished']
dtypes = ['SERIAL', 'text', 'text',
'text', 'text', 'timestamp', 'text',
'boolean', 'boolean']
drop_table(table_id, conn)
create_table(table_id, conn, column_names, dtypes,
primary_key='operation_id',
foreign_key={'log_sensor_file_id': 'log_sensor_file'})
yield src_dirname, dest_dirname
# cleanup
shutil.rmtree(src_dirname)
shutil.rmtree(dest_dirname)
def test_rename_subject_ids():
"""Test renaming of subject."""
table_id = 'subject'
column_names = [
'subject_id', 'redcap_event_name', 'first_name_birth',
'last_name_birth', 'date_of_birth_subject', 'old_subject_id'
]
dtypes = [
'VARCHAR (255)', 'VARCHAR (255)', 'VARCHAR (255)', 'VARCHAR (255)',
'date', 'VARCHAR (255)'
]
rows = [('1001', 'arm1', 'anoopum', 'gupta', '1985-11-28', None),
('1002', 'arm1', 'adonay', 'nunes', '1987-09-13', None)]
index = {
'subject_identifier':
['first_name_birth', 'last_name_birth', 'date_of_birth_subject']
}
with psycopg2.connect(connect_str) as conn:
drop_table(table_id, conn)
table_subject = create_table(table_id,
conn=conn,
column_names=column_names,
dtypes=dtypes,
index=index)
table_subject.insert_rows(rows, cols=column_names)
# Simulate changing subject_id in redcap and updating old_subject_id.
redcap_df = table_subject.query().reset_index()
row_idx = redcap_df['subject_id'] == '1001'
redcap_df.loc[row_idx, 'subject_id'] = '901'
redcap_df.loc[row_idx, 'old_subject_id'] = '1001'
# Add subject in redcap
redcap_df = redcap_df.append(pd.DataFrame({
'subject_id': ['1003'],
'first_name_birth': ['sheraz'],
'last_name_birth': ['khan'],
'date_of_birth_subject': ['1980-05-15'],
'old_subject_id': [None]
}),
ignore_index=True)
rename_subject_ids(table_subject, redcap_df)
# test renaming
table_df_updated = table_subject.query().reset_index()
assert '1001' in table_df_updated['old_subject_id'].values
assert '1001' not in table_df_updated['subject_id'].values
assert '901' in table_df_updated['subject_id'].values
# only rename, don't add rows
assert '1003' not in table_df_updated['subject_id'].values
def test_compound_primary_key():
"""Test compound primary key."""
conn = psycopg2.connect(connect_str)
table_id = 'test'
drop_table(table_id, conn)
column_names = ['subject_id', 'first_name_birth', 'last_name_birth', 'Age']
dtypes = ['VARCHAR (255)', 'VARCHAR (255)', 'VARCHAR (255)', 'INTEGER']
primary_key = ['first_name_birth', 'last_name_birth']
table_subject = create_table(table_id,
conn=conn,
column_names=column_names,
dtypes=dtypes,
primary_key=primary_key)
assert set(table_subject.primary_key) == set(primary_key)
def test_delete():
"""Test deleting rows"""
conn = psycopg2.connect(connect_str)
table_id = 'test'
drop_table(table_id, conn)
column_names = ['subject_id', 'first_name_birth', 'last_name_birth', 'Age']
dtypes = ['VARCHAR (255)', 'VARCHAR (255)', 'VARCHAR (255)', 'INTEGER']
table_subject = create_table(table_id,
conn=conn,
column_names=column_names,
dtypes=dtypes)
table_subject.insert_rows([('x5dc', 'mainak', 'jas', 21),
('y5d3', 'anoopum', 'gupta', 25),
('abcd', 'mayank', 'jas', 25)],
cols=column_names)
table_subject.delete_row("first_name_birth LIKE 'ma%'")
df = table_subject.query()
assert len(df['first_name_birth']) == 1
from neurobooth_terra import Table, create_table
from neurobooth_terra.fixes import OptionalSSHTunnelForwarder
from neurobooth_terra.dataflow import write_files
from config import ssh_args, db_args
do_create_table = False
write_table = True
dest_dir = '/autofs/nas/neurobooth/data_test/'
with OptionalSSHTunnelForwarder(**ssh_args) as tunnel:
with psycopg2.connect(port=tunnel.local_bind_port,
host=tunnel.local_bind_host, **db_args) as conn:
if do_create_table:
table_id = 'log_file'
column_names = ['operation_id', 'log_sensor_file_id', 'src_dirname',
'dest_dirname', 'fname', 'time_verified',
'rsync_operation', 'is_deleted']
dtypes = ['SERIAL', 'text', 'text',
'text', 'text', 'timestamp', 'text',
'boolean']
create_table(table_id, conn, column_names, dtypes,
primary_key='operation_id',
foreign_key={'log_sensor_file_id': 'log_sensor_file'})
if write_table:
sensor_file_table = Table('log_sensor_file', conn)
db_table = Table('log_file', conn)
write_files(sensor_file_table, db_table, dest_dir)
def test_upsert():
"""Test upsert operation."""
conn = psycopg2.connect(connect_str)
table_id = 'test'
drop_table(table_id, conn)
column_names = [
'subject_id', 'first_name_birth', 'last_name_birth', 'Age',
'attributes'
]
dtypes = [
'VARCHAR (255)', 'VARCHAR (255)', 'VARCHAR (255)', 'INTEGER', 'JSONB'
]
table_subject = create_table(table_id,
conn=conn,
column_names=column_names,
dtypes=dtypes)
table_subject.insert_rows([('x5dc', 'mainak', 'jas', 21, {
'a': 1
}), ('y5d3', 'anoopum', 'gupta', 25, {
'b': 2
}), ('abcd', 'mayank', 'jas', 25, {
'a': 1
})],
cols=column_names)
table_subject.insert_rows([('x5dc', 'mainak_new', 'jas_new', 21, {
'a': 1
}), ('zzzz', 'deepak', 'singh', 32, {
'd': 1
})],
cols=column_names,
on_conflict='nothing')
df = table_subject.query()
assert 'x5dc' in df.index
assert 'zzzz' in df.index
assert df.loc['x5dc']['first_name_birth'] == 'mainak' # not mainak_new
df = table_subject.query(where="subject_id = 'x5dc'")
assert 'x5dc' in df.index
assert 'zzzz' not in df.index
# smoke test include_columns
df = table_subject.query(include_columns='first_name_birth',
where="subject_id = 'x5dc'")
df = table_subject.query(include_columns=['first_name_birth'],
where="subject_id = 'x5dc'")
df = query(conn,
"SELECT first_name_birth FROM test where subject_id = 'x5dc' ",
column_names='first_name_birth')
table_subject.insert_rows([('x5dc', 'mainak', 'jazz', 32, {
'a': 1
})],
cols=column_names,
on_conflict='update')
df = table_subject.query(where="subject_id = 'x5dc'")
assert df.loc['x5dc']['last_name_birth'] == 'jazz'
# create index for conflicts
cmd = 'DROP INDEX IF EXISTS subject_identifier'
execute(table_subject.conn, table_subject.cursor, cmd)
cmd = (f'CREATE UNIQUE INDEX subject_identifier ON '
f'{table_id} (first_name_birth, last_name_birth);')
execute(table_subject.conn, table_subject.cursor, cmd)
table_subject.insert_rows(
[('x5de', 'mainak', 'jazz', 32, {
'a': 1
})],
cols=column_names,
on_conflict='update',
conflict_cols=['first_name_birth', 'last_name_birth'],
update_cols='subject_id')
df = table_subject.query()
assert 'x5de' in df.index
table_subject.insert_rows(
[('blah', 'mainak', 'jazz', 32, {
'a': 1
}), ('yyyy', 'mainak', 'jazz', 31, {
'a': 1
})],
cols=column_names,
on_conflict='update',
conflict_cols=['first_name_birth', 'last_name_birth'],
update_cols='subject_id',
where='test."Age" = excluded."Age"')
df = table_subject.query()
assert 'yyyy' not in df.index
assert 'blah' in df.index
def test_psql_connection():
"""Test that we can connect to the database"""
conn = psycopg2.connect(connect_str)
table_id = 'test'
drop_table(table_id, conn)
assert 'test' not in list_tables(conn)
column_names = ['subject_id', 'first_name_birth', 'last_name_birth', 'Age']
dtypes = ['VARCHAR (255)', 'VARCHAR (255)', 'VARCHAR (255)', 'INTEGER']
table_subject = create_table(table_id,
conn=conn,
column_names=column_names,
dtypes=dtypes)
with pytest.raises(Exception, match='already exists'):
table_subject = create_table(table_id,
conn=conn,
column_names=column_names,
dtypes=dtypes)
conn.rollback() # XXX: needed for the error testing above
table_subject.insert_rows([('x5dc', 'mainak', 'jas', 21),
('y5d3', 'anoopum', 'gupta', 25)],
cols=column_names)
with pytest.raises(ValueError, match='vals must be a list of tuple'):
table_subject.insert_rows('blah', ['subject_id'])
with pytest.raises(ValueError, match='entries in vals must be tuples'):
table_subject.insert_rows(['blah'], ['subject_id'])
with pytest.raises(ValueError, match='tuple length must match'):
table_subject.insert_rows([('x5dc', 'mainak', 'jas')], ['subject_id'])
assert sorted(table_subject.data_types) == sorted(dtypes)
table_subject.close()
table_test = Table('test', conn)
assert table_test.primary_key == ['subject_id']
# test updating row
table_test.update_row('y5d3', ('blah', 'anupum', 'gupta', 32),
cols=column_names)
df = table_test.query()
assert 'blah' in df.index
with pytest.raises(ValueError, match='vals must be a tuple'):
table_test.update_row('blah', 'mainak', ['first_name_birth'])
# test updating row partially
table_test.update_row('blah', ('mainak', ), ['first_name_birth'])
df = table_test.query()
assert df.loc['blah']['first_name_birth'] == 'mainak'
with pytest.raises(ValueError, match='column blah is not present'):
table_test.update_row('blah', ('mainak', ), ['blah'])
with pytest.raises(ValueError, match='length of vals'):
table_test.update_row('blah', ('mainak', ),
['first_name_birth', 'last_name_birth'])
# test dropping a column
table_test.drop_column(col='subject_id')
assert 'subject_id' not in table_test.column_names
# test adding an auto-incrementing default value to the column
table_test.add_column(col='subject_id', dtype='VARCHAR')
table_test.alter_column(col='subject_id', default=dict(prefix='SUBJ'))
pk_val = table_test.insert_rows(
[('mainak', 'jas', 21)],
cols=['first_name_birth', 'last_name_birth', 'Age'])
assert pk_val == 'SUBJ1'
df = table_test.query()
assert 'SUBJ1' in df.index
# test insertion of date
table_id = 'test_consent'
drop_table(table_id, conn)
column_names = ['subject_id', 'consent_date']
dtypes = ['VARCHAR (255)', 'date']
table_consent = create_table(table_id, conn, column_names, dtypes)
date = datetime.datetime.today().strftime('%Y-%m-%d')
table_consent.insert_rows([('x5dc', date)], cols=column_names)
conn.close()
with OptionalSSHTunnelForwarder(**ssh_args) as tunnel:
with psycopg2.connect(port=tunnel.local_bind_port,
host=tunnel.local_bind_host, **db_args) as conn:
table_metadata = Table('rc_data_dictionary', conn)
table_metadata.insert_rows(rows_metadata, cols_metadata,
on_conflict='update')
# for table_id, table_info in table_infos.items():
for table_id, _ in survey_ids.items():
table_info = table_infos[table_id]
print(f'Overwriting table {table_id}')
drop_table('rc_' + table_id, conn)
table = create_table('rc_' + table_id, conn, table_info['columns'],
table_info['dtypes'],
primary_key=['subject_id', 'redcap_event_name'])
df = fetch_survey(project, survey_name=table_id,
survey_id=survey_ids[table_id])
df = df.rename(columns={'record_id': 'subject_id'})
# XXX: not consistent.
complete_col = [col for col in df.columns if
col.endswith('complete')]
if len(complete_col) == 0:
warn(f'Skipping {table_id} because of missing complete col')
continue
df = df[df[complete_col[0]] == 2]
report_cols = set([col.split('___')[0] for col in df.columns])
extra_cols = report_cols - (set(table_info['columns']) |