diff --git a/IMPLEMENTATION_COVERAGE.md b/IMPLEMENTATION_COVERAGE.md
index 26da79baaf06..65005dede18e 100644
--- a/IMPLEMENTATION_COVERAGE.md
+++ b/IMPLEMENTATION_COVERAGE.md
@@ -8481,6 +8481,27 @@
- [ ] update_adapter
+## timestream-query
+
+40% implemented
+
+- [ ] cancel_query
+- [X] create_scheduled_query
+- [X] delete_scheduled_query
+- [ ] describe_account_settings
+- [X] describe_endpoints
+- [X] describe_scheduled_query
+- [ ] execute_scheduled_query
+- [ ] list_scheduled_queries
+- [ ] list_tags_for_resource
+- [ ] prepare_query
+- [X] query
+- [ ] tag_resource
+- [ ] untag_resource
+- [ ] update_account_settings
+- [X] update_scheduled_query
+
+
## timestream-write
78% implemented
@@ -9010,7 +9031,6 @@
- synthetics
- taxsettings
- timestream-influxdb
-- timestream-query
- tnb
- translate
- trustedadvisor
diff --git a/docs/docs/services/timestream-query.rst b/docs/docs/services/timestream-query.rst
new file mode 100644
index 000000000000..52ed9e8d8cde
--- /dev/null
+++ b/docs/docs/services/timestream-query.rst
@@ -0,0 +1,80 @@
+.. _implementedservice_timestream-query:
+
+.. |start-h3| raw:: html
+
+
+
+.. |end-h3| raw:: html
+
+
+
+================
+timestream-query
+================
+
+.. autoclass:: moto.timestreamquery.models.TimestreamQueryBackend
+
+|start-h3| Implemented features for this service |end-h3|
+
+- [ ] cancel_query
+- [X] create_scheduled_query
+- [X] delete_scheduled_query
+- [ ] describe_account_settings
+- [X] describe_endpoints
+- [X] describe_scheduled_query
+- [ ] execute_scheduled_query
+- [ ] list_scheduled_queries
+- [ ] list_tags_for_resource
+- [ ] prepare_query
+- [X] query
+
+ Moto does not have a builtin time-series Database, so calling this endpoint will return zero results by default.
+
+ You can use a dedicated API to configuring a queue of expected results.
+
+ An example invocation looks like this:
+
+ .. sourcecode:: python
+
+ first_result = {
+ 'QueryId': 'some_id',
+ 'Rows': [...],
+ 'ColumnInfo': [...],
+ 'QueryStatus': ...
+ }
+ result_for_unknown_query_string = {
+ 'QueryId': 'unknown',
+ 'Rows': [...],
+ 'ColumnInfo': [...],
+ 'QueryStatus': ...
+ }
+ expected_results = {
+ "account_id": "123456789012", # This is the default - can be omitted
+ "region": "us-east-1", # This is the default - can be omitted
+ "results": {
+ # Use the exact querystring, and a list of results for it
+ # For example
+ "SELECT data FROM mytable": [first_result, ...],
+ # Use None if the exact querystring is unknown/irrelevant
+ None: [result_for_unknown_query_string, ...],
+ }
+ }
+ requests.post(
+ "http://motoapi.amazonaws.com/moto-api/static/timestream/query-results",
+ json=expected_results,
+ )
+
+ When calling `query(QueryString='SELECT data FROM mytable')`, the `first_result` will be returned.
+ Call the query again for the second result, and so on.
+
+ If you don't know the exact query strings, use the `None`-key. In the above example, when calling `SELECT something FROM unknown`, there are no results for that specific query, so `result_for_unknown_query_string` will be returned.
+
+ Results for unknown queries are cached, so calling `SELECT something FROM unknown` will return the same result.
+
+
+
+- [ ] tag_resource
+- [ ] untag_resource
+- [ ] update_account_settings
+- [X] update_scheduled_query
+
diff --git a/moto/backend_index.py b/moto/backend_index.py
index b7d561b888e9..027edb94bd72 100644
--- a/moto/backend_index.py
+++ b/moto/backend_index.py
@@ -196,6 +196,10 @@
("support", re.compile("https?://support\\.(.+)\\.amazonaws\\.com")),
("swf", re.compile("https?://swf\\.(.+)\\.amazonaws\\.com")),
("textract", re.compile("https?://textract\\.(.+)\\.amazonaws\\.com")),
+ (
+ "timestreamquery",
+ re.compile("https?://query.timestream\\.(.+)\\.amazonaws\\.com"),
+ ),
(
"timestreamwrite",
re.compile("https?://ingest\\.timestream\\.(.+)\\.amazonaws\\.com"),
diff --git a/moto/backends.py b/moto/backends.py
index 2db397485d0f..15d8242aea34 100644
--- a/moto/backends.py
+++ b/moto/backends.py
@@ -141,6 +141,7 @@
from moto.support.models import SupportBackend
from moto.swf.models import SWFBackend
from moto.textract.models import TextractBackend
+ from moto.timestreamquery.models import TimestreamQueryBackend
from moto.timestreamwrite.models import TimestreamWriteBackend
from moto.transcribe.models import TranscribeBackend
from moto.transfer.models import TransferBackend
@@ -318,6 +319,7 @@ def get_service_from_url(url: str) -> Optional[str]:
"Literal['support']",
"Literal['swf']",
"Literal['textract']",
+ "Literal['timestream-query']",
"Literal['timestream-write']",
"Literal['transcribe']",
"Literal['transfer']",
@@ -710,6 +712,10 @@ def get_backend(name: "Literal['swf']") -> "BackendDict[SWFBackend]": ...
@overload
def get_backend(name: "Literal['textract']") -> "BackendDict[TextractBackend]": ...
@overload
+def get_backend(
+ name: "Literal['timestream-query']",
+) -> "BackendDict[TimestreamQueryBackend]": ...
+@overload
def get_backend(
name: "Literal['timestream-write']",
) -> "BackendDict[TimestreamWriteBackend]": ...
diff --git a/moto/moto_api/_internal/models.py b/moto/moto_api/_internal/models.py
index 70dc10f2a6b7..ba5241aa44e2 100644
--- a/moto/moto_api/_internal/models.py
+++ b/moto/moto_api/_internal/models.py
@@ -118,6 +118,23 @@ def set_inspector2_findings_result(
backend = inspector2_backends[account_id][region]
backend.findings_queue.append(results)
+ def set_timestream_result(
+ self,
+ query: Optional[str],
+ query_results: List[Dict[str, Any]],
+ account_id: str,
+ region: str,
+ ) -> None:
+ from moto.timestreamquery.models import (
+ TimestreamQueryBackend,
+ timestreamquery_backends,
+ )
+
+ backend: TimestreamQueryBackend = timestreamquery_backends[account_id][region]
+ if query not in backend.query_result_queue:
+ backend.query_result_queue[query] = []
+ backend.query_result_queue[query].extend(query_results)
+
def get_proxy_passthrough(self) -> Tuple[Set[str], Set[str]]:
return self.proxy_urls_to_passthrough, self.proxy_hosts_to_passthrough
diff --git a/moto/moto_api/_internal/responses.py b/moto/moto_api/_internal/responses.py
index f89899b44fbe..fdd8735771bd 100644
--- a/moto/moto_api/_internal/responses.py
+++ b/moto/moto_api/_internal/responses.py
@@ -291,6 +291,28 @@ def set_inspector2_findings_result(
)
return 201, {}, ""
+ def set_timestream_result(
+ self,
+ request: Any,
+ full_url: str, # pylint: disable=unused-argument
+ headers: Any,
+ ) -> TYPE_RESPONSE:
+ from .models import moto_api_backend
+
+ body = self._get_body(headers, request)
+ account_id = body.get("account_id", DEFAULT_ACCOUNT_ID)
+ region = body.get("region", "us-east-1")
+ results = body.get("results", {})
+
+ for query in results:
+ moto_api_backend.set_timestream_result(
+ query=None if query == "null" else query,
+ query_results=results[query],
+ account_id=account_id,
+ region=region,
+ )
+ return 201, {}, ""
+
def set_proxy_passthrough(
self,
request: Any,
diff --git a/moto/moto_api/_internal/urls.py b/moto/moto_api/_internal/urls.py
index 2195f2ee5057..7821f11b5c08 100644
--- a/moto/moto_api/_internal/urls.py
+++ b/moto/moto_api/_internal/urls.py
@@ -20,6 +20,7 @@
"{0}/moto-api/static/lambda-simple/response": response_instance.set_lambda_simple_result,
"{0}/moto-api/static/resilience-hub-assessments/response": response_instance.set_resilience_result,
"{0}/moto-api/static/sagemaker/endpoint-results": response_instance.set_sagemaker_result,
+ "{0}/moto-api/static/timestream/query-results": response_instance.set_timestream_result,
"{0}/moto-api/static/rds-data/statement-results": response_instance.set_rds_data_result,
"{0}/moto-api/state-manager/get-transition": response_instance.get_transition,
"{0}/moto-api/state-manager/set-transition": response_instance.set_transition,
diff --git a/moto/timestreamquery/__init__.py b/moto/timestreamquery/__init__.py
new file mode 100644
index 000000000000..05ba2ffefb2a
--- /dev/null
+++ b/moto/timestreamquery/__init__.py
@@ -0,0 +1,3 @@
+from .models import timestreamquery_backends # noqa: F401
+
+# Responses are handled by TimestreamWrite
diff --git a/moto/timestreamquery/exceptions.py b/moto/timestreamquery/exceptions.py
new file mode 100644
index 000000000000..d897e82c8766
--- /dev/null
+++ b/moto/timestreamquery/exceptions.py
@@ -0,0 +1,11 @@
+"""Exceptions raised by the timestreamquery service."""
+
+from moto.core.exceptions import JsonRESTError
+
+
+class ResourceNotFound(JsonRESTError):
+ def __init__(self, arn: str):
+ super().__init__(
+ error_type="ResourceNotFoundException",
+ message=f"The resource with arn {arn} does not exist.",
+ )
diff --git a/moto/timestreamquery/models.py b/moto/timestreamquery/models.py
new file mode 100644
index 000000000000..cb8847eb7fec
--- /dev/null
+++ b/moto/timestreamquery/models.py
@@ -0,0 +1,194 @@
+"""TimestreamQueryBackend class with methods for supported APIs."""
+
+from typing import Any, Dict, List, Optional, Union
+from uuid import uuid4
+
+from moto.core.base_backend import BackendDict, BaseBackend
+from moto.core.common_models import BaseModel
+from moto.core.utils import unix_time
+from moto.utilities.utils import get_partition
+
+from .exceptions import ResourceNotFound
+
+
+class ScheduledQuery(BaseModel):
+ def __init__(
+ self,
+ account_id: str,
+ region_name: str,
+ name: str,
+ query_string: str,
+ schedule_configuration: Dict[str, str],
+ notification_configuration: Dict[str, Dict[str, str]],
+ target_configuration: Optional[Dict[str, Any]],
+ scheduled_query_execution_role_arn: str,
+ tags: Optional[List[Dict[str, str]]],
+ kms_key_id: Optional[str],
+ error_report_configuration: Optional[Dict[str, Dict[str, str]]],
+ ):
+ self.account_id = account_id
+ self.region_name = region_name
+ self.name = name
+ self.query_string = query_string
+ self.schedule_configuration = schedule_configuration
+ self.notification_configuration = notification_configuration
+ self.target_configuration = target_configuration
+ self.scheduled_query_execution_role_arn = scheduled_query_execution_role_arn
+ self.tags = tags
+ self.kms_key_id = kms_key_id
+ self.error_report_configuration = error_report_configuration
+
+ self.created_on = unix_time()
+ self.updated_on = unix_time()
+
+ self.arn = f"arn:{get_partition(region_name)}:timestream:{region_name}:{account_id}:scheduled-query/{name}"
+ self.state = "ENABLED"
+
+ def description(self) -> Dict[str, Any]:
+ return {
+ "Arn": self.arn,
+ "Name": self.name,
+ "CreationTime": self.created_on,
+ "State": self.state,
+ "QueryString": self.query_string,
+ "ScheduleConfiguration": self.schedule_configuration,
+ "NotificationConfiguration": self.notification_configuration,
+ "TargetConfiguration": self.target_configuration,
+ "ScheduledQueryExecutionRoleArn": self.scheduled_query_execution_role_arn,
+ "KmsKeyId": self.kms_key_id,
+ "ErrorReportConfiguration": self.error_report_configuration,
+ }
+
+
+class TimestreamQueryBackend(BaseBackend):
+ """Implementation of TimestreamQuery APIs."""
+
+ def __init__(self, region_name: str, account_id: str):
+ super().__init__(region_name, account_id)
+ self.scheduled_queries: Dict[str, ScheduledQuery] = {}
+
+ self.query_result_queue: Dict[Optional[str], List[Dict[str, Any]]] = {}
+ self.query_results: Dict[str, Dict[str, Any]] = {}
+
+ def create_scheduled_query(
+ self,
+ name: str,
+ query_string: str,
+ schedule_configuration: Dict[str, str],
+ notification_configuration: Dict[str, Dict[str, str]],
+ target_configuration: Optional[Dict[str, Any]],
+ scheduled_query_execution_role_arn: str,
+ tags: Optional[List[Dict[str, str]]],
+ kms_key_id: Optional[str],
+ error_report_configuration: Dict[str, Dict[str, str]],
+ ) -> ScheduledQuery:
+ query = ScheduledQuery(
+ account_id=self.account_id,
+ region_name=self.region_name,
+ name=name,
+ query_string=query_string,
+ schedule_configuration=schedule_configuration,
+ notification_configuration=notification_configuration,
+ target_configuration=target_configuration,
+ scheduled_query_execution_role_arn=scheduled_query_execution_role_arn,
+ tags=tags,
+ kms_key_id=kms_key_id,
+ error_report_configuration=error_report_configuration,
+ )
+ self.scheduled_queries[query.arn] = query
+ return query
+
+ def delete_scheduled_query(self, scheduled_query_arn: str) -> None:
+ self.scheduled_queries.pop(scheduled_query_arn, None)
+
+ def describe_scheduled_query(self, scheduled_query_arn: str) -> ScheduledQuery:
+ if scheduled_query_arn not in self.scheduled_queries:
+ raise ResourceNotFound(scheduled_query_arn)
+ return self.scheduled_queries[scheduled_query_arn]
+
+ def update_scheduled_query(self, scheduled_query_arn: str, state: str) -> None:
+ query = self.scheduled_queries[scheduled_query_arn]
+ query.state = state
+
+ def query(self, query_string: str) -> Dict[str, Any]:
+ """
+ Moto does not have a builtin time-series Database, so calling this endpoint will return zero results by default.
+
+ You can use a dedicated API to configuring a queue of expected results.
+
+ An example invocation looks like this:
+
+ .. sourcecode:: python
+
+ first_result = {
+ 'QueryId': 'some_id',
+ 'Rows': [...],
+ 'ColumnInfo': [...],
+ 'QueryStatus': ...
+ }
+ result_for_unknown_query_string = {
+ 'QueryId': 'unknown',
+ 'Rows': [...],
+ 'ColumnInfo': [...],
+ 'QueryStatus': ...
+ }
+ expected_results = {
+ "account_id": "123456789012", # This is the default - can be omitted
+ "region": "us-east-1", # This is the default - can be omitted
+ "results": {
+ # Use the exact querystring, and a list of results for it
+ # For example
+ "SELECT data FROM mytable": [first_result, ...],
+ # Use None if the exact querystring is unknown/irrelevant
+ None: [result_for_unknown_query_string, ...],
+ }
+ }
+ requests.post(
+ "http://motoapi.amazonaws.com/moto-api/static/timestream/query-results",
+ json=expected_results,
+ )
+
+ When calling `query(QueryString='SELECT data FROM mytable')`, the `first_result` will be returned.
+ Call the query again for the second result, and so on.
+
+ If you don't know the exact query strings, use the `None`-key. In the above example, when calling `SELECT something FROM unknown`, there are no results for that specific query, so `result_for_unknown_query_string` will be returned.
+
+ Results for unknown queries are cached, so calling `SELECT something FROM unknown` will return the same result.
+
+ """
+ if self.query_result_queue.get(query_string):
+ return self.query_result_queue[query_string].pop(0)
+ if result := self.query_results.get(query_string):
+ return result
+ if self.query_result_queue.get(None):
+ self.query_results[query_string] = self.query_result_queue[None].pop(0)
+ return self.query_results[query_string]
+ return {"QueryId": str(uuid4()), "Rows": [], "ColumnInfo": []}
+
+ def describe_endpoints(self) -> List[Dict[str, Union[str, int]]]:
+ # https://docs.aws.amazon.com/timestream/latest/developerguide/Using-API.endpoint-discovery.how-it-works.html
+ # Usually, the address look like this:
+ # query-cell1.timestream.us-east-1.amazonaws.com
+ # Where 'cell1' can be any number, 'cell2', 'cell3', etc - whichever endpoint happens to be available for that particular account
+ # We don't implement a cellular architecture in Moto though, so let's keep it simple
+ return [
+ {
+ "Address": f"query.timestream.{self.region_name}.amazonaws.com",
+ "CachePeriodInMinutes": 1440,
+ }
+ ]
+
+
+timestreamquery_backends = BackendDict(
+ TimestreamQueryBackend,
+ "timestream-query",
+ additional_regions=[
+ "us-east-1",
+ "us-east-2",
+ "us-west-2",
+ "eu-central-1",
+ "eu-west-1",
+ "ap-southeast-2",
+ "ap-northeast-1",
+ ],
+)
diff --git a/moto/timestreamquery/urls.py b/moto/timestreamquery/urls.py
new file mode 100644
index 000000000000..9053d6c32365
--- /dev/null
+++ b/moto/timestreamquery/urls.py
@@ -0,0 +1,11 @@
+"""timestreamquery base URL and path."""
+
+from moto.timestreamwrite.responses import TimestreamWriteResponse
+
+url_bases = [
+ r"https?://query.timestream\.(.+)\.amazonaws\.com",
+]
+
+url_paths = {
+ "{0}/?$": TimestreamWriteResponse.dispatch,
+}
diff --git a/moto/timestreamwrite/responses.py b/moto/timestreamwrite/responses.py
index 179140d2032b..2c4ff9ed8017 100644
--- a/moto/timestreamwrite/responses.py
+++ b/moto/timestreamwrite/responses.py
@@ -1,6 +1,7 @@
import json
from moto.core.responses import BaseResponse
+from moto.timestreamquery.models import TimestreamQueryBackend, timestreamquery_backends
from .models import TimestreamWriteBackend, timestreamwrite_backends
@@ -9,6 +10,10 @@ class TimestreamWriteResponse(BaseResponse):
def __init__(self) -> None:
super().__init__(service_name="timestream-write")
+ @property
+ def timestreamquery_backend(self) -> TimestreamQueryBackend:
+ return timestreamquery_backends[self.current_account][self.region]
+
@property
def timestreamwrite_backend(self) -> TimestreamWriteBackend:
"""Return backend instance specific for this region."""
@@ -134,3 +139,61 @@ def untag_resource(self) -> str:
tag_keys = self._get_param("TagKeys")
self.timestreamwrite_backend.untag_resource(resource_arn, tag_keys)
return "{}"
+
+ # AWS uses the same path/headers for TimestreamWrite and TimestreamQuery
+ # The only difference is the host, but we don't have access to that in ServerMode
+ # Keep the `Query`-responses here
+ # So we don't have to jump through all kinds of hoops to explain the difference between Query and Write to MotoServer
+
+ def create_scheduled_query(self) -> str:
+ name = self._get_param("Name")
+ query_string = self._get_param("QueryString")
+ schedule_configuration = self._get_param("ScheduleConfiguration")
+ notification_configuration = self._get_param("NotificationConfiguration")
+ target_configuration = self._get_param("TargetConfiguration")
+ scheduled_query_execution_role_arn = self._get_param(
+ "ScheduledQueryExecutionRoleArn"
+ )
+ tags = self._get_param("Tags")
+ kms_key_id = self._get_param("KmsKeyId")
+ error_report_configuration = self._get_param("ErrorReportConfiguration")
+ scheduled_query = self.timestreamquery_backend.create_scheduled_query(
+ name=name,
+ query_string=query_string,
+ schedule_configuration=schedule_configuration,
+ notification_configuration=notification_configuration,
+ target_configuration=target_configuration,
+ scheduled_query_execution_role_arn=scheduled_query_execution_role_arn,
+ tags=tags,
+ kms_key_id=kms_key_id,
+ error_report_configuration=error_report_configuration,
+ )
+ return json.dumps(dict(Arn=scheduled_query.arn))
+
+ def delete_scheduled_query(self) -> str:
+ scheduled_query_arn = self._get_param("ScheduledQueryArn")
+ self.timestreamquery_backend.delete_scheduled_query(
+ scheduled_query_arn=scheduled_query_arn,
+ )
+ return "{}"
+
+ def update_scheduled_query(self) -> str:
+ scheduled_query_arn = self._get_param("ScheduledQueryArn")
+ state = self._get_param("State")
+ self.timestreamquery_backend.update_scheduled_query(
+ scheduled_query_arn=scheduled_query_arn,
+ state=state,
+ )
+ return "{}"
+
+ def query(self) -> str:
+ query_string = self._get_param("QueryString")
+ result = self.timestreamquery_backend.query(query_string=query_string)
+ return json.dumps(result)
+
+ def describe_scheduled_query(self) -> str:
+ scheduled_query_arn = self._get_param("ScheduledQueryArn")
+ scheduled_query = self.timestreamquery_backend.describe_scheduled_query(
+ scheduled_query_arn=scheduled_query_arn,
+ )
+ return json.dumps(dict(ScheduledQuery=scheduled_query.description()))
diff --git a/tests/test_timestreamquery/__init__.py b/tests/test_timestreamquery/__init__.py
new file mode 100644
index 000000000000..e69de29bb2d1
diff --git a/tests/test_timestreamquery/test_timestreamquery.py b/tests/test_timestreamquery/test_timestreamquery.py
new file mode 100644
index 000000000000..a330956f5297
--- /dev/null
+++ b/tests/test_timestreamquery/test_timestreamquery.py
@@ -0,0 +1,200 @@
+"""Unit tests for timestreamquery-supported APIs."""
+
+import boto3
+import pytest
+import requests
+from botocore.exceptions import ClientError
+
+from moto import mock_aws, settings
+from tests import DEFAULT_ACCOUNT_ID
+
+# See our Development Tips on writing tests for hints on how to write good tests:
+# http://docs.getmoto.org/en/latest/docs/contributing/development_tips/tests.html
+
+
+@mock_aws
+def test_create_scheduled_query():
+ target_config = {
+ "TimestreamConfiguration": {
+ "DatabaseName": "mydb",
+ "TableName": "mytab",
+ "TimeColumn": "tc",
+ "DimensionMappings": [],
+ }
+ }
+
+ client = boto3.client("timestream-query", region_name="us-east-2")
+ arn = client.create_scheduled_query(
+ Name="myquery",
+ QueryString="SELECT *",
+ ScheduleConfiguration={"ScheduleExpression": "* * * * * 1"},
+ NotificationConfiguration={"SnsConfiguration": {"TopicArn": "arn:some:topic"}},
+ TargetConfiguration=target_config,
+ ScheduledQueryExecutionRoleArn="some role",
+ ErrorReportConfiguration=get_error_config(),
+ KmsKeyId="arn:kms:key",
+ )["Arn"]
+
+ assert (
+ arn
+ == f"arn:aws:timestream:us-east-2:{DEFAULT_ACCOUNT_ID}:scheduled-query/myquery"
+ )
+
+ query = client.describe_scheduled_query(ScheduledQueryArn=arn)["ScheduledQuery"]
+ assert query["Arn"] == arn
+ assert query["Name"] == "myquery"
+ assert query["QueryString"] == "SELECT *"
+ assert query["CreationTime"]
+ assert query["State"] == "ENABLED"
+ assert query["ScheduleConfiguration"] == {"ScheduleExpression": "* * * * * 1"}
+ assert query["NotificationConfiguration"] == {
+ "SnsConfiguration": {"TopicArn": "arn:some:topic"}
+ }
+ assert query["TargetConfiguration"] == target_config
+ assert query["ScheduledQueryExecutionRoleArn"] == "some role"
+ assert query["ErrorReportConfiguration"] == get_error_config()
+ assert query["KmsKeyId"] == "arn:kms:key"
+
+
+@mock_aws
+def test_delete_scheduled_query():
+ client = boto3.client("timestream-query", region_name="us-east-2")
+
+ arn = client.create_scheduled_query(
+ Name="myquery",
+ QueryString="SELECT *",
+ ScheduleConfiguration={"ScheduleExpression": "* * * * * 1"},
+ NotificationConfiguration={"SnsConfiguration": {"TopicArn": "arn:some:topic"}},
+ ScheduledQueryExecutionRoleArn="some role",
+ ErrorReportConfiguration=get_error_config(),
+ )["Arn"]
+ client.delete_scheduled_query(ScheduledQueryArn=arn)
+
+ with pytest.raises(ClientError) as exc:
+ client.describe_scheduled_query(ScheduledQueryArn=arn)
+ err = exc.value.response["Error"]
+ assert err["Code"] == "ResourceNotFoundException"
+ assert err["Message"] == f"The resource with arn {arn} does not exist."
+
+
+@mock_aws
+def test_update_scheduled_query():
+ client = boto3.client("timestream-query", region_name="eu-west-1")
+ arn = client.create_scheduled_query(
+ Name="myquery",
+ QueryString="SELECT *",
+ ScheduleConfiguration={"ScheduleExpression": "* * * * * 1"},
+ NotificationConfiguration={"SnsConfiguration": {"TopicArn": "arn:some:topic"}},
+ ScheduledQueryExecutionRoleArn="some role",
+ ErrorReportConfiguration=get_error_config(),
+ )["Arn"]
+
+ client.update_scheduled_query(
+ ScheduledQueryArn=arn,
+ State="DISABLED",
+ )
+
+ query = client.describe_scheduled_query(ScheduledQueryArn=arn)["ScheduledQuery"]
+ assert query["State"] == "DISABLED"
+
+
+@mock_aws
+def test_query_default_results():
+ client = boto3.client("timestream-query", region_name="us-east-1")
+ resp = client.query(QueryString="SELECT *")
+
+ assert resp["QueryId"]
+ assert resp["Rows"] == []
+ assert resp["ColumnInfo"] == []
+
+
+@mock_aws
+def test_query__configured_results():
+ base_url = (
+ "localhost:5000" if settings.TEST_SERVER_MODE else "motoapi.amazonaws.com"
+ )
+
+ first_result = {
+ "QueryId": "some_id",
+ "Rows": [{"Data": [{"ScalarValue": "1"}]}],
+ "ColumnInfo": [
+ {"Name": "c", "Type": {"ScalarType": "VARCHAR"}},
+ ],
+ "QueryStatus": {
+ "ProgressPercentage": 50,
+ "CumulativeBytesScanned": 5,
+ "CumulativeBytesMetered": 5,
+ },
+ }
+ second_result = {
+ "QueryId": "some_id",
+ "Rows": [{"Data": [{"ScalarValue": "1"}]}, {"Data": [{"ScalarValue": "2"}]}],
+ "ColumnInfo": [
+ {"Name": "c", "Type": {"ScalarType": "VARCHAR"}},
+ {"Name": "c", "Type": {"ScalarType": "VARCHAR"}},
+ ],
+ "QueryStatus": {
+ "ProgressPercentage": 100,
+ "CumulativeBytesScanned": 10,
+ "CumulativeBytesMetered": 10,
+ },
+ }
+ third_result = {
+ "Rows": [{"Data": [{"ScalarValue": "5"}]}],
+ "ColumnInfo": [{"Name": "c", "Type": {"ScalarType": "VARCHAR"}}],
+ "QueryStatus": {
+ "ProgressPercentage": 100,
+ "CumulativeBytesScanned": 10,
+ "CumulativeBytesMetered": 10,
+ },
+ }
+
+ result = {
+ "results": {
+ "SELECT *": [first_result, second_result],
+ None: [third_result],
+ }
+ }
+ requests.post(
+ f"http://{base_url}/moto-api/static/timestream/query-results",
+ json=result,
+ )
+
+ client = boto3.client("timestream-query", region_name="us-east-1")
+
+ # Unknown QUERY returns third result
+ resp = client.query(QueryString="SELECT unknown")
+
+ assert resp["Rows"] == third_result["Rows"]
+ assert resp["ColumnInfo"] == third_result["ColumnInfo"]
+
+ # Known QUERY returns first result
+ resp = client.query(QueryString="SELECT *")
+
+ assert resp["QueryId"] == "some_id"
+ assert resp["Rows"] == first_result["Rows"]
+ assert resp["ColumnInfo"] == first_result["ColumnInfo"]
+
+ # Querying the same thing returns the second result
+ resp = client.query(QueryString="SELECT *")
+
+ assert resp["QueryId"] == "some_id"
+ assert resp["Rows"] == second_result["Rows"]
+ assert resp["ColumnInfo"] == second_result["ColumnInfo"]
+
+ # Unknown QUERY returns third result again
+ resp = client.query(QueryString="SELECT unknown")
+
+ assert resp["Rows"] == third_result["Rows"]
+ assert resp["ColumnInfo"] == third_result["ColumnInfo"]
+
+ # Known QUERY returns nothing - we've already returned all possible results
+ resp = client.query(QueryString="SELECT *")
+
+ assert resp["QueryId"]
+ assert resp["Rows"] == []
+ assert resp["ColumnInfo"] == []
+
+
+def get_error_config():
+ return {"S3Configuration": {"BucketName": "mybucket", "ObjectKeyPrefix": "prefix"}}