feat: batch STRM tree generation

This commit is contained in:
ShukeBta
2026-07-02 21:35:50 +08:00
parent d11f784438
commit 87928cfa4d
6 changed files with 117 additions and 4 deletions
+2
View File
@@ -119,6 +119,7 @@ type generateSTRMTreeReq struct {
Overwrite bool `json:"overwrite"`
Cleanup bool `json:"cleanup"`
DryRun bool `json:"dry_run"`
BatchLimit int `json:"batch_limit"`
}
func generateSTRMHandler(svc *service.Container) gin.HandlerFunc {
@@ -191,6 +192,7 @@ func generateSTRMFromTreeHandler(svc *service.Container) gin.HandlerFunc {
Overwrite: req.Overwrite,
Cleanup: req.Cleanup,
DryRun: req.DryRun,
BatchLimit: req.BatchLimit,
})
if err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
+12 -1
View File
@@ -20,6 +20,7 @@ type GenerateSTRMTreeOptions struct {
Overwrite bool `json:"overwrite"`
Cleanup bool `json:"cleanup"`
DryRun bool `json:"dry_run"`
BatchLimit int `json:"batch_limit,omitempty"`
}
type strmTreeSource struct {
@@ -55,8 +56,11 @@ func (s *STRMService) GenerateFromTree(ctx context.Context, opts GenerateSTRMTre
expectedFiles[filepath.Clean(item.FilePath)] = struct{}{}
}
result.addItem(item)
if strmTreeBatchLimitReached(result, opts.BatchLimit) {
break
}
}
if opts.Cleanup && !opts.DryRun && len(expectedFiles) > 0 {
if opts.Cleanup && !opts.DryRun && opts.BatchLimit <= 0 && len(expectedFiles) > 0 {
cleanupDir := outputDir
if prefix, err := strmTreeOutputPrefixPath(opts.OutputPrefix); err == nil && prefix != "" {
cleanupDir = filepath.Join(outputDir, prefix)
@@ -70,6 +74,13 @@ func (s *STRMService) GenerateFromTree(ctx context.Context, opts GenerateSTRMTre
return result, nil
}
func strmTreeBatchLimitReached(result *GenerateSTRMResult, limit int) bool {
if result == nil || limit <= 0 {
return false
}
return result.Generated+result.Updated+result.Previewed >= limit
}
func generateTreeSTRMItem(outputDir string, source strmTreeSource, opts GenerateSTRMTreeOptions) GenerateSTRMItem {
relSource := strmTreeRelativeSource(source.Path, opts.SourceRoot)
relPath, err := strmTreeOutputRelativePath(relSource)
+62
View File
@@ -137,3 +137,65 @@ func TestGenerateSTRMFromTreeDryRunDoesNotCreateOutputDir(t *testing.T) {
t.Fatalf("dry run should not create output dir, stat err=%v", err)
}
}
func TestGenerateSTRMFromTreeBatchLimitContinuesAfterExistingFiles(t *testing.T) {
outDir := filepath.Join(t.TempDir(), "strm")
svc := NewSTRMService(zap.NewNop(), nil, nil)
opts := GenerateSTRMTreeOptions{
Provider: "openlist",
Paths: []string{"/Movies/A.mkv", "/Movies/B.mkv", "/Movies/C.mkv"},
OutputDir: outDir,
BatchLimit: 2,
}
first, err := svc.GenerateFromTree(t.Context(), opts)
if err != nil {
t.Fatal(err)
}
if first.Generated != 2 || first.Skipped != 0 {
t.Fatalf("first batch = %#v, want two generated", first)
}
if _, err := os.Stat(filepath.Join(outDir, "Movies", "C.strm")); !os.IsNotExist(err) {
t.Fatalf("first batch should not write third item, stat err=%v", err)
}
second, err := svc.GenerateFromTree(t.Context(), opts)
if err != nil {
t.Fatal(err)
}
if second.Generated != 1 || second.Skipped != 2 {
t.Fatalf("second batch = %#v, want two existing skips then next generated", second)
}
if got := readSTRM(t, filepath.Join(outDir, "Movies", "C.strm")); !strings.Contains(got, "C.mkv") {
t.Fatalf("second batch C.strm = %q, want generated third item", got)
}
}
func TestGenerateSTRMFromTreeBatchLimitSkipsCleanup(t *testing.T) {
outDir := filepath.Join(t.TempDir(), "strm")
stale := filepath.Join(outDir, "Movies", "stale.strm")
if err := os.MkdirAll(filepath.Dir(stale), 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(stale, []byte("keep\n"), 0o644); err != nil {
t.Fatal(err)
}
svc := NewSTRMService(zap.NewNop(), nil, nil)
res, err := svc.GenerateFromTree(t.Context(), GenerateSTRMTreeOptions{
Provider: "openlist",
Paths: []string{"/Movies/A.mkv", "/Movies/B.mkv"},
OutputDir: outDir,
BatchLimit: 1,
Cleanup: true,
})
if err != nil {
t.Fatal(err)
}
if res.Generated != 1 || res.Cleaned != 0 {
t.Fatalf("batch result = %#v, want one generated and no cleanup", res)
}
if got := readSTRM(t, stale); got != "keep" {
t.Fatalf("batch cleanup should not touch stale file: %q", got)
}
}
+1
View File
@@ -40,6 +40,7 @@ export type GenerateSTRMTreeInput = {
overwrite?: boolean
cleanup?: boolean
dry_run?: boolean
batch_limit?: number
}
export const strmAPI = {
+21 -2
View File
@@ -29,6 +29,7 @@ const outputPrefixPresets = [
export function StrmTreeGenerateSection({
baseURL,
batchLimit,
cleanup,
generating,
onGenerate,
@@ -42,6 +43,7 @@ export function StrmTreeGenerateSection({
result,
runningMode,
setBaseURL,
setBatchLimit,
setCleanup,
setOutputDir,
setOutputPrefix,
@@ -53,6 +55,8 @@ export function StrmTreeGenerateSection({
sourceRoot,
treeText,
}: StrmTreeGenerateSectionProps) {
const batchLimitEnabled = Number.parseInt(batchLimit, 10) > 0
return (
<section className="glass-panel space-y-4">
<div>
@@ -112,11 +116,21 @@ export function StrmTreeGenerateSection({
<option value="成人">成人</option>
</select>
<input
className="input-base md:col-span-2"
className="input-base"
placeholder="输出分类,如 电影/欧美电影"
value={outputPrefix}
onChange={(e) => setOutputPrefix(e.target.value)}
/>
<input
className="input-base"
inputMode="numeric"
min={0}
placeholder="每批数量"
step={1}
type="number"
value={batchLimit}
onChange={(e) => setBatchLimit(e.target.value)}
/>
<input
className="input-base md:col-span-3"
placeholder="http://NAS-IP:18080 或 https://media.example.com"
@@ -166,7 +180,12 @@ export function StrmTreeGenerateSection({
覆盖已存在
</label>
<label className="flex min-h-10 items-center gap-2 rounded-2xl border border-gray-200 bg-white/70 px-3 py-2 text-sm text-ink-50 md:col-span-4">
<input type="checkbox" checked={cleanup} onChange={(e) => setCleanup(e.target.checked)} />
<input
type="checkbox"
checked={!batchLimitEnabled && cleanup}
disabled={batchLimitEnabled}
onChange={(e) => setCleanup(e.target.checked)}
/>
清理不在当前目录树中的旧 STRM
</label>
<button
+19 -1
View File
@@ -12,6 +12,7 @@ export function useStrmTreeGenerateForm() {
const [outputPrefix, setOutputPrefix] = useState('')
const [baseURL, setBaseURL] = useState(currentOrigin())
const [outputDir, setOutputDir] = useState('data/strm/tree')
const [batchLimit, setBatchLimit] = useState('')
const [treeText, setTreeText] = useState('')
const [pathsText, setPathsText] = useState('')
const [overwrite, setOverwrite] = useState(false)
@@ -49,6 +50,11 @@ export function useStrmTreeGenerateForm() {
toast.error('域名必须以 http:// 或 https:// 开头')
return
}
const parsedBatchLimit = parseBatchLimit(batchLimit)
if (parsedBatchLimit === null) {
toast.error('每批数量必须为非负整数')
return
}
setRunningMode(dryRun ? 'preview' : 'generate')
try {
const next = await strmAPI.generateFromTree({
@@ -60,8 +66,9 @@ export function useStrmTreeGenerateForm() {
tree_text: treeText.trim() || undefined,
paths: paths.length > 0 ? paths : undefined,
overwrite,
cleanup,
cleanup: parsedBatchLimit > 0 ? false : cleanup,
dry_run: dryRun,
batch_limit: parsedBatchLimit || undefined,
})
setResult(next)
if (dryRun) {
@@ -87,6 +94,7 @@ export function useStrmTreeGenerateForm() {
return {
baseURL,
batchLimit,
cleanup,
generating: runningMode !== null,
onImportTreeFile,
@@ -100,6 +108,7 @@ export function useStrmTreeGenerateForm() {
result,
runningMode,
setBaseURL,
setBatchLimit,
setCleanup,
setOutputDir,
setOutputPrefix,
@@ -128,3 +137,12 @@ function parsePathList(raw: string): string[] {
})
return out
}
function parseBatchLimit(raw: string): number | null {
const value = raw.trim()
if (value === '') return 0
if (!/^\d+$/.test(value)) return null
const parsed = Number.parseInt(value, 10)
if (!Number.isSafeInteger(parsed)) return null
return parsed
}