SERVER-107364 Flag Sync S3 Version (#36152)

GitOrigin-RevId: 83b8a408f600aaff50472a2544a90538cc761d2c

.bazelrc

@@ -518,3 +518,6 @@ try-import %workspace%/.bazelrc.local
 
 # Flag as built with bazelisk
 try-import %workspace%/.bazelrc.bazelisk
+
+# Flags synced from Flag Sync
+try-import %workspace%/.bazelrc.sync

.github/CODEOWNERS

@@ -3150,6 +3150,9 @@ WORKSPACE.bazel @10gen/devprod-build @svc-auto-approve-bot
 # The following patterns are parsed from ./src/third_party/libmongocrypt/OWNERS.yml
 /src/third_party/libmongocrypt/**/* @10gen/server-security @svc-auto-approve-bot
 
+# The following patterns are parsed from ./tools/flag_sync/OWNERS.yml
+/tools/flag_sync/**/* @10gen/devprod-build @svc-auto-approve-bot
+
 # The following patterns are parsed from ./tools/lint/OWNERS.yml
 /tools/lint/**/* @10gen/devprod-build @svc-auto-approve-bot
 

.gitignore

@@ -296,6 +296,7 @@ buildozer
 .bazelrc.workstation
 .bazelrc.common_bes
 .bazelrc.compiledb
+.bazelrc.sync
 .compiledb
 .bazelrc.xcode
 .bazelrc.bazelisk

bazel/wrapper_hook/BUILD.bazel

@@ -15,3 +15,22 @@ py_library(
     ],
     visibility = ["//visibility:public"],
 )
+
+py_library(
+    name = "post_bazel_hook",
+    srcs = [
+        "post_bazel_hook.py",
+    ],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":flag_sync",
+    ],
+)
+
+py_library(
+    name = "flag_sync",
+    srcs = [
+        "flag_sync.py",
+    ],
+    visibility = ["//visibility:public"],
+)

bazel/wrapper_hook/flag_sync.py

@@ -0,0 +1,58 @@
+import os
+import pathlib
+import sys
+import time
+from typing import Dict
+
+REPO_ROOT = pathlib.Path(__file__).parent.parent.parent
+sys.path.append(str(REPO_ROOT))
+
+from bazel.wrapper_hook.wrapper_debug import wrapper_debug
+from tools.flag_sync.util import get_flags
+
+# Allowed .bazelrc lines. Attempt to remove flag setting as attack vector.
+ALLOW_LINES = ["common --config=local"]
+
+
+def update_bazelrc(flags: Dict[str, Dict], verbose: bool):
+    bazelrc_path = f"{REPO_ROOT}/.bazelrc.sync"
+    if verbose:
+        print(f"Updating {bazelrc_path}")
+
+    enabled = set()
+    for flag in flags.values():
+        if flag["enabled"]:
+            enabled.add(flag["value"])
+
+    changed = True
+    if os.path.exists(bazelrc_path):
+        with open(bazelrc_path, "r") as bazelrc:
+            lines = bazelrc.readlines()
+            bazelrc = set([l.strip() for l in lines])
+            changed = bazelrc != enabled
+
+    if not changed:
+        return
+
+    with open(bazelrc_path, "w+") as bazelrc:
+        for line in enabled:
+            if line in ALLOW_LINES:
+                bazelrc.write(line + "\n")
+            else:
+                print("Tried to write unallowed line. Skipping...")
+
+
+def sync_and_update(namespace: str):
+    flags = get_flags(namespace)
+    update_bazelrc(flags, False)
+
+
+def sync_flags(namespace: str) -> bool:
+    start = time.time()
+    try:
+        sync_and_update(namespace)
+    except Exception:
+        print("Failed to sync bazel flags. Skipping...")
+        return False
+    wrapper_debug(f"flag sync time: {time.time() - start}")
+    return True

bazel/wrapper_hook/install_modules.py

@@ -105,7 +105,7 @@ def install_modules(bazel):
         with open(lockfile_hash_file, "w") as f:
             f.write(current_hash)
 
-    deps = ["retry", "gitpython"]
+    deps = ["retry", "gitpython", "requests", "timeout-decorator"]
     deps_installed = []
     deps_needed = search_for_modules(
         deps, deps_installed, lockfile_changed=old_hash != current_hash
@@ -122,7 +122,7 @@ def install_modules(bazel):
         cmd = [
             bazel,
             "build",
-        ] + ["@poetry//:library_" + dep for dep in deps_needed]
+        ] + ["@poetry//:library_" + dep.replace("-", "_") for dep in deps_needed]
         proc = subprocess.run(
             cmd
             + [

bazel/wrapper_hook/post_bazel_hook.py

@@ -0,0 +1,28 @@
+# Hook to be called around bazel invocation time. Does not run on Windows.
+import os
+import pathlib
+import sys
+
+REPO_ROOT = pathlib.Path(__file__).parent.parent.parent
+sys.path.append(str(REPO_ROOT))
+
+from bazel.wrapper_hook.install_modules import install_modules
+
+BAZEL_USER_NAMESPACE = "user-prod"
+BAZEL_CI_NAMESPACE = "ci-prod"
+
+
+def main():
+    install_modules(sys.argv[1])
+
+    from bazel.wrapper_hook.flag_sync import sync_flags
+
+    if os.environ.get("NO_FLAG_SYNC") is None:
+        if os.environ.get("CI") is None:
+            sync_flags(BAZEL_USER_NAMESPACE)
+        else:
+            sync_flags(BAZEL_CI_NAMESPACE)
+
+
+if __name__ == "__main__":
+    main()

poetry.lock

@@ -4293,6 +4293,18 @@ typing-extensions = ">=4.4.0,<5.0.0"
 [package.extras]
 syntax = ["tree-sitter (>=0.23.0)", "tree-sitter-bash (>=0.23.0)", "tree-sitter-css (>=0.23.0)", "tree-sitter-go (>=0.23.0)", "tree-sitter-html (>=0.23.0)", "tree-sitter-java (>=0.23.0)", "tree-sitter-javascript (>=0.23.0)", "tree-sitter-json (>=0.24.0)", "tree-sitter-markdown (>=0.3.0)", "tree-sitter-python (>=0.23.0)", "tree-sitter-regex (>=0.24.0)", "tree-sitter-rust (>=0.23.0,<=0.23.2)", "tree-sitter-sql (>=0.3.0,<0.3.8)", "tree-sitter-toml (>=0.6.0)", "tree-sitter-xml (>=0.7.0)", "tree-sitter-yaml (>=0.6.0)"]
 
+[[package]]
+name = "timeout-decorator"
+version = "0.5.0"
+description = "Timeout decorator"
+optional = false
+python-versions = "*"
+groups = ["testing"]
+markers = "platform_machine != \"s390x\" and platform_machine != \"ppc64le\" or platform_machine == \"s390x\" or platform_machine == \"ppc64le\""
+files = [
+    {file = "timeout-decorator-0.5.0.tar.gz", hash = "sha256:6a2f2f58db1c5b24a2cc79de6345760377ad8bdc13813f5265f6c3e63d16b3d7"},
+]
+
 [[package]]
 name = "toml"
 version = "0.10.2"
@@ -5447,4 +5459,4 @@ libdeps = ["cxxfilt", "eventlet", "flask", "flask-cors", "gevent", "lxml", "prog
 [metadata]
 lock-version = "2.1"
 python-versions = ">=3.10,<4.0"
-content-hash = "0f66bc1087116663ea12e384ed3266d07d495738b0e57eaddcc0b6b029a4b66e"
+content-hash = "622c4368619483bbf23d0e5d482c6905d2e947952135e3b13f688fa9b748f825"

pyproject.toml

@@ -173,6 +173,7 @@ opentelemetry-api = "*"
 opentelemetry-sdk = "*"
 opentelemetry-exporter-otlp-proto-common = "*"
 opentelemetry-exporter-otlp-proto-grpc = { version = "*", markers = "platform_machine != 's390x' and platform_machine != 'ppc64le'" }
+timeout-decorator = "0.5.0"
 
 # This can be installed with "poetry install -E libdeps"
 [project.optional-dependencies]
@@ -218,7 +219,7 @@ ignore = [
     ]
 
 [tool.ruff.lint.isort]
-known-first-party = ["buildscripts", "buildscripts/idl", "evergreen"]
+known-first-party = ["buildscripts", "buildscripts/idl", "evergreen", "tools"]
 
 [tool.pyright]
 include = [
@@ -283,4 +284,4 @@ reportConstantRedefinition = "none"
 
 [build-system]
 requires = ["poetry-core"]
-build-backend = "poetry.core.masonry.api"
+build-backend = "poetry.core.masonry.api"

tools/bazel

@@ -184,7 +184,8 @@ MONGO_BAZEL_WRAPPER_ARGS=$(mktemp)
 MONGO_BAZEL_WRAPPER_ARGS=$MONGO_BAZEL_WRAPPER_ARGS \
 MONGO_AUTOCOMPLETE_QUERY=$autocomplete_query \
 >&2 $python $REPO_ROOT/bazel/wrapper_hook/wrapper_hook.py $bazel_real "$@"
-if [[ $? != 0 ]]; then
+exit_code=$?
+if [[ $exit_code != 0 ]]; then
   if [[ ! -z "$CI" ]] || [[ $MONGO_BAZEL_WRAPPER_FALLBACK == 1 ]]; then
     >&2 echo "wrapper script failed! falling back to normal bazel call..."
     exec "$bazel_real" $@
@@ -212,5 +213,9 @@ if [[ $autocomplete_query == 1 ]]; then
   query_output=$("${new_args[@]@Q}") 
   echo $query_output $plus_targets | tr " " "\n"
 else
-  exec "$bazel_real" "${new_args[@]}"
+  $bazel_real "${new_args[@]}"
+  bazel_exit_code=$?
+  # Run post-bazel hook in subshell
+  (>&2 $python $REPO_ROOT/bazel/wrapper_hook/post_bazel_hook.py $bazel_real)
+  exit $bazel_exit_code
 fi

tools/flag_sync/BUILD.bazel

@@ -0,0 +1,45 @@
+load("@poetry//:dependencies.bzl", "dependency")
+
+py_binary(
+    name = "client",
+    srcs = [
+        "client.py",
+        "flag.py",
+        "namespace.py",
+    ],
+    main = "client.py",
+    visibility = ["//visibility:public"],
+    deps = [
+        ":util",
+        dependency(
+            "typer",
+            group = "core",
+        ),
+        dependency(
+            "requests",
+            group = "core",
+        ),
+        dependency(
+            "boto3",
+            group = "aws",
+        ),
+    ],
+)
+
+py_library(
+    name = "util",
+    srcs = [
+        "util.py",
+    ],
+    visibility = ["//visibility:public"],
+    deps = [
+        dependency(
+            "retry",
+            group = "testing",
+        ),
+        dependency(
+            "requests",
+            group = "core",
+        ),
+    ],
+)

tools/flag_sync/OWNERS.yml

@@ -0,0 +1,5 @@
+version: 1.0.0
+filters:
+  - "*":
+    approvers:
+      - 10gen/devprod-build

tools/flag_sync/client.py

@@ -0,0 +1,10 @@
+import typer
+
+from tools.flag_sync import flag, namespace
+
+app = typer.Typer()
+app.add_typer(flag.app, name="flag")
+app.add_typer(namespace.app, name="namespace")
+
+if __name__ == "__main__":
+    app()

tools/flag_sync/flag.py

@@ -0,0 +1,99 @@
+import io
+import json
+from typing import Optional
+
+import boto3
+import typer
+
+from tools.flag_sync import util
+
+app = typer.Typer()
+
+
+@app.command()
+def create(
+    namespace: str,
+    name: str,
+    value: str,
+    enabled: Optional[bool] = True,
+    validate: Optional[bool] = True,
+):
+    if validate:
+        if not util.validate_bazel_flag(value):
+            print("Failed to create flag. Flag failed bazel verification.")
+            return
+    s3 = boto3.client("s3")
+    flags = util.get_flags(namespace)
+    flags[name] = {"value": value, "enabled": enabled}
+    print(f"Created flag in namespace {namespace}:")
+    print_flag(name, value, enabled)
+    flags_json = json.dumps(flags)
+    f = io.BytesIO(flags_json.encode("utf-8"))
+    s3.upload_fileobj(f, util.S3_BUCKET, f"flag_sync/{namespace}.json")
+
+
+@app.command()
+def get(namespace: str, name: Optional[str] = None):
+    flags = util.get_flags(namespace)
+    for cur_name, flag in flags.items():
+        if name and cur_name != name:
+            continue
+        print_flag(cur_name, flag["value"], flag["enabled"])
+
+
+@app.command()
+def update(
+    namespace: str,
+    name: str,
+    value: Optional[str] = None,
+    enabled: Optional[bool] = None,
+    validate: Optional[bool] = True,
+):
+    if validate and value:
+        if not util.validate_bazel_flag(value):
+            print("Failed to update flag. Flag failed bazel verification.")
+            return
+    s3 = boto3.client("s3")
+    flags = util.get_flags(namespace)
+    if name not in flags:
+        print(f"Flag with name {name} does not exist.")
+        return
+    if value:
+        flags[name]["value"] = value
+        print(f"Updated {name} to {value} in namespace {namespace}")
+    if enabled is not None:
+        flags[name]["enabled"] = enabled
+        print(f"Updated enabled status of {name} to {enabled} in namespace {namespace}")
+    print(f"Updated flag in namespace {namespace}:")
+    print_flag(name, flags[name]["value"], flags[name]["enabled"])
+    flags_json = json.dumps(flags)
+    f = io.BytesIO(flags_json.encode("utf-8"))
+    s3.upload_fileobj(f, util.S3_BUCKET, f"flag_sync/{namespace}.json")
+
+
+@app.command()
+def toggle_on(namespace: str, name: str):
+    update(namespace=namespace, name=name, enabled=True)
+
+
+@app.command()
+def toggle_off(namespace: str, name: str):
+    update(namespace=namespace, name=name, enabled=False)
+
+
+@app.command()
+def delete(namespace: str, name: str):
+    s3 = boto3.client("s3")
+    flags = util.get_flags(namespace)
+    if name not in flags:
+        print(f"Flag with name {name} does not exist.")
+        return
+    del flags[name]
+    print(f"Deleted flag {name} in namespace {namespace}")
+    flags_json = json.dumps(flags)
+    f = io.BytesIO(flags_json.encode("utf-8"))
+    s3.upload_fileobj(f, util.S3_BUCKET, f"flag_sync/{namespace}.json")
+
+
+def print_flag(name: str, value: str, enabled: bool):
+    print(f"{name}[{'+' if enabled else '-'}]: {value}")

tools/flag_sync/namespace.py

@@ -0,0 +1,25 @@
+import io
+
+import boto3
+import typer
+
+from tools.flag_sync import util
+
+app = typer.Typer()
+
+
+@app.command()
+def create(namespace: str):
+    s3 = boto3.client("s3")
+    f = io.BytesIO(b"{}")
+    s3.upload_fileobj(f, util.S3_BUCKET, f"flag_sync/{namespace}.json")
+    print(f"Created namespace {namespace}")
+
+
+@app.command()
+def list():
+    s3 = boto3.client("s3")
+    res = s3.list_objects_v2(Bucket=util.S3_BUCKET, Prefix="flag_sync/")
+    print("====Namespaces====")
+    for f in res["Contents"]:
+        print(f["Key"].replace(".json", "").split("/")[-1])

tools/flag_sync/util.py

@@ -0,0 +1,31 @@
+import json
+import os
+import subprocess
+import tempfile
+
+import requests
+
+S3_BUCKET = "mdb-build-public"
+
+
+def get_flags(namespace: str):
+    S3_URL = f"https://mdb-build-public.s3.us-east-1.amazonaws.com/flag_sync/{namespace}.json"
+    r = requests.get(S3_URL)
+    if r.status_code != 200:
+        raise Exception("Namespace doesn't exist.")
+    try:
+        return json.loads(r.text)
+    except:
+        raise Exception("Couldn't parse flag json.")
+
+
+def validate_bazel_flag(line: str):
+    workspace_root = os.environ.get("BUILD_WORKSPACE_DIRECTORY")
+    with tempfile.NamedTemporaryFile(mode="w+") as tf:
+        tf.write(line)
+        tf.flush()
+        cmd = ["bazel", f"--bazelrc={tf.name}", "build", "//tools/flag_sync:client"]
+        p = subprocess.run(cmd, capture_output=True, cwd=workspace_root)
+        if p.returncode != 0:
+            return False
+    return True