feat: federation compose directive

Author: n9te9

docs/federation-compose-directive/ja/designdoc.md

@@ -332,34 +332,55 @@ sequenceDiagram
 
 ### Process
 
-1. **SubGraph V2 拡張**
-   1.1. `subgraph_v2.go` を修正
+**重要:** 以下のプロセスは TDD（テスト駆動開発）を厳守すること。各機能の実装前に必ずテストを書き、Red → Green → Refactor のサイクルを回すこと。
+
+1. **SubGraph V2 拡張 (TDD)**
+   1.1. **RED: テストを先に書く** - `subgraph_v2_test.go`
+        - @composeDirective のパース（既存）
+        - カスタムディレクティブ定義の抽出テスト
+        - 複数のカスタムディレクティブの処理テスト
+        - テストを実行して失敗することを確認
+   1.2. **GREEN: 最小限の実装** - `subgraph_v2.go`
         - SubGraphV2 構造に DirectiveDefinitions フィールドを追加
         - extractDirectiveDefinitions() 関数を追加
         - NewSubGraphV2() で extractDirectiveDefinitions() を呼び出し
-   1.2. `subgraph_v2_test.go` にテストを追加
-        - @composeDirective のパース（既存）
-        - カスタムディレクティブ定義の抽出
-        - 複数のカスタムディレクティブの処理
-
-2. **SuperGraph V2 拡張**
-   2.1. `super_graph_v2.go` を修正
+        - テストを実行して成功することを確認
+   1.3. **REFACTOR: リファクタリング**
+        - コードの重複を排除
+        - 可読性を向上
+        - テストが引き続き成功することを確認
+
+2. **SuperGraph V2 拡張 (TDD)**
+   2.1. **RED: テストを先に書く** - `super_graph_v2_test.go`
+        - カスタムディレクティブ定義のマージテスト
+        - ディレクティブ定義の一貫性検証（一致する場合）
+        - ディレクティブ定義の一貫性検証（不一致の場合、エラー）
+        - 複数サブグラフからの異なるカスタムディレクティブのマージ
+        - テストを実行して失敗することを確認
+   2.2. **GREEN: 最小限の実装** - `super_graph_v2.go`
         - SuperGraphV2 構造に DirectiveDefinitions フィールドを追加
         - mergeDirectiveDefinitions() 関数を追加
         - isDirectiveDefinitionEqual() ヘルパー関数を追加
         - isTypeEqual() ヘルパー関数を追加
         - NewSuperGraphV2() で mergeDirectiveDefinitions() を呼び出し
-   2.2. schema 定義への @composeDirective 追加（オプション）
-        - addSchemaDefinition() 関数を追加
-   2.3. `super_graph_v2_test.go` にテストを追加
-        - カスタムディレクティブ定義のマージ
-        - ディレクティブ定義の一貫性検証（一致する場合）
-        - ディレクティブ定義の一貫性検証（不一致の場合、エラー）
-        - 複数サブグラフからの異なるカスタムディレクティブのマージ
-
-3. **統合テスト**
-   3.1. カスタムディレクティブを使用するサンプルスキーマの作成
-   3.2. スーパーグラフにカスタムディレクティブが含まれることの確認
+        - テストを実行して成功することを確認
+   2.3. **REFACTOR: リファクタリング**
+        - コードの重複を排除
+        - 可読性を向上
+        - テストが引き続き成功することを確認
+   2.4. **[オプション] schema 定義への @composeDirective 追加 (TDD)**
+        - RED: テストを書く
+        - GREEN: addSchemaDefinition() 関数を実装
+        - REFACTOR: 改善
+
+3. **統合テスト (TDD)**
+   3.1. **RED: 統合テストを先に書く**
+        - カスタムディレクティブを使用するサンプルスキーマの作成
+        - スーパーグラフにカスタムディレクティブが含まれることを確認するテスト
+        - テストを実行して現状を確認
+   3.2. **GREEN: 必要な調整**
+        - テストが成功するように実装を調整
+   3.3. **REFACTOR: 統合レベルのリファクタリング**
 
 4. **ドキュメント**
    4.1. @composeDirective の使用方法を README に追加
@@ -369,6 +390,13 @@ sequenceDiagram
    5.1. `make test-all` で全ドメインのテストが通ることを確認
    5.2. カスタムディレクティブ機能が既存の動作を壊していないことを確認
 
+**TDD チェックリスト:**
+- [ ] 各機能について、実装前にテストを書いたか？
+- [ ] テストが最初は失敗することを確認したか？（RED）
+- [ ] テストが成功する最小限のコードを書いたか？（GREEN）
+- [ ] リファクタリング後もテストが成功することを確認したか？（REFACTOR）
+- [ ] 全てのテストが通ることを確認したか？
+
 ### Expected Outcomes
 
 - `@composeDirective` で指定されたカスタムディレクティブがスーパーグラフに含まれる

federation/graph/subgraph_v2.go

@@ -52,7 +52,8 @@ type SubGraphV2 struct {
 	entities map[string]*Entity // Entity map with entity name as key
 
 	// Federation v2 directives
-	ComposeDirectives []string // @composeDirective directives
+	ComposeDirectives    []string                            // @composeDirective directives
+	DirectiveDefinitions map[string]*ast.DirectiveDefinition // Custom directive definitions to compose
 }
 
 // NewSubGraphV2 initializes a SubGraphV2 by parsing the schema and extracting entities.
@@ -67,12 +68,14 @@ func NewSubGraphV2(name string, src []byte, host string) (*SubGraphV2, error) {
 	}
 
 	// Initialize SubGraph structure
+	composeDirectives := extractSchemaComposeDirectives(doc)
 	sg := &SubGraphV2{
-		Name:              name,
-		Host:              host,
-		Schema:            doc,
-		entities:          make(map[string]*Entity),
-		ComposeDirectives: extractSchemaComposeDirectives(doc),
+		Name:                 name,
+		Host:                 host,
+		Schema:               doc,
+		entities:             make(map[string]*Entity),
+		ComposeDirectives:    composeDirectives,
+		DirectiveDefinitions: extractDirectiveDefinitions(doc, composeDirectives),
 	}
 
 	// Traverse all type definitions
@@ -290,3 +293,33 @@ func extractSchemaComposeDirectives(doc *ast.Document) []string {
 func (sg *SubGraphV2) GetComposeDirectives() []string {
 	return sg.ComposeDirectives
 }
+
+// GetDirectiveDefinitions returns the custom directive definitions to compose.
+func (sg *SubGraphV2) GetDirectiveDefinitions() map[string]*ast.DirectiveDefinition {
+	return sg.DirectiveDefinitions
+}
+
+// extractDirectiveDefinitions extracts directive definitions that are listed in composeDirectives.
+func extractDirectiveDefinitions(
+	doc *ast.Document,
+	composeDirectives []string,
+) map[string]*ast.DirectiveDefinition {
+	definitions := make(map[string]*ast.DirectiveDefinition)
+
+	// Build a set of directive names to compose (strip leading "@")
+	composeSet := make(map[string]bool)
+	for _, name := range composeDirectives {
+		cleanName := strings.TrimPrefix(name, "@")
+		composeSet[cleanName] = true
+	}
+
+	for _, def := range doc.Definitions {
+		if directiveDef, ok := def.(*ast.DirectiveDefinition); ok {
+			if composeSet[directiveDef.Name.String()] {
+				definitions[directiveDef.Name.String()] = directiveDef
+			}
+		}
+	}
+
+	return definitions
+}

federation/graph/subgraph_v2_test.go

@@ -363,3 +363,126 @@ func TestNewSubGraphV2_WithComposeDirective(t *testing.T) {
 		t.Errorf("expected compose directive '@custom', got '%s'", composeDirectives[0])
 	}
 }
+
+func TestNewSubGraphV2_ExtractsDirectiveDefinitions(t *testing.T) {
+	schema := `
+		schema @composeDirective(name: "@rateLimit") {
+			query: Query
+		}
+
+		directive @rateLimit(limit: Int!, duration: Int!) on FIELD_DEFINITION
+
+		type Query {
+			expensiveQuery: String! @rateLimit(limit: 10, duration: 60)
+		}
+	`
+
+	sg, err := graph.NewSubGraphV2("product", []byte(schema), "http://product.example.com")
+	if err != nil {
+		t.Fatalf("NewSubGraphV2 failed: %v", err)
+	}
+
+	defs := sg.GetDirectiveDefinitions()
+	if len(defs) != 1 {
+		t.Fatalf("expected 1 directive definition, got %d", len(defs))
+	}
+
+	def, ok := defs["rateLimit"]
+	if !ok {
+		t.Fatal("expected 'rateLimit' directive definition, not found")
+	}
+
+	if def.Name.String() != "rateLimit" {
+		t.Errorf("expected directive name 'rateLimit', got '%s'", def.Name.String())
+	}
+}
+
+func TestNewSubGraphV2_ExtractsMultipleDirectiveDefinitions(t *testing.T) {
+	schema := `
+		schema
+			@composeDirective(name: "@rateLimit")
+			@composeDirective(name: "@cacheControl") {
+			query: Query
+		}
+
+		directive @rateLimit(limit: Int!) on FIELD_DEFINITION
+		directive @cacheControl(maxAge: Int!) on FIELD_DEFINITION | OBJECT
+
+		type Query {
+			cachedData: String! @cacheControl(maxAge: 3600) @rateLimit(limit: 100)
+		}
+	`
+
+	sg, err := graph.NewSubGraphV2("product", []byte(schema), "http://product.example.com")
+	if err != nil {
+		t.Fatalf("NewSubGraphV2 failed: %v", err)
+	}
+
+	defs := sg.GetDirectiveDefinitions()
+	if len(defs) != 2 {
+		t.Fatalf("expected 2 directive definitions, got %d", len(defs))
+	}
+
+	if _, ok := defs["rateLimit"]; !ok {
+		t.Error("expected 'rateLimit' directive definition, not found")
+	}
+	if _, ok := defs["cacheControl"]; !ok {
+		t.Error("expected 'cacheControl' directive definition, not found")
+	}
+}
+
+func TestNewSubGraphV2_DirectiveDefinitions_OnlyComposed(t *testing.T) {
+	// @internal is defined but NOT listed in @composeDirective, so should not be extracted
+	schema := `
+		schema @composeDirective(name: "@rateLimit") {
+			query: Query
+		}
+
+		directive @rateLimit(limit: Int!) on FIELD_DEFINITION
+		directive @internal on FIELD_DEFINITION
+
+		type Query {
+			data: String! @rateLimit(limit: 10) @internal
+		}
+	`
+
+	sg, err := graph.NewSubGraphV2("product", []byte(schema), "http://product.example.com")
+	if err != nil {
+		t.Fatalf("NewSubGraphV2 failed: %v", err)
+	}
+
+	defs := sg.GetDirectiveDefinitions()
+	if len(defs) != 1 {
+		t.Fatalf("expected 1 directive definition (only composed ones), got %d", len(defs))
+	}
+
+	if _, ok := defs["rateLimit"]; !ok {
+		t.Error("expected 'rateLimit' directive definition, not found")
+	}
+	if _, ok := defs["internal"]; ok {
+		t.Error("'internal' directive should not be extracted (not in @composeDirective)")
+	}
+}
+
+func TestNewSubGraphV2_NoComposeDirective_EmptyDefinitions(t *testing.T) {
+	schema := `
+		type Product @key(fields: "id") {
+			id: ID!
+			name: String!
+		}
+
+		type Query {
+			product(id: ID!): Product
+		}
+	`
+
+	sg, err := graph.NewSubGraphV2("product", []byte(schema), "http://product.example.com")
+	if err != nil {
+		t.Fatalf("NewSubGraphV2 failed: %v", err)
+	}
+
+	defs := sg.GetDirectiveDefinitions()
+	if len(defs) != 0 {
+		t.Errorf("expected 0 directive definitions, got %d", len(defs))
+	}
+}

federation/graph/super_graph_v2.go

@@ -8,10 +8,11 @@ import (
 
 // SuperGraphV2 represents an aggregated super graph composed of multiple subgraphs.
 type SuperGraphV2 struct {
-	SubGraphs []*SubGraphV2            // List of subgraphs
-	Schema    *ast.Document            // Composed schema
-	Ownership map[string][]*SubGraphV2 // Field ownership map (e.g., "Product.id" -> [SubGraph])
-	Graph     *WeightedDirectedGraph   // Weighted directed graph for Dijkstra-based plan optimization
+	SubGraphs            []*SubGraphV2                       // List of subgraphs
+	Schema               *ast.Document                       // Composed schema
+	Ownership            map[string][]*SubGraphV2            // Field ownership map (e.g., "Product.id" -> [SubGraph])
+	Graph                *WeightedDirectedGraph              // Weighted directed graph for Dijkstra-based plan optimization
+	DirectiveDefinitions map[string]*ast.DirectiveDefinition // Custom directive definitions merged from @composeDirective
 }
 
 // NewSuperGraphV2 creates a super graph from a list of SubGraphV2 instances.
@@ -26,6 +27,11 @@ func NewSuperGraphV2(subGraphs []*SubGraphV2) (*SuperGraphV2, error) {
 		return nil, err
 	}
 
+	// Merge and validate custom directive definitions from @composeDirective
+	if err := sg.mergeComposeDirectiveDefinitions(); err != nil {
+		return nil, err
+	}
+
 	// Build ownership map
 	if err := sg.buildOwnershipMap(); err != nil {
 		return nil, err
@@ -347,6 +353,74 @@ func (sg *SuperGraphV2) mergeDirectiveDefinition(newDef *ast.DirectiveDefinition
 	}
 }
 
+// mergeComposeDirectiveDefinitions validates and merges custom directive definitions
+// listed via @composeDirective across all subgraphs into the super graph.
+func (sg *SuperGraphV2) mergeComposeDirectiveDefinitions() error {
+	sg.DirectiveDefinitions = make(map[string]*ast.DirectiveDefinition)
+
+	for _, subGraph := range sg.SubGraphs {
+		for name, directiveDef := range subGraph.DirectiveDefinitions {
+			if existing, ok := sg.DirectiveDefinitions[name]; ok {
+				// Already seen this directive - validate consistency
+				if !isDirectiveDefinitionEqual(existing, directiveDef) {
+					return fmt.Errorf(
+						"inconsistent directive definition for '@%s' between subgraphs",
+						name,
+					)
+				}
+				// Consistent - skip
+				continue
+			}
+			// New directive definition
+			sg.DirectiveDefinitions[name] = directiveDef
+		}
+	}
+
+	return nil
+}
+
+// isDirectiveDefinitionEqual checks whether two DirectiveDefinition nodes are equivalent
+// by comparing their argument names/types and applicable locations.
+func isDirectiveDefinitionEqual(a, b *ast.DirectiveDefinition) bool {
+	if a.Name.String() != b.Name.String() {
+		return false
+	}
+
+	if len(a.Arguments) != len(b.Arguments) {
+		return false
+	}
+
+	for i := range a.Arguments {
+		if a.Arguments[i].Name.String() != b.Arguments[i].Name.String() {
+			return false
+		}
+		if !isTypeEqual(a.Arguments[i].Type, b.Arguments[i].Type) {
+			return false
+		}
+	}
+
+	if len(a.Locations) != len(b.Locations) {
+		return false
+	}
+
+	locSet := make(map[string]bool)
+	for _, loc := range a.Locations {
+		locSet[loc.String()] = true
+	}
+	for _, loc := range b.Locations {
+		if !locSet[loc.String()] {
+			return false
+		}
+	}
+
+	return true
+}
+
+// isTypeEqual compares two ast.Type values by their string representation.
+func isTypeEqual(a, b ast.Type) bool {
+	return a.String() == b.String()
+}
+
 // buildOwnershipMap constructs the ownership map.
 // It determines which subgraphs can resolve each field in the composed schema.
 func (sg *SuperGraphV2) buildOwnershipMap() error {