fix: allow db driver distinction on enforced URI params (#23769)

This commit is contained in:
Daniel Vaz Gaspar
2023-04-23 15:44:21 +01:00
committed by Elizabeth Thompson
parent 2f3471a87e
commit b26901cb05
9 changed files with 102 additions and 33 deletions

View File

@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
from datetime import datetime
from typing import Any, Dict, Optional
from typing import Any, Dict, Optional, Tuple
from urllib import parse
from sqlalchemy import types
@@ -69,11 +69,16 @@ class DrillEngineSpec(BaseEngineSpec):
return None
@classmethod
def adjust_database_uri(cls, uri: URL, selected_schema: Optional[str]) -> URL:
def adjust_database_uri(
cls,
uri: URL,
connect_args: Dict[str, Any],
selected_schema: Optional[str] = None,
) -> Tuple[URL, Dict[str, Any]]:
if selected_schema:
uri = uri.set(database=parse.quote(selected_schema, safe=""))
return uri
return uri, connect_args
@classmethod
def get_url_for_impersonation(