update azure-storage-blob (#29411)
Co-authored-by: Cameron Clough <cameronjclough@gmail.com>pull/29633/head
Adeeb Shihadeh
2 years ago
committed by
GitHub
parent
f1b8a86464
commit
c9e227a9c0
4 changed files with 212 additions and 73 deletions
@ -1,43 +1,64 @@ |
||||
#!/usr/bin/env python3 |
||||
import os |
||||
import sys |
||||
import subprocess |
||||
from datetime import datetime, timedelta |
||||
from functools import lru_cache |
||||
from pathlib import Path |
||||
from typing import IO, Union |
||||
|
||||
BASE_URL = "https://commadataci.blob.core.windows.net/openpilotci/" |
||||
TOKEN_PATH = "/data/azure_token" |
||||
DATA_CI_ACCOUNT = "commadataci" |
||||
DATA_CI_ACCOUNT_URL = f"https://{DATA_CI_ACCOUNT}.blob.core.windows.net" |
||||
DATA_CI_CONTAINER = "openpilotci" |
||||
BASE_URL = f"{DATA_CI_ACCOUNT_URL}/{DATA_CI_CONTAINER}/" |
||||
|
||||
TOKEN_PATH = Path("/data/azure_token") |
||||
|
||||
def get_url(route_name, segment_num, log_type="rlog"): |
||||
|
||||
def get_url(route_name: str, segment_num, log_type="rlog") -> str: |
||||
ext = "hevc" if log_type.endswith('camera') else "bz2" |
||||
return BASE_URL + f"{route_name.replace('|', '/')}/{segment_num}/{log_type}.{ext}" |
||||
|
||||
def get_sas_token(): |
||||
sas_token = os.environ.get("AZURE_TOKEN", None) |
||||
if os.path.isfile(TOKEN_PATH): |
||||
sas_token = open(TOKEN_PATH).read().strip() |
||||
|
||||
if sas_token is None: |
||||
sas_token = subprocess.check_output("az storage container generate-sas --account-name commadataci --name openpilotci \ |
||||
--https-only --permissions lrw --expiry $(date -u '+%Y-%m-%dT%H:%M:%SZ' -d '+1 hour') \ |
||||
--auth-mode login --as-user --output tsv", shell=True).decode().strip("\n") |
||||
@lru_cache |
||||
def get_azure_credential(): |
||||
if "AZURE_TOKEN" in os.environ: |
||||
return os.environ["AZURE_TOKEN"] |
||||
elif TOKEN_PATH.is_file(): |
||||
return TOKEN_PATH.read_text().strip() |
||||
else: |
||||
from azure.identity import AzureCliCredential |
||||
return AzureCliCredential() |
||||
|
||||
|
||||
return sas_token |
||||
@lru_cache |
||||
def get_container_sas(account_name: str, container_name: str): |
||||
from azure.storage.blob import BlobServiceClient, ContainerSasPermissions, generate_container_sas |
||||
start_time = datetime.utcnow() |
||||
expiry_time = start_time + timedelta(hours=1) |
||||
blob_service = BlobServiceClient( |
||||
account_url=f"https://{account_name}.blob.core.windows.net", |
||||
credential=get_azure_credential(), |
||||
) |
||||
return generate_container_sas( |
||||
account_name, |
||||
container_name, |
||||
user_delegation_key=blob_service.get_user_delegation_key(start_time, expiry_time), |
||||
permission=ContainerSasPermissions(read=True, write=True, list=True), |
||||
expiry=expiry_time, |
||||
) |
||||
|
||||
def upload_bytes(data, name): |
||||
from azure.storage.blob import BlockBlobService |
||||
service = BlockBlobService(account_name="commadataci", sas_token=get_sas_token()) |
||||
service.create_blob_from_bytes("openpilotci", name, data) |
||||
return BASE_URL + name |
||||
|
||||
def upload_file(path, name): |
||||
from azure.storage.blob import BlockBlobService |
||||
service = BlockBlobService(account_name="commadataci", sas_token=get_sas_token()) |
||||
service.create_blob_from_path("openpilotci", name, path) |
||||
return BASE_URL + name |
||||
def upload_bytes(data: Union[bytes, IO], blob_name: str) -> str: |
||||
from azure.storage.blob import BlobClient |
||||
blob = BlobClient( |
||||
account_url=DATA_CI_ACCOUNT_URL, |
||||
container_name=DATA_CI_CONTAINER, |
||||
blob_name=blob_name, |
||||
credential=get_azure_credential(), |
||||
) |
||||
blob.upload_blob(data) |
||||
return BASE_URL + blob_name |
||||
|
||||
|
||||
if __name__ == "__main__": |
||||
for f in sys.argv[1:]: |
||||
name = os.path.basename(f) |
||||
url = upload_file(f, name) |
||||
print(url) |
||||
def upload_file(path: Union[str, os.PathLike], blob_name: str) -> str: |
||||
with open(path, "rb") as f: |
||||
return upload_bytes(f, blob_name) |
||||
|
||||
Loading…
Reference in new issue