...

Text file src/edge-infra.dev/hack/build/rules/google_storage/gsutil_file.bzl

Documentation: edge-infra.dev/hack/build/rules/google_storage

     1"""Fetches a file from a gs bucket"""
     2
     3_GSUTIL_FILE_BUILD = """\
     4package(default_visibility = ["//visibility:public"])
     5
     6filegroup(
     7    name = "file",
     8    srcs = ["{}"],
     9)
    10"""
    11
    12def run_command(ctx, cmd):
    13    """Runs a cli command.
    14
    15    Args:
    16        ctx: the repository context
    17        cmd: the command to run as a list of strings
    18
    19    Returns:
    20    returns the exec_result object returned by the
    21    repository_ctx.execute() method
    22    """
    23
    24    result = ctx.execute(cmd, timeout = 1800)
    25    if result.return_code != 0:
    26        fail("Failed to run {}: {}".format(cmd, result.stderr))
    27
    28    return result
    29
    30def gsutil_download_file(ctx, path):
    31    """Downloads a file from a gs bucket to 'path'.
    32
    33    Args:
    34        ctx: the repository context
    35        path: path to download the file to
    36    """
    37
    38    # Create download directory if it doesn't exist
    39    mkdir_path = ctx.which("mkdir")
    40    mkdir_cmd = [mkdir_path, "-p", ctx.path(path).dirname]
    41    run_command(ctx, mkdir_cmd)
    42
    43    # Download the bucket file
    44    gsutil_path = ctx.which("gsutil")
    45    src_uri = "gs://{}/{}".format(ctx.attr.bucket, ctx.attr.path)
    46    download_cmd = [gsutil_path, "cp", src_uri, path]
    47    ctx.report_progress("Downloading {}.".format(src_uri))
    48    run_command(ctx, download_cmd)
    49
    50def validate_checksum(ctx, path, expected_sha256):
    51    """Validates the checksum of the downloaded file.
    52
    53    Args:
    54        ctx: the repository context
    55        path: the downloaded file path
    56        expected_sha256: the expected sha256 checksum
    57    """
    58
    59    if ctx.attr.sha256 == "":
    60        return
    61
    62    sha256_path = ctx.which("sha256sum")
    63    ctx.report_progress("Checksumming {}.".format(path))
    64    sha256_result = run_command(ctx, [sha256_path, path])
    65
    66    sha256 = sha256_result.stdout.split(" ")[0]
    67    if sha256 != expected_sha256:
    68        fail("Checksum mismatch for {}, expected {}, got {}.".format(
    69            path,
    70            expected_sha256,
    71            sha256,
    72        ))
    73
    74def validate_download_path(ctx, path):
    75    """Validates the download path.
    76
    77    Args:
    78        ctx: the repository context
    79        path: the expected download file path
    80    """
    81
    82    repo_root = ctx.path(".")
    83    forbidden_files = [
    84        repo_root,
    85        ctx.path("WORKSPACE"),
    86        ctx.path("BUILD"),
    87        ctx.path("BUILD.bazel"),
    88        ctx.path("file/BUILD"),
    89        ctx.path("file/BUILD.bazel"),
    90    ]
    91    if path in forbidden_files or not str(path).startswith(str(repo_root)):
    92        fail("'%s' cannot be used as downloaded_file_path in gsutil_file" % ctx.attr.downloaded_file_path)
    93
    94def _gsutil_file_impl(ctx):
    95    """Implementation of the gsutil_file rule."""
    96
    97    # Prepare download path
    98    downloaded_file_path = ctx.attr.downloaded_file_path
    99    download_path = ctx.path("file/" + downloaded_file_path)
   100
   101    # Validate download path
   102    validate_download_path(ctx, download_path)
   103
   104    # Download
   105    gsutil_download_file(ctx, download_path)
   106
   107    # Verify
   108    validate_checksum(ctx, download_path, ctx.attr.sha256)
   109
   110    # Create filegroup
   111    ctx.file("file/BUILD", _GSUTIL_FILE_BUILD.format(downloaded_file_path))
   112
   113_gsutil_file_attrs = {
   114    "bucket": attr.string(
   115        mandatory = True,
   116        doc = "Google storage bucket name",
   117    ),
   118    "path": attr.string(
   119        mandatory = True,
   120        doc = "Relative path to the archive file within the bucket",
   121    ),
   122    "downloaded_file_path": attr.string(
   123        default = "downloaded",
   124        doc = "Path assigned to the downloaded file",
   125    ),
   126    "sha256": attr.string(
   127        doc = "The expected SHA-256 of the downloaded file",
   128    ),
   129}
   130
   131gsutil_file = repository_rule(
   132    implementation = _gsutil_file_impl,
   133    attrs = _gsutil_file_attrs,
   134    doc =
   135        """Downloads a file from a google storage bucket and makes it available 
   136to be used as a file group.
   137
   138Examples:
   139  To get my_deb.deb from gs://my_google_storage_bucket, add this to your
   140  WORKSPACE file:
   141
   142  ```python
   143  load("//hack/build/rules/google_storage:gsutil_file.bzl", "gsutil_file")
   144  
   145  gsutil_file(
   146      name = "my_deb",
   147      bucket = "my_google_storage_bucket",
   148      path = "my_deb.deb",
   149  )
   150  ```
   151
   152  Targets would specify `@my_deb//file` as a dependency to depend on this file.
   153""",
   154)

View as plain text