Move BookmarkManager at the session level

Author: fbiville

neo4j/bookmarks.go

@@ -49,6 +49,8 @@ type BookmarkManager interface {
 	Forget(databases ...string)
 }
 
+// BookmarkManagerConfig is an experimental API and may be changed or removed
+// without prior notice
 type BookmarkManagerConfig struct {
 	// Initial bookmarks per database
 	InitialBookmarks map[string]Bookmarks

neo4j/config.go

@@ -140,11 +140,6 @@ type Config struct {
 	// If a single large result is to be retrieved, this is the most performant
 	// setting.
 	FetchSize int
-	// BookmarkManager defines a central point to externally supply bookmarks
-	// and be notified of bookmark updates per database
-	// Since 5.0
-	// default: nil (no-op)
-	BookmarkManager BookmarkManager
 }
 
 func defaultConfig() *Config {
@@ -159,7 +154,6 @@ func defaultConfig() *Config {
 		RootCAs:                      nil,
 		UserAgent:                    UserAgent,
 		FetchSize:                    FetchDefault,
-		BookmarkManager:              nil,
 	}
 }
 

neo4j/session_with_context.go

@@ -107,12 +107,12 @@ type SessionConfig struct {
 	// to the correct cluster member (different databases may have different
 	// leaders).
 	ImpersonatedUser string
-	// IgnoreBookmarkManager allows specific sessions to ignore the
-	// globally-configured bookmark manager
-	// Sessions with this setting will handle bookmarks as before the bookmark
-	// manager was introduced
+	// BookmarkManager defines a central point to externally supply bookmarks
+	// and be notified of bookmark updates per database
+	// This is experimental and may be changed or removed without prior notice
 	// Since 5.0
-	IgnoreBookmarkManager bool
+	// default: nil (no-op)
+	BookmarkManager BookmarkManager
 }
 
 // FetchAll turns off fetching records in batches.
@@ -157,16 +157,12 @@ func newSessionWithContext(config *Config, sessConfig SessionConfig, router sess
 		fetchSize = sessConfig.FetchSize
 	}
 
-	configuredBookmarkManager := config.BookmarkManager
-	if sessConfig.IgnoreBookmarkManager {
-		configuredBookmarkManager = nil
-	}
 	return &sessionWithContext{
 		config:           config,
 		router:           router,
 		pool:             pool,
 		defaultMode:      idb.AccessMode(sessConfig.AccessMode),
-		bookmarks:        newSessionBookmarks(configuredBookmarkManager, sessConfig.Bookmarks),
+		bookmarks:        newSessionBookmarks(sessConfig.BookmarkManager, sessConfig.Bookmarks),
 		databaseName:     sessConfig.DatabaseName,
 		impersonatedUser: sessConfig.ImpersonatedUser,
 		resolveHomeDb:    sessConfig.DatabaseName == "",

testkit-backend/backend.go

@@ -53,6 +53,7 @@ type backend struct {
 	wrLock               sync.Mutex
 	suppliedBookmarks    map[string]neo4j.Bookmarks
 	consumedBookmarks    map[string]struct{}
+	bookmarkManagers     map[string]neo4j.BookmarkManager
 }
 
 // To implement transactional functions a bit of extra state is needed on the
@@ -83,6 +84,7 @@ func newBackend(rd *bufio.Reader, wr io.Writer) *backend {
 		recordedErrors:       make(map[string]error),
 		resolvedAddresses:    make(map[string][]any),
 		id:                   0,
+		bookmarkManagers:     make(map[string]neo4j.BookmarkManager),
 		suppliedBookmarks:    make(map[string]neo4j.Bookmarks),
 		consumedBookmarks:    make(map[string]struct{}),
 	}
@@ -448,9 +450,6 @@ func (b *backend) handleRequest(req map[string]any) {
 			if data["connectionTimeoutMs"] != nil {
 				c.SocketConnectTimeout = time.Millisecond * time.Duration(asInt64(data["connectionTimeoutMs"].(json.Number)))
 			}
-			if data["bookmarkManager"] != nil {
-				c.BookmarkManager = neo4j.NewBookmarkManager(b.bookmarkManagerConfig(data["bookmarkManager"].(map[string]any)))
-			}
 		})
 		if err != nil {
 			b.writeError(err)
@@ -516,14 +515,35 @@ func (b *backend) handleRequest(req map[string]any) {
 		if data["impersonatedUser"] != nil {
 			sessionConfig.ImpersonatedUser = data["impersonatedUser"].(string)
 		}
-		if data["ignoreBookmarkManager"] != nil {
-			sessionConfig.IgnoreBookmarkManager = data["ignoreBookmarkManager"].(bool)
+		if data["bookmarkManagerId"] != nil {
+			bmmId := data["bookmarkManagerId"].(string)
+			bookmarkManager := b.bookmarkManagers[bmmId]
+			if bookmarkManager == nil {
+				b.writeError(fmt.Errorf("could not find bookmark manager with ID %s", bmmId))
+				return
+			}
+			sessionConfig.BookmarkManager = bookmarkManager
 		}
 		session := driver.NewSession(ctx, sessionConfig)
 		idKey := b.nextId()
 		b.sessionStates[idKey] = &sessionState{session: session}
 		b.writeResponse("Session", map[string]any{"id": idKey})
 
+	case "NewBookmarkManager":
+		bookmarkManagerId := b.nextId()
+		b.bookmarkManagers[bookmarkManagerId] = neo4j.NewBookmarkManager(
+			b.bookmarkManagerConfig(bookmarkManagerId, data))
+		b.writeResponse("BookmarkManager", map[string]any{
+			"id": bookmarkManagerId,
+		})
+
+	case "BookmarkManagerClose":
+		bookmarkManagerId := data["id"].(string)
+		delete(b.bookmarkManagers, bookmarkManagerId)
+		b.writeResponse("BookmarkManager", map[string]any{
+			"id": bookmarkManagerId,
+		})
+
 	case "SessionClose":
 		sessionId := data["sessionId"].(string)
 		sessionState := b.sessionStates[sessionId]
@@ -1076,50 +1096,59 @@ func patchNumbersInMap(dictionary map[string]any) error {
 	return nil
 }
 
-func (b *backend) bookmarkManagerConfig(config map[string]any) neo4j.BookmarkManagerConfig {
+func (b *backend) bookmarkManagerConfig(bookmarkManagerId string,
+	config map[string]any) neo4j.BookmarkManagerConfig {
+
 	var initialBookmarks map[string]neo4j.Bookmarks
 	if config["initialBookmarks"] != nil {
 		initialBookmarks = convertInitialBookmarks(config["initialBookmarks"].(map[string]any))
 	}
 	result := neo4j.BookmarkManagerConfig{InitialBookmarks: initialBookmarks}
 	supplierRegistered := config["bookmarksSupplierRegistered"]
 	if supplierRegistered != nil && supplierRegistered.(bool) {
-		result.BookmarkSupplier = &testkitBookmarkSupplier{supplierFn: b.supplyBookmarks}
+		result.BookmarkSupplier = &testkitBookmarkSupplier{
+			supplierFn: b.supplyBookmarks(bookmarkManagerId),
+		}
 	}
 	consumerRegistered := config["bookmarksConsumerRegistered"]
 	if consumerRegistered != nil && consumerRegistered.(bool) {
-		result.BookmarkUpdateNotifier = b.consumeBookmarks
+		result.BookmarkUpdateNotifier = b.consumeBookmarks(bookmarkManagerId)
 	}
 	return result
 }
 
-func (b *backend) supplyBookmarks(databases ...string) neo4j.Bookmarks {
-	if len(databases) > 1 {
-		panic("at most 1 database should be specified")
-	}
-	id := b.nextId()
-	msg := map[string]any{"id": id}
-	if len(databases) == 1 {
-		msg["database"] = databases[0]
-	}
-	b.writeResponse("BookmarksSupplierRequest", msg)
-	for {
-		b.process()
-		return b.suppliedBookmarks[id]
+func (b *backend) supplyBookmarks(bookmarkManagerId string) func(...string) neo4j.Bookmarks {
+	return func(databases ...string) neo4j.Bookmarks {
+		if len(databases) > 1 {
+			panic("at most 1 database should be specified")
+		}
+		id := b.nextId()
+		msg := map[string]any{"id": id, "bookmarkManagerId": bookmarkManagerId}
+		if len(databases) == 1 {
+			msg["database"] = databases[0]
+		}
+		b.writeResponse("BookmarksSupplierRequest", msg)
+		for {
+			b.process()
+			return b.suppliedBookmarks[id]
+		}
 	}
 }
 
-func (b *backend) consumeBookmarks(database string, bookmarks neo4j.Bookmarks) {
-	id := b.nextId()
-	b.writeResponse("BookmarksConsumerRequest", map[string]any{
-		"id":        id,
-		"database":  database,
-		"bookmarks": bookmarks,
-	})
-	for {
-		b.process()
-		if _, found := b.consumedBookmarks[id]; found {
-			return
+func (b *backend) consumeBookmarks(bookmarkManagerId string) func(string, neo4j.Bookmarks) {
+	return func(database string, bookmarks neo4j.Bookmarks) {
+		id := b.nextId()
+		b.writeResponse("BookmarksConsumerRequest", map[string]any{
+			"id":                id,
+			"bookmarkManagerId": bookmarkManagerId,
+			"database":          database,
+			"bookmarks":         bookmarks,
+		})
+		for {
+			b.process()
+			if _, found := b.consumedBookmarks[id]; found {
+				return
+			}
 		}
 	}
 }