Skip to content

Commit

Permalink
Address review comments.
Browse files Browse the repository at this point in the history
  • Loading branch information
jnturton committed Mar 31, 2022
1 parent fd6ab00 commit 488bd6e
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 4 deletions.
8 changes: 4 additions & 4 deletions superset/db_engine_specs/drill.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,18 +14,16 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
import logging
from datetime import datetime
from typing import Any, Dict, Optional
from urllib import parse

from sqlalchemy.engine.url import URL

from superset.db_engine_specs.exceptions import SupersetDBAPIProgrammingError
from superset.db_engine_specs.base import BaseEngineSpec
from superset.utils import core as utils

logger = logging.getLogger(__name__)


class DrillEngineSpec(BaseEngineSpec):
"""Engine spec for Apache Drill"""
Expand Down Expand Up @@ -90,4 +88,6 @@ def modify_url_for_impersonation(
elif url.drivername in ["drill+sadrill", "drill+jdbc"]:
url.query["impersonation_target"] = username
else:
logger.warning("impersonation is not supported for %s", url.drivername)
raise SupersetDBAPIProgrammingError(
f"impersonation is not supported for {url.drivername}"
)
24 changes: 24 additions & 0 deletions tests/unit_tests/db_engine_specs/test_drill.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
# pylint: disable=unused-argument, import-outside-toplevel, protected-access

from flask.ctx import AppContext
from superset.db_engine_specs.exceptions import SupersetDBAPIProgrammingError


def test_odbc_impersonation(app_context: AppContext) -> None:
Expand Down Expand Up @@ -65,3 +66,26 @@ def test_sadrill_impersonation(app_context: AppContext) -> None:
username = "DoAsUser"
DrillEngineSpec.modify_url_for_impersonation(url, True, username)
assert url.query["impersonation_target"] == username


def test_invalid_impersonation(app_context: AppContext) -> None:
"""
Test ``modify_url_for_impersonation`` method when driver == foobar.
The method raises an exception because impersonation is not supported
for drill+foobar.
"""
from sqlalchemy.engine.url import URL

from superset.db_engine_specs.drill import DrillEngineSpec

url = URL("drill+foobar")
username = "DoAsUser"
exception_type = None

try:
DrillEngineSpec.modify_url_for_impersonation(url, True, username)
except Exception as e:
exception_type = type(e)

assert exception_type == SupersetDBAPIProgrammingError

0 comments on commit 488bd6e

Please sign in to comment.