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..34d949c 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("sample_datasets", "SELECT 1", 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")