pt
/
superset
mirror of https://github.com/apache/superset.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix(ssh): Editting Database w/ SSH Tunneling (#24552)

This commit is contained in:
Hugh A. Miles II 2023-07-06 15:43:55 -04:00 committed by GitHub
parent 64d728f80f
commit d041648ad4
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 29 additions and 27 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

56
superset/databases/commands/update.py
View File

@ -69,10 +69,38 @@ class UpdateDatabaseCommand(BaseCommand):
try:
database = DatabaseDAO.update(self._model, self._properties, commit=False)
database.set_sqlalchemy_uri(database.sqlalchemy_uri)
if ssh_tunnel_properties := self._properties.get("ssh_tunnel"):
if not is_feature_enabled("SSH_TUNNELING"):
db.session.rollback()
raise SSHTunnelingNotEnabledError()
existing_ssh_tunnel_model = DatabaseDAO.get_ssh_tunnel(database.id)
if existing_ssh_tunnel_model is None:
# We couldn't found an existing tunnel so we need to create one
try:
CreateSSHTunnelCommand(database.id, ssh_tunnel_properties).run()
except (SSHTunnelInvalidError, SSHTunnelCreateFailedError) as ex:
# So we can show the original message
raise ex
except Exception as ex:
raise DatabaseUpdateFailedError() from ex
else:
# We found an existing tunnel so we need to update it
try:
UpdateSSHTunnelCommand(
existing_ssh_tunnel_model.id, ssh_tunnel_properties
).run()
except (SSHTunnelInvalidError, SSHTunnelUpdateFailedError) as ex:
# So we can show the original message
raise ex
except Exception as ex:
raise DatabaseUpdateFailedError() from ex
# adding a new database we always want to force refresh schema list
# TODO Improve this simplistic implementation for catching DB conn fails
try:
schemas = database.get_all_schema_names()
ssh_tunnel = DatabaseDAO.get_ssh_tunnel(database.id)
schemas = database.get_all_schema_names(ssh_tunnel=ssh_tunnel)
except Exception as ex:
db.session.rollback()
raise DatabaseConnectionFailedError() from ex
@ -104,32 +132,6 @@ class UpdateDatabaseCommand(BaseCommand):
"schema_access", security_manager.get_schema_perm(database, schema)
)
if ssh_tunnel_properties := self._properties.get("ssh_tunnel"):
if not is_feature_enabled("SSH_TUNNELING"):
db.session.rollback()
raise SSHTunnelingNotEnabledError()
existing_ssh_tunnel_model = DatabaseDAO.get_ssh_tunnel(database.id)
if existing_ssh_tunnel_model is None:
# We couldn't found an existing tunnel so we need to create one
try:
CreateSSHTunnelCommand(database.id, ssh_tunnel_properties).run()
except (SSHTunnelInvalidError, SSHTunnelCreateFailedError) as ex:
# So we can show the original message
raise ex
except Exception as ex:
raise DatabaseUpdateFailedError() from ex
else:
# We found an existing tunnel so we need to update it
try:
UpdateSSHTunnelCommand(
existing_ssh_tunnel_model.id, ssh_tunnel_properties
).run()
except (SSHTunnelInvalidError, SSHTunnelUpdateFailedError) as ex:
# So we can show the original message
raise ex
except Exception as ex:
raise DatabaseUpdateFailedError() from ex
db.session.commit()
except (DAOUpdateFailedError, DAOCreateFailedError) as ex: