fix(db): replace oversized upload unique index with tuple hash

Author: Mmx233

internal/cache/service.go

@@ -3,8 +3,10 @@ package cache
 import (
 	"context"
 	"crypto/rand"
+	"crypto/sha256"
 	"encoding/base64"
 	"encoding/binary"
+	"encoding/hex"
 	"errors"
 	"fmt"
 	"io"
@@ -440,6 +442,7 @@ func ReadScopesByPermission(scope auth.CacheScope) []string {
 }
 
 func (s *Service) createUploadRecord(ctx context.Context, key, version, scope, repoID string) (int64, error) {
+	tupleHash := uploadTupleHash(key, version, scope, repoID)
 	for i := 0; i < 5; i++ {
 		id, err := randomPositiveInt64()
 		if err != nil {
@@ -455,6 +458,7 @@ func (s *Service) createUploadRecord(ctx context.Context, key, version, scope, r
 			SetVersion(version).
 			SetScope(scope).
 			SetRepoId(repoID).
+			SetTupleHash(tupleHash).
 			Save(ctx)
 		if err == nil {
 			return id, nil
@@ -544,6 +548,17 @@ func uploadTuple(key, version, scope, repoID string) []entpredicate.Upload {
 	}
 }
 
+func uploadTupleHash(key, version, scope, repoID string) string {
+	hash := sha256.New()
+	var encodedLength [8]byte
+	for _, value := range [...]string{key, version, scope, repoID} {
+		binary.BigEndian.PutUint64(encodedLength[:], uint64(len(value)))
+		_, _ = hash.Write(encodedLength[:])
+		_, _ = hash.Write([]byte(value))
+	}
+	return hex.EncodeToString(hash.Sum(nil))
+}
+
 func (s *Service) isUploadAbandoned(currentUpload *ent.Upload) bool {
 	lastActivity := currentUpload.CreatedAt
 	if currentUpload.LastPartUploadedAt != nil {

internal/cache/service_test.go

@@ -3,6 +3,7 @@ package cache
 import (
 	"bytes"
 	"context"
+	"crypto/sha256"
 	"errors"
 	"fmt"
 	"io"
@@ -303,6 +304,18 @@ func TestCreateUploadReservationIsAtomic(t *testing.T) {
 	count, err := client.Upload.Query().Count(ctx)
 	require.NoError(t, err)
 	require.Equal(t, 1, count)
+	currentUpload := client.Upload.Query().OnlyX(ctx)
+	require.NotNil(t, currentUpload.TupleHash)
+	require.Equal(t, uploadTupleHash("key", "version", scope.Scopes[0].Scope, scope.RepoID), *currentUpload.TupleHash)
+}
+
+func TestUploadTupleHashPreservesFieldBoundaries(t *testing.T) {
+	first := uploadTupleHash("ab", "c", "scope", "repo")
+	second := uploadTupleHash("a", "bc", "scope", "repo")
+
+	require.Len(t, first, sha256.Size*2)
+	require.Equal(t, first, uploadTupleHash("ab", "c", "scope", "repo"))
+	require.NotEqual(t, first, second)
 }
 
 func TestSinglePartUsesPartsObjectForDirectDownload(t *testing.T) {

internal/db/db_test.go

@@ -103,6 +103,12 @@ func TestOpenAndMigrateSQLiteFromOriginalSchema(t *testing.T) {
 		_, err := sqlDB.ExecContext(ctx, stmt)
 		require.NoError(t, err)
 	}
+	_, err = sqlDB.ExecContext(ctx, `insert into uploads (
+		id, key, version, createdAt, folderName, scope, repoId
+	) values
+		(1, 'legacy-key', 'version', 1, 'legacy-1', 'scope', 'repo'),
+		(2, 'legacy-key', 'version', 2, 'legacy-2', 'scope', 'repo')`)
+	require.NoError(t, err)
 	require.NoError(t, sqlDB.Close())
 
 	client, err := OpenAndMigrate(ctx, config.DBConfig{
@@ -126,6 +132,10 @@ func TestOpenAndMigrateSQLiteFromOriginalSchema(t *testing.T) {
 	require.True(t, sqliteColumnExists(ctx, t, sqlDB, "storage_locations", "mergeLeaseToken"))
 	require.True(t, sqliteColumnExists(ctx, t, sqlDB, "uploads", "lastPartUploadedAt"))
 	require.True(t, sqliteColumnExists(ctx, t, sqlDB, "uploads", "committedPartCount"))
+	require.True(t, sqliteColumnExists(ctx, t, sqlDB, "uploads", "tupleHash"))
+	var legacyUploads int
+	require.NoError(t, sqlDB.QueryRowContext(ctx, `select count(*) from uploads where tupleHash is null`).Scan(&legacyUploads))
+	require.Equal(t, 2, legacyUploads)
 	var storageDeletionsTable string
 	require.NoError(t, sqlDB.QueryRowContext(
 		ctx,
@@ -144,6 +154,7 @@ func TestOpenAndMigrateSQLiteFromOriginalSchema(t *testing.T) {
 	require.False(t, sqliteColumnExists(ctx, t, sqlDB, "uploads", "last_part_uploaded_at"))
 	require.True(t, sqliteIndexExists(ctx, t, sqlDB, "idx_cache_entries_repo_scope_version_key"))
 	require.True(t, sqliteIndexExists(ctx, t, sqlDB, "idx_cache_entries_location_updated_at"))
+	require.True(t, sqliteIndexExists(ctx, t, sqlDB, "idx_uploads_tuple_hash"))
 }
 
 func TestUploadIDDoesNotRequireDatabaseIdentity(t *testing.T) {
@@ -168,7 +179,7 @@ func TestGeneratedSchemaMatchesOriginalColumns(t *testing.T) {
 		"id", "scope", "expiresAt", "storageLocationId",
 	}, columnNames(migrate.StorageReaderLeasesColumns))
 	require.Equal(t, []string{
-		"id", "key", "version", "scope", "repoId", "createdAt", "lastPartUploadedAt", "startedPartUploadCount", "finishedPartUploadCount", "folderName", "committedPartCount",
+		"id", "key", "version", "scope", "repoId", "createdAt", "lastPartUploadedAt", "startedPartUploadCount", "finishedPartUploadCount", "folderName", "committedPartCount", "tupleHash",
 	}, columnNames(migrate.UploadsColumns))
 }
 
@@ -202,7 +213,15 @@ func TestGeneratedSchemaMatchesOriginalIndexNames(t *testing.T) {
 		"idx_uploads_key_version",
 		"idx_uploads_scope",
 		"idx_uploads_repoId",
-	}, indexNames(migrate.UploadsTable.Indexes[:3]))
+		"idx_uploads_tuple_hash",
+	}, indexNames(migrate.UploadsTable.Indexes))
+}
+
+func TestMySQLUploadUniqueIndexUsesFixedLengthTupleHash(t *testing.T) {
+	uniqueIndex := migrate.UploadsTable.Indexes[3]
+	require.True(t, uniqueIndex.Unique)
+	require.Equal(t, []string{"tupleHash"}, columnNames(uniqueIndex.Columns))
+	require.Equal(t, "varchar(64)", migrate.UploadsColumns[11].SchemaType[dialect.MySQL])
 }
 
 func TestCacheMatchIndexAnnotations(t *testing.T) {

internal/ent/migrate/schema.go

@@ -45,6 +45,12 @@ func (Upload) Fields() []ent.Field {
 			Nillable().
 			NonNegative().
 			StorageKey("committedPartCount"),
+		field.String("tupleHash").
+			MaxLen(64).
+			Optional().
+			Nillable().
+			StorageKey("tupleHash").
+			SchemaType(originalBoundedStringColumnType(64)),
 	}
 }
 
@@ -53,6 +59,6 @@ func (Upload) Indexes() []ent.Index {
 		index.Fields("key", "version").StorageKey("idx_uploads_key_version"),
 		index.Fields("scope").StorageKey("idx_uploads_scope"),
 		index.Fields("repoId").StorageKey("idx_uploads_repoId"),
-		index.Fields("key", "version", "scope", "repoId").Unique(),
+		index.Fields("tupleHash").Unique().StorageKey("idx_uploads_tuple_hash"),
 	}
 }