We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Add support for kms-key in timestream service
### Describe the feature

Add support for kms-key in timestream service

### Extra information or context

_No response_

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `c7n/resources/timestream.py`
Content:
```
1 from c7n.manager import resources
2 from c7n.actions import Action
3 from c7n.query import DescribeSource, QueryResourceManager, TypeInfo
4 from c7n.utils import local_session, type_schema
5 from c7n.tags import (
6     TagDelayedAction,
7     TagActionFilter,
8     Tag as TagAction,
9     RemoveTag as RemoveTagAction
10 )
11 
12 
13 class DescribeTimestream(DescribeSource):
14     def augment(self, resources):
15         for r in resources:
16             client = local_session(self.manager.session_factory).client('timestream-write')
17             r['Tags'] = client.list_tags_for_resource(ResourceARN=r['Arn'])['Tags']
18         return resources
19 
20 
21 @resources.register('timestream-database')
22 class TimestreamDatabase(QueryResourceManager):
23     class resource_type(TypeInfo):
24         service = 'timestream-write'
25         arn_type = ''
26         name = 'DatabaseName'
27         id = arn = 'Arn'
28         enum_spec = ('list_databases', 'Databases', {})
29         permission_prefix = 'timestream'
30         permissions = ('timestream:ListDatabases', )
31         permissions_augment = ("timestream:ListTagsForResource",)
32     source_mapping = {
33         'describe': DescribeTimestream,
34     }
35 
36 
37 @resources.register('timestream-table')
38 class TimestreamTable(QueryResourceManager):
39     class resource_type(TypeInfo):
40         service = 'timestream-write'
41         arn_type = ''
42         name = 'TableName'
43         id = arn = 'Arn'
44         enum_spec = ('list_tables', 'Tables', {})
45         permission_prefix = 'timestream'
46         permissions = ('timestream:ListTables', )
47 
48     source_mapping = {
49         'describe': DescribeTimestream,
50     }
51 
52 
53 @TimestreamDatabase.action_registry.register('tag')
54 @TimestreamTable.action_registry.register('tag')
55 class TimestreamTag(TagAction):
56 
57     permissions = ('timestream:TagResource', )
58 
59     def process_resource_set(self, client, resource_set, tags):
60         for r in resource_set:
61             client.tag_resource(ResourceARN=r['Arn'], Tags=tags)
62 
63 
64 @TimestreamDatabase.action_registry.register('remove-tag')
65 @TimestreamTable.action_registry.register('remove-tag')
66 class TimestreamRemoveTag(RemoveTagAction):
67 
68     permissions = ('timestream:UntagResource', )
69 
70     def process_resource_set(self, client, resource_set, tag_keys):
71         for r in resource_set:
72             client.untag_resource(ResourceARN=r['Arn'], TagKeys=tag_keys)
73 
74 
75 TimestreamDatabase.action_registry.register('mark-for-op', TagDelayedAction)
76 TimestreamTable.action_registry.register('mark-for-op', TagDelayedAction)
77 
78 TimestreamDatabase.filter_registry.register('marked-for-op', TagActionFilter)
79 TimestreamTable.filter_registry.register('marked-for-op', TagActionFilter)
80 
81 
82 @TimestreamTable.action_registry.register('delete')
83 class TimestreamTableDelete(Action):
84     """
85     Deletes a timestream table
86     """
87 
88     schema = type_schema('delete')
89     permissions = ('timestream:DeleteTable', )
90 
91     def process(self, resources):
92         client = local_session(self.manager.session_factory).client('timestream-write')
93         for r in resources:
94             try:
95                 client.delete_table(
96                     DatabaseName=r['DatabaseName'],
97                     TableName=r['TableName']
98                 )
99             except client.exceptions.ResourceNotFoundException:
100                 continue
101 
102 
103 @TimestreamDatabase.action_registry.register('delete')
104 class TimestreamDatabaseDelete(Action):
105     """
106     Deletes a timestream database
107     """
108 
109     schema = type_schema('delete', force={'type': 'boolean', 'default': False})
110     permissions = (
111         'timestream:DeleteDatabase',
112         'timestream:ListTables', 'timestream:DeleteTable', )
113 
114     def process(self, resources):
115         client = local_session(self.manager.session_factory).client('timestream-write')
116         for r in resources:
117             try:
118                 client.delete_database(
119                     DatabaseName=r['DatabaseName'],
120                 )
121             except client.exceptions.ResourceNotFoundException:
122                 continue
123             except client.exceptions.ValidationException:
124                 if not self.data.get('force', False):
125                     self.log.error(
126                         f'Unable to delete database:{r["DatabaseName"]}, '
127                         'tables must be deleted first')
128                     continue
129                 tables = client.list_tables(DatabaseName=r['DatabaseName'])['Tables']
130                 TimestreamTableDelete(
131                     data={'type': 'delete'},
132                     manager=self.manager,
133                     log_dir=self.log_dir
134                 ).process(tables)
135                 client.delete_database(
136                     DatabaseName=r['DatabaseName'],
137                 )
138
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/c7n/resources/timestream.py b/c7n/resources/timestream.py
--- a/c7n/resources/timestream.py
+++ b/c7n/resources/timestream.py
@@ -1,5 +1,6 @@
 from c7n.manager import resources
 from c7n.actions import Action
+from c7n.filters.kms import KmsRelatedFilter
 from c7n.query import DescribeSource, QueryResourceManager, TypeInfo
 from c7n.utils import local_session, type_schema
 from c7n.tags import (
@@ -135,3 +136,8 @@
                 client.delete_database(
                     DatabaseName=r['DatabaseName'],
                 )
+
+
+@TimestreamDatabase.filter_registry.register('kms-key')
+class KmsFilter(KmsRelatedFilter):
+    RelatedIdsExpression = 'KmsKeyId'

{"golden_diff": "diff --git a/c7n/resources/timestream.py b/c7n/resources/timestream.py\n--- a/c7n/resources/timestream.py\n+++ b/c7n/resources/timestream.py\n@@ -1,5 +1,6 @@\n from c7n.manager import resources\n from c7n.actions import Action\n+from c7n.filters.kms import KmsRelatedFilter\n from c7n.query import DescribeSource, QueryResourceManager, TypeInfo\n from c7n.utils import local_session, type_schema\n from c7n.tags import (\n@@ -135,3 +136,8 @@\n                 client.delete_database(\n                     DatabaseName=r['DatabaseName'],\n                 )\n+\n+\n+@TimestreamDatabase.filter_registry.register('kms-key')\n+class KmsFilter(KmsRelatedFilter):\n+    RelatedIdsExpression = 'KmsKeyId'\n", "issue": "Add support for kms-key in timestream service\n### Describe the feature\n\nAdd support for kms-key in timestream service\n\n### Extra information or context\n\n_No response_\n", "before_files": [{"content": "from c7n.manager import resources\nfrom c7n.actions import Action\nfrom c7n.query import DescribeSource, QueryResourceManager, TypeInfo\nfrom c7n.utils import local_session, type_schema\nfrom c7n.tags import (\n    TagDelayedAction,\n    TagActionFilter,\n    Tag as TagAction,\n    RemoveTag as RemoveTagAction\n)\n\n\nclass DescribeTimestream(DescribeSource):\n    def augment(self, resources):\n        for r in resources:\n            client = local_session(self.manager.session_factory).client('timestream-write')\n            r['Tags'] = client.list_tags_for_resource(ResourceARN=r['Arn'])['Tags']\n        return resources\n\n\n@resources.register('timestream-database')\nclass TimestreamDatabase(QueryResourceManager):\n    class resource_type(TypeInfo):\n        service = 'timestream-write'\n        arn_type = ''\n        name = 'DatabaseName'\n        id = arn = 'Arn'\n        enum_spec = ('list_databases', 'Databases', {})\n        permission_prefix = 'timestream'\n        permissions = ('timestream:ListDatabases', )\n        permissions_augment = (\"timestream:ListTagsForResource\",)\n    source_mapping = {\n        'describe': DescribeTimestream,\n    }\n\n\n@resources.register('timestream-table')\nclass TimestreamTable(QueryResourceManager):\n    class resource_type(TypeInfo):\n        service = 'timestream-write'\n        arn_type = ''\n        name = 'TableName'\n        id = arn = 'Arn'\n        enum_spec = ('list_tables', 'Tables', {})\n        permission_prefix = 'timestream'\n        permissions = ('timestream:ListTables', )\n\n    source_mapping = {\n        'describe': DescribeTimestream,\n    }\n\n\n@TimestreamDatabase.action_registry.register('tag')\n@TimestreamTable.action_registry.register('tag')\nclass TimestreamTag(TagAction):\n\n    permissions = ('timestream:TagResource', )\n\n    def process_resource_set(self, client, resource_set, tags):\n        for r in resource_set:\n            client.tag_resource(ResourceARN=r['Arn'], Tags=tags)\n\n\n@TimestreamDatabase.action_registry.register('remove-tag')\n@TimestreamTable.action_registry.register('remove-tag')\nclass TimestreamRemoveTag(RemoveTagAction):\n\n    permissions = ('timestream:UntagResource', )\n\n    def process_resource_set(self, client, resource_set, tag_keys):\n        for r in resource_set:\n            client.untag_resource(ResourceARN=r['Arn'], TagKeys=tag_keys)\n\n\nTimestreamDatabase.action_registry.register('mark-for-op', TagDelayedAction)\nTimestreamTable.action_registry.register('mark-for-op', TagDelayedAction)\n\nTimestreamDatabase.filter_registry.register('marked-for-op', TagActionFilter)\nTimestreamTable.filter_registry.register('marked-for-op', TagActionFilter)\n\n\n@TimestreamTable.action_registry.register('delete')\nclass TimestreamTableDelete(Action):\n    \"\"\"\n    Deletes a timestream table\n    \"\"\"\n\n    schema = type_schema('delete')\n    permissions = ('timestream:DeleteTable', )\n\n    def process(self, resources):\n        client = local_session(self.manager.session_factory).client('timestream-write')\n        for r in resources:\n            try:\n                client.delete_table(\n                    DatabaseName=r['DatabaseName'],\n                    TableName=r['TableName']\n                )\n            except client.exceptions.ResourceNotFoundException:\n                continue\n\n\n@TimestreamDatabase.action_registry.register('delete')\nclass TimestreamDatabaseDelete(Action):\n    \"\"\"\n    Deletes a timestream database\n    \"\"\"\n\n    schema = type_schema('delete', force={'type': 'boolean', 'default': False})\n    permissions = (\n        'timestream:DeleteDatabase',\n        'timestream:ListTables', 'timestream:DeleteTable', )\n\n    def process(self, resources):\n        client = local_session(self.manager.session_factory).client('timestream-write')\n        for r in resources:\n            try:\n                client.delete_database(\n                    DatabaseName=r['DatabaseName'],\n                )\n            except client.exceptions.ResourceNotFoundException:\n                continue\n            except client.exceptions.ValidationException:\n                if not self.data.get('force', False):\n                    self.log.error(\n                        f'Unable to delete database:{r[\"DatabaseName\"]}, '\n                        'tables must be deleted first')\n                    continue\n                tables = client.list_tables(DatabaseName=r['DatabaseName'])['Tables']\n                TimestreamTableDelete(\n                    data={'type': 'delete'},\n                    manager=self.manager,\n                    log_dir=self.log_dir\n                ).process(tables)\n                client.delete_database(\n                    DatabaseName=r['DatabaseName'],\n                )\n", "path": "c7n/resources/timestream.py"}], "after_files": [{"content": "from c7n.manager import resources\nfrom c7n.actions import Action\nfrom c7n.filters.kms import KmsRelatedFilter\nfrom c7n.query import DescribeSource, QueryResourceManager, TypeInfo\nfrom c7n.utils import local_session, type_schema\nfrom c7n.tags import (\n    TagDelayedAction,\n    TagActionFilter,\n    Tag as TagAction,\n    RemoveTag as RemoveTagAction\n)\n\n\nclass DescribeTimestream(DescribeSource):\n    def augment(self, resources):\n        for r in resources:\n            client = local_session(self.manager.session_factory).client('timestream-write')\n            r['Tags'] = client.list_tags_for_resource(ResourceARN=r['Arn'])['Tags']\n        return resources\n\n\n@resources.register('timestream-database')\nclass TimestreamDatabase(QueryResourceManager):\n    class resource_type(TypeInfo):\n        service = 'timestream-write'\n        arn_type = ''\n        name = 'DatabaseName'\n        id = arn = 'Arn'\n        enum_spec = ('list_databases', 'Databases', {})\n        permission_prefix = 'timestream'\n        permissions = ('timestream:ListDatabases', )\n        permissions_augment = (\"timestream:ListTagsForResource\",)\n    source_mapping = {\n        'describe': DescribeTimestream,\n    }\n\n\n@resources.register('timestream-table')\nclass TimestreamTable(QueryResourceManager):\n    class resource_type(TypeInfo):\n        service = 'timestream-write'\n        arn_type = ''\n        name = 'TableName'\n        id = arn = 'Arn'\n        enum_spec = ('list_tables', 'Tables', {})\n        permission_prefix = 'timestream'\n        permissions = ('timestream:ListTables', )\n\n    source_mapping = {\n        'describe': DescribeTimestream,\n    }\n\n\n@TimestreamDatabase.action_registry.register('tag')\n@TimestreamTable.action_registry.register('tag')\nclass TimestreamTag(TagAction):\n\n    permissions = ('timestream:TagResource', )\n\n    def process_resource_set(self, client, resource_set, tags):\n        for r in resource_set:\n            client.tag_resource(ResourceARN=r['Arn'], Tags=tags)\n\n\n@TimestreamDatabase.action_registry.register('remove-tag')\n@TimestreamTable.action_registry.register('remove-tag')\nclass TimestreamRemoveTag(RemoveTagAction):\n\n    permissions = ('timestream:UntagResource', )\n\n    def process_resource_set(self, client, resource_set, tag_keys):\n        for r in resource_set:\n            client.untag_resource(ResourceARN=r['Arn'], TagKeys=tag_keys)\n\n\nTimestreamDatabase.action_registry.register('mark-for-op', TagDelayedAction)\nTimestreamTable.action_registry.register('mark-for-op', TagDelayedAction)\n\nTimestreamDatabase.filter_registry.register('marked-for-op', TagActionFilter)\nTimestreamTable.filter_registry.register('marked-for-op', TagActionFilter)\n\n\n@TimestreamTable.action_registry.register('delete')\nclass TimestreamTableDelete(Action):\n    \"\"\"\n    Deletes a timestream table\n    \"\"\"\n\n    schema = type_schema('delete')\n    permissions = ('timestream:DeleteTable', )\n\n    def process(self, resources):\n        client = local_session(self.manager.session_factory).client('timestream-write')\n        for r in resources:\n            try:\n                client.delete_table(\n                    DatabaseName=r['DatabaseName'],\n                    TableName=r['TableName']\n                )\n            except client.exceptions.ResourceNotFoundException:\n                continue\n\n\n@TimestreamDatabase.action_registry.register('delete')\nclass TimestreamDatabaseDelete(Action):\n    \"\"\"\n    Deletes a timestream database\n    \"\"\"\n\n    schema = type_schema('delete', force={'type': 'boolean', 'default': False})\n    permissions = (\n        'timestream:DeleteDatabase',\n        'timestream:ListTables', 'timestream:DeleteTable', )\n\n    def process(self, resources):\n        client = local_session(self.manager.session_factory).client('timestream-write')\n        for r in resources:\n            try:\n                client.delete_database(\n                    DatabaseName=r['DatabaseName'],\n                )\n            except client.exceptions.ResourceNotFoundException:\n                continue\n            except client.exceptions.ValidationException:\n                if not self.data.get('force', False):\n                    self.log.error(\n                        f'Unable to delete database:{r[\"DatabaseName\"]}, '\n                        'tables must be deleted first')\n                    continue\n                tables = client.list_tables(DatabaseName=r['DatabaseName'])['Tables']\n                TimestreamTableDelete(\n                    data={'type': 'delete'},\n                    manager=self.manager,\n                    log_dir=self.log_dir\n                ).process(tables)\n                client.delete_database(\n                    DatabaseName=r['DatabaseName'],\n                )\n\n\n@TimestreamDatabase.filter_registry.register('kms-key')\nclass KmsFilter(KmsRelatedFilter):\n    RelatedIdsExpression = 'KmsKeyId'\n", "path": "c7n/resources/timestream.py"}]}