prompt: 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()
  
  ```
  
golden_diff: 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'
  
verification_info: {"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"}]}