context-efficient-ai/database/postgres.ts

// import { Pool } from "pg";
import { neon } from "@neondatabase/serverless";
import { Kysely /*PostgresDialect*/ } from "kysely";
import { NeonDialect } from "kysely-neon";
import type Database from "./generated/Database";
import type {
  ConversationEntity,
  FactEntity,
  FactTriggerEntity,
  MessageEntity,
  UserEntity,
} from "./common.ts";
import type { CommittedMessage } from "../types";

// export const pool = new Pool({
//   connectionString: env.POSTGRES_CONNECTION_STRING as string,
//   // channelBinding: require ?
// });

// const dialect = new PostgresDialect({
//   pool,
// });

export function getDbClient(POSTGRES_CONNECTION_STRING: string) {
  const dialect = new NeonDialect({
    neon: neon(POSTGRES_CONNECTION_STRING as string),
  });

  // Database interface is passed to Kysely's constructor, and from now on, Kysely
  // knows your database structure.
  // Dialect is passed to Kysely's constructor, and from now on, Kysely knows how
  // to communicate with your database.
  const dbClient = new Kysely<Database>({
    dialect,
  });

  return dbClient;
}

export function getDb(POSTGRES_CONNECTION_STRING: string) {
  const dbClient = getDbClient(POSTGRES_CONNECTION_STRING);

  const conversations: ConversationEntity = {
    construct: (conversation) => conversation,
    create: async (conversation) => {
      const insertedRows = await dbClient
        .insertInto("conversations")
        .values(conversation)
        .returningAll()
        .execute();
      return insertedRows[0];
    },
    createMany: async (conversations) => {
      const insertedRows = await dbClient
        .insertInto("conversations")
        .values(conversations)
        .returningAll()
        .execute();
      return insertedRows;
    },
    findAll: async () => {
      const rows = await dbClient
        .selectFrom("conversations")
        .selectAll()
        .execute();
      return rows;
    },
    findById: async (id) => {
      const row = await dbClient
        .selectFrom("conversations")
        .selectAll()
        .where("id", "=", id)
        .execute();
      return row[0];
    },
    update: async (id, data) => {
      await dbClient
        .updateTable("conversations")
        .set(data)
        .where("id", "=", id)
        .execute();
    },
    delete: async (id) => {
      await dbClient.deleteFrom("conversations").where("id", "=", id).execute();
    },
    fetchMessages: async (conversationId) => {
      const rows = await dbClient
        .selectFrom("messages")
        .selectAll()
        .where("conversationId", "=", conversationId)
        .execute();
      return rows as Array<CommittedMessage>;
    },
  };

  const facts: FactEntity = {
    construct: (fact) => fact,
    create: async (fact) => {
      const insertedRows = await dbClient
        .insertInto("facts")
        .values(fact)
        .returningAll()
        .execute();
      return insertedRows[0];
    },
    createMany: async (facts) => {
      const insertedRows = await dbClient
        .insertInto("facts")
        .values(facts)
        .returningAll()
        .execute();
      return insertedRows;
    },
    findAll: async () => {
      const rows = await dbClient.selectFrom("facts").selectAll().execute();
      return rows;
    },
    findById: async (id) => {
      const row = await dbClient
        .selectFrom("facts")
        .selectAll()
        .where("id", "=", id)
        .execute();
      return row[0];
    },
    update: async (id, data) => {
      await dbClient
        .updateTable("facts")
        .set(data)
        .where("id", "=", id)
        .execute();
    },
    delete: async (id) => {
      await dbClient.deleteFrom("facts").where("id", "=", id).execute();
    },
    findByConversationId: async (conversationId) => {
      const rows = await dbClient
        .selectFrom("facts")
        .innerJoin("messages", "messages.id", "facts.sourceMessageId")
        .selectAll("facts")
        .where("conversationId", "=", conversationId)
        .execute();
      return rows;
    },
  };

  const factTriggers: FactTriggerEntity = {
    construct: (factTrigger) => factTrigger,
    create: async (factTrigger) => {
      const insertedRows = await dbClient
        .insertInto("fact_triggers")
        .values(factTrigger)
        .returningAll()
        .execute();
      return insertedRows[0];
    },
    createMany: async (factTriggers) => {
      const insertedRows = await dbClient
        .insertInto("fact_triggers")
        .values(factTriggers)
        .returningAll()
        .execute();
      return insertedRows;
    },
    findAll: async () => {
      const rows = await dbClient
        .selectFrom("fact_triggers")
        .selectAll()
        .execute();
      return rows;
    },
    findById: async (id) => {
      const row = await dbClient
        .selectFrom("fact_triggers")
        .selectAll()
        .where("id", "=", id)
        .execute();
      return row[0];
    },
    update: async (id, data) => {
      await dbClient
        .updateTable("fact_triggers")
        .set(data)
        .where("id", "=", id)
        .execute();
    },
    delete: async (id) => {
      await dbClient.deleteFrom("fact_triggers").where("id", "=", id).execute();
    },
    findByFactId: async (factId) => {
      const rows = await dbClient
        .selectFrom("fact_triggers")
        .innerJoin("facts", "facts.id", "fact_triggers.sourceFactId")
        .selectAll("fact_triggers")
        .where("sourceFactId", "=", factId)
        .execute();
      return rows;
    },
    findByConversationId: async (conversationId) => {
      const rows = await dbClient
        .selectFrom("fact_triggers")
        .innerJoin("facts", "facts.id", "fact_triggers.sourceFactId")
        .innerJoin("messages", "messages.id", "facts.sourceMessageId")
        .selectAll("fact_triggers")
        .where("messages.conversationId", "=", conversationId)
        .execute();
      return rows;
    },
  };

  const messages: MessageEntity = {
    construct: (message) => message,
    create: async (message) => {
      const insertedRows = await dbClient
        .insertInto("messages")
        .values({ ...message, parts: JSON.stringify(message.parts) })
        .returningAll()
        .execute();
      return insertedRows[0] as CommittedMessage;
    },
    createMany: async (messages) => {
      const insertedRows = await dbClient
        .insertInto("messages")
        .values(
          messages.map((message) => ({
            ...message,
            parts: JSON.stringify(message.parts),
          }))
        )
        .returningAll()
        .execute();
      return insertedRows as Array<CommittedMessage>;
    },
    findAll: async () => {
      const rows = await dbClient.selectFrom("messages").selectAll().execute();
      return rows as Array<CommittedMessage>;
    },
    findById: async (id) => {
      const row = await dbClient
        .selectFrom("messages")
        .selectAll()
        .where("id", "=", id)
        .execute();
      return row[0] as CommittedMessage;
    },
    update: async (id, data) => {
      await dbClient
        .updateTable("messages")
        .set(data)
        .where("id", "=", id)
        .execute();
    },
    delete: async (id) => {
      await dbClient.deleteFrom("messages").where("id", "=", id).execute();
    },
    findByConversationId: async (conversationId) => {
      const rows = (await dbClient
        .selectFrom("messages")
        .selectAll()
        .where("conversationId", "=", conversationId)
        .execute()) as Array<CommittedMessage>;
      return rows;
    },
  };

  const users: UserEntity = {
    construct: (user) => user,
    create: async (user) => {
      const insertedRows = await dbClient
        .insertInto("users")
        .values(user)
        .returningAll()
        .execute();
      return insertedRows[0];
    },
    createMany: async (users) => {
      const insertedRows = await dbClient
        .insertInto("users")
        .values(users)
        .returningAll()
        .execute();
      return insertedRows;
    },
    findAll: async () => {
      const rows = await dbClient.selectFrom("users").selectAll().execute();
      return rows;
    },
    findById: async (id) => {
      const row = await dbClient
        .selectFrom("users")
        .selectAll()
        .where("id", "=", id)
        .execute();
      return row[0];
    },
    update: async (id, data) => {
      await dbClient
        .updateTable("users")
        .set(data)
        .where("id", "=", id)
        .execute();
    },
    delete: async (id) => {
      await dbClient.deleteFrom("users").where("id", "=", id).execute();
    },
    findByEmailAddress: async (emailAddress) => {
      const row = await dbClient
        .selectFrom("users")
        .selectAll()
        .where("email", "=", emailAddress)
        .executeTakeFirst();
      return row;
    },
  };

  const db = {
    conversations,
    facts,
    factTriggers,
    messages,
    users,
  };

  return db;
}