Deploy dags from Google Cloud Storage to Astro

1
import os
2
import tarfile
3
import subprocess
4
from pathlib import Path
5
from google.cloud import storage
6
BUCKET = os.environ.get("BUCKET", "my-demo-bucket")
7
deploymentId = os.environ.get("ASTRO_DEPLOYMENT_ID", "missing-deployment-id")
8
9
def untar(filename: str, destination: str) -> None:
10
    with tarfile.open(filename) as file:
11
        file.extractall(destination)
12
13
def run_command(cmd: str) -> None:
14
    p = subprocess.Popen("set -x; " + cmd, shell=True)
15
    p.communicate()
16
17
def download_to_local(bucket_name: str, gcs_folder: str, local_dir: str = None) -> None:
18
    """Download the contents of a folder directory
19
    :param bucket_name: the name of the gcs bucket
20
    :param gcs_folder: the folder path in the gcs bucket
21
    :param local_dir: a relative or absolute directory path in the local file system
22
    """
23
24
    ## create a storage client to access GCS objects
25
    storage_client = storage.Client()
26
    source_bucket = storage_client.bucket(bucket_name)
27
28
    ## get a list of all the files in the bucket folder
29
    blobs = source_bucket.list_blobs(prefix=gcs_folder)
30
31
    ## download each of the dag to local
32
    for blob in blobs:
33
        if blob.name.endswith('/'):
34
            continue
35
36
    target = blob.name if local_dir is None \
37
        else os.path.join(local_dir, os.path.relpath(blob.name, gcs_folder))
38
    print(target)
39
    if not os.path.exists(os.path.dirname(target)):
40
        os.makedirs(os.path.dirname(target))
41
42
    blob.download_to_filename(target)
43
    print("downloaded file")
44
45
def astro_deploy(event, context) -> None:
46
    """Triggered by a change to a Cloud Storage bucket.
47
    :param event: Event payload.
48
    :param context: Metadata for the event.
49
    """
50
51
    base_dir = '/tmp/astro'
52
    ## download dag files to temp local storage
53
    download_to_local(BUCKET, 'dags', f'{base_dir}/dags')
54
55
    ## download astro cli binary and move to /tmp/astro
56
    download_to_local(BUCKET, 'cli_binary', base_dir)
57
58
    ## deploy to astro
59
    os.chdir(base_dir)
60
    untar('./astro_cli.tar.gz', '.')
61
    run_command('echo y | ./astro dev init')
62
    run_command(f'./astro deploy {deploymentId} --dags')