gga_init.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import argparse
import os
import re
import subprocess
import logging
import sys
import yaml
import shutil

from pathlib import Path
from jinja2 import Template, Environment, FileSystemLoader

import utilities
import speciesData


""" 
gga_init.py

Usage: $ python3 gga_init.py -i input_example.yml [OPTIONS]

"""


class DeploySpeciesStack(speciesData.SpeciesData):
    """
    Child of SpeciesData

    Contains methods and attributes to deploy a stack of services for a given organism, from creating/updating
    the organism's directory tree to create the required docker-compose files
    """

    def make_directory_tree(self):
        """
        Generate the directory tree for an organism and move datasets into src_data

        :return:
        """

        os.chdir(self.main_dir)

        try:
            os.mkdir(self.species_dir)
            logging.info("Making directory tree for %s" % self.full_name)
        except FileExistsError:
            logging.info("Updating directory tree for %s" % self.full_name)
        try:
            os.chdir(self.species_dir)
        except OSError:
            logging.critical("Cannot access " + self.species_dir + ", run with higher privileges")
            sys.exit()

        # Copy the custom banner to the species dir (banner used in tripal pages)
        # if "custom_banner" not in self.config.keys() or not self.config["custom_banner"] == "/path/to/banner" or not self.config["custom_banner"] == "":
        #     try:
        #         logging.debug("Custom banner path: %s" % self.config["custom_banner"])
        #         if os.path.isfile(os.path.abspath(self.config["custom_banner"])):
        #             shutil.copy(os.path.abspath(self.config["custom_banner"]), "%s/banner.png" % self.species_dir)
        #     except FileNotFoundError:
        #         logging.warning("Specified banner not found (%s), skipping" % self.config["custom_banner"])

        # Write nginx conf
        try:
            os.mkdir("./nginx/")
            os.mkdir("./nginx/conf")
            with open(os.path.abspath("./nginx/conf/default.conf"), 'w') as conf:
                conf.write("server {\n\tlisten 80;\n\tserver_name ~.;\n\tlocation /download/ {\n\t\talias /project_data/; \n\t\tautoindex on;\n\t}\n}")  # The species nginx conf
        except FileExistsError:
            logging.debug("Nginx conf already exists, skipping")


        organism_annotation_dir, organism_genome_dir = None, None

        # Creation (or updating) of the src_data directory tree
        # Depth 0
        try:
            os.mkdir("./src_data")
        except FileExistsError:
            logging.debug("src_data folder already exist for %s" % self.full_name)
        except PermissionError:
            logging.critical("Insufficient permission to create src_data directory tree")
            sys.exit()

        # Depth 1
        try:
            os.mkdir("./src_data/annotation")
            os.mkdir("./src_data/genome")
            os.mkdir("./src_data/tracks")
        except FileExistsError:
            logging.debug("Depth 1 src_data folder(s) already exist for %s" % self.full_name)
        except PermissionError as exc:
            logging.critical("Insufficient permission to create src_data directory tree")
            sys.exit(exc)

        # Depth 2
        try:
            os.mkdir("./src_data/annotation/" + self.species_folder_name)
            os.mkdir("./src_data/genome/" + self.species_folder_name)
        except FileExistsError:
            logging.debug("Depth 2 src_data folder(s) already exist for %s" % self.full_name)
        except PermissionError as exc:
            logging.critical("Insufficient permission to create src_data directory tree")
            sys.exit(exc)

        # Depth 3
        try:
            os.mkdir("./src_data/annotation/" + self.species_folder_name + "/OGS" + self.ogs_version)
            os.mkdir("./src_data/genome/" + self.species_folder_name + "/v" + self.genome_version)
            organism_annotation_dir = os.path.abspath("./src_data/annotation/" + self.species_folder_name + "/OGS" + self.genome_version)
            organism_genome_dir = os.path.abspath("./src_data/genome/" + self.species_folder_name + "/v" + self.genome_version)
        except FileExistsError:
            logging.debug("Depth 3 src_data folder(s) already exist for %s" % self.full_name)
        except PermissionError as exc:
            logging.critical("Insufficient permission to create src_data directory tree")
            sys.exit(exc)

        # Return to main_dir
        os.chdir(self.main_dir)

        logging.info("Directory tree generated for %s" % self.full_name)


    def make_compose_files(self):
        """
        Create a formatted copy of the template compose file inside a species directory tree

        :return:
        """

        os.chdir(self.main_dir)
        try:
            os.chdir(self.species_dir)
        except OSError:
            logging.critical("Cannot access " + self.species_dir)
            sys.exit(0)

        # Path to the templates used to generate the custom docker-compose files for an input species and the traefik+authelia services
        gspecies_template_path = self.script_dir + "/templates/gspecies_compose_template.yml.j2"  # Jinja template path
        traefik_template_path = self.script_dir + "/templates/traefik_compose_template.yml"
        # authelia_config_path = self.script_dir + "/templates/authelia_config_example.yml"  # Do not copy the authelia config!
        authelia_users_path = self.script_dir + "/templates/authelia_users_template.yml"

        # Set the genus_species_strain_sex value for replacing template
        genus_species_strain_sex = "{0}_{1}".format(self.genus.lower(), self.specieslower())
        if self.sex and self.strain:
            genus_species_strain_sex = "_".join([self.genus.lower(), self.species.lower(), self.strain, self.sex])
        elif self.sex and not self.strain:
            genus_species_strain_sex = "_".join([self.genus.lower(), self.specieslower(), self.sex])
        elif self.genus and not self.species and not self.sex and not self.strain:
            genus_species_strain_sex = self.genus.lower()
        else:
            genus_species_strain_sex = "{0}_{1}".format(self.genus.lower(), self.specieslower())

        # # Method without Jinja templating (deprecated)
        # with open(stack_template_path, 'r') as infile:
        #     organism_content = list()
        #     for line in infile:
        #         # Replace placeholders in the compose file and append line to output
        #         organism_content.append(
        #             line.replace("genus_species",
        #                          str(self.genus.lower() + "_" + self.species)).replace("Genus species",
        #                          str(self.genus_uppercase + " " + self.species)).replace("Genus/species",
        #                          str(self.genus_uppercase + "/" + self.species)).replace("gspecies",
        #                          str(self.genus.lower()[0] + self.species)).replace("genus_species_strain_sex",
        #                          genus_species_strain_sex))
        #     # Write/format the output compose file
        #     with open("./docker-compose.yml", 'w') as outfile:
        #         outfile.truncate(0)  # Delete file content
        #         for line in organism_content:  # Replace env variables by those in the config file
        #             for env_variable, value in self.config.items():  # env variables are stored in this dict
        #                 # print("ENV VARIABLE: " + env_variable + "\t VALUE: " + value)
        #                 if env_variable in line:
        #                     line = line.replace(env_variable, value)
        #                     break
        #             # Write the new line in the docker-compose
        #             outfile.write(line)

        #     # Create mounts for the current docker-compose
        #     self.create_mounts(working_dir=self.species_dir)


        # jinja templating, handled using the python jinja module
        # TODO: use ansible to handle the templating in production
        file_loader = FileSystemLoader(self.script_dir + "/templates")
        env = Environment(loader=file_loader)
        template = env.get_template("gspecies_compose_template.yml.j2")

        # We need a dict holding all key (variables) - values that needs to be replaced in the template as our rendering dict
        # To do so we need both input file vars and config vars
        # Create input file vars dict
        input_vars = {"genus": self.genus_lowercase, "Genus": self.genus_uppercase, "species": self.species,
                      "genus_species": self.genus_species, "genus_species_strain_sex": self.full_name,
                      "strain": self.strain, "sex": self.sex, "Genus_species": self.genus_species[0].upper() + self.genus_species[1:]}
        # Merge the two dicts
        render_vars = {**self.config, **input_vars}
        output = template.render(render_vars)

        with open(os.path.join(self.species_dir, "docker-compose.yml", "w")) as gspecies_compose_file:
            print("Writing gspecies compose yml file")

        try:
            os.chdir(os.path.abspath(self.main_dir))
            os.mkdir("./traefik")
            os.mkdir("./traefik/authelia")
            if self.config["custom_authelia_config_path"]:
                if os.path.isfile(os.path.abspath(self.config["custom_authelia_config_path"])):
                    try:
                        shutil.copy(os.path.abspath(self.config["custom_authelia_config_path"]), "./traefik/authelia")
                    except FileNotFoundError:
                        logging.critical("Cannot copy custom Authelia config file (%s)" % self.config["custom_authelia_config_path"])
                        sys.exit()
                else:
                    logging.critical("Custom Authelia config file not found (%s)" % self.config["custom_authelia_config_path"])
            if not os.path.isfile("./traefik/authelia/users.yml"):
                shutil.copy(authelia_users_path, "./traefik/authelia/users.yml")
        except FileExistsError:
            logging.debug("Traefik directory already exists: %s" % os.path.abspath("../traefik"))
        try:
            if not os.path.isfile("./traefik/docker-compose.yml"):
                shutil.copy(traefik_template_path, "./traefik/docker-compose.yml")
            else:
                logging.debug("Traefik compose file already exists: %s" % os.path.abspath("./traefik/docker-compose.yml"))
        except FileExistsError:
            logging.debug("Traefik compose file already exists: %s" % os.path.abspath("./traefik/docker-compose.yml"))

        traefik_dir = os.path.abspath(os.path.join(self.main_dir, "traefik"))

        # Create the mounts for the traefik+authelia containers
        self.create_mounts(working_dir=traefik_dir)

        os.chdir(self.main_dir)

    def create_mounts(self, working_dir):
        """
        Create the folders (volumes) required by a container (to see required volumes, check their compose file)

        :return:
        """

        # Change directory to create mount points for the container
        try:
            os.chdir(os.path.abspath(working_dir))
        except Exception:
            logging.critical("Cannot access %s, exiting" % working_dir)
            sys.exit()
        compose_yml = os.path.abspath("./docker-compose.yml")
        if not os.path.isfile(compose_yml):
            raise Exception("Could not find docker-compose.yml at %s" % compose_yml)

        with open(compose_yml) as f:
            compose = yaml.safe_load(f)

        if 'services' not in compose:
            raise Exception("Could not find services tag in docker-compose.yml")

        # Iterate over all services to find the "volumes" we need to create
        for service in compose['services']:
            if 'volumes' in compose['services'][service]:
                for volume in compose['services'][service]['volumes']:
                    # regex to match the volumes of the services
                    reg = re.match(r"^(\./[^:]+/):[^:]+(:\w+)?$", volume)
                    if reg:
                        vol_dir = os.path.abspath('./' + reg.group(1))
                        if not os.path.exists(vol_dir):
                            os.makedirs(vol_dir, exist_ok=True)
                    else:
                        reg = re.match(r"^(\./[^:]+):[^:]+(:\w+)?$", volume)
                        if reg:
                            vol_file = os.path.abspath('./' + reg.group(1))
                            vol_dir = os.path.dirname(vol_file)
                            if not os.path.exists(vol_dir):
                                os.makedirs(vol_dir, exist_ok=True)
                            if not os.path.exists(vol_file):
                                Path(vol_file).touch()

        # Go back to the "main" directory
        try:
            os.chdir(os.path.abspath(self.main_dir))
        except Exception:
            logging.critical("Cannot access main directory (%s), exiting" % self.main_dir)
            sys.exit()

    def deploy_stack(self):
        """
        Call the script "deploy.sh" used to initiliaze the swarm cluster if needed and
        launch/update the current organism's stack

        This script first try to deploy the traefik stack, then deploy the organism stack, then update the traefik stack
        The stacks are updated if already deployed

        :return:
        """

        # # Create our swarm cluster if it doesn't exist
        # subprocess.Popen(["docker", "swarm", "init"],
        #                  stdout=subprocess.PIPE, stderr=subprocess.PIPE, cwd=self.main_dir)
        #
        # # Deploy/update the stack for the current species
        # subprocess.Popen(["docker", "stack", "deploy", "-c", "./docker-compose.yml", "{0}_{1}".format(self.genus_lowercase, self.species)],
        #                  stdout=subprocess.PIPE, stderr=subprocess.PIPE, cwd=self.main_dir)

        # Launch and update docker stacks
        # noinspection PyArgumentList
        deploy_stacks_popen = subprocess.Popen(["sh", self.script_dir + "/deploy.sh", self.genus_species,
                                                self.main_dir + "/traefik"],
                                               stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                                               universal_newlines=True)

        for stdout_line in iter(deploy_stacks_popen.stdout.readline, ""):
            if "daemon" in stdout_line:  # Ignore swarm init error output
                pass
            else:
                logging.info("\t%s" % stdout_line.strip())
        deploy_stacks_popen.stdout.close()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Automatic data loading in containers and interaction "
                                                 "with galaxy instances for GGA"
                                                 ", following the protocol @ "
                                                 "http://gitlab.sb-roscoff.fr/abims/e-infra/gga")

    parser.add_argument("input",
                        type=str,
                        help="Input file (yml)")

    parser.add_argument("-v", "--verbose",
                        help="Increase output verbosity",
                        action="store_true")

    parser.add_argument("--config",
                        type=str,
                        help="Config path, default to the 'config' file inside the script repository")

    parser.add_argument("--main-directory",
                        type=str,
                        help="Where the stack containers will be located, defaults to current directory")

    args = parser.parse_args()

    if args.verbose:
        logging.basicConfig(level=logging.DEBUG)
    else:
        logging.basicConfig(level=logging.INFO)

    # Parsing the config file if provided, using the default config otherwise
    if not args.config:
        args.config = os.path.join(os.path.dirname(os.path.realpath(sys.argv[0])), "config")
    else:
        args.config = os.path.abspath(args.config)

    main_dir = None
    if not args.main_directory:
        main_dir = os.getcwd()
    else:
        main_dir = os.path.abspath(args.main_directory)

    sp_dict_list = utilities.parse_input(os.path.abspath(args.input))

    logging.info("Deploying stacks for organisms in input file %s" % args.input)
    for sp_dict in sp_dict_list:

        # Init instance
        deploy_stack_for_current_organism = DeploySpeciesStack(parameters_dictionary=sp_dict)

        # Setting some of the instance attributes
        deploy_stack_for_current_organism.main_dir = main_dir
        deploy_stack_for_current_organism.species_dir = os.path.join(deploy_stack_for_current_organism.main_dir,
                                                                     deploy_stack_for_current_organism.genus_species +
                                                                     "/")

        # Parse the config yaml file
        deploy_stack_for_current_organism.config = utilities.parse_config(args.config)

        # Set the instance url attribute
        for env_variable, value in deploy_stack_for_current_organism.config.items():
            if env_variable == "custom_host":
                deploy_stack_for_current_organism.instance_url = value + \
                                                                 deploy_stack_for_current_organism.genus_lowercase + \
                                                                 "_" + deploy_stack_for_current_organism.species + \
                                                                 "/galaxy/"
                break
            else:
                deploy_stack_for_current_organism.instance_url = "http://localhost:8888/sp/{0}_{1}/galaxy/".format(
                    deploy_stack_for_current_organism.genus_lowercase,
                    deploy_stack_for_current_organism.species)
                                                                    
        # Starting
        logging.info("gga_init.py called for %s" % deploy_stack_for_current_organism.full_name)

        # Make/update directory tree
        deploy_stack_for_current_organism.make_directory_tree()
        logging.info("Successfully generated the directory tree for %s" % deploy_stack_for_current_organism.full_name)

        # Make compose files
        deploy_stack_for_current_organism.make_compose_files()
        logging.info("Successfully generated the docker-compose files for %s" % deploy_stack_for_current_organism.full_name)

        # Deploy the stack
        logging.info("Deploying stack for %s..." % deploy_stack_for_current_organism.full_name)
        # deploy_stack_for_current_organism.deploy_stack()
        logging.info("Successfully deployed stack for %s" % deploy_stack_for_current_organism.full_name)

        logging.info("Stack deployed for %s" % deploy_stack_for_current_organism.full_name)

        # TODO: if GENUS°1 == GENUS°2 AND SP°1 == SP°2 --> skip init, continue to next item and only deploy once the loop is done

    # TODO: reload traefik outside loop
    logging.info("All stacks deployed for organisms in input file %s" % args.input)