separate migrations for pgvecto.rs

move pgvecto.rs migrations to separate folder
This commit is contained in:
mertalev 2023-11-17 16:28:39 -05:00
parent 71388eaec2
commit 8a7b00bdab
No known key found for this signature in database
GPG key ID: 9181CD92C0A1C5E3
5 changed files with 87 additions and 48 deletions

View file

@ -23,5 +23,30 @@ export const databaseConfig: PostgresConnectionOptions = {
...urlOrParts,
};
export const databaseConfigVector: PostgresConnectionOptions = {
...databaseConfig,
migrationsRun: false,
migrations: [__dirname + '/migrations/vector/*.{js,ts}'],
};
export async function initDataSource(): Promise<DataSource> {
const dataSource = await new DataSource(databaseConfig).initialize();
const hasVectorExtension = (await dataSource.query(
`SELECT * FROM pg_extension WHERE name = 'vectors'`,
)).length > 0;
if (hasVectorExtension) {
const dataSourceVector = await new DataSource(databaseConfigVector).initialize();
await dataSourceVector.runMigrations();
await dataSourceVector.destroy();
}
return dataSource;
}
// this export is used by TypeORM commands in package.json#scripts
export const dataSource = new DataSource(databaseConfig);
export let dataSource: DataSource;
(async () => {
dataSource = await initDataSource();
})();

View file

@ -1,47 +0,0 @@
import { MigrationInterface, QueryRunner } from 'typeorm';
export class UsePgVector1699746198141 implements MigrationInterface {
name = 'UsePgVector1699746198141';
public async up(queryRunner: QueryRunner): Promise<void> {
await queryRunner.query('DROP EXTENSION IF EXISTS vectors');
await queryRunner.query('CREATE EXTENSION vectors');
const faceDimQuery = await queryRunner.query(`
SELECT CARDINALITY(embedding) as dimsize
FROM asset_faces
LIMIT 1`);
const clipDimQuery = await queryRunner.query(`
SELECT CARDINALITY("clipEmbedding") as dimsize
FROM smart_info
LIMIT 1`);
const faceDimSize = faceDimQuery?.[0]?.['dimsize'] ?? 512;
const clipDimSize = clipDimQuery?.[0]?.['dimsize'] ?? 512;
await queryRunner.query(`ALTER TABLE asset_faces ALTER COLUMN embedding TYPE vector(${faceDimSize})`);
await queryRunner.query(`CREATE TABLE smart_search (
"assetId" uuid PRIMARY KEY NOT NULL REFERENCES assets(id) ON DELETE CASCADE,
embedding vector(${clipDimSize}) NOT NULL )`);
await queryRunner.query(`
INSERT INTO smart_search("assetId", embedding)
SELECT si."assetId", si."clipEmbedding"
FROM smart_info si
WHERE "clipEmbedding" IS NOT NULL
`);
await queryRunner.query(`ALTER TABLE smart_info DROP COLUMN "clipEmbedding"`);
}
public async down(queryRunner: QueryRunner): Promise<void> {
await queryRunner.query(`ALTER TABLE asset_faces ALTER COLUMN embedding TYPE real array`);
await queryRunner.query(`ALTER TABLE smart_info ADD COLUMN IF NOT EXISTS "clipEmbedding" TYPE real array`);
await queryRunner.query(`
INSERT INTO smart_info
("assetId", "clipEmbedding")
SELECT s."assetId", s.embedding
FROM smart_search s
ON CONFLICT (s."assetId") DO UPDATE SET "clipEmbedding" = s.embedding
`);
await queryRunner.query(`DROP TABLE IF EXISTS smart_search`);
}
}

View file

@ -0,0 +1,61 @@
import { MigrationInterface, QueryRunner } from 'typeorm';
export class UsePgVectors1699746198141 implements MigrationInterface {
name = 'UsePgVectors1699746198141';
public async up(queryRunner: QueryRunner): Promise<void> {
const faceDimQuery = await queryRunner.query(`
SELECT CARDINALITY(embedding) as dimsize
FROM asset_faces
LIMIT 1`);
const clipDimQuery = await queryRunner.query(`
SELECT CARDINALITY("clipEmbedding") as dimsize
FROM smart_info
LIMIT 1`);
const faceDimSize = faceDimQuery?.[0]?.['dimsize'] ?? 512;
const clipDimSize = clipDimQuery?.[0]?.['dimsize'] ?? 512;
await queryRunner.query('DROP EXTENSION IF EXISTS vectors');
await queryRunner.query('CREATE EXTENSION vectors');
await queryRunner.query(`
BEGIN;
ALTER TABLE asset_faces ALTER COLUMN embedding TYPE vector(${faceDimSize});
CREATE TABLE smart_search (
"assetId" uuid PRIMARY KEY NOT NULL REFERENCES assets(id) ON DELETE CASCADE,
embedding vector(${clipDimSize}) NOT NULL );
INSERT INTO smart_search("assetId", embedding)
SELECT si."assetId", si."clipEmbedding"
FROM smart_info si
WHERE "clipEmbedding" IS NOT NULL;
ALTER TABLE smart_info DROP COLUMN "clipEmbedding";
COMMIT;
`);
}
public async down(queryRunner: QueryRunner): Promise<void> {
await queryRunner.query(`
BEGIN;
ALTER TABLE asset_faces ALTER COLUMN embedding TYPE real array;
ALTER TABLE smart_info ADD COLUMN IF NOT EXISTS "clipEmbedding" TYPE real array;
INSERT INTO smart_info
("assetId", "clipEmbedding")
SELECT s."assetId", s.embedding
FROM smart_search s
ON CONFLICT (s."assetId") DO UPDATE SET "clipEmbedding" = s.embedding;
DROP TABLE IF EXISTS smart_search;
COMMIT;
`);
}
}