Skip to content

Assign port, http schema automatically based on passed hostname #340

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Mar 3, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,8 @@ $ pip install trino

Use the DBAPI interface to query Trino:

if `host` is a valid url, the port and http schema will be automatically determined. For example `https://my-trino-server:9999` will assign the `http_schema` property to `https` and port to `9999`.

```python
from trino.dbapi import connect

Expand Down
39 changes: 38 additions & 1 deletion tests/unit/test_dbapi.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@
)
from trino import constants
from trino.auth import OAuth2Authentication
from trino.dbapi import connect
from trino.dbapi import Connection, connect


@patch("trino.dbapi.trino.client")
Expand Down Expand Up @@ -272,3 +272,40 @@ def test_role_is_set_when_specified(mock_client):

_, passed_role = mock_client.ClientSession.call_args
assert passed_role["roles"] == roles


def test_hostname_parsing():
https_server_with_port = Connection("https://mytrinoserver.domain:9999")
assert https_server_with_port.host == "mytrinoserver.domain"
assert https_server_with_port.port == 9999
assert https_server_with_port.http_scheme == constants.HTTPS

https_server_without_port = Connection("https://mytrinoserver.domain")
assert https_server_without_port.host == "mytrinoserver.domain"
assert https_server_without_port.port == 8080
assert https_server_without_port.http_scheme == constants.HTTPS

http_server_with_port = Connection("http://mytrinoserver.domain:9999")
assert http_server_with_port.host == "mytrinoserver.domain"
assert http_server_with_port.port == 9999
assert http_server_with_port.http_scheme == constants.HTTP

http_server_without_port = Connection("http://mytrinoserver.domain")
assert http_server_without_port.host == "mytrinoserver.domain"
assert http_server_without_port.port == 8080
assert http_server_without_port.http_scheme == constants.HTTP

http_server_with_path = Connection("http://mytrinoserver.domain/some_path")
assert http_server_with_path.host == "mytrinoserver.domain/some_path"
assert http_server_with_path.port == 8080
assert http_server_with_path.http_scheme == constants.HTTP

only_hostname = Connection("mytrinoserver.domain")
assert only_hostname.host == "mytrinoserver.domain"
assert only_hostname.port == 8080
assert only_hostname.http_scheme == constants.HTTP

only_hostname_with_path = Connection("mytrinoserver.domain/some_path")
assert only_hostname_with_path.host == "mytrinoserver.domain/some_path"
assert only_hostname_with_path.port == 8080
assert only_hostname_with_path.http_scheme == constants.HTTP
12 changes: 8 additions & 4 deletions trino/dbapi.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
import uuid
from decimal import Decimal
from typing import Any, Dict, List, NamedTuple, Optional # NOQA for mypy types
from urllib.parse import urlparse

import trino.client
import trino.exceptions
Expand Down Expand Up @@ -92,7 +93,7 @@ class Connection(object):

def __init__(
self,
host,
host: str,
port=constants.DEFAULT_PORT,
user=None,
source=constants.DEFAULT_SOURCE,
Expand All @@ -114,8 +115,11 @@ def __init__(
roles=None,
timezone=None,
):
self.host = host
self.port = port
# Automatically assign http_schema, port based on hostname
parsed_host = urlparse(host, allow_fragments=False)

self.host = host if parsed_host.hostname is None else parsed_host.hostname + parsed_host.path
self.port = port if parsed_host.port is None else parsed_host.port
self.user = user
self.source = source
self.catalog = catalog
Expand All @@ -141,7 +145,7 @@ def __init__(
else:
self._http_session = http_session
self.http_headers = http_headers
self.http_scheme = http_scheme
self.http_scheme = http_scheme if not parsed_host.scheme else parsed_host.scheme
self.auth = auth
self.extra_credential = extra_credential
self.redirect_handler = redirect_handler
Expand Down