bedrock/bin/run-db-download.py

94 строки
2.5 KiB
Python
Executable File

#!/usr/bin/env python3
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at https://mozilla.org/MPL/2.0/.
import os
import sys
import requests
from db_s3_utils import (
DATA_PATH,
DB_FILE,
JSON_DATA_FILE_NAME,
get_db_checksum,
get_git_sha,
get_prev_db_data,
set_db_data,
)
BUCKET_NAME = os.getenv("AWS_DB_S3_BUCKET", "bedrock-db-dev")
REGION_NAME = os.getenv("AWS_DB_REGION", "us-west-2")
S3_BASE_URL = f"https://s3-{REGION_NAME}.amazonaws.com/{BUCKET_NAME}"
def get_file_url(filename):
return "/".join([S3_BASE_URL, filename])
def download_db_info():
try:
resp = requests.get(get_file_url(JSON_DATA_FILE_NAME))
resp.raise_for_status()
except requests.RequestException:
return None
try:
return resp.json()
except ValueError:
# not JSON
return None
def download_db_file(filename):
resp = requests.get(get_file_url(os.path.basename(filename)), stream=True)
with open(filename, "wb") as fp:
for chunk in resp.iter_content(chunk_size=128):
fp.write(chunk)
def update_live_db_file(filename):
os.rename(filename, DB_FILE)
def main(args):
force = "--force" in args
ignore_git = "--ignore-git" in args
db_info = download_db_info()
if not db_info:
return "ERROR: Could not get database info"
if not force:
prev_data = get_prev_db_data()
if prev_data and prev_data["checksum"] == db_info["checksum"]:
print("Checksums match. No update required.")
return 0
if prev_data and prev_data["updated"] > db_info["updated"]:
print("Remote database older than local. No update required.")
return 0
if not ignore_git:
git_sha = get_git_sha()
if git_sha != db_info["git_sha"]:
print("Git hashes do not match. No update required.")
return 0
new_db_file = db_info["file_name"]
new_db_file = f"{DATA_PATH}/{new_db_file}"
download_db_file(new_db_file)
checksum = get_db_checksum(new_db_file)
if checksum == db_info["checksum"]:
update_live_db_file(new_db_file)
set_db_data(db_info)
print("Database successfully updated")
return 0
os.remove(new_db_file)
return "ERROR: Checksums do not match. Bad db download. Aborting."
if __name__ == "__main__":
sys.exit(main(sys.argv[1:]))