Skip to content

Commit 0584c93

Browse files
aksakallihashhar
authored andcommitted
add Trino Python Client version to user agent
1 parent 31febef commit 0584c93

File tree

3 files changed

+9
-3
lines changed

3 files changed

+9
-3
lines changed

tests/unit/test_client.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@
3838
_get_token_requests,
3939
_post_statement_requests,
4040
)
41-
from trino import constants
41+
from trino import __version__, constants
4242
from trino.auth import KerberosAuthentication, _OAuth2TokenBearer
4343
from trino.client import (
4444
ClientSession,
@@ -125,6 +125,7 @@ def assert_headers(headers):
125125
assert headers[constants.HEADER_SOURCE] == source
126126
assert headers[constants.HEADER_USER] == user
127127
assert headers[constants.HEADER_SESSION] == ""
128+
assert headers[constants.HEADER_TRANSACTION] is None
128129
assert headers[constants.HEADER_TIMEZONE] == timezone
129130
assert headers[constants.HEADER_CLIENT_CAPABILITIES] == "PARAMETRIC_DATETIME"
130131
assert headers[accept_encoding_header] == accept_encoding_value
@@ -135,7 +136,8 @@ def assert_headers(headers):
135136
"catalog1=NONE,"
136137
"catalog2=" + urllib.parse.quote("ROLE{catalog2_role}")
137138
)
138-
assert len(headers.keys()) == 11
139+
assert headers["User-Agent"] == f"{constants.CLIENT_NAME}/{__version__}"
140+
assert len(headers.keys()) == 12
139141

140142
req.post("URL")
141143
_, post_kwargs = post.call_args

trino/client.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,7 @@
6262

6363
import trino.logging
6464
from trino import constants, exceptions
65+
from trino._version import __version__
6566

6667
__all__ = ["ClientSession", "TrinoQuery", "TrinoRequest", "PROXIES"]
6768

@@ -447,14 +448,15 @@ def transaction_id(self, value):
447448

448449
@property
449450
def http_headers(self) -> Dict[str, str]:
450-
headers = {}
451+
headers = requests.structures.CaseInsensitiveDict()
451452

452453
headers[constants.HEADER_CATALOG] = self._client_session.catalog
453454
headers[constants.HEADER_SCHEMA] = self._client_session.schema
454455
headers[constants.HEADER_SOURCE] = self._client_session.source
455456
headers[constants.HEADER_USER] = self._client_session.user
456457
headers[constants.HEADER_TIMEZONE] = self._client_session.timezone
457458
headers[constants.HEADER_CLIENT_CAPABILITIES] = 'PARAMETRIC_DATETIME'
459+
headers["user-agent"] = f"{constants.CLIENT_NAME}/{__version__}"
458460
if len(self._client_session.roles.values()):
459461
headers[constants.HEADER_ROLE] = ",".join(
460462
# ``name`` must not contain ``=``

trino/constants.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,8 @@
2626

2727
URL_STATEMENT_PATH = "/v1/statement"
2828

29+
CLIENT_NAME = "Trino Python Client"
30+
2931
HEADER_CATALOG = "X-Trino-Catalog"
3032
HEADER_SCHEMA = "X-Trino-Schema"
3133
HEADER_SOURCE = "X-Trino-Source"

0 commit comments

Comments
 (0)