Skip to content

Commit d510d08

Browse files
committed
refactor: use Schema.from_dict in task transforms
1 parent 2014a29 commit d510d08

1 file changed

Lines changed: 184 additions & 146 deletions

File tree

‎src/taskgraph/transforms/task.py‎

Lines changed: 184 additions & 146 deletions
Original file line numberDiff line numberDiff line change
@@ -50,9 +50,13 @@ def run_task_suffix():
5050
return hash_path(RUN_TASK)[0:20]
5151

5252

53-
class WorkerSchema(Schema, forbid_unknown_fields=False, kw_only=True):
54-
# The worker implementation type.
55-
implementation: str
53+
WORKER_SCHEMA = Schema.from_dict(
54+
{
55+
# The worker implementation type.
56+
"implementation": str,
57+
},
58+
forbid_unknown_fields=False,
59+
)
5660

5761

5862
#: Schema for the task transforms
@@ -134,7 +138,7 @@ class TaskDescriptionSchema(Schema):
134138
# the release promotion phase that this task belongs to.
135139
shipping_phase: Optional[Literal["build", "promote", "push", "ship"]] = None
136140
# Information specific to the worker implementation that will run this task.
137-
worker: Optional[WorkerSchema] = None
141+
worker: Optional[WORKER_SCHEMA] = None
138142

139143
def __post_init__(self):
140144
if self.dependencies:
@@ -241,57 +245,65 @@ def verify_index(config, index):
241245
DockerImage = Union[str, dict[str, str]]
242246

243247

244-
class DockerWorkerCacheEntry(Schema):
245-
# only one type is supported by any of the workers right now
246-
type: Literal["persistent"] = "persistent"
247-
# name of the cache, allowing reuse by subsequent tasks naming the same cache
248-
name: Optional[str] = None
249-
# location in the task image where the cache will be mounted
250-
mount_point: Optional[str] = None
251-
# Whether the cache is not used in untrusted environments (like the Try repo).
252-
skip_untrusted: Optional[bool] = None
253-
254-
255-
class DockerWorkerArtifact(Schema):
256-
# type of artifact -- simple file, or recursive directory, or a volume mounted directory.
257-
type: Optional[Literal["file", "directory", "volume"]] = None
258-
# task image path from which to read artifact
259-
path: Optional[str] = None
260-
# name of the produced artifact (root of the names for type=directory)
261-
name: Optional[str] = None
262-
263-
264-
class DockerWorkerPayloadSchema(Schema, forbid_unknown_fields=False, kw_only=True):
265-
implementation: Literal["docker-worker"]
266-
os: Literal["linux"]
267-
# For tasks that will run in docker-worker, this is the name of the docker
268-
# image or in-tree docker image to run the task in.
269-
docker_image: DockerImage
270-
# worker features that should be enabled
271-
chain_of_trust: bool
272-
taskcluster_proxy: bool
273-
allow_ptrace: bool
274-
loopback_video: bool
275-
# environment variables
276-
env: dict[str, taskref_or_string_msgspec]
277-
# the maximum time to run, in seconds
278-
max_run_time: int
279-
# Paths to Docker volumes.
280-
volumes: Optional[list[str]] = None
281-
# caches to set up for the task
282-
caches: Optional[list[DockerWorkerCacheEntry]] = None
283-
# artifacts to extract from the task image after completion
284-
artifacts: Optional[list[DockerWorkerArtifact]] = None
285-
# the command to run; if not given, docker-worker will default to the
286-
# command in the docker image
287-
command: Optional[list[taskref_or_string_msgspec]] = None
288-
# the exit status code(s) that indicates the task should be retried
289-
retry_exit_status: Optional[list[int]] = None
290-
# the exit status code(s) that indicates the caches used by the task
291-
# should be purged
292-
purge_caches_exit_status: Optional[list[int]] = None
293-
# Whether any artifacts are assigned to this worker
294-
skip_artifacts: Optional[bool] = None
248+
DOCKER_WORKER_CACHE_ENTRY = Schema.from_dict(
249+
{
250+
# only one type is supported by any of the workers right now
251+
"type": (Literal["persistent"], "persistent"),
252+
# name of the cache, allowing reuse by subsequent tasks naming the same cache
253+
"name": (str, None),
254+
# location in the task image where the cache will be mounted
255+
"mount-point": (str, None),
256+
# Whether the cache is not used in untrusted environments (like the Try repo).
257+
"skip-untrusted": (bool, None),
258+
},
259+
)
260+
261+
DOCKER_WORKER_ARTIFACT = Schema.from_dict(
262+
{
263+
# type of artifact -- simple file, or recursive directory, or a volume mounted directory.
264+
"type": (Optional[Literal["file", "directory", "volume"]], None),
265+
# task image path from which to read artifact
266+
"path": (str, None),
267+
# name of the produced artifact (root of the names for type=directory)
268+
"name": (str, None),
269+
},
270+
)
271+
272+
DockerWorkerPayloadSchema = Schema.from_dict(
273+
{
274+
"implementation": Literal["docker-worker"],
275+
"os": Literal["linux"],
276+
# For tasks that will run in docker-worker, this is the name of the docker
277+
# image or in-tree docker image to run the task in.
278+
"docker-image": DockerImage,
279+
# worker features that should be enabled
280+
"chain-of-trust": bool,
281+
"taskcluster-proxy": bool,
282+
"allow-ptrace": bool,
283+
"loopback-video": bool,
284+
# environment variables
285+
"env": dict[str, taskref_or_string_msgspec],
286+
# the maximum time to run, in seconds
287+
"max-run-time": int,
288+
# Paths to Docker volumes.
289+
"volumes": (list[str], None),
290+
# caches to set up for the task
291+
"caches": (list[DOCKER_WORKER_CACHE_ENTRY], None),
292+
# artifacts to extract from the task image after completion
293+
"artifacts": (list[DOCKER_WORKER_ARTIFACT], None),
294+
# the command to run; if not given, docker-worker will default to the
295+
# command in the docker image
296+
"command": (list[taskref_or_string_msgspec], None),
297+
# the exit status code(s) that indicates the task should be retried
298+
"retry-exit-status": (list[int], None),
299+
# the exit status code(s) that indicates the caches used by the task
300+
# should be purged
301+
"purge-caches-exit-status": (list[int], None),
302+
# Whether any artifacts are assigned to this worker
303+
"skip-artifacts": (bool, None),
304+
},
305+
forbid_unknown_fields=False,
306+
)
295307

296308

297309
@payload_builder("docker-worker", schema=DockerWorkerPayloadSchema)
@@ -493,69 +505,78 @@ def build_docker_worker_payload(config, task, task_def):
493505
payload["capabilities"] = capabilities
494506

495507

496-
class GenericWorkerArtifact(Schema):
497-
# type of artifact -- simple file, or recursive directory
498-
type: Literal["file", "directory"]
499-
# filesystem path from which to read artifact
500-
path: str
501-
# if not specified, path is used for artifact name
502-
name: Optional[str] = None
503-
504-
505-
class MountContentSchema(Schema):
506-
# Artifact name that contains the content.
507-
artifact: Optional[str] = None
508-
# Task ID that has the artifact that contains the content.
509-
task_id: Optional[taskref_or_string_msgspec] = None
510-
# URL that supplies the content in response to an unauthenticated GET request.
511-
url: Optional[str] = None
512-
513-
514-
class MountSchema(Schema):
515-
# A unique name for the cache volume.
516-
cache_name: Optional[str] = None
517-
# Optional content for pre-loading cache, or mandatory content for
518-
# read-only file or directory.
519-
content: Optional[MountContentSchema] = None
520-
# If mounting a cache or read-only directory.
521-
directory: Optional[str] = None
522-
# If mounting a file.
523-
file: Optional[str] = None
524-
# Archive format of the content.
525-
format: Optional[Literal["rar", "tar.bz2", "tar.gz", "zip"]] = None
526-
527-
528-
class GenericWorkerPayloadSchema(Schema, forbid_unknown_fields=False, kw_only=True):
529-
implementation: Literal["generic-worker"]
530-
os: Literal["windows", "macosx", "linux", "linux-bitbar"]
531-
# command is a list of commands to run, sequentially
532-
# On Windows, each command is a string; on Linux/OS X, each command is a string array
533-
command: list
534-
# environment variables
535-
env: dict[str, taskref_or_string_msgspec]
536-
# the maximum time to run, in seconds
537-
max_run_time: int
538-
# optional features
539-
chain_of_trust: bool
540-
# artifacts to extract from the task image after completion
541-
artifacts: Optional[list[GenericWorkerArtifact]] = None
542-
# Directories and/or files to be mounted.
543-
mounts: Optional[list[MountSchema]] = None
544-
# the exit status code(s) that indicates the task should be retried
545-
retry_exit_status: Optional[list[int]] = None
546-
# the exit status code(s) that indicates the caches used by the task
547-
# should be purged
548-
purge_caches_exit_status: Optional[list[int]] = None
549-
# os user groups for test task workers
550-
os_groups: Optional[list[str]] = None
551-
# feature for test task to run as administrator
552-
run_as_administrator: Optional[bool] = None
553-
# feature for task to run as current OS user
554-
run_task_as_current_user: Optional[bool] = None
555-
taskcluster_proxy: Optional[bool] = None
556-
hide_cmd_window: Optional[bool] = None
557-
# Whether any artifacts are assigned to this worker
558-
skip_artifacts: Optional[bool] = None
508+
GENERIC_WORKER_ARTIFACT = Schema.from_dict(
509+
{
510+
# type of artifact -- simple file, or recursive directory
511+
"type": Literal["file", "directory"],
512+
# filesystem path from which to read artifact
513+
"path": str,
514+
# if not specified, path is used for artifact name
515+
"name": (str, None),
516+
},
517+
)
518+
519+
MOUNT_SCHEMA = Schema.from_dict(
520+
{
521+
# A unique name for the cache volume.
522+
"cache-name": (str, None),
523+
# Optional content for pre-loading cache, or mandatory content for
524+
# read-only file or directory.
525+
"content": Schema.from_dict(
526+
{
527+
# Artifact name that contains the content.
528+
"artifact": (str, None),
529+
# Task ID that has the artifact that contains the content.
530+
"task-id": (taskref_or_string_msgspec, None),
531+
# URL that supplies the content in response to an unauthenticated GET request.
532+
"url": (str, None),
533+
},
534+
optional=True,
535+
),
536+
# If mounting a cache or read-only directory.
537+
"directory": (str, None),
538+
# If mounting a file.
539+
"file": (str, None),
540+
# Archive format of the content.
541+
"format": (Optional[Literal["rar", "tar.bz2", "tar.gz", "zip"]], None),
542+
},
543+
)
544+
545+
GenericWorkerPayloadSchema = Schema.from_dict(
546+
{
547+
"implementation": Literal["generic-worker"],
548+
"os": Literal["windows", "macosx", "linux", "linux-bitbar"],
549+
# command is a list of commands to run, sequentially
550+
# On Windows, each command is a string; on Linux/OS X, each command is a string array
551+
"command": list,
552+
# environment variables
553+
"env": dict[str, taskref_or_string_msgspec],
554+
# the maximum time to run, in seconds
555+
"max-run-time": int,
556+
# optional features
557+
"chain-of-trust": bool,
558+
# artifacts to extract from the task image after completion
559+
"artifacts": (list[GENERIC_WORKER_ARTIFACT], None),
560+
# Directories and/or files to be mounted.
561+
"mounts": (list[MOUNT_SCHEMA], None),
562+
# the exit status code(s) that indicates the task should be retried
563+
"retry-exit-status": (list[int], None),
564+
# the exit status code(s) that indicates the caches used by the task
565+
# should be purged
566+
"purge-caches-exit-status": (list[int], None),
567+
# os user groups for test task workers
568+
"os-groups": (list[str], None),
569+
# feature for test task to run as administrator
570+
"run-as-administrator": (bool, None),
571+
# feature for task to run as current OS user
572+
"run-task-as-current-user": (bool, None),
573+
"taskcluster-proxy": (bool, None),
574+
"hide-cmd-window": (bool, None),
575+
# Whether any artifacts are assigned to this worker
576+
"skip-artifacts": (bool, None),
577+
},
578+
forbid_unknown_fields=False,
579+
)
559580

560581

561582
@payload_builder("generic-worker", schema=GenericWorkerPayloadSchema)
@@ -673,13 +694,16 @@ def build_generic_worker_payload(config, task, task_def):
673694
task_def["payload"]["features"] = features
674695

675696

676-
class ReleaseProperties(Schema):
677-
app_name: Optional[str] = None
678-
app_version: Optional[str] = None
679-
branch: Optional[str] = None
680-
build_id: Optional[str] = None
681-
hash_type: Optional[str] = None
682-
platform: Optional[str] = None
697+
RELEASE_PROPERTIES = Schema.from_dict(
698+
{
699+
"app-name": (str, None),
700+
"app-version": (str, None),
701+
"branch": (str, None),
702+
"build-id": (str, None),
703+
"hash-type": (str, None),
704+
"platform": (str, None),
705+
},
706+
)
683707

684708

685709
class UpstreamArtifact(Schema, rename="camel"):
@@ -693,18 +717,22 @@ class UpstreamArtifact(Schema, rename="camel"):
693717
locale: str
694718

695719

696-
class BeetmoverPayloadSchema(Schema, forbid_unknown_fields=False, kw_only=True):
697-
implementation: Literal["beetmover"]
698-
# the maximum time to run, in seconds
699-
max_run_time: int
700-
# release properties
701-
release_properties: ReleaseProperties
702-
# list of artifact URLs for the artifacts that should be beetmoved
703-
upstream_artifacts: list[UpstreamArtifact]
704-
# locale key, if this is a locale beetmover task
705-
locale: Optional[str] = None
706-
partner_public: Optional[bool] = None
707-
artifact_map: Optional[object] = None
720+
BeetmoverPayloadSchema = Schema.from_dict(
721+
{
722+
"implementation": Literal["beetmover"],
723+
# the maximum time to run, in seconds
724+
"max-run-time": int,
725+
# release properties
726+
"release-properties": RELEASE_PROPERTIES,
727+
# list of artifact URLs for the artifacts that should be beetmoved
728+
"upstream-artifacts": list[UpstreamArtifact],
729+
# locale key, if this is a locale beetmover task
730+
"locale": (str, None),
731+
"partner-public": (bool, None),
732+
"artifact-map": (object, None),
733+
},
734+
forbid_unknown_fields=False,
735+
)
708736

709737

710738
@payload_builder("beetmover", schema=BeetmoverPayloadSchema)
@@ -733,23 +761,33 @@ def build_beetmover_payload(config, task, task_def):
733761
task_def["payload"]["is_partner_repack_public"] = worker["partner-public"]
734762

735763

736-
class InvalidPayloadSchema(Schema, forbid_unknown_fields=False, kw_only=True):
737-
# an invalid task is one which should never actually be created; this is used in
738-
# release automation on branches where the task just doesn't make sense
739-
implementation: Literal["invalid"]
764+
InvalidPayloadSchema = Schema.from_dict(
765+
{
766+
# an invalid task is one which should never actually be created; this is used in
767+
# release automation on branches where the task just doesn't make sense
768+
"implementation": Literal["invalid"],
769+
},
770+
forbid_unknown_fields=False,
771+
)
740772

741773

742774
@payload_builder("invalid", schema=InvalidPayloadSchema)
743775
def build_invalid_payload(config, task, task_def):
744776
task_def["payload"] = "invalid task - should never be created"
745777

746778

747-
class AlwaysOptimizedPayloadSchema(Schema, forbid_unknown_fields=False, kw_only=True):
748-
implementation: Literal["always-optimized"]
749-
779+
AlwaysOptimizedPayloadSchema = Schema.from_dict(
780+
{
781+
"implementation": Literal["always-optimized"],
782+
},
783+
forbid_unknown_fields=False,
784+
)
750785

751-
class SucceedPayloadSchema(Schema):
752-
implementation: Literal["succeed"]
786+
SucceedPayloadSchema = Schema.from_dict(
787+
{
788+
"implementation": Literal["succeed"],
789+
},
790+
)
753791

754792

755793
@payload_builder("always-optimized", schema=AlwaysOptimizedPayloadSchema)

0 commit comments

Comments
 (0)