snapshot.py

from concurrent.futures import ThreadPoolExecutor, Future
import os
from settings import Settings, parse_arguments
from google.cloud import storage
from google.oauth2 import service_account
import containers
import docker
from storage import create_tar, upload_file_to_bucket

def take_snapshot(
    docker_client: docker.DockerClient,
    storage_client: storage.Client,
    settings: Settings,
    container_config: dict[str, str]
) -> None:
    """Take a snapshot of the specified Docker containers and upload it to Google Cloud Storage."""
    
    container_name = container_config['container_name']
    bucket_path = container_config['bucket_path']
    tar_name = container_config['tar_name']
    data_path = container_config['data_path']

    # Step 1: Stop containers
    containers.stop(docker_client, [container_name])

    # Step 2: Create tarball from data_path
    tar_path = f"{tar_name}"
    create_tar([data_path], tar_path)

    # Step 3: Upload tarball to Google Cloud Storage
    upload_file_to_bucket(storage_client, settings.bucket_name, tar_path, bucket_path)

    # Step 4: Start containers back
    containers.start(docker_client, [container_name])
    
    # Step 5: Cleanup
    os.remove(tar_path)


def main() -> None:
    """Main function to parse arguments, initialize clients, and take snapshots in parallel."""
    args = parse_arguments()
    settings = Settings(config_path=args.config)

    credentials = service_account.Credentials.from_service_account_file(
        settings.google_credentials_json
    )
    print(settings)
    storage_client = storage.Client(credentials=credentials)
    docker_client = docker.DockerClient(base_url=settings.docker_host)
    
    containers_config = settings.docker_containers

    # Use ThreadPoolExecutor to take snapshots in parallel
    with ThreadPoolExecutor() as executor:
        futures: list[Future] = []
        for container_config in containers_config:
            futures.append(executor.submit(take_snapshot, docker_client, storage_client, settings, container_config))

        for future in futures:
            future.result()

    print("All snapshots taken and uploaded.")


if __name__ == "__main__":
    main()