feat: org agent management (#259)

* feat: org agent management

* feat: org agent management

* fix: adding migrations for legacy and refactoring.

* fix: refactoring

* fix: delete agent

* fix: refactoring

---------

Co-authored-by: charles-gauthereau <charles.gauthereau@soluce-technologies.com>

src/db/migrations/0049_chief_terrax.sql

@@ -0,0 +1,10 @@
+CREATE TABLE "organization_agents" (
+	"organization_id" uuid NOT NULL,
+	"agent_id" uuid NOT NULL,
+	CONSTRAINT "organization_agents_organization_id_agent_id_unique" UNIQUE("organization_id","agent_id")
+);
+--> statement-breakpoint
+ALTER TABLE "agents" ADD COLUMN "organization_id" uuid;--> statement-breakpoint
+ALTER TABLE "organization_agents" ADD CONSTRAINT "organization_agents_organization_id_organization_id_fk" FOREIGN KEY ("organization_id") REFERENCES "public"."organization"("id") ON DELETE cascade ON UPDATE no action;--> statement-breakpoint
+ALTER TABLE "organization_agents" ADD CONSTRAINT "organization_agents_agent_id_agents_id_fk" FOREIGN KEY ("agent_id") REFERENCES "public"."agents"("id") ON DELETE cascade ON UPDATE no action;--> statement-breakpoint
+ALTER TABLE "agents" ADD CONSTRAINT "agents_organization_id_organization_id_fk" FOREIGN KEY ("organization_id") REFERENCES "public"."organization"("id") ON DELETE cascade ON UPDATE no action;

src/db/migrations/0050_dark_saracen.sql

@@ -0,0 +1,31 @@
+-- Custom SQL migration file, put your code below! --
+
+
+DO $$
+    DECLARE
+        org RECORD;
+        proj RECORD;
+        db RECORD;
+    BEGIN
+        FOR org IN SELECT id FROM organization LOOP
+                FOR proj IN
+                    SELECT id FROM projects WHERE organization_id = org.id
+                    LOOP
+                        FOR db IN
+                            SELECT agent_id FROM databases WHERE project_id = proj.id
+                            LOOP
+                                IF db.agent_id IS NOT NULL THEN
+                                    INSERT INTO organization_agents (
+                                        organization_id,
+                                        agent_id
+                                    )
+                                    VALUES (
+                                               org.id,
+                                               db.agent_id
+                                           )
+                                    ON CONFLICT (organization_id, agent_id) DO NOTHING;
+                                END IF;
+                            END LOOP;
+                    END LOOP;
+            END LOOP;
+    END $$;

src/db/migrations/0051_young_senator_kelly.sql

@@ -0,0 +1,3 @@
+ALTER TABLE "organization_agents" ADD COLUMN "updated_at" timestamp;--> statement-breakpoint
+ALTER TABLE "organization_agents" ADD COLUMN "created_at" timestamp DEFAULT now() NOT NULL;--> statement-breakpoint
+ALTER TABLE "organization_agents" ADD COLUMN "deleted_at" timestamp;

src/db/migrations/0052_cute_punisher.sql

@@ -0,0 +1 @@
+ALTER TABLE "two_factor" ADD COLUMN "verified" boolean NOT NULL;

src/db/migrations/meta/_journal.json

@@ -344,6 +344,34 @@
       "when": 1774886139855,
       "tag": "0048_yellow_eddie_brock",
       "breakpoints": true
+    },
+    {
+      "idx": 49,
+      "version": "7",
+      "when": 1775585355435,
+      "tag": "0049_chief_terrax",
+      "breakpoints": true
+    },
+    {
+      "idx": 50,
+      "version": "7",
+      "when": 1775761395545,
+      "tag": "0050_dark_saracen",
+      "breakpoints": true
+    },
+    {
+      "idx": 51,
+      "version": "7",
+      "when": 1775762288699,
+      "tag": "0051_young_senator_kelly",
+      "breakpoints": true
+    },
+    {
+      "idx": 52,
+      "version": "7",
+      "when": 1775803959723,
+      "tag": "0052_cute_punisher",
+      "breakpoints": true
     }
   ]
 }

src/db/schema/02_user.ts

@@ -97,6 +97,7 @@ export const twoFactor = pgTable("two_factor", {
   id: uuid().defaultRandom().primaryKey(),
   secret: text("secret").notNull(),
   backupCodes: text("backup_codes").notNull(),
+  verified: boolean("verified").notNull(),
   userId: uuid("user_id")
     .notNull()
     .references(() => user.id, { onDelete: "cascade" }),

src/db/schema/07_database.ts

@@ -1,5 +1,5 @@
 import {pgTable, text, boolean, timestamp, uuid, integer, pgEnum} from "drizzle-orm/pg-core";
-import {Agent, agent} from "./08_agent";
+import {Agent, agent, AgentWith} from "./08_agent";
 import {Project, project} from "./06_project";
 import {relations} from "drizzle-orm";
 import {dbmsEnum, statusEnum} from "./types";
@@ -123,7 +123,7 @@ export type RetentionPolicy = z.infer<typeof retentionPolicySchema>;
 
 
 export type DatabaseWith = Database & {
-    agent?: Agent | null;
+    agent?: Agent | AgentWith | null;
     project?: Project | null;
     backups?: Backup[] | null;
     restorations?: Restoration[] | null;

src/db/schema/08_agent.ts

@@ -1,9 +1,10 @@
-import {boolean, pgTable, text, timestamp, uuid, integer} from "drizzle-orm/pg-core";
+import {boolean, pgTable, text, timestamp, uuid, integer, unique} from "drizzle-orm/pg-core";
 import {createSelectSchema} from "drizzle-zod";
 import {z} from "zod";
 import {Database, database} from "@/db/schema/07_database";
 import {relations} from "drizzle-orm";
 import {timestamps} from "@/db/schema/00_common";
+import {organization} from "@/db/schema/03_organization";
 
 export const agent = pgTable("agents", {
     id: uuid("id").primaryKey().defaultRandom(),
@@ -14,22 +15,56 @@ export const agent = pgTable("agents", {
     description: text("description").notNull(),
     isArchived: boolean("is_archived").default(false),
     lastContact: timestamp("last_contact"),
+    organizationId: uuid("organization_id").references(() => organization.id, {onDelete: "cascade"}),
     ...timestamps
 });
 
+
+export const organizationAgent = pgTable(
+    "organization_agents",
+    {
+        organizationId: uuid('organization_id')
+            .notNull()
+            .references(() => organization.id, {onDelete: 'cascade'}),
+        agentId: uuid('agent_id')
+            .notNull()
+            .references(() => agent.id, {onDelete: 'cascade'}),
+        ...timestamps
+    },
+    (t) => [unique().on(t.organizationId, t.agentId)]
+
+);
+
 export const agentSchema = createSelectSchema(agent);
 export type Agent = z.infer<typeof agentSchema>;
 
 
 export const agentRelations = relations(agent, ({many}) => ({
     databases: many(database),
+    organizations: many(organizationAgent),
+}));
+
+export const organizationAgentRelations = relations(organizationAgent, ({one}) => ({
+    organization: one(organization, {
+        fields: [organizationAgent.organizationId],
+        references: [organization.id],
+    }),
+    agent: one(agent, {
+        fields: [organizationAgent.agentId],
+        references: [agent.id],
+    }),
 }));
 
 
 export type AgentWith = Agent & {
     databases?: Database[] | null;
+    organizations: {
+        organizationId: string;
+        agentId: string;
+    }[];
 };
 
 export type AgentWithDatabases = Agent & {
     databases: Database[] | [];
 };
+

src/db/schema/12_storage-channel.ts

@@ -4,8 +4,6 @@ import {organization} from "@/db/schema/03_organization";
 import {relations} from "drizzle-orm";
 import {createSelectSchema} from "drizzle-zod";
 import {z} from "zod";
-import {database} from "@/db/schema/07_database";
-
 
 export const providerStorageKindEnum = pgEnum('provider_storage_kind', ['local', 's3', 'google-drive']);
 
@@ -32,7 +30,6 @@ export const organizationStorageChannel = pgTable(
     (t) => [unique().on(t.organizationId, t.storageChannelId)]
 );
 
-
 export const storageChannelRelations = relations(storageChannel, ({many}) => ({
     organizations: many(organizationStorageChannel),
 }));

src/db/schema/13_storage-policy.ts

@@ -1,13 +1,10 @@
 import {boolean, pgTable, uuid} from "drizzle-orm/pg-core";
 import {timestamps} from "@/db/schema/00_common";
 import {relations} from "drizzle-orm";
-import {Backup, Database, database, Restoration, RetentionPolicy} from "@/db/schema/07_database";
+import {database} from "@/db/schema/07_database";
 import {createSelectSchema} from "drizzle-zod";
 import {z} from "zod";
 import {StorageChannel, storageChannel} from "@/db/schema/12_storage-channel";
-import {Agent} from "@/db/schema/08_agent";
-import {Project} from "@/db/schema/06_project";
-import {AlertPolicy} from "@/db/schema/10_alert-policy";
 
 export const storagePolicy = pgTable('storage_policy', {
     id: uuid('id').defaultRandom().primaryKey(),

src/db/services/agent.ts

@@ -0,0 +1,43 @@
+import {and, desc, eq, sql} from "drizzle-orm";
+import {db} from "@/db";
+import {Agent, agent, organizationAgent} from "@/db/schema/08_agent";
+import {Database, database} from "@/db/schema/07_database";
+
+export async function getOrganizationAgents(organizationId: string) {
+
+    return await db
+        .select({
+            id: agent.id,
+            name: agent.name,
+            organizationId: agent.organizationId,
+            slug: agent.slug,
+            healthErrorCount: agent.healthErrorCount,
+            description: agent.description,
+            isArchived: agent.isArchived,
+            lastContact: agent.lastContact,
+            version: agent.version,
+            updatedAt: agent.updatedAt,
+            createdAt: agent.createdAt,
+            deletedAt: agent.deletedAt,
+            databases: sql<Database[]>`
+                COALESCE(
+                    json_agg(${database}.*) FILTER (WHERE ${database}.id IS NOT NULL),
+                    '[]'
+                )
+            `,
+        })
+        .from(organizationAgent)
+        .innerJoin(
+            agent,
+            eq(organizationAgent.agentId, agent.id)
+        )
+        .leftJoin(database, eq(database.agentId, agent.id))
+        .groupBy(agent.id)
+        .orderBy(desc(agent.createdAt))
+        .where(
+            and(
+                eq(organizationAgent.organizationId, organizationId),
+                eq(agent.isArchived, false)
+            )
+        ) as unknown as Agent[];
+}

src/db/services/database.ts

@@ -0,0 +1,39 @@
+"use server"
+import {db} from "@/db";
+import {DatabaseWith} from "@/db/schema/07_database";
+import {AgentWith} from "@/db/schema/08_agent";
+
+export async function getOrganizationAvailableDatabases(
+    organizationId: string,
+    projectId?: string
+) {
+
+    const availableDatabases = (
+        await db.query.database.findMany({
+            where: (db, { eq, or, isNull }) =>
+                projectId
+                    ? or(isNull(db.projectId), eq(db.projectId, projectId))
+                    : isNull(db.projectId),
+            with: {
+                agent: {
+                    with: {
+                        organizations: true
+                    }
+                },
+                project: true,
+                backups: true,
+                restorations: true,
+            },
+            orderBy: (db, {desc}) => [desc(db.createdAt)],
+        })
+    ) as DatabaseWith[];
+
+    return availableDatabases.filter(db => {
+        const agent = db.agent as AgentWith;
+        if (agent?.isArchived) return false;
+        return (
+            agent?.organizationId === organizationId ||
+            agent?.organizations?.some(org => org.organizationId === organizationId)
+        );
+    })
+}