Skip to content

Commit

Permalink
feat: Adding option to set_database_uri CLI command (#12740)
Browse files Browse the repository at this point in the history
* Adding option to set_database_uri CLI command

* Fixing flag logic
  • Loading branch information
craig-rueda committed Jan 25, 2021
1 parent 51d1b66 commit 0fed1e0
Show file tree
Hide file tree
Showing 2 changed files with 16 additions and 7 deletions.
11 changes: 9 additions & 2 deletions superset/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -191,9 +191,16 @@ def load_examples(
@superset.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: str, uri: str) -> None:
@click.option(
"--skip_create",
"-s",
is_flag=True,
default=False,
help="Create the DB if it doesn't exist",
)
def set_database_uri(database_name: str, uri: str, skip_create: bool) -> None:
"""Updates a database connection URI """
utils.get_or_create_db(database_name, uri)
utils.get_or_create_db(database_name, uri, not skip_create)


@superset.command()
Expand Down
12 changes: 7 additions & 5 deletions superset/utils/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -1095,7 +1095,7 @@ def user_label(user: User) -> Optional[str]:


def get_or_create_db(
database_name: str, sqlalchemy_uri: str, *args: Any, **kwargs: Any
database_name: str, sqlalchemy_uri: str, always_create: Optional[bool] = True
) -> "Database":
from superset import db
from superset.models import core as models
Expand All @@ -1104,13 +1104,15 @@ def get_or_create_db(
db.session.query(models.Database).filter_by(database_name=database_name).first()
)

if not database:
if not database and always_create:
logger.info("Creating database reference for %s", database_name)
database = models.Database(database_name=database_name, *args, **kwargs)
database = models.Database(database_name=database_name)
db.session.add(database)

database.set_sqlalchemy_uri(sqlalchemy_uri)
db.session.commit()
if database:
database.set_sqlalchemy_uri(sqlalchemy_uri)
db.session.commit()

return database


Expand Down

0 comments on commit 0fed1e0

Please sign in to comment.