From 6d6533d45fcf231780f53eb9022b76660daca31d Mon Sep 17 00:00:00 2001 From: Timi Fasubaa Date: Fri, 26 Jan 2018 01:26:11 -0800 Subject: [PATCH 1/5] move helpers to utils --- superset/utils.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/superset/utils.py b/superset/utils.py index 8224843213d2a..cef20098ba443 100644 --- a/superset/utils.py +++ b/superset/utils.py @@ -465,6 +465,15 @@ def table_has_constraint(table, name, db): return True return False +def upload_file(file): + """Takes in the name of the file to be uploaded and uploads that file""" + + from superset import app + config = app.config + if not file or not file.filename: + raise Exception + file.save(os.path.join(config['UPLOAD_FOLDER'], file.filename)) + assert file.filename in os.listdir(config['UPLOAD_FOLDER']) class timeout(object): """ From 6da85638f677b2552c5615e01c364132c0d21979 Mon Sep 17 00:00:00 2001 From: Timi Fasubaa Date: Fri, 26 Jan 2018 01:26:46 -0800 Subject: [PATCH 2/5] make form use queryselector --- superset/forms.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/superset/forms.py b/superset/forms.py index a07790440ffde..1759499f6b59b 100644 --- a/superset/forms.py +++ b/superset/forms.py @@ -11,13 +11,18 @@ from wtforms import ( BooleanField, IntegerField, SelectField, StringField) from wtforms.validators import DataRequired, NumberRange, Optional +from wtforms.ext.sqlalchemy.fields import QuerySelectField -from superset import app +from superset import app, db config = app.config class CsvToDatabaseForm(DynamicForm): + def all_db_items(): + from superset.models import core as models + return db.session.query(models.Database) + name = StringField( _('Table Name'), description=_('Name of table to be created from csv data.'), @@ -28,12 +33,9 @@ class CsvToDatabaseForm(DynamicForm): description=_('Select a CSV file to be uploaded to a database.'), validators=[ FileRequired(), FileAllowed(['csv'], _('CSV Files Only!'))]) - - con = SelectField( - _('Database'), - description=_('database in which to add above table.'), - validators=[DataRequired()], - choices=[]) + con = QuerySelectField( + query_factory=all_db_items, + get_pk=lambda a: a.id, get_label=lambda a: a.database_name) sep = StringField( _('Delimiter'), description=_('Delimiter used by CSV file (for whitespace use \s+).'), From 40fa488dae3d3408c43d7cefc10eef75594413f6 Mon Sep 17 00:00:00 2001 From: Timi Fasubaa Date: Fri, 26 Jan 2018 01:27:31 -0800 Subject: [PATCH 3/5] refactor exception throwing and handling --- superset/views/core.py | 41 +++++++++++++++++------------------------ 1 file changed, 17 insertions(+), 24 deletions(-) diff --git a/superset/views/core.py b/superset/views/core.py index ec4cce1fb1568..804937a5a3310 100755 --- a/superset/views/core.py +++ b/superset/views/core.py @@ -26,6 +26,7 @@ from flask_babel import lazy_gettext as _ import pandas as pd import sqlalchemy as sqla +import sqlalchemy.exc as sqla_exceptions from sqlalchemy import create_engine from sqlalchemy.engine.url import make_url from sqlalchemy.exc import OperationalError @@ -164,8 +165,6 @@ def apply(self, query, func): # noqa return query - - class DatabaseView(SupersetModelView, DeleteMixin, YamlExportMixin): # noqa datamodel = SQLAInterface(models.Database) @@ -301,7 +300,6 @@ class DatabaseAsync(DatabaseView): 'allow_run_async', 'allow_run_sync', 'allow_dml', ] - appbuilder.add_view_no_menu(DatabaseAsync) @@ -320,47 +318,42 @@ def form_get(self, form): form.infer_datetime_format.data = True form.decimal.data = '.' form.if_exists.data = 'append' - all_datasources = ( - db.session.query( - models.Database.sqlalchemy_uri, - models.Database.database_name) - .all() - ) - form.con.choices += all_datasources + def form_post(self, form): - def _upload_file(csv_file): - if csv_file and csv_file.filename: - filename = secure_filename(csv_file.filename) - csv_file.save(os.path.join(config['UPLOAD_FOLDER'], filename)) - return filename - csv_file = form.csv_file.data - _upload_file(csv_file) + form.csv_file.data.filename = secure_filename(form.csv_file.data.filename) + csv_filename = form.csv_file.data.filename + utils.upload_file(csv_file) table = SqlaTable(table_name=form.name.data) database = ( db.session.query(models.Database) - .filter_by(sqlalchemy_uri=form.data.get('con')) + .filter_by(id=form.data.get('con').id) .one() ) table.database = database table.database_id = database.id try: database.db_engine_spec.create_table_from_csv(form, table) + except sqla_exceptions.IntegrityError as e1: + os.remove(os.path.join(config['UPLOAD_FOLDER'], csv_filename)) # move this to a finally? + flash("Table name {} already exists. Please pick another".format(form.name.data), 'warning') + return redirect('/csvtodatabaseview/form') + except Exception as e: - os.remove(os.path.join(config['UPLOAD_FOLDER'], csv_file.filename)) - flash(e, 'error') - return redirect('/tablemodelview/list/') + os.remove(os.path.join(config['UPLOAD_FOLDER'], csv_filename)) + flash(str(e), 'error') + return redirect('csvtodatabaseview/form') - os.remove(os.path.join(config['UPLOAD_FOLDER'], csv_file.filename)) + os.remove(os.path.join(config['UPLOAD_FOLDER'], csv_filename)) # Go back to welcome page / splash screen db_name = ( db.session.query(models.Database.database_name) - .filter_by(sqlalchemy_uri=form.data.get('con')) + .filter_by(id=form.data.get('con').id) .one() ) message = _('CSV file "{0}" uploaded to table "{1}" in ' - 'database "{2}"'.format(form.csv_file.data.filename, + 'database "{2}"'.format(csv_filename, form.name.data, db_name[0])) flash(message, 'info') From 509b9a21ebac9fa0de606df9880b8ec287914259 Mon Sep 17 00:00:00 2001 From: Timi Fasubaa Date: Fri, 26 Jan 2018 01:30:36 -0800 Subject: [PATCH 4/5] update db_connection access point --- superset/db_engine_specs.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/superset/db_engine_specs.py b/superset/db_engine_specs.py index d26f633bbd266..e55fa9482902a 100644 --- a/superset/db_engine_specs.py +++ b/superset/db_engine_specs.py @@ -133,13 +133,14 @@ def _allowed_file(filename): 'table': table, 'df': df, 'name': form.name.data, - 'con': create_engine(form.con.data, echo=False), + 'con': create_engine(form.con.data.sqlalchemy_uri, echo=False), 'schema': form.schema.data, 'if_exists': form.if_exists.data, 'index': form.index.data, 'index_label': form.index_label.data, 'chunksize': 10000, } + BaseEngineSpec.df_to_db(**df_to_db_kwargs) @classmethod From 5f9055450c95603f65741d1e4fb35f67813e99a3 Mon Sep 17 00:00:00 2001 From: Timi Fasubaa Date: Fri, 26 Jan 2018 02:11:46 -0800 Subject: [PATCH 5/5] nits --- superset/forms.py | 2 +- superset/utils.py | 11 ----------- superset/views/core.py | 33 ++++++++++++++++----------------- tests/core_tests.py | 14 ++++++++------ 4 files changed, 25 insertions(+), 35 deletions(-) diff --git a/superset/forms.py b/superset/forms.py index 893b209c9e565..cacb9067eb81b 100644 --- a/superset/forms.py +++ b/superset/forms.py @@ -14,6 +14,7 @@ from wtforms.validators import DataRequired, NumberRange, Optional from superset import app, db +from superset.models import core as models config = app.config @@ -21,7 +22,6 @@ class CsvToDatabaseForm(DynamicForm): # pylint: disable=E0211 def all_db_items(): - from superset.models import core as models return db.session.query(models.Database) name = StringField( diff --git a/superset/utils.py b/superset/utils.py index f71a9fb31eebe..a5058b7522bfd 100644 --- a/superset/utils.py +++ b/superset/utils.py @@ -513,17 +513,6 @@ def table_has_constraint(table, name, db): return False -def upload_file(file): - """Takes in the name of the file to be uploaded and uploads that file""" - - from superset import app - config = app.config - if not file or not file.filename: - raise Exception - file.save(os.path.join(config['UPLOAD_FOLDER'], file.filename)) - assert file.filename in os.listdir(config['UPLOAD_FOLDER']) - - class timeout(object): """ To be used in a ``with`` block and timeout its content. diff --git a/superset/views/core.py b/superset/views/core.py index bf49b235af3a6..48767b8433455 100755 --- a/superset/views/core.py +++ b/superset/views/core.py @@ -24,11 +24,11 @@ from flask_babel import gettext as __ from flask_babel import lazy_gettext as _ import pandas as pd +from six import text_type import sqlalchemy as sqla -import sqlalchemy.exc as sqla_exceptions from sqlalchemy import create_engine from sqlalchemy.engine.url import make_url -from sqlalchemy.exc import OperationalError +from sqlalchemy.exc import IntegrityError, OperationalError from unidecode import unidecode from werkzeug.routing import BaseConverter from werkzeug.utils import secure_filename @@ -323,29 +323,28 @@ def form_post(self, form): csv_file = form.csv_file.data form.csv_file.data.filename = secure_filename(form.csv_file.data.filename) csv_filename = form.csv_file.data.filename - utils.upload_file(csv_file) - table = SqlaTable(table_name=form.name.data) - table.database = form.data.get('con') - table.database_id = table.database.id try: + csv_file.save(os.path.join(config['UPLOAD_FOLDER'], csv_filename)) + table = SqlaTable(table_name=form.name.data) + table.database = form.data.get('con') + table.database_id = table.database.id table.database.db_engine_spec.create_table_from_csv(form, table) - except sqla_exceptions.IntegrityError: - os.remove(os.path.join(config['UPLOAD_FOLDER'], csv_filename)) + except Exception as e: + try: + os.remove(os.path.join(config['UPLOAD_FOLDER'], csv_filename)) + except OSError: + pass + message = u'Table name {} already exists. Please pick another'.format( + form.name.data) if isinstance(e, IntegrityError) else text_type(e) flash( - 'Table name {} already exists. Please pick another'.format( - form.name.data), - 'warning') + message, + 'danger') return redirect('/csvtodatabaseview/form') - except Exception as e: - os.remove(os.path.join(config['UPLOAD_FOLDER'], csv_filename)) - flash(str(e), 'error') - return redirect('csvtodatabaseview/form') - os.remove(os.path.join(config['UPLOAD_FOLDER'], csv_filename)) # Go back to welcome page / splash screen db_name = table.database.database_name - message = _('CSV file "{0}" uploaded to table "{1}" in ' + message = _(u'CSV file "{0}" uploaded to table "{1}" in ' 'database "{2}"'.format(csv_filename, form.name.data, db_name)) diff --git a/tests/core_tests.py b/tests/core_tests.py index e36c7a998d534..ca8f2460126a9 100644 --- a/tests/core_tests.py +++ b/tests/core_tests.py @@ -803,20 +803,22 @@ def test_import_csv(self): test_file.write('john,1\n') test_file.write('paul,2\n') test_file.close() - main_db_uri = db.session.query( - models.Database)\ - .filter_by(database_name='main').all() + main_db_uri = ( + db.session.query(models.Database) + .filter_by(database_name='main') + .all() + ) test_file = open(filename, 'rb') form_data = { 'csv_file': test_file, 'sep': ',', 'name': table_name, - 'con': main_db_uri, + 'con': main_db_uri[0].id, 'if_exists': 'append', 'index_label': 'test_label', - 'mangle_dupe_cols': False} - + 'mangle_dupe_cols': False, + } url = '/databaseview/list/' add_datasource_page = self.get_resp(url) assert 'Upload a CSV' in add_datasource_page