feat: reader for postgres (#1813)

Co-authored-by: Marcus Schiesser <mail@marcusschiesser.de>

.changeset/beige-seals-live.md

@@ -0,0 +1,5 @@
+---
+"@llamaindex/postgres": patch
+---
+
+Added reader for the postgres package

examples/vector-store/pg/pg-reader.ts

@@ -0,0 +1,54 @@
+import { OpenAIEmbedding } from "@llamaindex/openai";
+import { PGVectorStore, SimplePostgresReader } from "@llamaindex/postgres";
+import { Document, Settings } from "llamaindex";
+
+const containerName = "llamaindex-postgres-example";
+const dbConfig = {
+  host: "localhost",
+  port: 5432,
+  database: "test",
+  user: "postgres",
+  password: "postgres",
+};
+Settings.embedModel = new OpenAIEmbedding();
+async function main() {
+  try {
+    // Initialize PGVectorStore and add a document
+    // This is just to generate test data in the DB for the SimplePostgresReader
+    // Usually it's purpose is to read documents from a Postgres DB
+    const vectorStore = new PGVectorStore({
+      clientConfig: dbConfig,
+      dimensions: 3,
+      tableName: "llamaindex_vector",
+    });
+
+    // Add a test document with embedding
+    await vectorStore.add([
+      new Document({
+        text: "This is a test document about AI technology.",
+        metadata: {
+          author: "Test Author",
+          category: "technology",
+        },
+        embedding: [1, 2, 3],
+      }),
+    ]);
+    console.log("Added document to vector store");
+    // Initialize PostgresReader to read the document
+    const reader = new SimplePostgresReader({ clientConfig: dbConfig });
+
+    // Read documents using PostgresReader
+    const documents = await reader.loadData({
+      query: "SELECT document as content, metadata FROM llamaindex_vector",
+      fields: ["content"],
+      metadataFields: ["metadata"],
+    });
+
+    console.log("Read documents:");
+    console.log(JSON.stringify(documents, null, 2));
+  } catch (error) {
+    console.error("Error:", error);
+  }
+}
+
+main().catch(console.error);

packages/providers/storage/postgres/package.json

@@ -35,7 +35,8 @@
   },
   "scripts": {
     "build": "bunchee",
-    "dev": "bunchee --watch"
+    "dev": "bunchee --watch",
+    "test": "vitest"
   },
   "devDependencies": {
     "@types/pg": "^8.11.8",

packages/providers/storage/postgres/src/SimplePostgresReader.ts

@@ -0,0 +1,147 @@
+import { type BaseReader, Document } from "@llamaindex/core/schema";
+import pg from "pg";
+
+export type PostgresReaderConfig = {
+  /**
+   * The query to execute. Must return at least one column.
+   */
+  query: string;
+  /**
+   * An array of field names to retrieve from each row. Defaults to first column.
+   */
+  fields?: string[];
+  /**
+   * The separator to join multiple field values. Defaults to an empty string.
+   */
+  fieldSeparator?: string;
+  /**
+   * An optional array of metadata field names to extract as metadata.
+   */
+  metadataFields?: string[];
+};
+
+export type PostgresReaderClientConfig =
+  | {
+      clientConfig?: pg.ClientConfig | undefined;
+    }
+  | {
+      shouldConnect?: boolean | undefined;
+      client?: pg.Client | pg.PoolClient;
+    };
+
+export class SimplePostgresReader implements BaseReader<Document> {
+  private isDBConnected: boolean = false;
+  private clientConfig: pg.ClientConfig | undefined = undefined;
+  private db: pg.ClientBase | undefined = undefined;
+
+  constructor(config: PostgresReaderClientConfig) {
+    if (this.hasClientConfig(config)) {
+      this.clientConfig = config.clientConfig;
+    } else if (this.hasClient(config)) {
+      this.isDBConnected =
+        config?.shouldConnect !== undefined ? !config.shouldConnect : false;
+      this.db = config.client;
+    }
+  }
+
+  private hasClientConfig(
+    config: PostgresReaderClientConfig,
+  ): config is { clientConfig: pg.ClientConfig } {
+    return "clientConfig" in config;
+  }
+
+  private hasClient(
+    config: PostgresReaderClientConfig,
+  ): config is { client: pg.Client | pg.PoolClient; shouldConnect?: boolean } {
+    return "client" in config;
+  }
+
+  private async importPg() {
+    const pg = await import("pg");
+    return pg;
+  }
+
+  private async initializeConnection(): Promise<pg.ClientBase> {
+    const pg = await this.importPg();
+    const { Client } = pg.default ? pg.default : pg;
+    const client = new Client({ ...this.clientConfig });
+    await client.connect();
+    return client;
+  }
+
+  private checkDBIsNotConnected(): boolean {
+    return this.db !== null && this.isDBConnected === false;
+  }
+
+  private async getDb(): Promise<pg.ClientBase> {
+    if (!this.db) {
+      this.db = await this.initializeConnection();
+      this.isDBConnected = true;
+    }
+    if (this.checkDBIsNotConnected()) {
+      await this.db.connect();
+      this.isDBConnected = true;
+    }
+    this.db.on("end", () => {
+      this.isDBConnected = false;
+    });
+
+    return this.db;
+  }
+
+  async loadData(config: PostgresReaderConfig): Promise<Document[]> {
+    const db = await this.getDb();
+    try {
+      const result = await db.query(config.query);
+      return this.processResult(result, config);
+    } catch (err) {
+      throw new Error(`Error executing query: ${(err as Error).message}`);
+    }
+  }
+
+  private processResult(
+    result: pg.QueryResult,
+    config: PostgresReaderConfig,
+  ): Document[] {
+    const documents: Document[] = [];
+    for (const row of result.rows) {
+      const text = this.extractText(row, config);
+      const metadata = this.extractMetadata(row, config);
+
+      documents.push(new Document({ text, metadata }));
+    }
+    return documents;
+  }
+
+  private extractFieldsFromRow(
+    row: Record<string, unknown>,
+    config: PostgresReaderConfig,
+  ): (string | undefined)[] {
+    const fields = config.fields || [Object.keys(row)[0]];
+    return fields;
+  }
+
+  private extractText(
+    row: Record<string, unknown>,
+    config: PostgresReaderConfig,
+  ): string {
+    const fields = this.extractFieldsFromRow(row, config);
+    const texts = fields
+      .filter((f): f is string => typeof f === "string")
+      .map((name) => String(row[name] ?? ""));
+    return texts.join(config.fieldSeparator || "");
+  }
+
+  private extractMetadata(
+    row: Record<string, unknown>,
+    config: PostgresReaderConfig,
+  ): Record<string, unknown> {
+    let metadata = {};
+    if (config.metadataFields) {
+      metadata = Object.fromEntries(
+        config.metadataFields.map((name) => [name, row[name]]),
+      );
+    }
+    return metadata;
+  }
+}

packages/providers/storage/postgres/src/index.ts

@@ -2,3 +2,4 @@ export * from "./PGVectorStore";
 export * from "./PostgresDocumentStore";
 export * from "./PostgresIndexStore";
 export * from "./PostgresKVStore";
+export * from "./SimplePostgresReader";

packages/providers/storage/postgres/tests/SimplePostgresReader.test.ts

@@ -0,0 +1,71 @@
+import { Client } from "pg";
+import { describe, expect, it, vi } from "vitest";
+import { SimplePostgresReader } from "../src/SimplePostgresReader";
+
+vi.mock("pg", () => {
+  return {
+    Client: vi.fn().mockImplementation(() => ({
+      connect: vi.fn(),
+      query: vi.fn().mockResolvedValue({
+        rows: [
+          {
+            title: "Test Document",
+            content: "This is a test document content.",
+            author: "Test Author",
+            created_at: new Date("2024-01-01"),
+          },
+        ],
+      }),
+      end: vi.fn(),
+      on: vi.fn(),
+    })),
+  };
+});
+
+describe("SimplePostgresReader", () => {
+  it("should read documents with fields and metadata", async () => {
+    const client = new Client();
+    const reader = new SimplePostgresReader({ client });
+
+    const documents = await reader.loadData({
+      query: "SELECT * FROM test_documents",
+      fields: ["title", "content"],
+      fieldSeparator: "\n",
+      metadataFields: ["author", "created_at"],
+    });
+
+    expect(documents).toHaveLength(1);
+    expect(documents[0].text).toBe(
+      "Test Document\nThis is a test document content.",
+    );
+    expect(documents[0].metadata).toEqual({
+      author: "Test Author",
+      created_at: new Date("2024-01-01"),
+    });
+  });
+
+  it("should read the column as specified in fields", async () => {
+    const client = new Client();
+    const reader = new SimplePostgresReader({ client });
+
+    const documents = await reader.loadData({
+      query: "SELECT content FROM test_documents",
+      fields: ["content"],
+    });
+
+    expect(documents).toHaveLength(1);
+    expect(documents[0].text).toBe("This is a test document content.");
+  });
+
+  it("should default to first column when no fields specified", async () => {
+    const client = new Client();
+    const reader = new SimplePostgresReader({ client });
+
+    const documents = await reader.loadData({
+      query: "SELECT content FROM test_documents",
+    });
+
+    expect(documents).toHaveLength(1);
+    expect(documents[0].text).toBe("Test Document");
+  });
+});