[WiP] work on Superset example improvement

requirements.txt

@@ -76,6 +76,7 @@ six==1.11.0               # via bleach, cryptography, flask-jwt-extended, flask-
 sqlalchemy-utils==0.33.11
 sqlalchemy==1.3.5
 sqlparse==0.2.4
+tabulate==0.8.3
 urllib3==1.24.3           # via requests, selenium
 vine==1.1.4               # via amqp
 webencodings==0.5.1       # via bleach

setup.py

@@ -105,6 +105,7 @@ def get_git_sha():
         'sqlalchemy>=1.3.5,<2.0',
         'sqlalchemy-utils>=0.33.2',
         'sqlparse',
+        'tabulate>=0.8.3',
         'wtforms-json',
     ],
     extras_require={

superset/cli.py

@@ -17,20 +17,32 @@
 # under the License.
 # pylint: disable=C,R,W
 from datetime import datetime
+import json
 import logging
 from subprocess import Popen
-from sys import stdout
+from sys import exit, stdout
+import tarfile
+import tempfile
 
 import click
 from colorama import Fore, Style
 from pathlib2 import Path
+import requests
 import yaml
 
 from superset import (
     app, appbuilder, data, db, security_manager,
 )
+from superset.data.helpers import (
+    download_url_to_blob_url, get_examples_file_list, get_examples_uris,
+    list_examples_table,
+)
+from superset.exceptions import DashboardNotFoundException, ExampleNotFoundException
 from superset.utils import (
-    core as utils, dashboard_import_export, dict_import_export)
+    core as utils, dashboard_import_export, dict_import_export,
+)
+
+logging.getLogger('urllib3').setLevel(logging.WARNING)
 
 config = app.config
 celery_app = utils.get_celery_app(config)
@@ -48,7 +60,7 @@ def make_shell_context():
 @app.cli.command()
 def init():
     """Inits the Superset application"""
-    utils.get_or_create_main_db()
+    utils.get_or_create_db_by_name(db_name='main')
     appbuilder.add_permissions(update_perms=True)
     security_manager.sync_role_definitions()
 
@@ -124,10 +136,254 @@ def load_examples_run(load_test_data):
 @app.cli.command()
 @click.option('--load-test-data', '-t', is_flag=True, help='Load additional test data')
 def load_examples(load_test_data):
-    """Loads a set of Slices and Dashboards and a supporting dataset """
+    """Loads a set of charts and dashboards and a supporting dataset"""
     load_examples_run(load_test_data)
 
 
+def exclusive(ctx_params, exclusive_params, error_message):
+    """Provide exclusive option grouping"""
+    if sum([1 if ctx_params[p] else 0 for p in exclusive_params]) > 1:
+        raise click.UsageError(error_message)
+
+
+@app.cli.group()
+def examples():
+    """Manages example dashboards/datasets"""
+    pass
+
+
+@examples.command('export')
+@click.option(
+    '--dashboard-id', '-i', default=None, type=int,
+    help='Specify a single dashboard id to export')
+@click.option(
+    '--dashboard-title', '-t', default=None,
+    help='Specify a single dashboard title to export')
+@click.option(
+    '--description', '-d', help='Description of new example', required=True)
+@click.option(
+    '--example-title', '-e', help='Title for new example', required=False)
+@click.option(
+    '--file-name', '-f', default='dashboard.tar.gz',
+    help='Specify export file name. Defaults to dashboard.tar.gz')
+@click.option(
+    '--license', '-l', '_license', default='Apache 2.0',
+    help='License of the example dashboard')
+@click.option(
+    '--url', '-u', default=None, help='URL of dataset home page')
+def export_example(dashboard_id, dashboard_title, description, example_title,
+                   file_name, _license, url):
+    """Export example dashboard/datasets tarball"""
+    if not (dashboard_id or dashboard_title):
+        raise click.UsageError('must supply --dashboard-id/-i or --dashboard-title/-t')
+    exclusive(
+        click.get_current_context().params,
+        ['dashboard_id', 'dashboard_title'],
+        'options --dashboard-id/-i and --dashboard-title/-t mutually exclusive')
+
+    # Export into a temporary directory and then tarball that directory
+    with tempfile.TemporaryDirectory() as tmp_dir_name:
+
+        try:
+            data = dashboard_import_export.export_dashboards(
+                db.session,
+                dashboard_ids=[dashboard_id],
+                dashboard_titles=[dashboard_title],
+                export_data=True,
+                export_data_dir=tmp_dir_name,
+                description=description,
+                export_title=example_title or dashboard_title,
+                _license=_license,
+                url=url,
+                strip_database=True)
+
+            dashboard_slug = dashboard_import_export.get_slug(
+                db.session,
+                dashboard_id=dashboard_id,
+                dashboard_title=dashboard_title)
+
+            out_path = f'{tmp_dir_name}/dashboard.json'
+
+            with open(out_path, 'w') as data_stream:
+                data_stream.write(data)
+
+            with tarfile.open(file_name, 'w:gz') as tar:
+                tar.add(tmp_dir_name, arcname=f'{dashboard_slug}')
+
+            click.secho(str(f'Exported example to {file_name}'), fg='blue')
+
+        except DashboardNotFoundException as e:
+            click.secho(str(e), fg='red')
+            exit(1)
+
+
+@examples.command('list')
+@click.option(
+    '--examples-repo', '-r',
+    help='Full name of Github repository containing examples, ex: ' +
+         "'apache-superset/examples-data'",
+    default=None)
+@click.option(
+    '--examples-tag', '-r',
+    help="Tag or branch of Github repository containing examples. Defaults to 'master'",
+    default='master')
+@click.option(
+    '--full-fields', '-ff', is_flag=True, default=False, help='Print full length fields')
+def _list_examples(examples_repo, examples_tag, full_fields):
+    """List example dashboards/datasets"""
+    click.echo(
+        list_examples_table(examples_repo, examples_tag=examples_tag,
+                            full_fields=full_fields))
+
+
+@examples.command('import')
+@click.option(
+    '--database-uri', '-d', help='Database URI to import example to',
+    default=config.get('SQLALCHEMY_EXAMPLES_URI'))
+@click.option(
+    '--examples-repo', '-r',
+    help='Full name of Github repository containing examples, ex: ' +
+         "'apache-superset/examples-data'",
+    default=None)
+@click.option(
+    '--examples-tag', '-r',
+    help="Tag or branch of Github repository containing examples. Defaults to 'master'",
+    default='master')
+@click.option(
+    '--example-title', '-e', help='Title of example to import', required=True)
+def import_example(example_title, examples_repo, examples_tag, database_uri):
+    """Import an example dashboard/dataset"""
+
+    # First fetch the example information from Github
+    examples_repos = [(examples_repo, examples_tag)] \
+        if examples_repo else config.get('EXAMPLE_REPOS_TAGS')
+    examples_repos_uris = [(r[0], r[1]) + get_examples_uris(r[0], r[1])
+                           for r in examples_repos]
+    examples_files = get_examples_file_list(examples_repos_uris)
+
+    # Github authentication via a Personal Access Token for rate limit problems
+    headers = None
+    token = config.get('GITHUB_AUTH_TOKEN')
+    if token:
+        headers = {'Authorization': 'token %s' % token}
+
+    import_example_json = None
+    import_data_info = None
+    for example_file in examples_files:
+
+        metadata_download_url = example_file['metadata_file']['download_url']
+        example_metadata_json = requests.get(metadata_download_url,
+                                             headers=headers).content
+        # Cheaply load json without generating objects
+        example_metadata = json.loads(example_metadata_json)
+        if example_metadata['description']['title'] == example_title:
+            import_example_json = example_metadata_json
+            import_data_info = example_file['data_files']
+            logging.info(
+                f"Will import example '{example_title}' from {metadata_download_url}")
+            break
+
+    if not import_example_json:
+        e = ExampleNotFoundException(f'Example {example_title} not found!')
+        click.secho(str(e), fg='red')
+        exit(1)
+
+    # Parse data to get file download_urls -> blob_urls
+    example_metadata = json.loads(import_example_json,
+                                  object_hook=utils.decode_dashboards)
+
+    # The given download url won't work for data files, need a blob url
+    data_blob_urls = {}
+    for ex_file in example_metadata['files']:
+        github_info = [t for t in import_data_info
+                       if t['name'] == ex_file['file_name']][0]
+        blob_url = download_url_to_blob_url(github_info['download_url'])
+        data_blob_urls[github_info['name']] = blob_url
+
+    try:
+        dashboard_import_export.import_dashboards(
+            db.session,
+            import_example_json,
+            is_example=True,
+            data_blob_urls=data_blob_urls,
+            database_uri=database_uri)
+    except Exception as e:
+        logging.error(f"Error importing example dashboard '{example_title}'!")
+        logging.exception(e)
+
+
+@examples.command('remove')
+@click.option(
+    '--example-title', '-e', help='Title of example to remove', required=True)
+@click.option(
+    '--database-uri', '-d', help='Database URI to remove example from',
+    default=config.get('SQLALCHEMY_EXAMPLES_URI'))
+@click.option(
+    '--examples-repo', '-r',
+    help='Full name of Github repository containing examples, ex: ' +
+         "'apache-superset/examples-data'",
+    default=None)
+@click.option(
+    '--examples-tag', '-r',
+    help="Tag or branch of Github repository containing examples. Defaults to 'master'",
+    default='master')
+def remove_example(example_title, database_uri, examples_repo, examples_tag):
+    """Remove an example dashboard/dataset"""
+
+    # First fetch the example information from Github
+    examples_repos = [(examples_repo, examples_tag)] \
+        if examples_repo else config.get('EXAMPLE_REPOS_TAGS')
+    examples_repos_uris = [(r[0], r[1]) + get_examples_uris(r[0], r[1])
+                           for r in examples_repos]
+    examples_files = get_examples_file_list(examples_repos_uris)
+
+    # Github authentication via a Personal Access Token for rate limit problems
+    headers = None
+    token = config.get('GITHUB_AUTH_TOKEN')
+    if token:
+        headers = {'Authorization': 'token %s' % token}
+
+    # temporary - substitute url provided
+    db_name = 'superset'
+
+    import_example_data = None
+    for example_file in examples_files:
+
+        metadata_download_url = example_file['metadata_file']['download_url']
+        example_metadata_json = requests.get(metadata_download_url,
+                                             headers=headers).content
+        # Cheaply load json without generating objects
+        example_metadata = json.loads(example_metadata_json)
+        if example_metadata['description']['title'] == example_title:
+            import_example_data = json.loads(example_metadata_json)
+            logging.info(
+                f"Will remove example '{example_title}' from '{db_name}'")
+            break
+
+    logging.info(import_example_data['files'])
+
+    # Get the dashboard and associated records
+    dashboard_title = \
+        import_example_data['dashboards'][0]['__Dashboard__']['dashboard_title']
+    logging.info(f'Got dashboard title {dashboard_title} for removal...')
+
+    utils.get_or_create_db_by_name(db_name='main')
+    session = db.session()
+
+    try:
+        dashboard_import_export.remove_dashboard(
+            session,
+            import_example_data,
+            dashboard_title,
+            database_uri=database_uri,
+        )
+    except DashboardNotFoundException as e:
+        logging.exception(e)
+        click.secho(
+            f'Example {example_title} associated dashboard {dashboard_title} not found!',
+            fg='red')
+
+
 @app.cli.command()
 @click.option('--datasource', '-d', help='Specify which datasource name to load, if '
                                          'omitted, all datasources will be refreshed')
@@ -156,7 +412,7 @@ def refresh_druid(datasource, merge):
 
 @app.cli.command()
 @click.option(
-    '--path', '-p',
+    '--path', '-p', required=True,
     help='Path to a single JSON file or path containing multiple JSON files'
          'files to import (*.json)')
 @click.option(
@@ -177,7 +433,7 @@ def import_dashboards(path, recursive):
         try:
             with f.open() as data_stream:
                 dashboard_import_export.import_dashboards(
-                    db.session, data_stream)
+                    db.session, data_stream.read())
         except Exception as e:
             logging.error('Error when importing dashboard from file %s', f)
             logging.error(e)
@@ -190,9 +446,31 @@ def import_dashboards(path, recursive):
 @click.option(
     '--print_stdout', '-p', is_flag=True, default=False,
     help='Print JSON to stdout')
-def export_dashboards(print_stdout, dashboard_file):
-    """Export dashboards to JSON"""
-    data = dashboard_import_export.export_dashboards(db.session)
+@click.option(
+    '--dashboard-ids', '-i', default=None, type=int, multiple=True,
+    help='Specify dashboard id to export')
+@click.option(
+    '--dashboard-titles', '-t', default=None, multiple=True,
+    help='Specify dashboard title to export')
+@click.option(
+    '--export-data', '-x', default=None, is_flag=True,
+    help="Export the dashboard's data tables as CSV files.")
+@click.option(
+    '--export-data-dir', '-d', default=config.get('DASHBOARD_EXPORT_DIR'),
+    help="Specify export directory path. Defaults to '/tmp'")
+def export_dashboards(print_stdout, dashboard_file, dashboard_ids,
+                      dashboard_titles, export_data, export_data_dir):
+    """Export dashboards to JSON and optionally tables to CSV"""
+    try:
+        data = dashboard_import_export.export_dashboards(
+            db.session,
+            dashboard_ids=dashboard_ids,
+            dashboard_titles=dashboard_titles,
+            export_data=export_data,
+            export_data_dir=export_data_dir)
+    except DashboardNotFoundException as e:
+        click.secho(str(e), fg='red')
+        exit(1)
     if print_stdout or not dashboard_file:
         print(data)
     if dashboard_file:
@@ -369,7 +647,7 @@ def load_test_users_run():
         gamma_sqllab_role = security_manager.add_role('gamma_sqllab')
         for perm in security_manager.find_role('Gamma').permissions:
             security_manager.add_permission_role(gamma_sqllab_role, perm)
-        utils.get_or_create_main_db()
+        utils.get_or_create_db_by_name(db_name='main')
         db_perm = utils.get_main_database(security_manager.get_session).perm
         security_manager.add_permission_view_menu('database_access', db_perm)
         db_pvm = security_manager.find_permission_view_menu(