mirror of
https://github.com/apache/superset.git
synced 2026-04-11 12:26:05 +00:00
ebaad10d6c
Co-authored-by: Antonio Rivero Martinez <38889534+Antonio-RiveroMartnez@users.noreply.github.com> Co-authored-by: Elizabeth Thompson <eschutho@gmail.com>
93 lines
4.0 KiB
Python
93 lines
4.0 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
import logging
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
from flask_appbuilder.models.sqla import Model
|
|
from marshmallow import ValidationError
|
|
|
|
from superset.commands.base import BaseCommand
|
|
from superset.dao.exceptions import DAOCreateFailedError
|
|
from superset.databases.ssh_tunnel.commands.exceptions import (
|
|
SSHTunnelCreateFailedError,
|
|
SSHTunnelInvalidError,
|
|
SSHTunnelRequiredFieldValidationError,
|
|
)
|
|
from superset.databases.ssh_tunnel.dao import SSHTunnelDAO
|
|
from superset.extensions import db, event_logger
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class CreateSSHTunnelCommand(BaseCommand):
|
|
def __init__(self, database_id: int, data: Dict[str, Any]):
|
|
self._properties = data.copy()
|
|
self._properties["database_id"] = database_id
|
|
|
|
def run(self) -> Model:
|
|
try:
|
|
# Start nested transaction since we are always creating the tunnel
|
|
# through a DB command (Create or Update). Without this, we cannot
|
|
# safely rollback changes to databases if any, i.e, things like
|
|
# test_do_not_create_database_if_ssh_tunnel_creation_fails test will fail
|
|
db.session.begin_nested()
|
|
self.validate()
|
|
tunnel = SSHTunnelDAO.create(self._properties, commit=False)
|
|
except DAOCreateFailedError as ex:
|
|
# Rollback nested transaction
|
|
db.session.rollback()
|
|
raise SSHTunnelCreateFailedError() from ex
|
|
except SSHTunnelInvalidError as ex:
|
|
# Rollback nested transaction
|
|
db.session.rollback()
|
|
raise ex
|
|
|
|
return tunnel
|
|
|
|
def validate(self) -> None:
|
|
# TODO(hughhh): check to make sure the server port is not localhost
|
|
# using the config.SSH_TUNNEL_MANAGER
|
|
exceptions: List[ValidationError] = []
|
|
database_id: Optional[int] = self._properties.get("database_id")
|
|
server_address: Optional[str] = self._properties.get("server_address")
|
|
server_port: Optional[int] = self._properties.get("server_port")
|
|
username: Optional[str] = self._properties.get("username")
|
|
private_key: Optional[str] = self._properties.get("private_key")
|
|
private_key_password: Optional[str] = self._properties.get(
|
|
"private_key_password"
|
|
)
|
|
if not database_id:
|
|
exceptions.append(SSHTunnelRequiredFieldValidationError("database_id"))
|
|
if not server_address:
|
|
exceptions.append(SSHTunnelRequiredFieldValidationError("server_address"))
|
|
if not server_port:
|
|
exceptions.append(SSHTunnelRequiredFieldValidationError("server_port"))
|
|
if not username:
|
|
exceptions.append(SSHTunnelRequiredFieldValidationError("username"))
|
|
if private_key_password and private_key is None:
|
|
exceptions.append(SSHTunnelRequiredFieldValidationError("private_key"))
|
|
if exceptions:
|
|
exception = SSHTunnelInvalidError()
|
|
exception.add_list(exceptions)
|
|
event_logger.log_with_context(
|
|
action="ssh_tunnel_creation_failed.{}.{}".format(
|
|
exception.__class__.__name__,
|
|
".".join(exception.get_list_classnames()),
|
|
)
|
|
)
|
|
raise exception
|