def copy_to_csv(self, file: TextIO, select: Select) -> None:
"Executes query and write the rows into a file object in CSV format."
postgres_copy.copy_to(select,
file,
self.engine,
format="csv",
header=True)
def coordinates_csv(survey, active_users, start, end):
def _tz_formatters(table, cols):
formatters = []
for col_name in cols:
col = getattr(table, col_name)
timestamp = db.func.timezone('UTC', col).label('timestamp_UTC')
timestamp_epoch = db.func.extract('epoch', col).cast(
db.Integer).label('timestamp_epoch')
formatters.append(timestamp)
formatters.append(timestamp_epoch)
return formatters
coordinates_csv = io.BytesIO()
coordinates_csv.write(codecs.BOM_UTF8)
active_mobile_ids = [u.id for u in active_users]
formatters = _tz_formatters(MobileCoordinate, ['timestamp'])
coordinates = (db.session.query(
MobileUser, MobileCoordinate,
*formatters).join(MobileCoordinate).filter(
db.and_(MobileCoordinate.mobile_id.in_(active_mobile_ids),
MobileCoordinate.timestamp >= start,
MobileCoordinate.timestamp <= end)).options(
db.Load(MobileUser).defer('id').load_only('uuid'),
db.Load(MobileCoordinate).defer('mobile_id').defer(
'survey_id').defer('timestamp')).order_by(
MobileCoordinate.id))
postgres_copy.copy_to(coordinates,
coordinates_csv,
db.engine,
format='csv',
header=True)
return coordinates_csv
def test_copy_csv(self, session, objects):
sio = io.StringIO()
flags = {'format': 'csv', 'header': True}
copy_to(session.query(Album), sio, session.connection().engine, **flags)
lines = sio.getvalue().strip().split('\n')
assert len(lines) == 4
assert lines[0].split(',') == ['aid', 'name']
assert lines[1].split(',') == [str(objects[0].id), objects[0].name]
def make_bundle(resource):
s3_key = task_utils.get_s3_key(resource['name'])
with smart_open(s3_key, "wb") as fp:
query = query_with_labels(resource['query'], resource['schema'])
copy_to(query,
fp,
db.session.connection().engine,
format='csv',
header=True)
def test_rename_columns(self, session, objects):
sio = io.StringIO()
flags = {'format': 'csv', 'header': True}
query = relabel_query(session.query(Album.id, Album.name.label('title')))
copy_to(query, sio, session.connection().engine, **flags)
lines = sio.getvalue().strip().split('\n')
assert len(lines) == 4
assert lines[0].split(',') == ['id', 'title']
assert lines[1].split(',') == [str(objects[0].id), objects[0].name]
def test_rename_model(self, session, objects):
sio = io.StringIO()
flags = {'format': 'csv', 'header': True}
query = relabel_query(session.query(Album))
copy_to(query, sio, session.connection().engine, **flags)
lines = sio.getvalue().strip().split('\n')
assert len(lines) == 4
assert lines[0].split(',') == ['id', 'name']
assert lines[1].split(',') == [str(objects[0].id), objects[0].name]
def dump(self, table: str, filepath: str) -> None:
self.log.info(f"dumping content of table {table!r} to {filepath!r}")
with self.sqlalchemy_session() as session:
with open(filepath, "w") as fp:
postgres_copy.copy_to(
session.query(self.get_sqlalchemy_table(table)),
fp,
self.engine,
format="csv",
header=True,
)
def make_bundle(resource):
with tempfile.TemporaryDirectory(dir=os.getenv('TMPDIR')) as tmpdir:
csv_path = os.path.join(tmpdir, 'data.csv')
with open(csv_path, 'w') as fp:
query = query_with_labels(resource['query'], resource['schema'])
copy_to(query, db.session.connection().engine, fp, format='csv', header=True)
row_count = wc(csv_path) - 1
make_manifest(resource, row_count, tmpdir)
with tempfile.TemporaryFile(mode='w+b', dir=os.getenv('TMPDIR')) as tmpfile:
archive = zipfile.ZipFile(tmpfile, 'w')
for path in os.listdir(tmpdir):
_, arcname = os.path.split(path)
archive.write(os.path.join(tmpdir, path), arcname=arcname)
archive.close()
tmpfile.seek(0)
upload_s3(resource['name'], tmpfile)
def make_bundle(resource):
with tempfile.TemporaryDirectory(dir=os.getenv('TMPDIR')) as tmpdir:
csv_path = os.path.join(tmpdir, 'data.csv')
with open(csv_path, 'w') as fp:
query = query_with_labels(resource['query'], resource['schema'])
copy_to(query,
db.session.connection().engine,
fp,
format='csv',
header=True)
row_count = wc(csv_path) - 1
make_manifest(resource, row_count, tmpdir)
with tempfile.TemporaryFile(mode='w+b',
dir=os.getenv('TMPDIR')) as tmpfile:
archive = zipfile.ZipFile(tmpfile, 'w')
for path in os.listdir(tmpdir):
_, arcname = os.path.split(path)
archive.write(os.path.join(tmpdir, path), arcname=arcname)
archive.close()
tmpfile.seek(0)
upload_s3(resource['name'], tmpfile)
def test_copy_table(self, session, objects):
sio = io.StringIO()
copy_to(Album.__table__.select(), sio, session.connection().engine)
lines = sio.getvalue().strip().split('\n')
assert len(lines) == 3
assert lines[0].split('\t') == [str(objects[0].id), objects[0].name]
def test_copy_table(self, session, objects):
sio = io.StringIO()
copy_to(Album.__table__.select(), sio, session.connection().engine)
lines = sio.getvalue().strip().split('\n')
assert len(lines) == 3
assert lines[0].split('\t') == [str(objects[0].id), objects[0].name]
