superset2/superset/cli/update.py

# 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 json
import logging
import os
import sys
from typing import Optional

import click
from apispec import APISpec
from apispec.ext.marshmallow import MarshmallowPlugin
from flask import current_app
from flask.cli import with_appcontext
from flask_appbuilder import Model
from flask_appbuilder.api import BaseApi
from flask_appbuilder.api.manager import resolver

import superset.utils.database as database_utils
from superset.utils.decorators import transaction
from superset.utils.encrypt import ENCRYPTION_ENGINES, SecretsMigrator

logger = logging.getLogger(__name__)


@click.command()
@with_appcontext
@transaction()
@click.option("--database_name", "-d", required=True, help="Database name to change")
@click.option("--uri", "-u", required=True, help="Database URI to change")
@click.option(
    "--skip_create",
    "-s",
    is_flag=True,
    default=False,
    help="Don't 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"""
    database_utils.get_or_create_db(database_name, uri, not skip_create)


@click.command()
@with_appcontext
@transaction()
def sync_tags() -> None:
    """Rebuilds special tags (owner, type, favorited by)."""
    # pylint: disable=no-member
    metadata = Model.metadata

    # pylint: disable=import-outside-toplevel
    from superset.common.tags import add_favorites, add_owners, add_types

    add_types(metadata)
    add_owners(metadata)
    add_favorites(metadata)


@click.command()
@with_appcontext
def update_api_docs() -> None:
    """Regenerate the openapi.json file in docs"""
    superset_dir = os.path.abspath(os.path.dirname(__file__))
    openapi_json = os.path.join(
        superset_dir, "..", "..", "docs", "static", "resources", "openapi.json"
    )
    api_version = "v1"

    version_found = False
    api_spec = APISpec(
        title=current_app.appbuilder.app_name,
        version=api_version,
        openapi_version="3.0.2",
        info={"description": current_app.appbuilder.app_name},
        plugins=[MarshmallowPlugin(schema_name_resolver=resolver)],
        servers=[{"url": "http://localhost:8088"}],
    )
    for base_api in current_app.appbuilder.baseviews:
        if isinstance(base_api, BaseApi) and base_api.version == api_version:
            base_api.add_api_spec(api_spec)
            version_found = True
    if version_found:
        click.secho("Generating openapi.json", fg="green")
        with open(openapi_json, "w") as outfile:
            json.dump(api_spec.to_dict(), outfile, sort_keys=True, indent=2)
            outfile.write("\n")
    else:
        click.secho("API version not found", err=True)


@click.command()
@with_appcontext
@click.option(
    "--previous_secret_key",
    "-a",
    required=False,
    help="An optional previous secret key, if PREVIOUS_SECRET_KEY "
    "is not set on the config",
)
@click.option(
    "--engine",
    "-e",
    "target_engine_name",
    required=False,
    type=click.Choice(sorted(ENCRYPTION_ENGINES), case_sensitive=False),
    help="Re-encrypt all app-encrypted fields with this encryption engine "
    "(e.g. 'aes-gcm' for authenticated encryption). The SECRET_KEY is "
    "unchanged. Take a metadata-DB backup first, then set "
    "SQLALCHEMY_ENCRYPTED_FIELD_ENGINE to the same value and restart.",
)
def re_encrypt_secrets(
    previous_secret_key: Optional[str] = None,
    target_engine_name: Optional[str] = None,
) -> None:
    """Re-encrypt every app-encrypted field via :class:`SecretsMigrator`.

    Supports key rotation (``previous_secret_key``, falling back to the
    ``PREVIOUS_SECRET_KEY`` config) and engine migration (``target_engine_name``,
    a case-insensitive ``ENCRYPTION_ENGINES`` key such as ``aes-gcm``); the two
    can combine. With neither provided the command is a no-op. Exits non-zero on
    failure.
    """
    previous_secret_key = previous_secret_key or current_app.config.get(
        "PREVIOUS_SECRET_KEY"
    )
    target_engine = (
        ENCRYPTION_ENGINES[target_engine_name] if target_engine_name else None
    )
    if previous_secret_key is None and target_engine is None:
        click.secho(
            "No previous secret key or target engine provided; nothing to re-encrypt.",
            fg="yellow",
        )
        return
    secrets_migrator = SecretsMigrator(
        previous_secret_key=previous_secret_key,
        target_engine=target_engine,
    )
    try:
        stats = secrets_migrator.run()
    except Exception as exc:  # pylint: disable=broad-except
        click.secho(f"Re-encryption failed: {exc}", err=True)
        sys.exit(1)
    click.secho(
        f"Re-encryption complete: {stats.re_encrypted} re-encrypted, "
        f"{stats.skipped} skipped, {stats.null} null, {stats.failed} failed.",
        fg="green",
    )