From cf9bc410107c1e828ec13f2dcf4805a99833f25a Mon Sep 17 00:00:00 2001 From: Vu Tung Date: Mon, 16 Jun 2025 18:14:46 +0700 Subject: [PATCH 1/2] Support Trino query type --- .gitignore | 1 + tdclient/client.py | 2 +- tdclient/job_api.py | 2 +- tdclient/test/client_test.py | 14 ++++++++++++++ 4 files changed, 17 insertions(+), 2 deletions(-) diff --git a/.gitignore b/.gitignore index 648ca83..3022a30 100644 --- a/.gitignore +++ b/.gitignore @@ -8,6 +8,7 @@ /.eggs /.envrc /.tox +/.venv /build /dist /tmp diff --git a/tdclient/client.py b/tdclient/client.py index 54280e5..30b93cc 100644 --- a/tdclient/client.py +++ b/tdclient/client.py @@ -248,7 +248,7 @@ def query( ValueError: if unknown query type has been specified """ # for compatibility, assume type is hive unless specifically specified - if type not in ["hive", "pig", "impala", "presto"]: + if type not in ["hive", "pig", "impala", "presto", "trino"]: raise ValueError("The specified query type is not supported: %s" % (type)) job_id = self.api.query( q, diff --git a/tdclient/job_api.py b/tdclient/job_api.py index 94dd85c..bb06e10 100644 --- a/tdclient/job_api.py +++ b/tdclient/job_api.py @@ -390,7 +390,7 @@ def query( Args: q (str): Query string. - type (str): Query type. `hive`, `presto`, `bulkload`. Default: `hive` + type (str): Query type. `hive`, `presto`, `trino`, `bulkload`. Default: `hive` db (str): Database name. result_url (str): Result output URL. e.g., ``postgresql://:@://`` diff --git a/tdclient/test/client_test.py b/tdclient/test/client_test.py index a407b04..b3b6332 100644 --- a/tdclient/test/client_test.py +++ b/tdclient/test/client_test.py @@ -197,6 +197,20 @@ def test_query(): ) assert job.job_id == "12345" +def test_trino_query(): + td = client.Client("APIKEY") + td._api = mock.MagicMock() + td._api.query = mock.MagicMock(return_value=("12345")) + job = td.query("tung_db", "SELECT * LIMIT 1 from tung_table", type="trino") + td.api.query.assert_called_with( + "SELECT 1", + db="sample_datasets", + type="trino", + retry_limit=None, + priority=None, + result_url=None, + ) + assert job.job_id == "12345" def test_jobs(): td = client.Client("APIKEY") From d99ac638997a9849042588568729c9e4156addcf Mon Sep 17 00:00:00 2001 From: Vu Tung Date: Wed, 18 Jun 2025 10:11:52 +0700 Subject: [PATCH 2/2] Correct mock query call --- tdclient/test/client_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tdclient/test/client_test.py b/tdclient/test/client_test.py index b3b6332..34d949c 100644 --- a/tdclient/test/client_test.py +++ b/tdclient/test/client_test.py @@ -201,7 +201,7 @@ def test_trino_query(): td = client.Client("APIKEY") td._api = mock.MagicMock() td._api.query = mock.MagicMock(return_value=("12345")) - job = td.query("tung_db", "SELECT * LIMIT 1 from tung_table", type="trino") + job = td.query("sample_datasets", "SELECT 1", type="trino") td.api.query.assert_called_with( "SELECT 1", db="sample_datasets",