fixes from comments

https://github.com/chrislusf/seaweedfs/pull/2996

4 files changed, 64 insertions, 60 deletions

diff --git a/weed/filer/ydb/ydb_queries.go b/weed/filer/ydb/ydb_queries.go
index c8876e004..f1db9a143 100644
--- a/weed/filer/ydb/ydb_queries.go
+++ b/weed/filer/ydb/ydb_queries.go
@@ -6,7 +6,7 @@ package ydb
 import asql "github.com/chrislusf/seaweedfs/weed/filer/abstract_sql"
 
 const (
-	insertQuery = `
+	upsertQuery = `
 		PRAGMA TablePathPrefix("%v");
 		DECLARE $dir_hash AS int64;
 		DECLARE $directory AS Utf8;
@@ -19,19 +19,6 @@ const (
 		VALUES
 			($dir_hash, $name, $directory, $meta, $expire_at);`
 
-	updateQuery = `
-		PRAGMA TablePathPrefix("%v");
-		DECLARE $dir_hash AS int64;
-		DECLARE $directory AS Utf8;
-		DECLARE $name AS Utf8;
-		DECLARE $meta AS String;
-		DECLARE $expire_at AS Optional<uint32>;
-
-		REPLACE INTO ` + asql.DEFAULT_TABLE + `
-			(dir_hash, name, directory, meta, expire_at)
-		VALUES
-			($dir_hash, $name, $directory, $meta, $expire_at);`
-
 	deleteQuery = `
 		PRAGMA TablePathPrefix("%v");
 		DECLARE $dir_hash AS int64;
diff --git a/weed/filer/ydb/ydb_store.go b/weed/filer/ydb/ydb_store.go
index 5b0e4e764..b06b35ac2 100644
--- a/weed/filer/ydb/ydb_store.go
+++ b/weed/filer/ydb/ydb_store.go
@@ -28,6 +28,7 @@ import (
 
 const (
 	defaultDialTimeOut = 10
+	maxRowsInQuery     = 1000 // Limit number of rows in query results https://cloud.yandex.com/en-ru/docs/ydb/concepts/limits-ydb
 )
 
 var (
@@ -89,7 +90,7 @@ func (store *YdbStore) initialize(dirBuckets string, dsn string, tablePathPrefix
 		dsn = os.Getenv("YDB_CONNECTION_STRING")
 	}
 	store.DB, err = ydb.Open(ctx, dsn, opts...)
-	if err != nil || store.DB == nil {
+	if err != nil {
 		if store.DB != nil {
 			_ = store.DB.Close(ctx)
 			store.DB = nil
@@ -139,7 +140,7 @@ func (store *YdbStore) doTxOrDB(ctx context.Context, query *string, params *tabl
 	return err
 }
 
-func (store *YdbStore) insertOrUpdateEntry(ctx context.Context, entry *filer.Entry, isUpdate bool) (err error) {
+func (store *YdbStore) insertOrUpdateEntry(ctx context.Context, entry *filer.Entry) (err error) {
 	dir, name := entry.FullPath.DirAndName()
 	meta, err := entry.EncodeAttributesAndChunks()
 	if err != nil {
@@ -151,21 +152,15 @@ func (store *YdbStore) insertOrUpdateEntry(ctx context.Context, entry *filer.Ent
 	}
 	tablePathPrefix, shortDir := store.getPrefix(ctx, &dir)
 	fileMeta := FileMeta{util.HashStringToLong(dir), name, *shortDir, meta}
-	var query *string
-	if isUpdate {
-		query = withPragma(tablePathPrefix, updateQuery)
-	} else {
-		query = withPragma(tablePathPrefix, insertQuery)
-	}
-	return store.doTxOrDB(ctx, query, fileMeta.queryParameters(entry.TtlSec), rwTX, nil)
+	return store.doTxOrDB(ctx, withPragma(tablePathPrefix, upsertQuery), fileMeta.queryParameters(entry.TtlSec), rwTX, nil)
 }
 
 func (store *YdbStore) InsertEntry(ctx context.Context, entry *filer.Entry) (err error) {
-	return store.insertOrUpdateEntry(ctx, entry, false)
+	return store.insertOrUpdateEntry(ctx, entry)
 }
 
 func (store *YdbStore) UpdateEntry(ctx context.Context, entry *filer.Entry) (err error) {
-	return store.insertOrUpdateEntry(ctx, entry, true)
+	return store.insertOrUpdateEntry(ctx, entry)
 }
 
 func (store *YdbStore) FindEntry(ctx context.Context, fullpath util.FullPath) (entry *filer.Entry, err error) {
@@ -179,14 +174,15 @@ func (store *YdbStore) FindEntry(ctx context.Context, fullpath util.FullPath) (e
 		table.ValueParam("$name", types.UTF8Value(name)))
 
 	err = store.doTxOrDB(ctx, query, queryParams, roTX, func(res result.Result) error {
-		for res.NextResultSet(ctx) {
-			for res.NextRow() {
-				if err = res.ScanNamed(named.OptionalWithDefault("meta", &data)); err != nil {
-					return fmt.Errorf("scanNamed %s : %v", fullpath, err)
-				}
-				entryFound = true
-				return nil
+		if !res.NextResultSet(ctx) || !res.HasNextRow() {
+			return nil
+		}
+		for res.NextRow() {
+			if err = res.ScanNamed(named.OptionalWithDefault("meta", &data)); err != nil {
+				return fmt.Errorf("scanNamed %s : %v", fullpath, err)
 			}
+			entryFound = true
+			return nil
 		}
 		return res.Err()
 	})
@@ -230,7 +226,7 @@ func (store *YdbStore) DeleteFolderChildren(ctx context.Context, fullpath util.F
 }
 
 func (store *YdbStore) ListDirectoryEntries(ctx context.Context, dirPath util.FullPath, startFileName string, includeStartFile bool, limit int64, eachEntryFunc filer.ListEachEntryFunc) (lastFileName string, err error) {
-	return store.ListDirectoryPrefixedEntries(ctx, dirPath, startFileName, includeStartFile, limit, "", nil)
+	return store.ListDirectoryPrefixedEntries(ctx, dirPath, startFileName, includeStartFile, limit, "", eachEntryFunc)
 }
 
 func (store *YdbStore) ListDirectoryPrefixedEntries(ctx context.Context, dirPath util.FullPath, startFileName string, includeStartFile bool, limit int64, prefix string, eachEntryFunc filer.ListEachEntryFunc) (lastFileName string, err error) {
@@ -242,17 +238,39 @@ func (store *YdbStore) ListDirectoryPrefixedEntries(ctx context.Context, dirPath
 	} else {
 		query = withPragma(tablePathPrefix, listDirectoryQuery)
 	}
-	queryParams := table.NewQueryParameters(
-		table.ValueParam("$dir_hash", types.Int64Value(util.HashStringToLong(*shortDir))),
-		table.ValueParam("$directory", types.UTF8Value(*shortDir)),
-		table.ValueParam("$start_name", types.UTF8Value(startFileName)),
-		table.ValueParam("$prefix", types.UTF8Value(prefix+"%")),
-		table.ValueParam("$limit", types.Uint64Value(uint64(limit))),
-	)
-	err = store.doTxOrDB(ctx, query, queryParams, roTX, func(res result.Result) error {
-		var name string
-		var data []byte
-		for res.NextResultSet(ctx) {
+	truncated := true
+	eachEntryFuncIsNotBreake := true
+	shortLimit := limit
+	if limit > maxRowsInQuery {
+		shortLimit = maxRowsInQuery * 2
+	}
+	entryCount := int64(0)
+	for truncated && eachEntryFuncIsNotBreake {
+		if lastFileName != "" {
+			startFileName = lastFileName
+			if includeStartFile {
+				query = withPragma(tablePathPrefix, listDirectoryQuery)
+			}
+		}
+		restLimit := limit - entryCount
+		if maxRowsInQuery > restLimit {
+			shortLimit = restLimit
+		}
+		queryParams := table.NewQueryParameters(
+			table.ValueParam("$dir_hash", types.Int64Value(util.HashStringToLong(*shortDir))),
+			table.ValueParam("$directory", types.UTF8Value(*shortDir)),
+			table.ValueParam("$start_name", types.UTF8Value(startFileName)),
+			table.ValueParam("$prefix", types.UTF8Value(prefix+"%")),
+			table.ValueParam("$limit", types.Uint64Value(uint64(shortLimit))),
+		)
+		err = store.doTxOrDB(ctx, query, queryParams, roTX, func(res result.Result) error {
+			var name string
+			var data []byte
+			if !res.NextResultSet(ctx) || !res.HasNextRow() {
+				truncated = false
+				return nil
+			}
+			truncated = res.CurrentResultSet().Truncated()
 			for res.NextRow() {
 				if err := res.ScanNamed(
 					named.OptionalWithDefault("name", &name),
@@ -267,12 +285,14 @@ func (store *YdbStore) ListDirectoryPrefixedEntries(ctx context.Context, dirPath
 					return fmt.Errorf("scan decode %s : %v", entry.FullPath, err)
 				}
 				if !eachEntryFunc(entry) {
+					eachEntryFuncIsNotBreake = false
 					break
 				}
+				entryCount += 1
 			}
-		}
-		return res.Err()
-	})
+			return res.Err()
+		})
+	}
 	if err != nil {
 		return lastFileName, err
 	}
diff --git a/weed/filer/ydb/ydb_store_kv.go b/weed/filer/ydb/ydb_store_kv.go
index 61c3a273c..72bbfff42 100644
--- a/weed/filer/ydb/ydb_store_kv.go
+++ b/weed/filer/ydb/ydb_store_kv.go
@@ -18,7 +18,7 @@ func (store *YdbStore) KvPut(ctx context.Context, key []byte, value []byte) (err
 	dirStr, dirHash, name := abstract_sql.GenDirAndName(key)
 	fileMeta := FileMeta{dirHash, name, dirStr, value}
 	return store.DB.Table().Do(ctx, func(ctx context.Context, s table.Session) (err error) {
-		_, _, err = s.Execute(ctx, rwTX, *withPragma(&store.tablePathPrefix, insertQuery),
+		_, _, err = s.Execute(ctx, rwTX, *withPragma(&store.tablePathPrefix, upsertQuery),
 			fileMeta.queryParameters(0))
 		if err != nil {
 			return fmt.Errorf("kv put execute %s: %v", util.NewFullPath(dirStr, name).Name(), err)
@@ -39,14 +39,15 @@ func (store *YdbStore) KvGet(ctx context.Context, key []byte) (value []byte, err
 			return fmt.Errorf("kv get execute %s: %v", util.NewFullPath(dirStr, name).Name(), err)
 		}
 		defer func() { _ = res.Close() }()
-		for res.NextResultSet(ctx) {
-			for res.NextRow() {
-				if err := res.ScanNamed(named.OptionalWithDefault("meta", &value)); err != nil {
-					return fmt.Errorf("scanNamed %s : %v", util.NewFullPath(dirStr, name).Name(), err)
-				}
-				valueFound = true
-				return nil
+		if !res.NextResultSet(ctx) || !res.HasNextRow() {
+			return nil
+		}
+		for res.NextRow() {
+			if err := res.ScanNamed(named.OptionalWithDefault("meta", &value)); err != nil {
+				return fmt.Errorf("scanNamed %s : %v", util.NewFullPath(dirStr, name).Name(), err)
 			}
+			valueFound = true
+			return nil
 		}
 		return res.Err()
 	})
@@ -61,7 +62,7 @@ func (store *YdbStore) KvGet(ctx context.Context, key []byte) (value []byte, err
 func (store *YdbStore) KvDelete(ctx context.Context, key []byte) (err error) {
 	dirStr, dirHash, name := abstract_sql.GenDirAndName(key)
 	return store.DB.Table().Do(ctx, func(ctx context.Context, s table.Session) (err error) {
-		_, _, err = s.Execute(ctx, rwTX, *withPragma(&store.tablePathPrefix, insertQuery),
+		_, _, err = s.Execute(ctx, rwTX, *withPragma(&store.tablePathPrefix, deleteQuery),
 			table.NewQueryParameters(
 				table.ValueParam("$dir_hash", types.Int64Value(dirHash)),
 				table.ValueParam("$name", types.UTF8Value(name))))
diff --git a/weed/filer/ydb/ydb_types.go b/weed/filer/ydb/ydb_types.go
index d256eaf50..4e5100236 100644
--- a/weed/filer/ydb/ydb_types.go
+++ b/weed/filer/ydb/ydb_types.go
@@ -10,9 +10,6 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/table/types"
 )
 
-//go:generate ydbgen
-
-//ydb:gen
 type FileMeta struct {
 	DirHash   int64  `ydb:"type:int64"`
 	Name      string `ydb:"type:utf8"`
@@ -20,7 +17,6 @@ type FileMeta struct {
 	Meta      []byte `ydb:"type:string"`
 }
 
-//ydb:gen scan,value
 type FileMetas []FileMeta
 
 func (fm *FileMeta) queryParameters(ttlSec int32) *table.QueryParameters {