feat: add per-user tenant scoping
This commit is contained in:
+26
-15
@@ -31,10 +31,10 @@ func (db *DB) InsertThought(ctx context.Context, thought thoughttypes.Thought, e
|
||||
}()
|
||||
|
||||
row := tx.QueryRow(ctx, `
|
||||
insert into thoughts (content, metadata, project_id)
|
||||
values ($1, $2::jsonb, $3)
|
||||
insert into thoughts (content, metadata, project_id, tenant_key)
|
||||
values ($1, $2::jsonb, $3, $4)
|
||||
returning id, guid, created_at, updated_at
|
||||
`, thought.Content, metadata, thought.ProjectID)
|
||||
`, thought.Content, metadata, thought.ProjectID, tenantKeyPtr(ctx))
|
||||
|
||||
created := thought
|
||||
created.Embedding = nil
|
||||
@@ -107,6 +107,7 @@ func (db *DB) ListThoughts(ctx context.Context, filter thoughttypes.ListFilter)
|
||||
args := make([]any, 0, 6)
|
||||
conditions := []string{}
|
||||
|
||||
addTenantCondition(ctx, &args, &conditions, "tenant_key")
|
||||
if !filter.IncludeArchived {
|
||||
conditions = append(conditions, "archived_at is null")
|
||||
}
|
||||
@@ -170,11 +171,14 @@ func (db *DB) ListThoughts(ctx context.Context, filter thoughttypes.ListFilter)
|
||||
|
||||
func (db *DB) Stats(ctx context.Context) (thoughttypes.ThoughtStats, error) {
|
||||
var total int
|
||||
if err := db.pool.QueryRow(ctx, `select count(*) from thoughts where archived_at is null`).Scan(&total); err != nil {
|
||||
statsArgs := []any{}
|
||||
statsConditions := []string{"archived_at is null"}
|
||||
addTenantCondition(ctx, &statsArgs, &statsConditions, "tenant_key")
|
||||
if err := db.pool.QueryRow(ctx, `select count(*) from thoughts where `+strings.Join(statsConditions, " and "), statsArgs...).Scan(&total); err != nil {
|
||||
return thoughttypes.ThoughtStats{}, fmt.Errorf("count thoughts: %w", err)
|
||||
}
|
||||
|
||||
rows, err := db.pool.Query(ctx, `select metadata from thoughts where archived_at is null`)
|
||||
rows, err := db.pool.Query(ctx, `select metadata from thoughts where `+strings.Join(statsConditions, " and "), statsArgs...)
|
||||
if err != nil {
|
||||
return thoughttypes.ThoughtStats{}, fmt.Errorf("query stats metadata: %w", err)
|
||||
}
|
||||
@@ -217,11 +221,11 @@ func (db *DB) Stats(ctx context.Context) (thoughttypes.ThoughtStats, error) {
|
||||
}
|
||||
|
||||
func (db *DB) GetThought(ctx context.Context, id uuid.UUID) (thoughttypes.Thought, error) {
|
||||
args := []any{id}
|
||||
row := db.pool.QueryRow(ctx, `
|
||||
select id, guid, content, metadata, project_id, archived_at, created_at, updated_at
|
||||
from thoughts
|
||||
where guid = $1
|
||||
`, id)
|
||||
where guid = $1`+tenantSQL(ctx, &args, "tenant_key"), args...)
|
||||
|
||||
var model generatedmodels.ModelPublicThoughts
|
||||
if err := row.Scan(&model.ID, &model.GUID, &model.Content, &model.Metadata, &model.ProjectID, &model.ArchivedAt, &model.CreatedAt, &model.UpdatedAt); err != nil {
|
||||
@@ -240,11 +244,11 @@ func (db *DB) GetThought(ctx context.Context, id uuid.UUID) (thoughttypes.Though
|
||||
}
|
||||
|
||||
func (db *DB) GetThoughtByID(ctx context.Context, id int64) (thoughttypes.Thought, error) {
|
||||
args := []any{id}
|
||||
row := db.pool.QueryRow(ctx, `
|
||||
select id, guid, content, metadata, project_id, archived_at, created_at, updated_at
|
||||
from thoughts
|
||||
where id = $1
|
||||
`, id)
|
||||
where id = $1`+tenantSQL(ctx, &args, "tenant_key"), args...)
|
||||
|
||||
var model generatedmodels.ModelPublicThoughts
|
||||
if err := row.Scan(&model.ID, &model.GUID, &model.Content, &model.Metadata, &model.ProjectID, &model.ArchivedAt, &model.CreatedAt, &model.UpdatedAt); err != nil {
|
||||
@@ -276,14 +280,14 @@ func (db *DB) UpdateThought(ctx context.Context, id uuid.UUID, content string, e
|
||||
_ = tx.Rollback(ctx)
|
||||
}()
|
||||
|
||||
args := []any{id, content, metadataBytes, projectID}
|
||||
tag, err := tx.Exec(ctx, `
|
||||
update thoughts
|
||||
set content = $2,
|
||||
metadata = $3::jsonb,
|
||||
project_id = $4,
|
||||
updated_at = now()
|
||||
where guid = $1
|
||||
`, id, content, metadataBytes, projectID)
|
||||
where guid = $1`+tenantSQL(ctx, &args, "tenant_key"), args...)
|
||||
if err != nil {
|
||||
return thoughttypes.Thought{}, fmt.Errorf("update thought: %w", err)
|
||||
}
|
||||
@@ -317,12 +321,12 @@ func (db *DB) UpdateThoughtMetadata(ctx context.Context, id int64, metadata thou
|
||||
return thoughttypes.Thought{}, fmt.Errorf("marshal updated metadata: %w", err)
|
||||
}
|
||||
|
||||
args := []any{id, metadataBytes}
|
||||
tag, err := db.pool.Exec(ctx, `
|
||||
update thoughts
|
||||
set metadata = $2::jsonb,
|
||||
updated_at = now()
|
||||
where id = $1
|
||||
`, id, metadataBytes)
|
||||
where id = $1`+tenantSQL(ctx, &args, "tenant_key"), args...)
|
||||
if err != nil {
|
||||
return thoughttypes.Thought{}, fmt.Errorf("update thought metadata: %w", err)
|
||||
}
|
||||
@@ -334,7 +338,8 @@ func (db *DB) UpdateThoughtMetadata(ctx context.Context, id int64, metadata thou
|
||||
}
|
||||
|
||||
func (db *DB) DeleteThought(ctx context.Context, id uuid.UUID) error {
|
||||
tag, err := db.pool.Exec(ctx, `delete from thoughts where guid = $1`, id)
|
||||
args := []any{id}
|
||||
tag, err := db.pool.Exec(ctx, `delete from thoughts where guid = $1`+tenantSQL(ctx, &args, "tenant_key"), args...)
|
||||
if err != nil {
|
||||
return fmt.Errorf("delete thought: %w", err)
|
||||
}
|
||||
@@ -345,7 +350,8 @@ func (db *DB) DeleteThought(ctx context.Context, id uuid.UUID) error {
|
||||
}
|
||||
|
||||
func (db *DB) ArchiveThought(ctx context.Context, id uuid.UUID) error {
|
||||
tag, err := db.pool.Exec(ctx, `update thoughts set archived_at = now(), updated_at = now() where guid = $1`, id)
|
||||
args := []any{id}
|
||||
tag, err := db.pool.Exec(ctx, `update thoughts set archived_at = now(), updated_at = now() where guid = $1`+tenantSQL(ctx, &args, "tenant_key"), args...)
|
||||
if err != nil {
|
||||
return fmt.Errorf("archive thought: %w", err)
|
||||
}
|
||||
@@ -424,6 +430,7 @@ func (db *DB) SearchSimilarThoughts(ctx context.Context, embedding []float32, em
|
||||
"1 - (e.embedding <=> $1) > $2",
|
||||
"e.model = $3",
|
||||
}
|
||||
addTenantCondition(ctx, &args, &conditions, "t.tenant_key")
|
||||
if projectID != nil {
|
||||
args = append(args, *projectID)
|
||||
conditions = append(conditions, fmt.Sprintf("t.project_id = $%d", len(args)))
|
||||
@@ -472,6 +479,7 @@ func (db *DB) HasEmbeddingsForModel(ctx context.Context, model string, projectID
|
||||
"e.model = $1",
|
||||
"t.archived_at is null",
|
||||
}
|
||||
addTenantCondition(ctx, &args, &conditions, "t.tenant_key")
|
||||
if projectID != nil {
|
||||
args = append(args, *projectID)
|
||||
conditions = append(conditions, fmt.Sprintf("t.project_id = $%d", len(args)))
|
||||
@@ -490,6 +498,7 @@ func (db *DB) HasEmbeddingsForModel(ctx context.Context, model string, projectID
|
||||
func (db *DB) ListThoughtsMissingEmbedding(ctx context.Context, model string, limit int, projectID *int64, includeArchived bool, olderThanDays int) ([]thoughttypes.Thought, error) {
|
||||
args := []any{model}
|
||||
conditions := []string{"e.id is null"}
|
||||
addTenantCondition(ctx, &args, &conditions, "t.tenant_key")
|
||||
|
||||
if !includeArchived {
|
||||
conditions = append(conditions, "t.archived_at is null")
|
||||
@@ -539,6 +548,7 @@ func (db *DB) ListThoughtsMissingEmbedding(ctx context.Context, model string, li
|
||||
func (db *DB) ListThoughtsForMetadataReparse(ctx context.Context, limit int, projectID *int64, includeArchived bool, olderThanDays int) ([]thoughttypes.Thought, error) {
|
||||
args := make([]any, 0, 3)
|
||||
conditions := make([]string, 0, 4)
|
||||
addTenantCondition(ctx, &args, &conditions, "tenant_key")
|
||||
|
||||
if !includeArchived {
|
||||
conditions = append(conditions, "archived_at is null")
|
||||
@@ -608,6 +618,7 @@ func (db *DB) SearchThoughtsText(ctx context.Context, query string, limit int, p
|
||||
"t.archived_at is null",
|
||||
"(to_tsvector('simple', t.content) || to_tsvector('simple', coalesce(p.name, ''))) @@ websearch_to_tsquery('simple', $1)",
|
||||
}
|
||||
addTenantCondition(ctx, &args, &conditions, "t.tenant_key")
|
||||
if projectID != nil {
|
||||
args = append(args, *projectID)
|
||||
conditions = append(conditions, fmt.Sprintf("t.project_id = $%d", len(args)))
|
||||
|
||||
Reference in New Issue
Block a user