From 835b2f88220cfbf3e4110592e2940df7df126965 Mon Sep 17 00:00:00 2001
From: Pol Dellaiera <pol.dellaiera@protonmail.com>
Date: Fri, 16 Aug 2024 18:14:15 +0200
Subject: [PATCH] nixos/chromadb: init

---
 .../manual/release-notes/rl-2411.section.md   |   3 +
 nixos/modules/module-list.nix                 |   1 +
 nixos/modules/services/databases/chromadb.nix | 107 ++++++++++++++++++
 nixos/tests/all-tests.nix                     |   1 +
 nixos/tests/chromadb.nix                      |  26 +++++
 5 files changed, 138 insertions(+)
 create mode 100644 nixos/modules/services/databases/chromadb.nix
 create mode 100644 nixos/tests/chromadb.nix

diff --git a/nixos/doc/manual/release-notes/rl-2411.section.md b/nixos/doc/manual/release-notes/rl-2411.section.md
index 972e13852eca..547d458578e4 100644
--- a/nixos/doc/manual/release-notes/rl-2411.section.md
+++ b/nixos/doc/manual/release-notes/rl-2411.section.md
@@ -85,6 +85,9 @@
 
 - [Proton Mail bridge](https://proton.me/mail/bridge), a desktop application that runs in the background, encrypting and decrypting messages as they enter and leave your computer. It lets you add your Proton Mail account to your favorite email client via IMAP/SMTP by creating a local email server on your computer.
 
+- [chromadb](https://www.trychroma.com/), an open-source AI application
+  database. Batteries included. Available as [services.chromadb](options.html#opt-services.chromadb.enable).
+
 ## Backward Incompatibilities {#sec-release-24.11-incompatibilities}
 
 - `transmission` package has been aliased with a `trace` warning to `transmission_3`. Since [Transmission 4 has been released last year](https://github.com/transmission/transmission/releases/tag/4.0.0), and Transmission 3 will eventually go away, it was decided perform this warning alias to make people aware of the new version. The `services.transmission.package` defaults to `transmission_3` as well because the upgrade can cause data loss in certain specific usage patterns (examples: [#5153](https://github.com/transmission/transmission/issues/5153), [#6796](https://github.com/transmission/transmission/issues/6796)). Please make sure to back up to your data directory per your usage:
diff --git a/nixos/modules/module-list.nix b/nixos/modules/module-list.nix
index 210ca98e2f27..dbeae42298de 100644
--- a/nixos/modules/module-list.nix
+++ b/nixos/modules/module-list.nix
@@ -457,6 +457,7 @@
   ./services/continuous-integration/woodpecker/server.nix
   ./services/databases/aerospike.nix
   ./services/databases/cassandra.nix
+  ./services/databases/chromadb.nix
   ./services/databases/clickhouse.nix
   ./services/databases/cockroachdb.nix
   ./services/databases/couchdb.nix
diff --git a/nixos/modules/services/databases/chromadb.nix b/nixos/modules/services/databases/chromadb.nix
new file mode 100644
index 000000000000..d8d60078cf45
--- /dev/null
+++ b/nixos/modules/services/databases/chromadb.nix
@@ -0,0 +1,107 @@
+{
+  config,
+  pkgs,
+  lib,
+  ...
+}:
+
+let
+  cfg = config.services.chromadb;
+  inherit (lib)
+    mkEnableOption
+    mkOption
+    mkIf
+    types
+    literalExpression
+    ;
+in
+{
+
+  meta.maintainers = with lib.maintainers; [ drupol ];
+
+  options = {
+    services.chromadb = {
+      enable = mkEnableOption "ChromaDB, an open-source AI application database.";
+
+      package = mkOption {
+        type = types.package;
+        example = literalExpression "pkgs.python3Packages.chromadb";
+        default = pkgs.python3Packages.chromadb;
+        defaultText = "pkgs.python3Packages.chromadb";
+        description = "ChromaDB package to use.";
+      };
+
+      host = mkOption {
+        type = types.str;
+        default = "127.0.0.1";
+        description = ''
+          Defines the IP address by which ChromaDB will be accessible.
+        '';
+      };
+
+      port = mkOption {
+        type = types.port;
+        default = 8000;
+        description = ''
+          Defined the port number to listen.
+        '';
+      };
+
+      logFile = mkOption {
+        type = types.path;
+        default = "/var/log/chromadb/chromadb.log";
+        description = ''
+          Specifies the location of file for logging output.
+        '';
+      };
+
+      dbpath = mkOption {
+        type = types.str;
+        default = "/var/lib/chromadb";
+        description = "Location where ChromaDB stores its files";
+      };
+
+      openFirewall = mkOption {
+        type = types.bool;
+        default = false;
+        description = ''
+          Whether to automatically open the specified TCP port in the firewall.
+        '';
+      };
+    };
+  };
+
+  config = mkIf cfg.enable {
+    systemd.services.chromadb = {
+      description = "ChromaDB";
+      after = [ "network.target" ];
+      wantedBy = [ "multi-user.target" ];
+      serviceConfig = {
+        Type = "simple";
+        StateDirectory = "chromadb";
+        WorkingDirectory = "/var/lib/chromadb";
+        LogsDirectory = "chromadb";
+        ExecStart = "${lib.getExe cfg.package} run --path ${cfg.dbpath} --host ${cfg.host} --port ${toString cfg.port} --log-path ${cfg.logFile}";
+        Restart = "on-failure";
+        ProtectHome = true;
+        ProtectSystem = "strict";
+        PrivateTmp = true;
+        PrivateDevices = true;
+        ProtectHostname = true;
+        ProtectClock = true;
+        ProtectKernelTunables = true;
+        ProtectKernelModules = true;
+        ProtectKernelLogs = true;
+        ProtectControlGroups = true;
+        NoNewPrivileges = true;
+        RestrictRealtime = true;
+        RestrictSUIDSGID = true;
+        RemoveIPC = true;
+        PrivateMounts = true;
+        DynamicUser = true;
+      };
+    };
+
+    networking.firewall.allowedTCPPorts = lib.optionals cfg.openFirewall [ cfg.port ];
+  };
+}
diff --git a/nixos/tests/all-tests.nix b/nixos/tests/all-tests.nix
index f129a36e139e..d13e542808b0 100644
--- a/nixos/tests/all-tests.nix
+++ b/nixos/tests/all-tests.nix
@@ -192,6 +192,7 @@ in {
   cfssl = handleTestOn ["aarch64-linux" "x86_64-linux"] ./cfssl.nix {};
   cgit = handleTest ./cgit.nix {};
   charliecloud = handleTest ./charliecloud.nix {};
+  chromadb = runTest ./chromadb.nix;
   chromium = (handleTestOn ["aarch64-linux" "x86_64-linux"] ./chromium.nix {}).stable or {};
   chrony = handleTestOn ["aarch64-linux" "x86_64-linux"] ./chrony.nix {};
   chrony-ptp = handleTestOn ["aarch64-linux" "x86_64-linux"] ./chrony-ptp.nix {};
diff --git a/nixos/tests/chromadb.nix b/nixos/tests/chromadb.nix
new file mode 100644
index 000000000000..be04d10e74de
--- /dev/null
+++ b/nixos/tests/chromadb.nix
@@ -0,0 +1,26 @@
+{ lib, pkgs, ... }:
+
+let
+  lib = pkgs.lib;
+
+in
+{
+  name = "chromadb";
+  meta.maintainers = [ lib.maintainers.drupol ];
+
+  nodes = {
+    machine =
+      { pkgs, ... }:
+      {
+        services.chromadb = {
+          enable = true;
+        };
+      };
+  };
+
+  testScript = ''
+    machine.start()
+    machine.wait_for_unit("chromadb.service")
+    machine.wait_for_open_port(8000)
+  '';
+}