fix(Trino): create PrestoBaseEngineSpec base class to share common code between Trino and Presto (#21066)

* chore: create `PrestoBaseEngineSpec` class that share common functions between Presto and Trino

Signed-off-by: Đặng Minh Dũng <dungdm93@live.com>

* feat(Trino): support CertificateAuthentication

* chore(Presto): move `get_function_names` to `PrestoBaseEngineSpec`

Signed-off-by: Đặng Minh Dũng <dungdm93@live.com>

* chores(Presto): remove `is_readonly_query`

* feat(Trino): implement `extra_table_metadata`

* feat(Trino): specify `User-Agent`

Signed-off-by: Đặng Minh Dũng <dungdm93@live.com>

* fix: pylint

Signed-off-by: Đặng Minh Dũng <dungdm93@live.com>

* chores(Presto): move `PrestoBaseEngineSpec` to `presto.py`

Signed-off-by: Đặng Minh Dũng <dungdm93@live.com>

* fix(Presto): typing annotations

Signed-off-by: Đặng Minh Dũng <dungdm93@live.com>

Signed-off-by: Đặng Minh Dũng <dungdm93@live.com>
This commit is contained in:
Đặng Minh Dũng
2022-08-29 14:21:28 +07:00
committed by GitHub
parent 5f76ac9001
commit ccb293a083
7 changed files with 233 additions and 205 deletions

View File

@@ -19,9 +19,9 @@ from typing import Any, Dict
from unittest.mock import Mock, patch
import pytest
from sqlalchemy.engine.url import URL
import superset.config
from superset.constants import USER_AGENT
from superset.db_engine_specs.trino import TrinoEngineSpec
from tests.integration_tests.db_engine_specs.base_tests import TestDbEngineSpec
@@ -33,12 +33,15 @@ class TestTrinoDbEngineSpec(TestDbEngineSpec):
database.extra = json.dumps({})
database.server_cert = None
extra = TrinoEngineSpec.get_extra_params(database)
expected = {"engine_params": {"connect_args": {}}}
expected = {"engine_params": {"connect_args": {"source": USER_AGENT}}}
self.assertEqual(extra, expected)
expected = {
"first": 1,
"engine_params": {"second": "two", "connect_args": {"third": "three"}},
"engine_params": {
"second": "two",
"connect_args": {"source": "foobar", "third": "three"},
},
}
database.extra = json.dumps(expected)
database.server_cert = None
@@ -93,6 +96,21 @@ class TestTrinoDbEngineSpec(TestDbEngineSpec):
self.assertEqual(connect_args.get("http_scheme"), "https")
auth.assert_called_once_with(**auth_params)
@patch("trino.auth.CertificateAuthentication")
def test_auth_certificate(self, auth: Mock):
database = Mock()
auth_params = {"cert": "/path/to/cert.pem", "key": "/path/to/key.pem"}
database.encrypted_extra = json.dumps(
{"auth_method": "certificate", "auth_params": auth_params}
)
params: Dict[str, Any] = {}
TrinoEngineSpec.update_encrypted_extra_params(database, params)
connect_args = params.setdefault("connect_args", {})
self.assertEqual(connect_args.get("http_scheme"), "https")
auth.assert_called_once_with(**auth_params)
@patch("trino.auth.JWTAuthentication")
def test_auth_jwt(self, auth: Mock):
database = Mock()