Support Multiple S3 and Docker Endpoints, and Offering Vault Integration alongside SSM (#350)

umang8223 · robot-clickhouse · web-flow · commit a16f00e704f8 · 2023-06-16T12:24:43.000-04:00
* Parametrize Docker image registry

* Refactoring scripts for generalization

* Automatic style fix

* pr changes

* added import hvac in get_vault_robot_tokens

* pr-changes

---------

Co-authored-by: robot-clickhouse &lt;robot-clickhouse@users.noreply.github.com&gt;
diff --git a/docker/packager/binary/Dockerfile b/docker/packager/binary/Dockerfile
@@ -1,6 +1,7 @@
 # docker build -t clickhouse/binary-builder .
 ARG FROM_TAG=latest
-FROM clickhouse/test-util:latest AS cctools
+ARG DOCKER_REPO=docker.io
+FROM $DOCKER_REPO/clickhouse/test-util:latest AS cctools
 # The cctools are built always from the clickhouse/test-util:latest and cached inline
 # Theoretically, it should improve rebuild speed significantly
 ENV CC=clang-${LLVM_VERSION}
@@ -39,7 +40,7 @@ RUN git clone --depth 1 https://github.com/tpoechtrager/cctools-port.git \
 # END COMPILE
 # !!!!!!!!!!!
 
-FROM clickhouse/test-util:$FROM_TAG
+FROM $DOCKER_REPO/clickhouse/test-util:$FROM_TAG
 ENV CC=clang-${LLVM_VERSION}
 ENV CXX=clang++-${LLVM_VERSION}
 
diff --git a/docker/test/base/Dockerfile b/docker/test/base/Dockerfile
@@ -1,7 +1,8 @@
 # rebuild in #33610
 # docker build -t clickhouse/test-base .
 ARG FROM_TAG=latest
-FROM clickhouse/test-util:$FROM_TAG
+ARG DOCKER_REPO=docker.io
+FROM $DOCKER_REPO/clickhouse/test-util:$FROM_TAG
 
 RUN apt-get update \
     && apt-get install \
diff --git a/docker/test/codebrowser/Dockerfile b/docker/test/codebrowser/Dockerfile
@@ -2,7 +2,8 @@
 # docker build --network=host -t clickhouse/codebrowser .
 # docker run --volume=path_to_repo:/repo_folder --volume=path_to_result:/test_output clickhouse/codebrowser
 ARG FROM_TAG=latest
-FROM clickhouse/binary-builder:$FROM_TAG
+ARG DOCKER_REPO=docker.io
+FROM $DOCKER_REPO/clickhouse/binary-builder:$FROM_TAG
 
 # ARG for quick switch to a given ubuntu mirror
 ARG apt_archive="http://archive.ubuntu.com"
diff --git a/docker/test/fasttest/Dockerfile b/docker/test/fasttest/Dockerfile
@@ -1,7 +1,8 @@
 # rebuild in #33610
 #  docker build -t clickhouse/fasttest .
 ARG FROM_TAG=latest
-FROM clickhouse/test-util:$FROM_TAG
+ARG DOCKER_REPO=docker.io
+FROM $DOCKER_REPO/clickhouse/test-util:$FROM_TAG
 
 RUN apt-get update \
     && apt-get install \
diff --git a/docker/test/fuzzer/Dockerfile b/docker/test/fuzzer/Dockerfile
@@ -1,7 +1,8 @@
 # rebuild in #33610
 # docker build -t clickhouse/fuzzer .
 ARG FROM_TAG=latest
-FROM clickhouse/test-base:$FROM_TAG
+ARG DOCKER_REPO=docker.io
+FROM $DOCKER_REPO/clickhouse/test-base:$FROM_TAG
 
 # ARG for quick switch to a given ubuntu mirror
 ARG apt_archive="http://archive.ubuntu.com"
diff --git a/docker/test/integration/base/Dockerfile b/docker/test/integration/base/Dockerfile
@@ -1,7 +1,8 @@
 # rebuild in #33610
 # docker build -t clickhouse/integration-test .
 ARG FROM_TAG=latest
-FROM clickhouse/test-base:$FROM_TAG
+ARG DOCKER_REPO=docker.io
+FROM $DOCKER_REPO/clickhouse/test-base:$FROM_TAG
 
 SHELL ["/bin/bash", "-c"]
 
diff --git a/docker/test/keeper-jepsen/Dockerfile b/docker/test/keeper-jepsen/Dockerfile
@@ -1,7 +1,8 @@
 # rebuild in #33610
 # docker build -t clickhouse/keeper-jepsen-test .
 ARG FROM_TAG=latest
-FROM clickhouse/test-base:$FROM_TAG
+ARG DOCKER_REPO=docker.io
+FROM $DOCKER_REPO/clickhouse/test-base:$FROM_TAG
 
 ENV DEBIAN_FRONTEND=noninteractive
 ENV CLOJURE_VERSION=1.10.3.814
diff --git a/docker/test/server-jepsen/Dockerfile b/docker/test/server-jepsen/Dockerfile
@@ -1,7 +1,8 @@
 # rebuild in #33610
 # docker build -t clickhouse/server-jepsen-test .
 ARG FROM_TAG=latest
-FROM clickhouse/test-base:$FROM_TAG
+ARG DOCKER_REPO=docker.io
+FROM $DOCKER_REPO/clickhouse/test-base:$FROM_TAG
 
 ENV DEBIAN_FRONTEND=noninteractive
 ENV CLOJURE_VERSION=1.10.3.814
diff --git a/docker/test/sqllogic/Dockerfile b/docker/test/sqllogic/Dockerfile
@@ -1,6 +1,7 @@
 # docker build -t clickhouse/sqllogic-test .
 ARG FROM_TAG=latest
-FROM clickhouse/test-base:$FROM_TAG
+ARG DOCKER_REPO=docker.io
+FROM $DOCKER_REPO/clickhouse/test-base:$FROM_TAG
 
 RUN apt-get update --yes \
     && env DEBIAN_FRONTEND=noninteractive \
diff --git a/docker/test/stateful/Dockerfile b/docker/test/stateful/Dockerfile
@@ -1,7 +1,8 @@
 # rebuild in #47031
 # docker build -t clickhouse/stateful-test .
 ARG FROM_TAG=latest
-FROM clickhouse/stateless-test:$FROM_TAG
+ARG DOCKER_REPO=docker.io
+FROM $DOCKER_REPO/clickhouse/stateless-test:$FROM_TAG
 
 RUN apt-get update -y \
     && env DEBIAN_FRONTEND=noninteractive \
diff --git a/docker/test/stateless/Dockerfile b/docker/test/stateless/Dockerfile
@@ -1,7 +1,8 @@
 # rebuild in #33610
 # docker build -t clickhouse/stateless-test .
 ARG FROM_TAG=latest
-FROM clickhouse/test-base:$FROM_TAG
+ARG DOCKER_REPO=docker.io
+FROM $DOCKER_REPO/clickhouse/test-base:$FROM_TAG
 
 ARG odbc_driver_url="https://github.com/ClickHouse/clickhouse-odbc/releases/download/v1.1.4.20200302/clickhouse-odbc-1.1.4-Linux.tar.gz"
 
diff --git a/docker/test/stress/Dockerfile b/docker/test/stress/Dockerfile
@@ -1,7 +1,8 @@
 # rebuild in #33610
 # docker build -t clickhouse/stress-test .
 ARG FROM_TAG=latest
-FROM clickhouse/stateful-test:$FROM_TAG
+ARG DOCKER_REPO=docker.io
+FROM $DOCKER_REPO/clickhouse/stateful-test:$FROM_TAG
 
 RUN apt-get update -y \
     && env DEBIAN_FRONTEND=noninteractive \
diff --git a/docker/test/style/Dockerfile b/docker/test/style/Dockerfile
@@ -18,7 +18,12 @@ RUN apt-get update && env DEBIAN_FRONTEND=noninteractive apt-get install --yes \
     python3-pip \
     shellcheck \
     yamllint \
-    && pip3 install black==23.1.0 boto3 codespell==2.2.1 dohq-artifactory mypy PyGithub unidiff pylint==2.6.2 \
+    && if [ "$DOCKER_REPO" = "docker.io" ]; then \
+         pip3 install black==23.1.0 boto3 codespell==2.2.1 dohq-artifactory mypy PyGithub unidiff pylint==2.6.2; \
+       else \
+       # the `hvac` is an optional dependency for Hashicorp Vault as a parameters store
+         pip3 install black==23.1.0 boto3 hvac codespell==2.2.1 dohq-artifactory mypy PyGithub unidiff pylint==2.6.2; \
+       fi \
     && apt-get clean \
     && rm -rf /root/.cache/pip
 
diff --git a/docker/test/unit/Dockerfile b/docker/test/unit/Dockerfile
@@ -1,7 +1,8 @@
 # rebuild in #33610
 # docker build -t clickhouse/unit-test .
 ARG FROM_TAG=latest
-FROM clickhouse/stateless-test:$FROM_TAG
+ARG DOCKER_REPO=docker.io
+FROM $DOCKER_REPO/clickhouse/stateless-test:$FROM_TAG
 
 RUN apt-get install gdb
 
diff --git a/docker/test/upgrade/Dockerfile b/docker/test/upgrade/Dockerfile
@@ -1,7 +1,8 @@
 # rebuild in #33610
 # docker build -t clickhouse/upgrade-check .
 ARG FROM_TAG=latest
-FROM clickhouse/stateful-test:$FROM_TAG
+ARG DOCKER_REPO=docker.io
+FROM $DOCKER_REPO/clickhouse/stateless-test:$FROM_TAG
 
 RUN apt-get update -y \
     && env DEBIAN_FRONTEND=noninteractive \
diff --git a/tests/ci/clickhouse_helper.py b/tests/ci/clickhouse_helper.py
@@ -7,6 +7,7 @@
 import requests  # type: ignore
 
 from get_robot_token import get_parameter_from_ssm
+from env_helper import GITHUB_REPOSITORY
 from pr_info import PRInfo
 from report import TestResults
 
@@ -141,7 +142,7 @@ def prepare_tests_results_for_clickhouse(
     report_url: str,
     check_name: str,
 ) -> List[dict]:
-    pull_request_url = "https://github.com/ClickHouse/ClickHouse/commits/master"
+    pull_request_url = f"https://github.com/{GITHUB_REPOSITORY}/commits/master"
     base_ref = "master"
     head_ref = "master"
     base_repo = pr_info.repo_full_name
diff --git a/tests/ci/docker_images_check.py b/tests/ci/docker_images_check.py
@@ -15,7 +15,13 @@
 
 from clickhouse_helper import ClickHouseHelper, prepare_tests_results_for_clickhouse
 from commit_status_helper import format_description, get_commit, post_commit_status
-from env_helper import GITHUB_WORKSPACE, RUNNER_TEMP, GITHUB_RUN_URL
+from env_helper import (
+    GITHUB_WORKSPACE,
+    RUNNER_TEMP,
+    GITHUB_RUN_URL,
+    DOCKER_USER,
+    DOCKER_REPO,
+)
 from get_robot_token import get_best_robot_token, get_parameter_from_ssm
 from pr_info import PRInfo
 from report import TestResults, TestResult
@@ -94,7 +100,7 @@ def get_images_dict(repo_path: str, image_file_path: str) -> ImagesDict:
 
 
 def get_changed_docker_images(
-    pr_info: PRInfo, images_dict: ImagesDict
+    pr_info: PRInfo, images_dict: ImagesDict, DOCKER_REPO: str
 ) -> Set[DockerImage]:
     if not images_dict:
         return set()
@@ -113,7 +119,7 @@ def get_changed_docker_images(
     for dockerfile_dir, image_description in images_dict.items():
         for f in files_changed:
             if f.startswith(dockerfile_dir):
-                name = image_description["name"]
+                name = f"{DOCKER_REPO}/{image_description['name']}"
                 only_amd64 = image_description.get("only_amd64", False)
                 logging.info(
                     "Found changed file '%s' which affects "
@@ -137,7 +143,7 @@ def get_changed_docker_images(
                 dependent,
                 image,
             )
-            name = images_dict[dependent]["name"]
+            name = f"{DOCKER_REPO}/{images_dict[dependent]['name']}"
             only_amd64 = images_dict[dependent].get("only_amd64", False)
             changed_images.append(DockerImage(dependent, name, only_amd64, image))
         index += 1
@@ -248,6 +254,7 @@ def build_and_push_one_image(
         "docker buildx build --builder default "
         f"--label build-url={GITHUB_RUN_URL} "
         f"{from_tag_arg}"
+        f"--build-arg DOCKER_REPO={DOCKER_REPO} "
         # A hack to invalidate cache, grep for it in docker/ dir
         f"--build-arg CACHE_INVALIDATOR={GITHUB_RUN_URL} "
         f"--tag {image.repo}:{version_string} "
@@ -400,7 +407,7 @@ def main():
 
     if args.push:
         subprocess.check_output(  # pylint: disable=unexpected-keyword-arg
-            "docker login --username 'robotclickhouse' --password-stdin",
+            f"docker login {DOCKER_REPO} --username '{DOCKER_USER}' --password-stdin",
             input=get_parameter_from_ssm("dockerhub_robot_password"),
             encoding="utf-8",
             shell=True,
@@ -424,7 +431,7 @@ def main():
             # If the event does not contain diff, nothing will be built
             pass
 
-    changed_images = get_changed_docker_images(pr_info, images_dict)
+    changed_images = get_changed_docker_images(pr_info, images_dict, DOCKER_REPO)
     if changed_images:
         logging.info(
             "Has changed images: %s", ", ".join([im.path for im in changed_images])
diff --git a/tests/ci/docker_manifests_merge.py b/tests/ci/docker_manifests_merge.py
@@ -11,7 +11,7 @@
 
 from clickhouse_helper import ClickHouseHelper, prepare_tests_results_for_clickhouse
 from commit_status_helper import format_description, get_commit, post_commit_status
-from env_helper import RUNNER_TEMP
+from env_helper import RUNNER_TEMP, DOCKER_USER, DOCKER_REPO
 from get_robot_token import get_best_robot_token, get_parameter_from_ssm
 from pr_info import PRInfo
 from report import TestResults, TestResult
@@ -174,7 +174,7 @@ def main():
     args = parse_args()
     if args.push:
         subprocess.check_output(  # pylint: disable=unexpected-keyword-arg
-            "docker login --username 'robotclickhouse' --password-stdin",
+            f"docker login {DOCKER_REPO} --username '{DOCKER_USER}' --password-stdin",
             input=get_parameter_from_ssm("dockerhub_robot_password"),
             encoding="utf-8",
             shell=True,
diff --git a/tests/ci/env_helper.py b/tests/ci/env_helper.py
@@ -14,6 +14,8 @@
 CLOUDFLARE_TOKEN = os.getenv("CLOUDFLARE_TOKEN")
 GITHUB_EVENT_PATH = os.getenv("GITHUB_EVENT_PATH", "")
 GITHUB_JOB = os.getenv("GITHUB_JOB", "local")
+DOCKER_REPO = os.getenv("DOCKER_REPO", "docker.io")
+DOCKER_USER = os.getenv("DOCKER_USER", "robotclickhouse")
 GITHUB_REPOSITORY = os.getenv("GITHUB_REPOSITORY", "ClickHouse/ClickHouse")
 GITHUB_RUN_ID = os.getenv("GITHUB_RUN_ID", "0")
 GITHUB_SERVER_URL = os.getenv("GITHUB_SERVER_URL", "https://github.com")
@@ -23,6 +25,7 @@
 REPORTS_PATH = os.getenv("REPORTS_PATH", p.abspath(p.join(module_dir, "./reports")))
 REPO_COPY = os.getenv("REPO_COPY", git_root)
 RUNNER_TEMP = os.getenv("RUNNER_TEMP", p.abspath(p.join(module_dir, "./tmp")))
+S3_REGION = os.getenv("S3_REGION", "us-east-1")
 S3_BUILDS_BUCKET = os.getenv("S3_BUILDS_BUCKET", "clickhouse-builds")
 S3_TEST_REPORTS_BUCKET = os.getenv("S3_TEST_REPORTS_BUCKET", "clickhouse-test-reports")
 S3_URL = os.getenv("S3_URL", "https://s3.amazonaws.com")
@@ -31,6 +34,10 @@
     f"{S3_DOWNLOAD}/{S3_BUILDS_BUCKET}/"
     "{pr_or_release}/{commit}/{build_name}/{artifact}"
 )
+VAULT_PATH = os.getenv("VAULT_PATH")
+VAULT_TOKEN = os.getenv("VAULT_TOKEN")
+VAULT_URL = os.getenv("VAULT_URL")
+VAULT_MOUNT_POINT = os.getenv("VAULT_MOUNT_POINT", "secret")
 
 # These parameters are set only on demand, and only once
 _GITHUB_JOB_ID = ""
diff --git a/tests/ci/get_robot_token.py b/tests/ci/get_robot_token.py
@@ -7,6 +7,8 @@
 from github import Github
 from github.AuthenticatedUser import AuthenticatedUser
 
+from env_helper import VAULT_URL, VAULT_TOKEN, VAULT_PATH, VAULT_MOUNT_POINT, S3_REGION
+
 
 @dataclass
 class Token:
@@ -16,9 +18,20 @@ class Token:
 
 
 def get_parameter_from_ssm(name, decrypt=True, client=None):
-    if not client:
-        client = boto3.client("ssm", region_name="us-east-1")
-    return client.get_parameter(Name=name, WithDecryption=decrypt)["Parameter"]["Value"]
+    if VAULT_URL:
+        import hvac
+        if not client:
+            client = hvac.Client(url=VAULT_URL, token=VAULT_TOKEN)
+        parameter = client.secrets.kv.v2.read_secret_version(
+            mount_point=VAULT_MOUNT_POINT, path=VAULT_PATH
+        )["data"]["data"][name]
+    else:
+        if not client:
+            client = boto3.client("ssm", region_name=S3_REGION)
+        parameter = client.get_parameter(Name=name, WithDecryption=decrypt)[
+            "Parameter"
+        ]["Value"]
+    return parameter
 
 
 ROBOT_TOKEN = None  # type: Optional[Token]
@@ -28,18 +41,49 @@ def get_best_robot_token(token_prefix_env_name="github_robot_token_"):
     global ROBOT_TOKEN
     if ROBOT_TOKEN is not None:
         return ROBOT_TOKEN.value
-    client = boto3.client("ssm", region_name="us-east-1")
-    parameters = client.describe_parameters(
-        ParameterFilters=[
-            {"Key": "Name", "Option": "BeginsWith", "Values": [token_prefix_env_name]}
-        ]
-    )["Parameters"]
-    assert parameters
-
-    for token_name in [p["Name"] for p in parameters]:
-        value = get_parameter_from_ssm(token_name, True, client)
+
+    def get_vault_robot_tokens():
+        import hvac
+        client = hvac.Client(url=VAULT_URL, token=VAULT_TOKEN)
+        parameters = client.secrets.kv.v2.read_secret_version(
+            mount_point=VAULT_MOUNT_POINT, path=VAULT_PATH
+        )["data"]["data"]
+        parameters = {
+            key: value
+            for key, value in parameters.items()
+            if key.startswith(token_prefix_env_name)
+        }
+        assert parameters
+        return list(parameters.values())
+
+    def get_ssm_robot_tokens():
+        client = boto3.client("ssm", region_name=S3_REGION)
+        parameters = client.describe_parameters(
+            ParameterFilters=[
+                {
+                    "Key": "Name",
+                    "Option": "BeginsWith",
+                    "Values": [token_prefix_env_name],
+                }
+            ]
+        )["Parameters"]
+        assert parameters
+        for token_name in [p["Name"] for p in parameters]:
+            value = get_parameter_from_ssm(token_name, True, client)
+            values.append(value)
+        return values
+
+    client = None
+    values = []
+
+    if VAULT_URL:
+        values = get_vault_robot_tokens()
+    else:
+        values = get_ssm_robot_tokens()
+
+    for value in values:
         gh = Github(value, per_page=100)
-        # Do not spend additional request to API by accessin user.login unless
+        # Do not spend additional request to API by accessing user.login unless
         # the token is chosen by the remaining requests number
         user = gh.get_user()
         rest, _ = gh.rate_limiting
diff --git a/tests/ci/jepsen_check.py b/tests/ci/jepsen_check.py
diff --git a/tests/ci/s3_helper.py b/tests/ci/s3_helper.py
