nixpkgs/nixos/modules/services/scheduling/prefect.nix

{
  lib,
  pkgs,
  config,
  ...
}:

let
  cfg = config.services.prefect;
  inherit (lib.types)
    bool
    str
    enum
    path
    attrsOf
    nullOr
    submodule
    port
    ;

in
{
  options.services.prefect = {
    enable = lib.mkOption {
      type = bool;
      default = false;
      description = "enable prefect server and worker services";
    };

    package = lib.mkPackageOption pkgs "prefect" { };

    host = lib.mkOption {
      type = str;
      default = "127.0.0.1";
      example = "0.0.0.0";
      description = "Prefect server host";
    };

    port = lib.mkOption {
      type = port;
      default = 4200;
      description = "Prefect server port";
    };

    dataDir = lib.mkOption {
      type = path;
      default = "/var/lib/prefect-server";
      description = ''
        Specify the directory for Prefect.
      '';
    };

    database = lib.mkOption {
      type = enum [
        "sqlite"
        "postgres"
      ];
      default = "sqlite";
      description = "which database to use for prefect server: sqlite or postgres";
    };

    databaseHost = lib.mkOption {
      type = str;
      default = "localhost";
      description = "database host for postgres only";
    };

    databasePort = lib.mkOption {
      type = str;
      default = "5432";
      description = "database port for postgres only";
    };

    databaseName = lib.mkOption {
      type = str;
      default = "prefect";
      description = "database name for postgres only";
    };

    databaseUser = lib.mkOption {
      type = str;
      default = "postgres";
      description = "database user for postgres only";
    };

    databasePasswordFile = lib.mkOption {
      type = nullOr str;
      default = null;
      description = ''
        path to a file containing e.g.:
          DBPASSWORD=supersecret

        stored outside the nix store, read by systemd as EnvironmentFile.
      '';
    };

    # now define workerPools as an attribute set of submodules,
    # each key is the pool name, and the submodule has an installPolicy
    workerPools = lib.mkOption {
      type = attrsOf (submodule {
        options = {
          installPolicy = lib.mkOption {
            type = enum [
              "always"
              "if-not-present"
              "never"
              "prompt"
            ];
            default = "always";
            description = "install policy for the worker (always, if-not-present, never, prompt)";
          };
        };
      });
      default = { };
      description = ''
        define a set of worker pools with submodule config. example:
        workerPools.my-pool = {
          installPolicy = "never";
        };
      '';
    };

    baseUrl = lib.mkOption {
      type = nullOr str;
      default = null;
      description = "external url when served by a reverse proxy, e.g. https://example.com/prefect";
    };
  };

  config = lib.mkIf cfg.enable {
    # define systemd.services as the server plus any worker definitions
    systemd.services =
      {
        "prefect-server" = {
          description = "prefect server";
          wantedBy = [ "multi-user.target" ];
          after = [ "network.target" ];

          serviceConfig = {
            DynamicUser = true;
            StateDirectory = "prefect-server";
            # TODO all my efforts to setup the database url
            # have failed with some unable to open file
            Environment = [
              "PREFECT_HOME=%S/prefect-server"
              "PREFECT_UI_STATIC_DIRECTORY=%S/prefect-server"
              "PREFECT_SERVER_ANALYTICS_ENABLED=off"
              "PREFECT_UI_API_URL=${cfg.baseUrl}/api"
              "PREFECT_UI_URL=${cfg.baseUrl}"
            ];
            EnvironmentFile =
              if cfg.database == "postgres" && cfg.databasePasswordFile != null then
                [ cfg.databasePasswordFile ]
              else
                [ ];

            # ReadWritePaths = [ cfg.dataDir ];
            ProtectSystem = "strict";
            ProtectHome = true;
            PrivateTmp = true;
            NoNewPrivileges = true;
            MemoryDenyWriteExecute = true;
            LockPersonality = true;
            CapabilityBoundingSet = [ ];
            AmbientCapabilities = [ ];
            RestrictSUIDSGID = true;
            RestrictAddressFamilies = [
              "AF_INET"
              "AF_INET6"
              "AF_UNIX"
            ];
            ProtectKernelTunables = true;
            ProtectKernelModules = true;
            ProtectKernelLogs = true;
            ProtectControlGroups = true;
            MemoryAccounting = true;
            CPUAccounting = true;

            ExecStart = "${pkgs.prefect}/bin/prefect server start --host ${cfg.host} --port ${toString cfg.port}";
            Restart = "always";
            WorkingDirectory = cfg.dataDir;
          };
        };
      }
      // lib.concatMapAttrs (poolName: poolCfg: {
        # return a partial attr set with one key: "prefect-worker-..."
        "prefect-worker-${poolName}" = {
          description = "prefect worker for pool '${poolName}'";
          wantedBy = [ "multi-user.target" ];
          after = [ "network.target" ];

          environment.systemPackages = cfg.package;

          serviceConfig = {
            DynamicUser = true;
            StateDirectory = "prefect-worker-${poolName}";
            Environment = [
              "PREFECT_HOME=%S/prefect-worker-${poolName}"
              "PREFECT_API_URL=${cfg.baseUrl}/api"
            ];
            ProtectSystem = "strict";
            ProtectHome = true;
            PrivateTmp = true;
            NoNewPrivileges = true;
            MemoryDenyWriteExecute = true;
            LockPersonality = true;
            CapabilityBoundingSet = [ ];
            AmbientCapabilities = [ ];
            RestrictSUIDSGID = true;
            RestrictAddressFamilies = [
              "AF_INET"
              "AF_INET6"
              "AF_UNIX"
            ];
            ProtectKernelTunables = true;
            ProtectKernelModules = true;
            ProtectKernelLogs = true;
            ProtectControlGroups = true;
            MemoryAccounting = true;
            CPUAccounting = true;
            ExecStart = ''
              ${pkgs.prefect}/bin/prefect worker start \
                --pool ${poolName} \
                --type process \
                --install-policy ${poolCfg.installPolicy}
            '';
            Restart = "always";
          };
        };
      }) cfg.workerPools;
  };
}