support get all nodes by ListPodNodes API

projecteru2 · Nov 9, 2020 · 16bf528 · 16bf528
1 parent 600dc8b
commit 16bf528
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 3 deletions.
diff --git a/store/etcdv3/node.go b/store/etcdv3/node.go
@@ -113,12 +113,30 @@ func (m *Mercury) GetNodes(ctx context.Context, nodenames []string) ([]*types.No
 // GetNodesByPod get all nodes bound to pod
 // here we use podname instead of pod instance
 func (m *Mercury) GetNodesByPod(ctx context.Context, podname string, labels map[string]string, all bool) ([]*types.Node, error) {
-	key := fmt.Sprintf(nodePodKey, podname, "")
-	resp, err := m.Get(ctx, key, clientv3.WithPrefix())
+	do := func(podname string) ([]*types.Node, error) {
+		key := fmt.Sprintf(nodePodKey, podname, "")
+		resp, err := m.Get(ctx, key, clientv3.WithPrefix())
+		if err != nil {
+			return []*types.Node{}, err
+		}
+		return m.doGetNodes(ctx, resp.Kvs, labels, all)
+	}
+	if podname != "" {
+		return do(podname)
+	}
+	pods, err := m.GetAllPods(ctx)
 	if err != nil {
 		return []*types.Node{}, err
 	}
-	return m.doGetNodes(ctx, resp.Kvs, labels, all)
+	result := []*types.Node{}
+	for _, pod := range pods {
+		ns, err := do(pod.Name)
+		if err != nil {
+			return []*types.Node{}, err
+		}
+		result = append(result, ns...)
+	}
+	return result, nil
 }
 
 // UpdateNode update a node, save it to etcd

diff --git a/store/etcdv3/node_test.go b/store/etcdv3/node_test.go
@@ -181,6 +181,11 @@ func TestGetNodesByPod(t *testing.T) {
 	ns, err = m.GetNodesByPod(ctx, "testpod", nil, false)
 	assert.NoError(t, err)
 	assert.NotEmpty(t, ns)
+	_, err = m.AddPod(ctx, "testpod", "")
+	assert.NoError(t, err)
+	ns, err = m.GetNodesByPod(ctx, "", nil, false)
+	assert.NoError(t, err)
+	assert.NotEmpty(t, ns)
 }
 
 func TestUpdateNode(t *testing.T) {