dagster-aws 0.20.15__tar.gz → 0.28.5__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (84) hide show
  1. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/LICENSE +1 -1
  2. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/MANIFEST.in +2 -0
  3. dagster_aws-0.28.5/PKG-INFO +44 -0
  4. dagster_aws-0.28.5/README.md +4 -0
  5. dagster_aws-0.28.5/dagster_aws/__init__.py +5 -0
  6. dagster_aws-0.28.5/dagster_aws/_stubs.py +50 -0
  7. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/athena/__init__.py +1 -1
  8. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/athena/resources.py +2 -2
  9. dagster_aws-0.28.5/dagster_aws/cloudwatch/__init__.py +1 -0
  10. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/cloudwatch/loggers.py +18 -6
  11. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/ecr/__init__.py +1 -1
  12. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/ecr/resources.py +26 -2
  13. dagster_aws-0.28.5/dagster_aws/ecs/__init__.py +3 -0
  14. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/ecs/container_context.py +22 -5
  15. dagster_aws-0.28.5/dagster_aws/ecs/executor.py +450 -0
  16. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/ecs/launcher.py +263 -85
  17. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/ecs/tasks.py +125 -40
  18. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/ecs/test_utils.py +6 -5
  19. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/ecs/utils.py +75 -35
  20. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/emr/__init__.py +5 -3
  21. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/emr/configs.py +1 -1
  22. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/emr/configs_spark.py +974 -118
  23. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/emr/emr.py +31 -30
  24. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/emr/pyspark_step_launcher.py +14 -14
  25. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/emr/types.py +4 -0
  26. dagster_aws-0.28.5/dagster_aws/pipes/__init__.py +35 -0
  27. dagster_aws-0.28.5/dagster_aws/pipes/clients/__init__.py +15 -0
  28. dagster_aws-0.28.5/dagster_aws/pipes/clients/ecs.py +329 -0
  29. dagster_aws-0.28.5/dagster_aws/pipes/clients/emr.py +336 -0
  30. dagster_aws-0.28.5/dagster_aws/pipes/clients/emr_containers.py +242 -0
  31. dagster_aws-0.28.5/dagster_aws/pipes/clients/emr_serverless.py +367 -0
  32. dagster_aws-0.28.5/dagster_aws/pipes/clients/glue.py +203 -0
  33. dagster_aws-0.28.5/dagster_aws/pipes/clients/lambda_.py +103 -0
  34. dagster_aws-0.28.5/dagster_aws/pipes/clients/utils.py +159 -0
  35. dagster_aws-0.28.5/dagster_aws/pipes/context_injectors.py +61 -0
  36. dagster_aws-0.28.5/dagster_aws/pipes/message_readers.py +423 -0
  37. dagster_aws-0.28.5/dagster_aws/rds/__init__.py +1 -0
  38. dagster_aws-0.28.5/dagster_aws/rds/resources.py +78 -0
  39. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/redshift/__init__.py +1 -1
  40. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/redshift/resources.py +8 -8
  41. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/s3/__init__.py +7 -7
  42. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/s3/compute_log_manager.py +44 -21
  43. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/s3/file_manager.py +1 -1
  44. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/s3/io_manager.py +10 -14
  45. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/s3/ops.py +4 -3
  46. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/s3/resources.py +5 -7
  47. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/s3/s3_fake_resource.py +3 -3
  48. dagster_aws-0.28.5/dagster_aws/s3/sensor.py +101 -0
  49. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/s3/utils.py +13 -12
  50. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/secretsmanager/__init__.py +2 -2
  51. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/secretsmanager/resources.py +29 -14
  52. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/secretsmanager/secrets.py +27 -3
  53. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/ssm/__init__.py +1 -1
  54. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/ssm/parameters.py +29 -8
  55. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/ssm/resources.py +38 -23
  56. dagster_aws-0.28.5/dagster_aws/utils/__init__.py +71 -0
  57. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/utils/mrjob/log4j.py +2 -2
  58. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/utils/mrjob/retry.py +3 -2
  59. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/utils/mrjob/utils.py +3 -3
  60. dagster_aws-0.28.5/dagster_aws/version.py +1 -0
  61. dagster_aws-0.28.5/dagster_aws.egg-info/PKG-INFO +44 -0
  62. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws.egg-info/SOURCES.txt +15 -1
  63. dagster_aws-0.28.5/dagster_aws.egg-info/requires.txt +22 -0
  64. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/setup.py +13 -6
  65. dagster-aws-0.20.15/PKG-INFO +0 -18
  66. dagster-aws-0.20.15/README.md +0 -4
  67. dagster-aws-0.20.15/dagster_aws/__init__.py +0 -5
  68. dagster-aws-0.20.15/dagster_aws/cloudwatch/__init__.py +0 -1
  69. dagster-aws-0.20.15/dagster_aws/ecs/__init__.py +0 -2
  70. dagster-aws-0.20.15/dagster_aws/ext.py +0 -32
  71. dagster-aws-0.20.15/dagster_aws/s3/sensor.py +0 -41
  72. dagster-aws-0.20.15/dagster_aws/utils/__init__.py +0 -38
  73. dagster-aws-0.20.15/dagster_aws/version.py +0 -1
  74. dagster-aws-0.20.15/dagster_aws.egg-info/PKG-INFO +0 -18
  75. dagster-aws-0.20.15/dagster_aws.egg-info/requires.txt +0 -15
  76. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/emr/emr_step_main.py +0 -0
  77. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/emr/main.py.template +0 -0
  78. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/emr/utils.py +0 -0
  79. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/py.typed +0 -0
  80. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws/utils/mrjob/__init__.py +0 -0
  81. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws.egg-info/dependency_links.txt +0 -0
  82. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws.egg-info/not-zip-safe +0 -0
  83. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/dagster_aws.egg-info/top_level.txt +0 -0
  84. {dagster-aws-0.20.15 → dagster_aws-0.28.5}/setup.cfg +0 -0
@@ -186,7 +186,7 @@
186
186
  same "printed page" as the copyright notice for easier
187
187
  identification within third-party archives.
188
188
 
189
- Copyright 2023 Dagster Labs, Inc".
189
+ Copyright 2025 Dagster Labs, Inc.
190
190
 
191
191
  Licensed under the Apache License, Version 2.0 (the "License");
192
192
  you may not use this file except in compliance with the License.
@@ -4,3 +4,5 @@ recursive-include dagster_aws *.txt
4
4
  recursive-include dagster_aws *.template
5
5
  include LICENSE
6
6
  include dagster_aws/py.typed
7
+
8
+ exclude ruff.toml
@@ -0,0 +1,44 @@
1
+ Metadata-Version: 2.4
2
+ Name: dagster-aws
3
+ Version: 0.28.5
4
+ Summary: Package for AWS-specific Dagster framework solid and resource components.
5
+ Home-page: https://github.com/dagster-io/dagster/tree/master/python_modules/libraries/dagster-aws
6
+ Author: Dagster Labs
7
+ Author-email: hello@dagsterlabs.com
8
+ License: Apache-2.0
9
+ Classifier: Programming Language :: Python :: 3.10
10
+ Classifier: Programming Language :: Python :: 3.11
11
+ Classifier: Programming Language :: Python :: 3.12
12
+ Classifier: Programming Language :: Python :: 3.13
13
+ Classifier: License :: OSI Approved :: Apache Software License
14
+ Classifier: Operating System :: OS Independent
15
+ Requires-Python: >=3.10,<3.14
16
+ License-File: LICENSE
17
+ Requires-Dist: boto3
18
+ Requires-Dist: dagster==1.12.5
19
+ Requires-Dist: packaging
20
+ Requires-Dist: requests
21
+ Provides-Extra: redshift
22
+ Requires-Dist: psycopg2-binary; extra == "redshift"
23
+ Provides-Extra: pyspark
24
+ Requires-Dist: dagster-pyspark; extra == "pyspark"
25
+ Provides-Extra: stubs
26
+ Requires-Dist: boto3-stubs-lite[ecs,emr,emr-containers,emr-serverless,glue,logs,s3]; extra == "stubs"
27
+ Provides-Extra: test
28
+ Requires-Dist: botocore!=1.32.1; extra == "test"
29
+ Requires-Dist: moto[emrcontainers,emrserverless,glue,logs,s3,server]<5.0,>=2.2.8; extra == "test"
30
+ Requires-Dist: requests-mock; extra == "test"
31
+ Requires-Dist: xmltodict==0.12.0; extra == "test"
32
+ Requires-Dist: flaky; extra == "test"
33
+ Requires-Dist: pytest-cases; extra == "test"
34
+ Requires-Dist: s3fs; extra == "test"
35
+ Dynamic: author
36
+ Dynamic: author-email
37
+ Dynamic: classifier
38
+ Dynamic: home-page
39
+ Dynamic: license
40
+ Dynamic: license-file
41
+ Dynamic: provides-extra
42
+ Dynamic: requires-dist
43
+ Dynamic: requires-python
44
+ Dynamic: summary
@@ -0,0 +1,4 @@
1
+ # dagster-aws
2
+
3
+ The docs for `dagster-aws` can be found
4
+ [here](https://docs.dagster.io/api/python-api/libraries/dagster-aws).
@@ -0,0 +1,5 @@
1
+ from dagster_shared.libraries import DagsterLibraryRegistry
2
+
3
+ from dagster_aws.version import __version__
4
+
5
+ DagsterLibraryRegistry.register("dagster-aws", __version__)
@@ -0,0 +1,50 @@
1
+ # Stubs have been sourced from
2
+ #
3
+ # https://youtype.github.io/boto3_stubs_docs/mypy_boto3_s3/type_defs/#objecttypedef
4
+ #
5
+ # as to not require any additional dependencies
6
+
7
+ from datetime import datetime
8
+ from typing import Literal, Optional, TypedDict
9
+
10
+ ChecksumAlgorithmType = Literal[
11
+ "CRC32",
12
+ "CRC32C",
13
+ "SHA1",
14
+ "SHA256",
15
+ ]
16
+
17
+ ObjectStorageClassType = Literal[
18
+ "DEEP_ARCHIVE",
19
+ "EXPRESS_ONEZONE",
20
+ "GLACIER",
21
+ "GLACIER_IR",
22
+ "INTELLIGENT_TIERING",
23
+ "ONEZONE_IA",
24
+ "OUTPOSTS",
25
+ "REDUCED_REDUNDANCY",
26
+ "SNOW",
27
+ "STANDARD",
28
+ "STANDARD_IA",
29
+ ]
30
+
31
+
32
+ class OwnerTypeDef(TypedDict):
33
+ DisplayName: Optional[str]
34
+ ID: Optional[str]
35
+
36
+
37
+ class RestoreStatusTypeDef(TypedDict):
38
+ IsRestoreInProgress: Optional[bool]
39
+ RestoreExpiryDate: Optional[datetime]
40
+
41
+
42
+ class ObjectTypeDef(TypedDict):
43
+ Key: str
44
+ LastModified: datetime
45
+ ETag: Optional[str]
46
+ ChecksumAlgorithm: Optional[list[ChecksumAlgorithmType]]
47
+ Size: Optional[int]
48
+ StorageClass: Optional[ObjectStorageClassType]
49
+ Owner: Optional[OwnerTypeDef]
50
+ RestoreStatus: Optional[RestoreStatusTypeDef]
@@ -1,4 +1,4 @@
1
- from .resources import (
1
+ from dagster_aws.athena.resources import (
2
2
  AthenaClient as AthenaClient,
3
3
  AthenaClientResource as AthenaClientResource,
4
4
  AthenaError as AthenaError,
@@ -82,7 +82,7 @@ class AthenaClient:
82
82
  raise AthenaTimeout()
83
83
 
84
84
  if state != "SUCCEEDED":
85
- raise AthenaError(execution["Status"]["StateChangeReason"])
85
+ raise AthenaError(execution["Status"]["StateChangeReason"]) # pyright: ignore[reportPossiblyUnboundVariable]
86
86
 
87
87
  def _results(self, execution_id):
88
88
  execution = self.client.get_query_execution(QueryExecutionId=execution_id)["QueryExecution"]
@@ -254,7 +254,7 @@ class AthenaClientResource(ResourceWithAthenaConfig):
254
254
  def example_athena_asset(athena: AthenaClientResource):
255
255
  return athena.get_client().execute_query("SELECT 1", fetch_results=True)
256
256
 
257
- defs = Definitions(
257
+ Definitions(
258
258
  assets=[example_athena_asset],
259
259
  resources={"athena": AthenaClientResource()}
260
260
  )
@@ -0,0 +1 @@
1
+ from dagster_aws.cloudwatch.loggers import cloudwatch_logger as cloudwatch_logger
@@ -1,15 +1,17 @@
1
1
  import datetime
2
2
  import logging
3
+ from typing import Optional
3
4
 
4
5
  import boto3
5
6
  from dagster import (
6
7
  Field,
7
8
  StringSource,
8
9
  _check as check,
9
- _seven,
10
10
  logger,
11
11
  )
12
+ from dagster._annotations import deprecated
12
13
  from dagster._core.utils import coerce_valid_log_level
14
+ from dagster_shared import seven
13
15
 
14
16
  # The maximum batch size is 1,048,576 bytes, and this size is calculated as the sum of all event
15
17
  # messages in UTF-8, plus 26 bytes for each log event.
@@ -28,20 +30,29 @@ def millisecond_timestamp(dt):
28
30
  return int(microsecond_timestamp / 1000)
29
31
 
30
32
 
33
+ @deprecated(breaking_version="0.27", additional_warn_text="Use `PipesCloudWatchLogReader` instead.")
31
34
  class CloudwatchLogsHandler(logging.Handler):
32
35
  def __init__(
33
36
  self,
34
37
  log_group_name,
35
38
  log_stream_name,
36
39
  aws_region=None,
37
- aws_secret_access_key=None,
38
- aws_access_key_id=None,
40
+ aws_access_key_id: Optional[str] = None,
41
+ aws_secret_access_key: Optional[str] = None,
42
+ endpoint_url: Optional[str] = None,
43
+ use_ssl: bool = True,
44
+ aws_session_token: Optional[str] = None,
45
+ verify: Optional[bool] = None,
39
46
  ):
40
47
  self.client = boto3.client(
41
48
  "logs",
42
49
  region_name=aws_region,
43
50
  aws_access_key_id=aws_access_key_id,
44
51
  aws_secret_access_key=aws_secret_access_key,
52
+ endpoint_url=endpoint_url,
53
+ use_ssl=use_ssl,
54
+ aws_session_token=aws_session_token,
55
+ verify=verify,
45
56
  )
46
57
  self.log_group_name = check.str_param(log_group_name, "log_group_name")
47
58
  # Maybe we should make this optional, and default to the run_id
@@ -53,7 +64,7 @@ class CloudwatchLogsHandler(logging.Handler):
53
64
  self.check_log_group()
54
65
  self.check_log_stream()
55
66
 
56
- super(CloudwatchLogsHandler, self).__init__()
67
+ super().__init__()
57
68
 
58
69
  def check_log_group(self):
59
70
  # Check that log group exists
@@ -111,7 +122,7 @@ class CloudwatchLogsHandler(logging.Handler):
111
122
  def log_error(self, record, exc):
112
123
  logging.critical("Error while logging!")
113
124
  try:
114
- logging.error(f"Attempted to log: {_seven.json.dumps(record.__dict__)}")
125
+ logging.error(f"Attempted to log: {seven.json.dumps(record.__dict__)}")
115
126
  except Exception:
116
127
  pass
117
128
  logging.exception(str(exc))
@@ -123,7 +134,7 @@ class CloudwatchLogsHandler(logging.Handler):
123
134
  self._emit(record, retry=True)
124
135
 
125
136
  def _emit(self, record, retry=False):
126
- message = _seven.json.dumps(record.__dict__)
137
+ message = seven.json.dumps(record.__dict__)
127
138
  timestamp = millisecond_timestamp(
128
139
  datetime.datetime.strptime(record.dagster_meta["log_timestamp"], "%Y-%m-%dT%H:%M:%S.%f")
129
140
  )
@@ -191,6 +202,7 @@ class CloudwatchLogsHandler(logging.Handler):
191
202
  },
192
203
  description="The default colored console logger.",
193
204
  )
205
+ @deprecated(breaking_version="0.27", additional_warn_text="Use `PipesCloudWatchLogReader` instead.")
194
206
  def cloudwatch_logger(init_context):
195
207
  """This logger provides support for sending Dagster logs to AWS CloudWatch.
196
208
 
@@ -1,4 +1,4 @@
1
- from .resources import (
1
+ from dagster_aws.ecr.resources import (
2
2
  ECRPublicResource as ECRPublicResource,
3
3
  FakeECRPublicResource as FakeECRPublicResource,
4
4
  ecr_public_resource as ecr_public_resource,
@@ -1,14 +1,34 @@
1
1
  import datetime
2
+ from typing import Optional
2
3
 
3
4
  import boto3
4
5
  from botocore.stub import Stubber
5
6
  from dagster import ConfigurableResource, resource
7
+ from dagster._annotations import beta
6
8
  from dagster._core.definitions.resource_definition import dagster_maintained_resource
7
9
 
8
10
 
9
11
  class ECRPublicClient:
10
- def __init__(self):
11
- self.client = boto3.client("ecr-public")
12
+ def __init__(
13
+ self,
14
+ region_name: Optional[str] = None,
15
+ endpoint_url: Optional[str] = None,
16
+ use_ssl: bool = True,
17
+ aws_access_key_id: Optional[str] = None,
18
+ aws_secret_access_key: Optional[str] = None,
19
+ aws_session_token: Optional[str] = None,
20
+ verify: Optional[bool] = None,
21
+ ):
22
+ self.client = boto3.client(
23
+ "ecr-public",
24
+ region_name=region_name,
25
+ use_ssl=use_ssl,
26
+ verify=verify,
27
+ endpoint_url=endpoint_url,
28
+ aws_access_key_id=aws_access_key_id,
29
+ aws_secret_access_key=aws_secret_access_key,
30
+ aws_session_token=aws_session_token,
31
+ )
12
32
 
13
33
  def get_login_password(self):
14
34
  return self.client.get_authorization_token()["authorizationData"]["authorizationToken"]
@@ -34,6 +54,7 @@ class FakeECRPublicClient(ECRPublicClient):
34
54
  return result
35
55
 
36
56
 
57
+ @beta
37
58
  class ECRPublicResource(ConfigurableResource):
38
59
  """This resource enables connecting to AWS Public and getting a login password from it.
39
60
  Similar to the AWS CLI's `aws ecr-public get-login-password` command.
@@ -47,6 +68,7 @@ class ECRPublicResource(ConfigurableResource):
47
68
  return ECRPublicClient()
48
69
 
49
70
 
71
+ @beta
50
72
  class FakeECRPublicResource(ConfigurableResource):
51
73
  """This resource behaves like ecr_public_resource except it stubs out the real AWS API
52
74
  requests and always returns `'token'` as its login password.
@@ -60,6 +82,7 @@ class FakeECRPublicResource(ConfigurableResource):
60
82
  return FakeECRPublicClient()
61
83
 
62
84
 
85
+ @beta
63
86
  @dagster_maintained_resource
64
87
  @resource(
65
88
  description=(
@@ -71,6 +94,7 @@ def ecr_public_resource(context) -> ECRPublicClient:
71
94
  return ECRPublicResource.from_resource_context(context).get_client()
72
95
 
73
96
 
97
+ @beta
74
98
  @dagster_maintained_resource
75
99
  @resource(
76
100
  description=(
@@ -0,0 +1,3 @@
1
+ from dagster_aws.ecs.executor import ecs_executor as ecs_executor
2
+ from dagster_aws.ecs.launcher import EcsRunLauncher as EcsRunLauncher
3
+ from dagster_aws.ecs.tasks import EcsEventualConsistencyTimeout as EcsEventualConsistencyTimeout
@@ -1,4 +1,5 @@
1
- from typing import TYPE_CHECKING, Any, Mapping, NamedTuple, Optional, Sequence, cast
1
+ from collections.abc import Mapping, Sequence
2
+ from typing import TYPE_CHECKING, Any, NamedTuple, Optional, cast
2
3
 
3
4
  from dagster import (
4
5
  Array,
@@ -16,10 +17,10 @@ from dagster._core.errors import DagsterInvalidConfigError
16
17
  from dagster._core.storage.dagster_run import DagsterRun
17
18
  from dagster._core.utils import parse_env_var
18
19
 
19
- from ..secretsmanager import get_tagged_secrets
20
+ from dagster_aws.secretsmanager import get_tagged_secrets
20
21
 
21
22
  if TYPE_CHECKING:
22
- from . import EcsRunLauncher
23
+ from dagster_aws.ecs import EcsRunLauncher
23
24
 
24
25
  # Config shared between EcsRunLauncher and EcsContainerContext
25
26
  SHARED_ECS_SCHEMA = {
@@ -186,6 +187,12 @@ ECS_CONTAINER_CONTEXT_SCHEMA = {
186
187
  is_required=False,
187
188
  description="The ephemeral storage, in GiB, to use for the launched task.",
188
189
  ),
190
+ "replica_count": Field(
191
+ int,
192
+ default_value=1,
193
+ is_required=False,
194
+ description="The number of code server instances to launch.",
195
+ ),
189
196
  }
190
197
  )
191
198
  ),
@@ -209,6 +216,11 @@ ECS_CONTAINER_CONTEXT_SCHEMA = {
209
216
  is_required=False,
210
217
  description="Additional sidecar containers to include in run task definitions.",
211
218
  ),
219
+ "server_health_check": Field(
220
+ Permissive(),
221
+ is_required=False,
222
+ description="Health check to include in code server task definitions.",
223
+ ),
212
224
  **SHARED_TASK_DEFINITION_FIELDS,
213
225
  **SHARED_ECS_SCHEMA,
214
226
  }
@@ -235,6 +247,7 @@ class EcsContainerContext(
235
247
  ("server_ecs_tags", Sequence[Mapping[str, Optional[str]]]),
236
248
  ("run_ecs_tags", Sequence[Mapping[str, Optional[str]]]),
237
249
  ("repository_credentials", Optional[str]),
250
+ ("server_health_check", Optional[Mapping[str, Any]]),
238
251
  ],
239
252
  )
240
253
  ):
@@ -259,8 +272,9 @@ class EcsContainerContext(
259
272
  server_ecs_tags: Optional[Sequence[Mapping[str, Optional[str]]]] = None,
260
273
  run_ecs_tags: Optional[Sequence[Mapping[str, Optional[str]]]] = None,
261
274
  repository_credentials: Optional[str] = None,
275
+ server_health_check: Optional[Mapping[str, Any]] = None,
262
276
  ):
263
- return super(EcsContainerContext, cls).__new__(
277
+ return super().__new__(
264
278
  cls,
265
279
  secrets=check.opt_sequence_param(secrets, "secrets"),
266
280
  secrets_tags=check.opt_sequence_param(secrets_tags, "secrets_tags"),
@@ -287,6 +301,7 @@ class EcsContainerContext(
287
301
  repository_credentials=check.opt_str_param(
288
302
  repository_credentials, "repository_credentials"
289
303
  ),
304
+ server_health_check=check.opt_mapping_param(server_health_check, "server_health_check"),
290
305
  )
291
306
 
292
307
  def merge(self, other: "EcsContainerContext") -> "EcsContainerContext":
@@ -311,6 +326,7 @@ class EcsContainerContext(
311
326
  server_ecs_tags=[*other.server_ecs_tags, *self.server_ecs_tags],
312
327
  run_ecs_tags=[*other.run_ecs_tags, *self.run_ecs_tags],
313
328
  repository_credentials=other.repository_credentials or self.repository_credentials,
329
+ server_health_check=other.server_health_check or self.server_health_check,
314
330
  )
315
331
 
316
332
  def get_secrets_dict(self, secrets_manager) -> Mapping[str, str]:
@@ -383,7 +399,7 @@ class EcsContainerContext(
383
399
  run_ecs_container_context,
384
400
  )
385
401
 
386
- processed_context_value = cast(Mapping[str, Any], processed_container_context.value)
402
+ processed_context_value = cast("Mapping[str, Any]", processed_container_context.value)
387
403
 
388
404
  return shared_container_context.merge(
389
405
  EcsContainerContext(
@@ -404,5 +420,6 @@ class EcsContainerContext(
404
420
  server_ecs_tags=processed_context_value.get("server_ecs_tags"),
405
421
  run_ecs_tags=processed_context_value.get("run_ecs_tags"),
406
422
  repository_credentials=processed_context_value.get("repository_credentials"),
423
+ server_health_check=processed_context_value.get("server_health_check"),
407
424
  )
408
425
  )