diff --git a/rag/engine/postgres.go b/rag/engine/postgres.go index 43d6901..6c7eb61 100644 --- a/rag/engine/postgres.go +++ b/rag/engine/postgres.go @@ -97,10 +97,21 @@ func NewPostgresDBCollection(collectionName, databaseURL string, openaiClient *o } func sanitizeTableName(name string) string { - // Replace invalid characters with underscores - name = strings.ReplaceAll(name, "-", "_") - name = strings.ReplaceAll(name, ".", "_") - name = strings.ReplaceAll(name, " ", "_") + // Replace every character that is not a valid PostgreSQL identifier + // character with an underscore. Allowlisting (rather than stripping a + // hardcoded set such as '-', '.', ' ') guarantees the result is a legal + // unquoted identifier whatever the collection name contains: the ':' + // namespace separator used for per-user collections (e.g. the synthetic + // "legacy-api-key:" name) previously slipped through and produced + // "syntax error at or near ':'" on CREATE TABLE. + name = strings.Map(func(r rune) rune { + switch { + case r >= 'a' && r <= 'z', r >= 'A' && r <= 'Z', r >= '0' && r <= '9', r == '_': + return r + default: + return '_' + } + }, name) // Ensure it starts with a letter if len(name) > 0 && (name[0] < 'a' || name[0] > 'z') && (name[0] < 'A' || name[0] > 'Z') { name = "col_" + name diff --git a/rag/engine/sanitize_test.go b/rag/engine/sanitize_test.go new file mode 100644 index 0000000..77d0133 --- /dev/null +++ b/rag/engine/sanitize_test.go @@ -0,0 +1,42 @@ +package engine + +import ( + "regexp" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// validPostgresIdentifier matches an unquoted PostgreSQL identifier: it must +// start with a letter or underscore and contain only letters, digits and +// underscores. Anything else triggers a SQL syntax error when interpolated +// into DDL such as CREATE TABLE. +var validPostgresIdentifier = regexp.MustCompile(`^[A-Za-z_][A-Za-z0-9_]*$`) + +var _ = Describe("sanitizeTableName", func() { + It("strips the ':' namespace separator used for per-user collections", func() { + // Regression for LocalAI #10375: agents created under a legacy API key + // get the collection name "legacy-api-key:", and the colon broke + // PostgreSQL table creation (ERROR: syntax error at or near ":"). + got := sanitizeTableName("legacy-api-key:LiteraryResearch") + Expect(got).NotTo(ContainSubstring(":")) + Expect(got).To(MatchRegexp(validPostgresIdentifier.String())) + }) + + It("produces a valid identifier for any input", func() { + for _, in := range []string{ + "legacy-api-key:LiteraryResearch", + "user@example.com/notes", + "a b.c-d:e", + "123starts-with-digit", + "with$pecial%chars!", + } { + Expect(sanitizeTableName(in)).To(MatchRegexp(validPostgresIdentifier.String()), + "input %q should sanitize to a valid PostgreSQL identifier", in) + } + }) + + It("still maps the previously-handled characters to underscores", func() { + Expect(sanitizeTableName("my-collection.name here")).To(Equal("documents_my_collection_name_here")) + }) +})