Service: TimestreamQuery (#8019)

IMPLEMENTATION_COVERAGE.md

@@ -8481,6 +8481,27 @@
 - [ ] update_adapter
 </details>
 
+## timestream-query
+<details>
+<summary>40% implemented</summary>
+
+- [ ] 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
+</details>
+
 ## timestream-write
 <details>
 <summary>78% implemented</summary>
@@ -9010,7 +9031,6 @@
 - synthetics
 - taxsettings
 - timestream-influxdb
-- timestream-query
 - tnb
 - translate
 - trustedadvisor

docs/docs/services/timestream-query.rst

@@ -0,0 +1,80 @@
+.. _implementedservice_timestream-query:
+
+.. |start-h3| raw:: html
+
+    <h3>
+
+.. |end-h3| raw:: html
+
+    </h3>
+
+================
+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
+

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"),

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]": ...

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
 

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,

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,

moto/timestreamquery/__init__.py

@@ -0,0 +1,3 @@
+from .models import timestreamquery_backends  # noqa: F401
+
+# Responses are handled by TimestreamWrite

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.",
+        )