From 994ac04c1f9af9e11a69d9b8f9f2e56cdf4a3da0 Mon Sep 17 00:00:00 2001
From: Daniel Vaz Gaspar <danielvazgaspar@gmail.com>
Date: Sun, 28 Jul 2019 01:43:39 +0100
Subject: [PATCH] [cli] New, command line option to create or set a db URI
 (#7918)

* [cli] New, command line option to create or set a db URI

* [tests] New, get or create db on utils

* [tests] flake8 and black

* [tests] Fix, remove dummy comments

* [cli] dummy commit to re-trigger build

* [cli] dummy commit to re-trigger build
---
 superset/cli.py      |  8 ++++++++
 tests/utils_tests.py | 24 +++++++++++++++++++++++-
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/superset/cli.py b/superset/cli.py
index cc66f5884..65b08bfef 100755
--- a/superset/cli.py
+++ b/superset/cli.py
@@ -141,6 +141,14 @@ def load_examples(load_test_data, only_metadata=False, force=False):
     load_examples_run(load_test_data, only_metadata, force)
 
 
+@app.cli.command()
+@click.option("--database_name", "-d", help="Database name to change")
+@click.option("--uri", "-u", help="Database URI to change")
+def set_database_uri(database_name, uri):
+    """Updates a database connection URI """
+    utils.get_or_create_db(database_name, uri)
+
+
 @app.cli.command()
 @click.option(
     "--datasource",
diff --git a/tests/utils_tests.py b/tests/utils_tests.py
index 3094f669b..1df71a778 100644
--- a/tests/utils_tests.py
+++ b/tests/utils_tests.py
@@ -23,13 +23,16 @@ import uuid
 from flask import Flask
 from flask_caching import Cache
 import numpy
+from sqlalchemy.exc import ArgumentError
 
-from superset import app
+from superset import app, db, security_manager
 from superset.exceptions import SupersetException
+from superset.models.core import Database
 from superset.utils.core import (
     base_json_conv,
     convert_legacy_filters_into_adhoc,
     datetime_f,
+    get_or_create_db,
     get_since_until,
     get_stacktrace,
     json_int_dttm_ser,
@@ -813,3 +816,22 @@ class UtilsTestCase(unittest.TestCase):
             except Exception:
                 stacktrace = get_stacktrace()
                 assert stacktrace is None
+
+    def test_get_or_create_db(self):
+        get_or_create_db("test_db", "sqlite:///superset.db")
+        database = db.session.query(Database).filter_by(database_name="test_db").one()
+        self.assertIsNotNone(database)
+        self.assertEqual(database.sqlalchemy_uri, "sqlite:///superset.db")
+        self.assertIsNotNone(
+            security_manager.find_permission_view_menu(
+                "datasource_access", database.perm
+            )
+        )
+        # Test change URI
+        get_or_create_db("test_db", "sqlite:///changed.db")
+        database = db.session.query(Database).filter_by(database_name="test_db").one()
+        self.assertEqual(database.sqlalchemy_uri, "sqlite:///changed.db")
+
+    def test_get_or_create_db_invalid_uri(self):
+        with self.assertRaises(ArgumentError):
+            get_or_create_db("test_db", "yoursql:superset.db/()")