From 369b8af109a14b774e253c7d0450f061b885ef52 Mon Sep 17 00:00:00 2001
From: Lincoln Nogueira <lincolnthalles@users.noreply.github.com>
Date: Tue, 2 Jan 2024 21:31:59 -0300
Subject: [PATCH] chore: improve resource internal_path migrator (#2698)

* chore: improve internal path migrator
- handle mixed path styles
- handle Windows paths
- add tests

* chore: fix goimports error
---
 bin/memos/main.go           | 13 +++++----
 store/migrator.go           | 48 +++++++++++++++++++------------
 test/store/migrator_test.go | 56 +++++++++++++++++++++++++++++++++++++
 3 files changed, 94 insertions(+), 23 deletions(-)
 create mode 100644 test/store/migrator_test.go

diff --git a/bin/memos/main.go b/bin/memos/main.go
index 4c67e76b..5ffbfabf 100644
--- a/bin/memos/main.go
+++ b/bin/memos/main.go
@@ -59,11 +59,14 @@ var (
 			}
 
 			store := store.New(dbDriver, profile)
-			if err := store.MigrateResourceInternalPath(ctx); err != nil {
-				cancel()
-				log.Error("failed to migrate resource internal path", zap.Error(err))
-				return
-			}
+
+			go func() {
+				if err := store.MigrateResourceInternalPath(ctx); err != nil {
+					cancel()
+					log.Error("failed to migrate resource internal path", zap.Error(err))
+					return
+				}
+			}()
 
 			s, err := server.NewServer(ctx, profile, store)
 			if err != nil {
diff --git a/store/migrator.go b/store/migrator.go
index 9caa4615..9b9995af 100644
--- a/store/migrator.go
+++ b/store/migrator.go
@@ -2,10 +2,15 @@ package store
 
 import (
 	"context"
-	"path/filepath"
+	"fmt"
+	"os"
+	"time"
+
 	"strings"
 
 	"github.com/pkg/errors"
+
+	"github.com/usememos/memos/internal/log"
 )
 
 // MigrateResourceInternalPath migrates resource internal path from absolute path to relative path.
@@ -15,30 +20,37 @@ func (s *Store) MigrateResourceInternalPath(ctx context.Context) error {
 		return errors.Wrap(err, "failed to list resources")
 	}
 
+	dataPath := strings.ReplaceAll(s.Profile.Data, `\`, "/")
+	migrateStartTime := time.Now()
+	migratedCount := 0
 	for _, resource := range resources {
 		if resource.InternalPath == "" {
 			continue
 		}
 
-		internalPath := resource.InternalPath
-		if filepath.IsAbs(internalPath) {
-			if !strings.HasPrefix(internalPath, s.Profile.Data) {
-				// Invalid internal path, skip.
-				continue
-			}
-			internalPath = strings.TrimPrefix(internalPath, s.Profile.Data)
-			for strings.HasPrefix(internalPath, "/") {
-				internalPath = strings.TrimPrefix(internalPath, "/")
-			}
-			_, err := s.UpdateResource(ctx, &UpdateResource{
-				ID:           resource.ID,
-				InternalPath: &internalPath,
-			})
-			if err != nil {
-				return errors.Wrap(err, "failed to update resource")
-			}
+		internalPath := strings.ReplaceAll(resource.InternalPath, `\`, "/")
+		if !strings.HasPrefix(internalPath, dataPath) {
+			continue
 		}
+
+		internalPath = strings.TrimPrefix(internalPath, dataPath)
+
+		for os.IsPathSeparator(internalPath[0]) {
+			internalPath = internalPath[1:]
+		}
+
+		_, err := s.UpdateResource(ctx, &UpdateResource{
+			ID:           resource.ID,
+			InternalPath: &internalPath,
+		})
+		if err != nil {
+			return errors.Wrap(err, "failed to update local resource path")
+		}
+		migratedCount++
 	}
 
+	if migratedCount > 0 {
+		log.Info(fmt.Sprintf("migrated %d local resource paths in %s", migratedCount, time.Since(migrateStartTime)))
+	}
 	return nil
 }
diff --git a/test/store/migrator_test.go b/test/store/migrator_test.go
new file mode 100644
index 00000000..a971f28f
--- /dev/null
+++ b/test/store/migrator_test.go
@@ -0,0 +1,56 @@
+package teststore
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/usememos/memos/store"
+)
+
+func TestMigrateResourceInternalPath(t *testing.T) {
+	ctx := context.Background()
+	ts := NewTestingStore(ctx, t)
+	user, err := createTestingHostUser(ctx, ts)
+	require.NoError(t, err)
+
+	testCases := []map[string]string{
+		{
+			ts.Profile.Data + "/assets/test.jpg": "assets/test.jpg",
+		},
+		{
+			ts.Profile.Data + `\assets\test.jpg`: "assets/test.jpg",
+		},
+		{
+			"/unhandled/path/test.jpg": "/unhandled/path/test.jpg",
+		},
+		{
+			`C:\unhandled\path\assets\test.jpg`: `C:\unhandled\path\assets\test.jpg`,
+		},
+	}
+
+	for _, testCase := range testCases {
+		for input, expectedOutput := range testCase {
+			resourceCreate := &store.Resource{
+				CreatorID:    user.ID,
+				InternalPath: input,
+			}
+			createdResource, err := ts.CreateResource(ctx, resourceCreate)
+			require.NoError(t, err)
+
+			err = ts.MigrateResourceInternalPath(ctx)
+			require.NoError(t, err)
+
+			findResource := &store.FindResource{
+				ID: &createdResource.ID,
+			}
+			resource, err := ts.GetResource(ctx, findResource)
+			require.NoError(t, err)
+
+			require.Equal(t, expectedOutput, resource.InternalPath)
+		}
+	}
+
+	ts.Close()
+}