mirror of
https://github.com/apache/superset.git
synced 2026-09-01 13:01:33 +00:00
251 lines
7.9 KiB
Python
251 lines
7.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.
|
|
|
|
from __future__ import annotations
|
|
|
|
import sqlite3
|
|
from types import SimpleNamespace
|
|
from typing import Any
|
|
|
|
import pytest
|
|
from pytest_mock import MockerFixture
|
|
from sqlalchemy import create_engine, event, text
|
|
from sqlalchemy.engine import Engine
|
|
from sqlalchemy.exc import OperationalError
|
|
|
|
from superset.databases.error_provenance import (
|
|
is_database_engine_error,
|
|
mark_database_engine_error,
|
|
)
|
|
from superset.models.core import Database
|
|
|
|
|
|
def _database_engine(mocker: MockerFixture, **kwargs: Any) -> Engine:
|
|
real_create_engine = create_engine
|
|
mocker.patch(
|
|
"superset.models.core.create_engine",
|
|
side_effect=lambda url, **engine_kwargs: real_create_engine(
|
|
url, **engine_kwargs, **kwargs
|
|
),
|
|
)
|
|
return Database(
|
|
database_name="database",
|
|
sqlalchemy_uri="sqlite://",
|
|
)._get_sqla_engine(nullpool=False)
|
|
|
|
|
|
def test_connect_error_is_marked_without_changing_exception(
|
|
app_context: None,
|
|
mocker: MockerFixture,
|
|
) -> None:
|
|
original = sqlite3.OperationalError("connect failed")
|
|
|
|
def fail_to_connect() -> None:
|
|
raise original
|
|
|
|
engine = _database_engine(mocker, creator=fail_to_connect)
|
|
observed: list[BaseException] = []
|
|
|
|
def observe(context: Any) -> None:
|
|
observed.append(context.sqlalchemy_exception)
|
|
|
|
event.listen(engine, "handle_error", observe)
|
|
|
|
with pytest.raises(OperationalError) as exc_info:
|
|
engine.connect()
|
|
|
|
assert observed == [exc_info.value]
|
|
assert observed[0] is exc_info.value
|
|
assert is_database_engine_error(exc_info.value)
|
|
assert exc_info.value.orig is original
|
|
assert str(exc_info.value.orig) == "connect failed"
|
|
|
|
|
|
def test_sqlalchemy_core_error_is_marked_without_changing_exception(
|
|
app_context: None,
|
|
mocker: MockerFixture,
|
|
) -> None:
|
|
engine = _database_engine(mocker)
|
|
observed: list[BaseException] = []
|
|
|
|
def observe(context: Any) -> None:
|
|
observed.append(context.sqlalchemy_exception)
|
|
|
|
event.listen(engine, "handle_error", observe)
|
|
|
|
with engine.connect() as connection, pytest.raises(OperationalError) as exc_info:
|
|
connection.execute(text("SELECT * FROM table_that_does_not_exist"))
|
|
|
|
assert observed == [exc_info.value]
|
|
assert observed[0] is exc_info.value
|
|
assert is_database_engine_error(exc_info.value)
|
|
assert isinstance(exc_info.value.orig, sqlite3.OperationalError)
|
|
assert "table_that_does_not_exist" in str(exc_info.value)
|
|
|
|
|
|
def test_raw_dbapi_cursor_error_is_outside_listener_scope(
|
|
app_context: None,
|
|
mocker: MockerFixture,
|
|
) -> None:
|
|
database = Database(database_name="database", sqlalchemy_uri="sqlite://")
|
|
real_create_engine = create_engine
|
|
mocker.patch(
|
|
"superset.models.core.create_engine",
|
|
side_effect=lambda url, **engine_kwargs: real_create_engine(
|
|
url, **engine_kwargs
|
|
),
|
|
)
|
|
|
|
with (
|
|
database.get_raw_connection() as connection,
|
|
pytest.raises(sqlite3.OperationalError) as exc_info,
|
|
):
|
|
connection.cursor().execute("SELECT * FROM table_that_does_not_exist")
|
|
|
|
assert not is_database_engine_error(exc_info.value)
|
|
|
|
|
|
def test_unrelated_and_metadata_engine_errors_are_not_marked() -> None:
|
|
assert not is_database_engine_error(ValueError("unrelated"))
|
|
|
|
metadata_engine = create_engine("sqlite://")
|
|
with (
|
|
metadata_engine.connect() as connection,
|
|
pytest.raises(OperationalError) as exc_info,
|
|
):
|
|
connection.execute(text("SELECT * FROM table_that_does_not_exist"))
|
|
|
|
assert not is_database_engine_error(exc_info.value)
|
|
|
|
|
|
def test_public_engine_context_preserves_marker_for_private_engine(
|
|
app_context: None,
|
|
mocker: MockerFixture,
|
|
) -> None:
|
|
original = sqlite3.OperationalError("public path connect failed")
|
|
|
|
def fail_to_connect() -> None:
|
|
raise original
|
|
|
|
real_create_engine = create_engine
|
|
mocker.patch(
|
|
"superset.models.core.create_engine",
|
|
side_effect=lambda url, **engine_kwargs: real_create_engine(
|
|
url,
|
|
**engine_kwargs,
|
|
creator=fail_to_connect,
|
|
),
|
|
)
|
|
database = Database(database_name="database", sqlalchemy_uri="sqlite://")
|
|
mocker.patch.object(
|
|
database.db_engine_spec,
|
|
"get_prequeries",
|
|
return_value=["SELECT 1"],
|
|
)
|
|
|
|
with pytest.raises(OperationalError) as exc_info:
|
|
with database.get_sqla_engine() as engine:
|
|
engine.connect()
|
|
|
|
assert exc_info.value.orig is original
|
|
assert is_database_engine_error(exc_info.value)
|
|
|
|
|
|
def test_cached_engine_has_one_instance_listener_without_database_closure(
|
|
app_context: None,
|
|
mocker: MockerFixture,
|
|
) -> None:
|
|
from superset.models.core import _ENGINE_CACHE
|
|
|
|
_ENGINE_CACHE.clear()
|
|
try:
|
|
listen = mocker.spy(__import__("sqlalchemy").event, "listen")
|
|
database = Database(database_name="database", sqlalchemy_uri="sqlite://")
|
|
database.id = 1
|
|
|
|
first = database._get_sqla_engine(nullpool=False)
|
|
second = database._get_sqla_engine(nullpool=False)
|
|
|
|
assert first is second
|
|
handle_error_calls = [
|
|
call
|
|
for call in listen.call_args_list
|
|
if len(call.args) > 1 and call.args[1] == "handle_error"
|
|
]
|
|
assert len(handle_error_calls) == 1
|
|
callback = handle_error_calls[0].args[2]
|
|
assert callback is mark_database_engine_error
|
|
finally:
|
|
_ENGINE_CACHE.clear()
|
|
|
|
|
|
def test_marker_assignment_failure_does_not_mask_error() -> None:
|
|
class UnmarkableError(Exception):
|
|
def __setattr__(self, name: str, value: object) -> None:
|
|
raise AttributeError(name)
|
|
|
|
unmarkable = UnmarkableError("original")
|
|
context = SimpleNamespace(sqlalchemy_exception=unmarkable)
|
|
|
|
mark_database_engine_error(context)
|
|
assert context.sqlalchemy_exception is unmarkable
|
|
assert not is_database_engine_error(unmarkable)
|
|
|
|
|
|
def test_marker_helper_marks_exception() -> None:
|
|
exception = RuntimeError("database error")
|
|
context = SimpleNamespace(sqlalchemy_exception=exception)
|
|
|
|
mark_database_engine_error(context)
|
|
|
|
assert is_database_engine_error(exception)
|
|
|
|
|
|
def test_original_exception_is_marked_when_sqlalchemy_exception_is_none() -> None:
|
|
exception = UnicodeDecodeError("utf-8", b"\xff", 0, 1, "invalid byte")
|
|
context = SimpleNamespace(
|
|
sqlalchemy_exception=None,
|
|
original_exception=exception,
|
|
)
|
|
|
|
mark_database_engine_error(context)
|
|
|
|
assert is_database_engine_error(exception)
|
|
|
|
|
|
def test_later_listener_replacement_is_not_marked(
|
|
app_context: None,
|
|
mocker: MockerFixture,
|
|
) -> None:
|
|
engine = _database_engine(mocker)
|
|
original: list[BaseException] = []
|
|
|
|
def replace(context: Any) -> RuntimeError:
|
|
original.append(context.sqlalchemy_exception)
|
|
return RuntimeError("replacement")
|
|
|
|
event.listen(engine, "handle_error", replace, retval=True)
|
|
|
|
with (
|
|
engine.connect() as connection,
|
|
pytest.raises(RuntimeError, match="replacement") as exc_info,
|
|
):
|
|
connection.execute(text("SELECT * FROM table_that_does_not_exist"))
|
|
|
|
assert is_database_engine_error(original[0])
|
|
assert not is_database_engine_error(exc_info.value)
|