diff --git a/superset/views/core.py b/superset/views/core.py index 3bb6a215c8089..e30258f5a93c9 100755 --- a/superset/views/core.py +++ b/superset/views/core.py @@ -112,6 +112,7 @@ from clickhouse_driver import Client import os from time import sleep +from urllib.parse import urlparse config = app.config CACHE_DEFAULT_TIMEOUT = config.get("CACHE_DEFAULT_TIMEOUT", 0) @@ -2739,6 +2740,12 @@ def csv(self, client_id): logging.info("Exporting CSV file [{}]".format(client_id)) query = db.session.query(Query).filter_by(client_id=client_id).one() + # Handle the situations when schema is empty + if query.schema is None: + # Utilize the endpoint of the SQLAlchemy URI as our fallback schema + query.schema = urlparse(query.database.sqlalchemy_uri).path.strip('/') + logging.info(f'Empty query schema. Replaced it with the endpoint of sqlalchemy_uri: {query.schema}') + rejected_tables = security_manager.rejected_tables( query.sql, query.database, query.schema ) @@ -2791,17 +2798,22 @@ def csv(self, client_id): CLICKHOUSE_UNAME = os.environ.get('CLICKHOUSE_UNAME') CLICKHOUSE_PWD = os.environ.get('CLICKHOUSE_PWD') client = Client(host=CLICKHOUSE_HOST, database=query.schema, user=CLICKHOUSE_UNAME, password=CLICKHOUSE_PWD) - rows_gen = client.execute_iter(sql, settings={'max_block_size': 10000}) + rows_gen = client.execute_iter(sql, with_column_types=True, settings={'max_block_size': 10000}) # Utilize the generator pattern to stream CSV contents # ref: https://flask.palletsprojects.com/en/1.1.x/patterns/streaming/ # ref: https://clickhouse-driver.readthedocs.io/en/latest/quickstart.html#streaming-results def generate(): - cnt = 0 + # Determine whether this row is CSV header(columns) or not + isHeader = True for row in rows_gen: # We add sleep(0) between generator iterations to give the worker a break # to update the heartbeat file and keep the connection and worker process alive. sleep(0) s = '' + if isHeader: + # Transform headers from a list of tuples to a comma-seperated string + s = ','.join([col[0] for col in row]) + isHeader = False for item in row: # Remove extra commas item = str(item).replace(',', ' ') @@ -2821,7 +2833,12 @@ def generate(): s = s[:-1] + '\n' yield s - return Response(generate(), mimetype='text/csv'), extra_info + response = Response(generate(), mimetype='text/csv') + # Add the header to assign the filename and filename extension of the response + response.headers[ + "Content-Disposition" + ] = f"attachment; filename={query.name}.csv" + return response, extra_info @api @handle_api_exception