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
import functions_framework
7
BUCKET = os.environ.get("BUCKET", "missing-bucket")
8
deploymentId = os.environ.get("ASTRO_DEPLOYMENT_ID", "missing-deployment-id")
9
10
11
import shutil
12
13
def clear_dir(path: str) -> None:
14
    if os.path.exists(path):
15
        print(f"Clearing directory: {path}")
16
        shutil.rmtree(path)
17
    os.makedirs(path, exist_ok=True)
18
    print(f"Re-created directory: {path}")
19
20
def untar(filename: str, destination: str) -> None:
21
    with tarfile.open(filename) as file:
22
        file.extractall(destination)
23
24
def run_command(cmd: str) -> None:
25
    print(f'running command: {cmd}')
26
    p = subprocess.Popen("set -x; " + cmd, shell=True)
27
    p.communicate()
28
29
def download_to_local(bucket_name: str, gcs_folder: str, local_dir: str = None) -> None:
30
    """Download the contents of a folder directory
31
    :param bucket_name: the name of the gcs bucket
32
    :param gcs_folder: the folder path in the gcs bucket
33
    :param local_dir: a relative or absolute directory path in the local file system
34
    """
35
36
    ## create a storage client to access GCS objects
37
    storage_client = storage.Client()
38
    source_bucket = storage_client.bucket(bucket_name)
39
40
    ## get a list of all the files in the bucket folder
41
    blobs = source_bucket.list_blobs(prefix=gcs_folder)
42
43
    ## download each of the dag to local
44
    for blob in blobs:
45
        if blob.name.endswith('/'):
46
            continue
47
48
        target = blob.name if local_dir is None \
49
            else os.path.join(local_dir, os.path.relpath(blob.name, gcs_folder))
50
        print(target)
51
        if not os.path.exists(os.path.dirname(target)):
52
            os.makedirs(os.path.dirname(target))
53
54
        blob.download_to_filename(target)
55
        print("downloaded file")
56
57
@functions_framework.cloud_event
58
def astro_deploy(cloud_event) -> None:
59
    base_dir = '/tmp/astro'
60
    dags_dir = f'{base_dir}/dags'
61
62
    clear_dir(dags_dir)
63
64
    # --- Download DAGs ---
65
    print('downloading dags')
66
    download_to_local(BUCKET, 'dags/', f'{base_dir}/dags')  # NOTE: use "dags/" prefix
67
68
    # --- Download CLI ---
69
    print('downloading cli')
70
    download_to_local(BUCKET, 'cli_binary/', base_dir)
71
72
    # --- Initialize project ---
73
    os.chdir(base_dir)
74
    untar('./astro_cli.tar.gz', '.')
75
    run_command('echo y | ./astro dev init')
76
77
    # --- Remove generated example DAG(s) ---
78
    example_paths = [
79
        "dags/example_dag.py",
80
        "dags/exampledag.py",
81
    ]
82
83
    for path in example_paths:
84
        full_path = os.path.join(base_dir, path)
85
        if os.path.exists(full_path):
86
            print(f"Removing generated example DAG: {full_path}")
87
            os.remove(full_path)
88
        else:
89
            print(f"Example DAG not found: {full_path}")
90
91
    # --- Deploy ----
92
    run_command(f'./astro deploy {deploymentId} --dags')