datahub/docker/build.gradle

plugins {
    id 'java' // required by versioning
    id 'docker-compose'
}
import com.avast.gradle.dockercompose.tasks.ComposeUp
import com.avast.gradle.dockercompose.tasks.ComposeDownForced
import org.yaml.snakeyaml.Yaml

apply from: "../gradle/versioning/versioning.gradle"

ext {
    compose_base = "profiles/docker-compose.yml"
    project_name = "datahub"

    backend_profile_modules = [
            ':docker:elasticsearch-setup',
            ':docker:mysql-setup',
            ':docker:kafka-setup',
            ':datahub-upgrade',
            ':metadata-service:war',
    ]

    python_services_modules = [
    ]

    // Common configuration for all tasks
    common_config = [
            captureContainersOutput: true,
            captureContainersOutputToFiles: project.file('build/container-logs')
    ]

    // declarative task configuration
    quickstart_configs = [
            'quickstart': [
                    profile: 'quickstart-consumers',
                    modules: python_services_modules + backend_profile_modules + [
                            ':datahub-frontend',
                            ':metadata-jobs:mce-consumer-job',
                            ':metadata-jobs:mae-consumer-job',
                            ':datahub-actions',
                    ]
            ],
            'quickstartCLI': [
                    profile: 'quickstart',
                    modules: python_services_modules + backend_profile_modules + [
                            ':datahub-frontend',
                            ':datahub-actions',
                    ]
            ],
            'quickstartDebug': [
                    profile: 'debug',
                    modules: python_services_modules + backend_profile_modules + [':datahub-frontend', ':datahub-actions'],
                    isDebug: true,
                    additionalEnv: [
                            DATAHUB_LOCAL_ACTIONS_ENV: "${rootProject.project(':smoke-test').projectDir}/test_resources/actions/actions.env"
                    ]
            ],
            'quickstartDebugMin': [
                    profile: 'debug-min',
                    modules: backend_profile_modules + [':datahub-frontend'],
                    isDebug: true
            ],
            'quickstartDebugConsumers': [
                    profile: 'debug-consumers',
                    modules: python_services_modules + backend_profile_modules + [':datahub-frontend',
                                                                          ':metadata-jobs:mce-consumer-job',
                                                                          ':metadata-jobs:mae-consumer-job',
                                                                          ':datahub-actions'],
                    isDebug: true,
                    additionalEnv: [
                            DATAHUB_LOCAL_ACTIONS_ENV: "${rootProject.project(':smoke-test').projectDir}/test_resources/actions/actions.env"
                    ]
            ],
            'quickstartPg': [
                    profile: 'quickstart-postgres',
                    modules: (backend_profile_modules - [':docker:mysql-setup']) + [
                            ':docker:postgres-setup',
                            ':datahub-frontend',
                            ':datahub-actions',
                    ]
            ],
            'quickstartPgDebug': [
                    profile: 'debug-postgres',
                    modules: python_services_modules + (backend_profile_modules - [':docker:mysql-setup']) + [
                            ':docker:postgres-setup',
                            ':datahub-frontend'
                    ],
                    isDebug: true
            ],
            'quickstartSlim': [
                    profile: 'quickstart-backend',
                    modules: backend_profile_modules + [
                        ':datahub-actions',
                    ]
            ],
            'quickstartSpark': [
                    profile: 'quickstart-backend',
                    modules: backend_profile_modules + [
                            ':datahub-actions',
                    ],
                    additionalEnv: [
                            'DATAHUB_LOCAL_COMMON_ENV': "${rootProject.project(':metadata-integration:java:spark-lineage-legacy').projectDir}/spark-smoke-test/smoke-gms.env",
                            'METADATA_SERVICE_AUTH_ENABLED': 'false'
                    ]
            ],
            'quickstartStorage': [
                    profile: 'quickstart-storage',
                    preserveVolumes: true
            ],
            'quickstartBackendDebug': [
                profile: 'debug-backend-aws',
                modules: python_services_modules + backend_profile_modules + [':datahub-frontend', ':datahub-actions'],
                isDebug: true,
                additionalEnv: [
                        DATAHUB_LOCAL_ACTIONS_ENV: "${rootProject.project(':smoke-test').projectDir}/test_resources/actions/actions.env"
                ]
            ],
            'allImages': [ //This is a special task just to include all images as dependencies - and is useful when CI needs to publish all images
                    profile: 'quickstart-consumers',
                    modules: python_services_modules + backend_profile_modules + [
                        ':datahub-frontend',
                        ':metadata-jobs:mce-consumer-job',
                        ':metadata-jobs:mae-consumer-job',
                        ':datahub-actions',
                        ':docker:datahub-ingestion',
                        ':docker:postgres-setup'
                    ]
            ]
    ]

    // only for debug variants of quickstart to enable <variant>Reload tasks. 
    // The actual service name needs the profile to be appended,  <container-name>-<profile>
    // This list only contains modules that can be reloaded via the reloadTask. If other modules need to be reloaded, quickstart* needs to be used.
    moduleToContainer = [
        ':metadata-service:war': 'datahub-gms',
        ':datahub-frontend': 'frontend',
        ':metadata-jobs:mce-consumer-job': 'datahub-mce-consumer',
        ':metadata-jobs:mae-consumer-job': 'datahub-mae-consumer',
    ]
}

// Register all quickstart tasks
quickstart_configs.each { taskName, config ->
    tasks.register(taskName) {
        group = 'quickstart'
    }
    
}

// Dynamically create all quickstart tasks and configurations
dockerCompose {
    // Configure default settings that apply to all configurations
    useComposeFiles = [compose_base]
    projectName = project_name
    projectNamePrefix = ''
    buildBeforeUp = false
    buildBeforePull = false
    stopContainers = false
    removeVolumes = false

    quickstart_configs.each { taskName, config ->
        "${taskName}" {
            isRequiredBy(tasks.named(taskName))
            if (config.profile) {
                composeAdditionalArgs = ['--profile', config.profile]
            }

            // Common environment variables
            environment.put 'DATAHUB_VERSION', System.getenv("DATAHUB_VERSION") ?: "v${version}"
            environment.put 'DATAHUB_TELEMETRY_ENABLED', 'false'
            environment.put "METADATA_TESTS_ENABLED", "true"
            environment.put "DATAHUB_REPO", "${docker_registry}"


            // Additional environment variables if specified
            if (config.additionalEnv) {
                config.additionalEnv.each { key, value ->
                    environment.put key, value
                }
            }

            useComposeFiles = [compose_base]
            projectName = project_name
            projectNamePrefix = ''
            buildBeforeUp = false
            buildBeforePull = false
            stopContainers = false
            removeVolumes = false
            retainContainersOnStartupFailure = config.isDebug ? true: false //forcing nulls to bool

            // Apply common configuration
            common_config.each { key, value ->
                delegate."${key}" = value
            }

            // Apply additional task-specific configuration if specified
            if (config.additionalConfig) {
                config.additionalConfig.each { key, value ->
                    delegate."${key}" = value
                }
            }
        }
    }
}


task generateQuickstartComposeConfig {
    description = "Generates the docker-compose config output for quickstart configuration for use with oss quickstart cli"
    group = "docker"

    doLast {
        def composeUpTask = tasks.named("quickstartCLIComposeUp")

        // Some of these details task gleaned from the docker-compose plugin source code
        // https://github.com/avast/gradle-docker-compose-plugin/blob/main/src/main/groovy/com/avast/gradle/dockercompose/ComposeExecutor.groovy
        def composeExecutor = composeUpTask.get().getComposeExecutor().get()
        composeExecutor.parameters.environment.put("DATAHUB_VERSION", '__VERSION__')
        composeExecutor.parameters.environment.put("METADATA_SERVICE_AUTH_ENABLED", 'false')

        def configOutput = composeExecutor.execute('config')
        
        // The `config` returns a fully resolved compose file. We need to replace the DATAHUB_VERSION and HOME with values not available at build time
        configOutput = configOutput.replace('__VERSION__', '${DATAHUB_VERSION}')
        configOutput = configOutput.replace(System.getenv('HOME'), '${HOME}')

        // Parse configOutput as yaml and add UI_INGESTION_DEFAULT_CLI_VERSION as an env var under datahub-gms-quickstart.environment key
        // This needs to be matched with the version of the CLI used to launch quickstart
        def yaml = new Yaml()
        def config = yaml.load(configOutput)
        config['services']['datahub-gms-quickstart']['environment']['UI_INGESTION_DEFAULT_CLI_VERSION'] = '${UI_INGESTION_DEFAULT_CLI_VERSION}'

        // Configure YAML dump options to match Docker Compose format
        def dumpOptions = new org.yaml.snakeyaml.DumperOptions()
        dumpOptions.setDefaultFlowStyle(org.yaml.snakeyaml.DumperOptions.FlowStyle.BLOCK)
        dumpOptions.setIndent(2)
        dumpOptions.setWidth(120)
        yaml = new Yaml(dumpOptions)
        configOutput = yaml.dump(config)

        configOutput = "# This file is generated as part of build process. If any build changes cause this file to be modified, please check in the generated file\n" + configOutput
        
        def outputFile = project.file("${project.projectDir}/quickstart/docker-compose.quickstart-profile.yml")
        outputFile.getParentFile().mkdirs()

        outputFile.text = configOutput

        logger.lifecycle("quickstart Docker Compose config saved to: ${outputFile.absolutePath}")
    }
}

// Register all quickstart tasks
quickstart_configs.each { taskName, config ->
    tasks.register("prepareAll${taskName}"){
        group = 'quickstart-ci'
    }
}

quickstart_configs.each { taskName, config ->
    if (config.modules) {
        tasks.getByName("prepareAll${taskName}").dependsOn(
            config.modules.collect { it + ':dockerPrepare' }
        )
    }
}


quickstart_configs.each { taskName, config ->
    tasks.register("buildImages${taskName}", Exec) {
        ext{
            bakeSpec = [:]
        }
        group = 'quickstart-ci'
        def taskSuffix = config.isDebug? 'debug' : ''
        dependsOn(config.modules.collect { it + ":generateBakeSnippet${taskSuffix}" })
        dependsOn(tasks.getByName("prepareAll${taskName}"))
        dependsOn tasks.named("generateQuickstartComposeConfig")

        def jsonFile = new File(rootProject.buildDir, "bake-spec-${taskName}.json")

        def bakeCmdArgs = ["bake", "-f",  "${jsonFile.absolutePath}"]
        def buildCmd = []
        if (System.getenv("DOCKER_CACHE") == "DEPOT") {
            buildCmd << "depot"
            buildCmd += bakeCmdArgs
            buildCmd += ['--save', '--metadata-file', "${rootProject.buildDir}/build-metadata.json"]
            if (project.properties.getOrDefault("dockerPush", false)){
                buildCmd << "--push"
            }
        } else {
            buildCmd += ["docker",  "buildx" ]
            buildCmd +=bakeCmdArgs
        }
        commandLine buildCmd
        workingDir rootProject.projectDir

        doFirst {
            def bakeSnippets = [:]
            def targets = []

            config.modules.each { module ->
                def moduleProject = project.project(module)
                def generateBakeSnippetsTask = moduleProject.tasks.getByName("generateBakeSnippet${taskSuffix}")
                bakeSnippets.putAll(generateBakeSnippetsTask.bakeSpec.target)
                targets.addAll(generateBakeSnippetsTask.bakeSpec.target.keySet())
            }

            ext.bakeSpec.group = [ "default": ["targets": targets] ]
            ext.bakeSpec.target = bakeSnippets

            jsonFile.parentFile.mkdirs()
            jsonFile.text = groovy.json.JsonOutput.prettyPrint(groovy.json.JsonOutput.toJson(ext.bakeSpec))
        }
    }
}

// Configure dependencies for ComposeUp tasks
quickstart_configs.each { taskName, config ->
    if (config.modules) {
        tasks.getByName("${taskName}ComposeUp").dependsOn(
                tasks.getByName("buildImages${taskName}")
        )
    }
}


tasks.register('minDockerCompose2.20', Exec) {
    executable 'bash'
    args '-c', 'echo -e "$(docker compose version --short)\n2.20"|sort --version-sort --check=quiet --reverse'
}

tasks.register('quickstartNuke') {
    group = 'quickstart'
    doFirst {
        quickstart_configs.each { taskName, config ->
            dockerCompose."${taskName}".removeVolumes = !config.preserveVolumes
        }
    }
    finalizedBy(tasks.withType(ComposeDownForced))
}

tasks.register('quickstartDown') {
    group = 'quickstart'
    finalizedBy(tasks.withType(ComposeDownForced))
}

tasks.withType(ComposeUp).configureEach {
    shouldRunAfter('quickstartNuke')
    dependsOn tasks.named("minDockerCompose2.20")
}

// Register all quickstart Reload tasks. For quickstartDebug, the reload task is DebugReload. (Taskname without quickstart prefix)
quickstart_configs.each { taskName, config ->
    if (config.isDebug) {
        def reloadTaskName = taskName.replaceFirst(/^quickstart/, "")
        tasks.register("${reloadTaskName}Reload", Exec) {
            dependsOn tasks.named("prepareAll${taskName}")
            group = 'quickstart'
            description = "Build and reload only changed containers for the ${taskName} task"
            doFirst {
                def executedTasks = project.gradle.taskGraph.allTasks.findAll { it.state.executed }
                def containersToRestart = []

                moduleToContainer.each { modulePath, containerName ->
                    def moduleProject = project.project(modulePath)
                    def dockerPrepareTask = moduleProject.tasks.findByName('dockerPrepare')
                    
                    if (dockerPrepareTask && executedTasks.contains(dockerPrepareTask) && !dockerPrepareTask.state.upToDate) {
                        containersToRestart << "${containerName}-${config.profile}"
                    }
                }
                
                // Only restart containers that had their modules rebuilt
                if (containersToRestart) {
                    def composeFiles = dockerCompose."${taskName}".useComposeFiles.get()
                    def composeFileArgs = composeFiles.collectMany { ['-f', it] }

                    def cmd = ["docker compose -p ${project_name} --profile ${config.profile}"] + composeFileArgs + ['restart'] + containersToRestart
                    println(cmd.join(" "))
                    commandLine 'bash', '-c', cmd.join(" ")
                } else {
                    // If no containers need restart, make this a no-op
                    commandLine 'bash', '-c', 'echo "No containers need restarting - all modules are up to date"'
                }    
            }
        }
        tasks.register("${reloadTaskName}ReloadEnv", Exec) {
            dependsOn tasks.named("prepareAll${taskName}")
            group = 'quickstart'
            description = "Build changed containers but recreate all services for the ${taskName} task"
            doFirst {
                def containersToRestart = []
                moduleToContainer.each { modulePath, containerName ->
                    // Find which of of the reloadable modules are in used in this task
                    if (config.modules.contains(modulePath)) {
                        containersToRestart << "${containerName}-${config.profile}"
                    }
                }
                def composeFiles = dockerCompose."${taskName}".useComposeFiles.get()
                def composeFileArgs = composeFiles.collectMany { ['-f', it] }

                def cmd = ["docker compose -p ${project_name} --profile ${config.profile}"] + composeFileArgs + ['up', '-d', '--no-deps'] + containersToRestart
                println(cmd.join(" "))
                commandLine 'bash', '-c', cmd.join(" ")
            }
        }
    }
}