mirror of
https://github.com/apache/superset.git
synced 2026-04-12 12:47:53 +00:00
98 lines
3.9 KiB
Python
98 lines
3.9 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 functools import partial
|
|
from typing import Any, Optional
|
|
|
|
from flask_appbuilder.models.sqla import Model
|
|
from marshmallow import ValidationError
|
|
|
|
from superset.commands.base import BaseCommand
|
|
from superset.commands.database.ssh_tunnel.exceptions import (
|
|
SSHTunnelCreateFailedError,
|
|
SSHTunnelDatabasePortError,
|
|
SSHTunnelInvalidError,
|
|
SSHTunnelRequiredFieldValidationError,
|
|
)
|
|
from superset.daos.database import SSHTunnelDAO
|
|
from superset.databases.utils import make_url_safe
|
|
from superset.extensions import event_logger
|
|
from superset.models.core import Database
|
|
from superset.utils.decorators import on_error, transaction
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class CreateSSHTunnelCommand(BaseCommand):
|
|
_database: Database
|
|
|
|
def __init__(self, database: Database, data: dict[str, Any]):
|
|
self._properties = data.copy()
|
|
self._properties["database"] = database
|
|
self._database = database
|
|
|
|
@transaction(on_error=partial(on_error, reraise=SSHTunnelCreateFailedError))
|
|
def run(self) -> Model:
|
|
"""
|
|
Create an SSH tunnel.
|
|
|
|
:returns: The SSH tunnel model
|
|
:raises SSHTunnelCreateFailedError: If the model creation fails
|
|
:raises SSHTunnelInvalidError: If the configuration are invalid
|
|
"""
|
|
|
|
self.validate()
|
|
return SSHTunnelDAO.create(attributes=self._properties)
|
|
|
|
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] = []
|
|
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")
|
|
password: Optional[str] = self._properties.get("password")
|
|
private_key: Optional[str] = self._properties.get("private_key")
|
|
private_key_password: Optional[str] = self._properties.get(
|
|
"private_key_password"
|
|
)
|
|
url = make_url_safe(self._database.sqlalchemy_uri)
|
|
if not url.port:
|
|
raise SSHTunnelDatabasePortError()
|
|
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 not private_key and not password:
|
|
exceptions.append(SSHTunnelRequiredFieldValidationError("password"))
|
|
if private_key_password and private_key is None:
|
|
exceptions.append(SSHTunnelRequiredFieldValidationError("private_key"))
|
|
if exceptions:
|
|
exception = SSHTunnelInvalidError()
|
|
exception.extend(exceptions)
|
|
event_logger.log_with_context(
|
|
# pylint: disable=consider-using-f-string
|
|
action="ssh_tunnel_creation_failed.{}.{}".format(
|
|
exception.__class__.__name__,
|
|
".".join(exception.get_list_classnames()),
|
|
)
|
|
)
|
|
raise exception
|