protocol-builder-mock/example_data.py

52 lines
1.7 KiB
Python

import csv
from pb import db, session
from pb.models import Sponsor
class ExampleDataLoader:
@staticmethod
def clean_db():
session.flush() # Clear out any transactions before deleting it all to avoid spurious errors.
for table in reversed(db.metadata.sorted_tables):
session.execute(table.delete())
session.commit()
session.flush()
@staticmethod
def load_all():
ExampleDataLoader().load_sponsors()
# self.load_studies()
# self.load_investigators()
@staticmethod
def load_sponsors():
# Load sponsors from csv
with open('./pb/static/csv/sponsors.csv') as csv_file:
data = csv.reader(csv_file, delimiter=',')
first_line = True
sponsors = []
for row in data:
# Skip first line, which will be the column headings
if first_line:
# row[0]: SPONSOR_ID
# row[1]: SP_NAME
# row[2]: SP_MAILING_ADDRESS
# row[3]: SP_TYPE
first_line = False
elif int(row[0] or -1) != -1:
sponsor_id = int(row[0])
is_duplicate = session.query(Sponsor).filter(Sponsor.SPONSOR_ID == sponsor_id).count() > 0
# Make sure we're not creating duplicates
if not is_duplicate:
new_sponsor = Sponsor(SPONSOR_ID=sponsor_id, SP_NAME=row[1], SP_MAILING_ADDRESS=row[2], SP_TYPE=row[3])
new_sponsor.SP_TYPE_GROUP_NAME = Sponsor.get_type_group_name(new_sponsor.SP_TYPE)
sponsors.append(new_sponsor)
session.add_all(sponsors)
session.commit()