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 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]