AppFlowy-Cloud/migrations/20250318120849_departition_af_collab.sql
2025-04-01 10:37:11 +02:00

104 lines
No EOL
4 KiB
PL/PgSQL

set statement_timeout to 3600000;
-- create new uniform collab table
create table af_collab_temp
(
oid uuid not null primary key,
workspace_id uuid not null references public.af_workspace on delete cascade,
owner_uid bigint not null,
partition_key integer not null,
len integer,
blob bytea not null,
deleted_at timestamp with time zone,
created_at timestamp with time zone default CURRENT_TIMESTAMP,
updated_at timestamp with time zone default CURRENT_TIMESTAMP not null,
indexed_at timestamp with time zone
);
-- create a table for collabs that have non-UUID object ids
create table af_collab_conflicts
(
oid text not null primary key,
workspace_id uuid not null references public.af_workspace on delete cascade,
owner_uid bigint not null,
partition_key integer not null,
len integer,
blob bytea not null,
deleted_at timestamp with time zone,
created_at timestamp with time zone default CURRENT_TIMESTAMP,
updated_at timestamp with time zone default CURRENT_TIMESTAMP not null,
indexed_at timestamp with time zone
);
-- move non-UUID object ids to the new table (execution time: 31 secs)
insert into af_collab_conflicts(oid, workspace_id, owner_uid, partition_key, len, blob, deleted_at, created_at, updated_at, indexed_at)
select oid, workspace_id, owner_uid, partition_key, len, blob, deleted_at, created_at, updated_at, indexed_at
from af_collab
where oid !~ E'^[[:xdigit:]]{8}-([[:xdigit:]]{4}-){3}[[:xdigit:]]{12}$';
-- copy data from all collab partitions to new collab table (execution time: 7 mins)
insert into af_collab_temp(oid, workspace_id, owner_uid, partition_key, len, blob, deleted_at, created_at, updated_at, indexed_at)
select oid::uuid as oid, workspace_id, owner_uid, partition_key, len, blob, deleted_at, created_at, updated_at, indexed_at
from af_collab
where oid ~ E'^[[:xdigit:]]{8}-([[:xdigit:]]{4}-){3}[[:xdigit:]]{12}$';
-- prune embeddings table
truncate table af_collab_embeddings;
alter table af_collab_embeddings
drop constraint af_collab_embeddings_oid_partition_key_fkey;
-- replace af_collab table
drop table af_collab;
alter table af_collab_temp rename to af_collab;
-- modify embeddings to make use of new uuid columns
alter table af_collab_embeddings
alter column oid type uuid using oid::uuid;
create index if not exists ix_af_collab_embeddings_oid
on af_collab_embeddings(oid);
-- add foreign key constraint to af_collab_embeddings
alter table af_collab_embeddings
add constraint fk_af_collab_embeddings_oid foreign key (oid)
references af_collab (oid) on delete cascade;
alter table af_collab_embeddings drop partition_key;
-- add trigger for af_collab.updated_at
create trigger set_updated_at
before insert or update
on af_collab
for each row
execute procedure update_updated_at_column();
-- add remaining indexes to new af_collab table (execution time: 25 sec + 25sec)
create index if not exists idx_workspace_id_on_af_collab
on af_collab (workspace_id);
create index if not exists idx_af_collab_updated_at
on af_collab (updated_at);
create or replace procedure af_collab_embeddings_upsert(IN p_workspace_id uuid, IN p_oid uuid, IN p_tokens_used integer, IN p_fragments af_fragment_v3[])
language plpgsql
as
$$
BEGIN
DELETE FROM af_collab_embeddings WHERE oid = p_oid;
INSERT INTO af_collab_embeddings (fragment_id, oid, content_type, content, embedding, indexed_at, metadata, fragment_index, embedder_type)
SELECT
f.fragment_id,
p_oid,
f.content_type,
f.contents,
f.embedding,
NOW(),
f.metadata,
f.fragment_index,
f.embedder_type
FROM UNNEST(p_fragments) as f;
-- Update the usage tracking table
INSERT INTO af_workspace_ai_usage(created_at, workspace_id, search_requests, search_tokens_consumed, index_tokens_consumed)
VALUES (now()::date, p_workspace_id, 0, 0, p_tokens_used)
ON CONFLICT (created_at, workspace_id)
DO UPDATE SET index_tokens_consumed = af_workspace_ai_usage.index_tokens_consumed + p_tokens_used;
END
$$;