From 925dee002b92f1210456dce16c18c6b436e912f3 Mon Sep 17 00:00:00 2001
From: chen <wszgrcy@gmail.com>
Date: Mon, 29 Sep 2025 16:36:35 +0800
Subject: [PATCH] feat: entity schema support trees (#11606)

---
 src/entity-schema/EntitySchemaOptions.ts      |  3 +
 src/entity-schema/EntitySchemaTransformer.ts  | 10 +++
 .../trees/nested-set/entity/Category.ts       | 34 +++++++++
 .../trees/nested-set/nested-set.ts            | 69 +++++++++++++++++++
 4 files changed, 116 insertions(+)
 create mode 100644 test/functional/entity-schema/trees/nested-set/entity/Category.ts
 create mode 100644 test/functional/entity-schema/trees/nested-set/nested-set.ts

diff --git a/src/entity-schema/EntitySchemaOptions.ts b/src/entity-schema/EntitySchemaOptions.ts
index d8cdc969d..a7fb435f1 100644
--- a/src/entity-schema/EntitySchemaOptions.ts
+++ b/src/entity-schema/EntitySchemaOptions.ts
@@ -14,6 +14,7 @@ import { EntitySchemaExclusionOptions } from "./EntitySchemaExclusionOptions"
 import { EntitySchemaInheritanceOptions } from "./EntitySchemaInheritanceOptions"
 import { EntitySchemaRelationIdOptions } from "./EntitySchemaRelationIdOptions"
 import { EntitySchemaForeignKeyOptions } from "./EntitySchemaForeignKeyOptions"
+import { TreeMetadataArgs } from "../metadata-args/TreeMetadataArgs"
 
 /**
  * Interface for entity metadata mappings stored inside "schemas" instead of models decorated by decorators.
@@ -135,4 +136,6 @@ export class EntitySchemaOptions<T> {
      * Custom discriminator value for Single Table Inheritance.
      */
     discriminatorValue?: string
+
+    trees?: Omit<TreeMetadataArgs, "target">[]
 }
diff --git a/src/entity-schema/EntitySchemaTransformer.ts b/src/entity-schema/EntitySchemaTransformer.ts
index 59a6127af..ac7d7bb76 100644
--- a/src/entity-schema/EntitySchemaTransformer.ts
+++ b/src/entity-schema/EntitySchemaTransformer.ts
@@ -396,5 +396,15 @@ export class EntitySchemaTransformer {
                 )
             })
         }
+
+        if (options.trees) {
+            options.trees.forEach((tree) => {
+                metadataArgsStorage.trees.push({
+                    target: options.target || options.name,
+                    type: tree.type,
+                    options: tree.options,
+                })
+            })
+        }
     }
 }
diff --git a/test/functional/entity-schema/trees/nested-set/entity/Category.ts b/test/functional/entity-schema/trees/nested-set/entity/Category.ts
new file mode 100644
index 000000000..5678ecd18
--- /dev/null
+++ b/test/functional/entity-schema/trees/nested-set/entity/Category.ts
@@ -0,0 +1,34 @@
+import { EntitySchema } from "../../../../../../src"
+
+export const Category = new EntitySchema<{
+    id: number
+    name: string
+    parentCategory: any
+    childCategories: any[]
+}>({
+    name: "Category",
+    columns: {
+        id: {
+            primary: true,
+            type: Number,
+            generated: "increment",
+        },
+        name: {
+            type: String,
+        },
+    },
+    relations: {
+        parentCategory: {
+            type: "many-to-one",
+            treeParent: true,
+            target: "Category",
+        },
+        childCategories: {
+            type: "one-to-many",
+            treeChildren: true,
+            target: "Category",
+            cascade: true,
+        },
+    },
+    trees: [{ type: "nested-set" }],
+})
diff --git a/test/functional/entity-schema/trees/nested-set/nested-set.ts b/test/functional/entity-schema/trees/nested-set/nested-set.ts
new file mode 100644
index 000000000..6118169e4
--- /dev/null
+++ b/test/functional/entity-schema/trees/nested-set/nested-set.ts
@@ -0,0 +1,69 @@
+import "reflect-metadata"
+import { Category } from "./entity/Category"
+
+import { DataSource } from "../../../../../src"
+import {
+    createTestingConnections,
+    reloadTestingDatabases,
+    closeTestingConnections,
+} from "../../../../utils/test-utils"
+
+describe("entity-schema > tree tables > nested-set", () => {
+    let connections: DataSource[]
+    before(
+        async () =>
+            (connections = await createTestingConnections({
+                entities: [Category],
+                enabledDrivers: ["better-sqlite3"],
+            })),
+    )
+    beforeEach(() => reloadTestingDatabases(connections))
+    after(() => closeTestingConnections(connections))
+
+    it("attach should work properly", () =>
+        Promise.all(
+            connections.map(async (connection) => {
+                const categoryRepository =
+                    connection.getTreeRepository(Category)
+
+                const a1 = await categoryRepository.save({ name: "a1" })
+
+                const a11 = await categoryRepository.save({
+                    name: "a11",
+                    parentCategory: a1,
+                })
+
+                await categoryRepository.save({
+                    name: "a111",
+                    parentCategory: a11,
+                })
+
+                await categoryRepository.save({
+                    name: "a12",
+                    parentCategory: a1,
+                })
+
+                const rootCategories = await categoryRepository.findRoots()
+                rootCategories.should.be.eql([
+                    {
+                        id: 1,
+                        name: "a1",
+                    },
+                ])
+
+                const a11Parent = await categoryRepository.findAncestors(a11)
+                const a11ParentNames = a11Parent.map((i) => i.name)
+                a11ParentNames.length.should.be.equal(2)
+                a11ParentNames.should.deep.include("a1")
+                a11ParentNames.should.deep.include("a11")
+
+                const a1Children = await categoryRepository.findDescendants(a1)
+                const a1ChildrenNames = a1Children.map((i) => i.name)
+                a1ChildrenNames.length.should.be.equal(4)
+                a1ChildrenNames.should.deep.include("a1")
+                a1ChildrenNames.should.deep.include("a11")
+                a1ChildrenNames.should.deep.include("a111")
+                a1ChildrenNames.should.deep.include("a12")
+            }),
+        ))
+})