mirror of https://github.com/mongodb/mongo
89 lines
2.5 KiB
Python
89 lines
2.5 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
S3-based lock mechanism.
|
|
|
|
This script attempts to create a lock file in S3. If the file already exists,
|
|
it exits with failure (exit code 1). If the file doesn't exist and is successfully
|
|
created, it exits with success (exit code 0).
|
|
|
|
Usage:
|
|
python s3_lock.py --bucket BUCKET_NAME --key LOCK_KEY [--content CONTENT]
|
|
|
|
Example:
|
|
python s3_lock.py --bucket my-bucket --key locks/mylock.txt --content "locked by build 123"
|
|
"""
|
|
|
|
import sys
|
|
|
|
import typer
|
|
from botocore.exceptions import ClientError
|
|
|
|
from buildscripts.util.download_utils import get_s3_client
|
|
|
|
|
|
def acquire_s3_lock(bucket, key, content="locked"):
|
|
"""
|
|
Attempt to acquire a lock by uploading a file to S3 only if it doesn't exist.
|
|
|
|
Args:
|
|
bucket: S3 bucket name
|
|
key: S3 object key (file path)
|
|
content: Content to write to the lock file
|
|
|
|
Returns:
|
|
True if lock was acquired (file didn't exist), False if lock already exists
|
|
"""
|
|
s3_client = get_s3_client()
|
|
|
|
try:
|
|
# Try to upload the file only if it doesn't already exist
|
|
# using the IfNoneMatch condition with "*" (no existing ETag should match)
|
|
s3_client.put_object(
|
|
Bucket=bucket,
|
|
Key=key,
|
|
Body=content.encode("utf-8"),
|
|
IfNoneMatch="*", # Only succeed if object doesn't exist
|
|
)
|
|
return True
|
|
except ClientError as e:
|
|
error_code = e.response.get("Error", {}).get("Code", "")
|
|
|
|
# PreconditionFailed means the object already exists
|
|
if error_code == "PreconditionFailed":
|
|
return False
|
|
|
|
# Any other error should be raised
|
|
raise
|
|
|
|
|
|
app = typer.Typer()
|
|
|
|
|
|
@app.command()
|
|
def main(
|
|
bucket: str = typer.Option(..., help="S3 bucket name"),
|
|
key: str = typer.Option(..., help="S3 object key (path) for the lock file"),
|
|
content: str = typer.Option("locked", help="Content to write to the lock file"),
|
|
):
|
|
"""Acquire an S3-based lock by creating a file if it doesn't exist."""
|
|
try:
|
|
lock_acquired = acquire_s3_lock(bucket, key, content)
|
|
|
|
if lock_acquired:
|
|
print(f"Lock acquired: s3://{bucket}/{key}")
|
|
sys.exit(0)
|
|
else:
|
|
print(f"Lock already exists: s3://{bucket}/{key}", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
except ClientError as e:
|
|
print(f"S3 error: {e}", file=sys.stderr)
|
|
sys.exit(2)
|
|
except Exception as e:
|
|
print(f"Unexpected error: {e}", file=sys.stderr)
|
|
sys.exit(2)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
app()
|