From 3c9f08387f45f157ecbe588ccfa490c7fcee8c5f Mon Sep 17 00:00:00 2001
From: Roman Bodavskiy <roman.bodavskiy@metalamp.io>
Date: Mon, 5 Feb 2024 15:22:10 +0300
Subject: [PATCH] clickhouse init script

---
 nix/clickhouse.nix | 160 +++++++++++++++++++++++++++++++++------------
 1 file changed, 119 insertions(+), 41 deletions(-)

diff --git a/nix/clickhouse.nix b/nix/clickhouse.nix
index ea417c9b..e57620b2 100644
--- a/nix/clickhouse.nix
+++ b/nix/clickhouse.nix
@@ -32,55 +32,133 @@ in
       default = "";
     };
 
+    initialDatabases = lib.mkOption {
+      type = types.listOf (types.submodule {
+        options = {
+          name = lib.mkOption {
+            type = types.str;
+            description = ''
+              The name of the database to create.
+            '';
+          };
+          schemas = lib.mkOption {
+            type = types.nullOr (types.listOf types.path);
+            default = null;
+            description = ''
+              The initial list of schemas for the database; if null (the default),
+              an empty database is created.
+            '';
+          };
+        };
+      });
+      default = [ ];
+      description = ''
+        List of database names and their initial schemas that should be used to create databases on the first startup
+        of Postgres. The schema attribute is optional: If not specified, an empty database is created.
+      '';
+      example = lib.literalExpression ''
+        [
+          {
+            name = "foodatabase";
+            schemas = [ ./fooschemas ./bar.sql ];
+          }
+          { name = "bardatabase"; }
+        ]
+      '';
+    };
     outputs.settings = lib.mkOption {
       type = types.deferredModule;
       internal = true;
       readOnly = true;
       default = {
-        processes = {
-          "${name}" =
-            let
-              clickhouseConfig = pkgs.writeText "clickhouse-config.yaml" ''
-                logger:
-                  level: warning
-                  console: 1
-                tcp_port: ${toString config.port}
-                default_profile: default
-                default_database: default
-                path: ${config.dataDir}/clickhouse
-                tmp_path: ${config.dataDir}/clickhouse/tmp
-                user_files_path: ${config.dataDir}/clickhouse/user_files
-                format_schema_path: ${config.dataDir}/clickhouse/format_schemas
-                user_directories:
-                  users_xml:
-                    path: ${config.package}/etc/clickhouse-server/users.xml
-                ${config.extraConfig}
-              '';
-              startScript = pkgs.writeShellApplication {
-                name = "start-clickhouse";
-                runtimeInputs = [ config.package ];
-                text = ''
-                  clickhouse-server --config-file=${clickhouseConfig}
-                '';
+        processes =
+          let
+            clickhouseConfig = pkgs.writeText "clickhouse-config.yaml" ''
+              logger:
+                level: warning
+                console: 1
+              tcp_port: ${toString config.port}
+              default_profile: default
+              default_database: default
+              path: ${config.dataDir}/clickhouse
+              tmp_path: ${config.dataDir}/clickhouse/tmp
+              user_files_path: ${config.dataDir}/clickhouse/user_files
+              format_schema_path: ${config.dataDir}/clickhouse/format_schemas
+              user_directories:
+                users_xml:
+                  path: ${config.package}/etc/clickhouse-server/users.xml
+              ${config.extraConfig}
+            '';
+          in
+          {
+            # DB initialization
+            "${name}-init" =
+              let
+                # https://github.com/ClickHouse/ClickHouse/issues/4491
+                setupInitialSchema = schema: '' < ${schema} tr -s '\r\n' ' ' | clickhouse-client -mn --port ${builtins.toString config.port}; '';
+                setupInitialDatabases =
+                  lib.concatMapStrings
+                    (database: ''
+                      echo "Creating database: ${database.name}"
+                      clickhouse-client --port ${builtins.toString config.port} --query "CREATE DATABASE iF NOT EXISTS ${database.name}"
+                      echo "Database successfully created: ${database.name}"
+                      ${lib.optionalString (database.schemas != null)
+                        (lib.concatMapStrings (schema: setupInitialSchema schema) database.schemas)}
+                    '')
+                    config.initialDatabases;
+                setupScript = pkgs.writeShellApplication {
+                  name = "setup-clickhouse";
+                  runtimeInputs = with pkgs; [ config.package coreutils gnugrep gawk ];
+                  text = ''
+                    if test -d ${config.dataDir}
+                      then echo "Clickhouse database directory ${config.dataDir} appears to contain a database; Skipping initialization"
+                      else
+                        echo "Clickhouse is setting up the initial database."
+                        set -m
+                        clickhouse-server --config-file=${clickhouseConfig} &
+                        sleep 5s
+                        echo "Clickhouse server started."
+                        ${setupInitialDatabases}
+                        echo "Clickhouse db setting is done."
+                        kill %1
+                        echo "Clickhouse server stopped."
+                    fi
+                  '';
+                };
+              in
+              {
+                command = setupScript;
+                namespace = name;
               };
-            in
-            {
-              command = "${lib.getExe startScript}";
 
-              readiness_probe = {
-                exec.command = ''${config.package}/bin/clickhouse-client --query "SELECT 1" --port ${builtins.toString config.port}'';
-                initial_delay_seconds = 2;
-                period_seconds = 10;
-                timeout_seconds = 4;
-                success_threshold = 1;
-                failure_threshold = 5;
-              };
-              namespace = name;
+            # DB process
+            "${name}" =
+              let
+                startScript = pkgs.writeShellApplication {
+                  name = "start-clickhouse";
+                  runtimeInputs = [ config.package ];
+                  text = ''
+                    clickhouse-server --config-file=${clickhouseConfig}
+                  '';
+                };
+              in
+              {
+                command = "${lib.getExe startScript}";
 
-              # https://github.com/F1bonacc1/process-compose#-auto-restart-if-not-healthy
-              availability.restart = "on_failure";
-            };
-        };
+                readiness_probe = {
+                  exec.command = ''${config.package}/bin/clickhouse-client --query "SELECT 1" --port ${builtins.toString config.port}'';
+                  initial_delay_seconds = 2;
+                  period_seconds = 10;
+                  timeout_seconds = 4;
+                  success_threshold = 1;
+                  failure_threshold = 5;
+                };
+                namespace = name;
+                depends_on."${name}-init".condition = "process_completed_successfully";
+                # https://github.com/F1bonacc1/process-compose#-auto-restart-if-not-healthy
+                availability.restart = "on_failure";
+              };
+          };
       };
     };
   };