From 601467fd5158ce796a0dea2f49ed551847e953b1 Mon Sep 17 00:00:00 2001
From: Lauri Huotari <65590793+Mrflatt@users.noreply.github.com>
Date: Fri, 14 Feb 2025 00:46:26 +0200
Subject: [PATCH] feat(inputs.gnmi): Add support for "depth" extension (#16480)

---
 plugins/inputs/gnmi/README.md      |  7 +++++++
 plugins/inputs/gnmi/gnmi.go        | 16 ++++++++++++++++
 plugins/inputs/gnmi/sample.conf    |  7 +++++++
 plugins/inputs/gnmi/sample.conf.in |  7 +++++++
 4 files changed, 37 insertions(+)

diff --git a/plugins/inputs/gnmi/README.md b/plugins/inputs/gnmi/README.md
index 2c0fa4681..7e7a13f14 100644
--- a/plugins/inputs/gnmi/README.md
+++ b/plugins/inputs/gnmi/README.md
@@ -75,12 +75,19 @@ details on how to use them.
   ## gRPC Maximum Message Size
   # max_msg_size = "4MB"
 
+  ## Subtree depth for depth extension (disables if < 1)
+  ## see https://github.com/openconfig/reference/blob/master/rpc/gnmi/gnmi-depth.md
+  # depth = 0
+
   ## Enable to get the canonical path as field-name
   # canonical_field_names = false
 
   ## Remove leading slashes and dots in field-name
   # trim_field_names = false
 
+  ## Only receive updates for the state, also suppresses receiving the initial state
+  # updates_only = false
+
   ## Guess the path-tag if an update does not contain a prefix-path
   ## Supported values are
   ##   none         -- do not add a 'path' tag
diff --git a/plugins/inputs/gnmi/gnmi.go b/plugins/inputs/gnmi/gnmi.go
index dcf7101ff..64736338f 100644
--- a/plugins/inputs/gnmi/gnmi.go
+++ b/plugins/inputs/gnmi/gnmi.go
@@ -14,6 +14,7 @@ import (
 
 	"github.com/google/gnxi/utils/xpath"
 	"github.com/openconfig/gnmi/proto/gnmi"
+	"github.com/openconfig/gnmi/proto/gnmi_ext"
 	"google.golang.org/grpc/keepalive"
 	"google.golang.org/grpc/metadata"
 
@@ -54,6 +55,7 @@ type GNMI struct {
 	Password             config.Secret     `toml:"password"`
 	Redial               config.Duration   `toml:"redial"`
 	MaxMsgSize           config.Size       `toml:"max_msg_size"`
+	Depth                int32             `toml:"depth"`
 	Trace                bool              `toml:"dump_responses"`
 	CanonicalFieldNames  bool              `toml:"canonical_field_names"`
 	TrimFieldNames       bool              `toml:"trim_field_names"`
@@ -376,6 +378,19 @@ func (c *GNMI) newSubscribeRequest() (*gnmi.SubscribeRequest, error) {
 		return nil, fmt.Errorf("unsupported encoding %s", c.Encoding)
 	}
 
+	var extensions []*gnmi_ext.Extension
+	if c.Depth > 0 {
+		extensions = []*gnmi_ext.Extension{
+			{
+				Ext: &gnmi_ext.Extension_Depth{
+					Depth: &gnmi_ext.Depth{
+						Level: uint32(c.Depth),
+					},
+				},
+			},
+		}
+	}
+
 	return &gnmi.SubscribeRequest{
 		Request: &gnmi.SubscribeRequest_Subscribe{
 			Subscribe: &gnmi.SubscriptionList{
@@ -386,6 +401,7 @@ func (c *GNMI) newSubscribeRequest() (*gnmi.SubscribeRequest, error) {
 				UpdatesOnly:  c.UpdatesOnly,
 			},
 		},
+		Extension: extensions,
 	}, nil
 }
 
diff --git a/plugins/inputs/gnmi/sample.conf b/plugins/inputs/gnmi/sample.conf
index ac178431d..b7b6ef1fb 100644
--- a/plugins/inputs/gnmi/sample.conf
+++ b/plugins/inputs/gnmi/sample.conf
@@ -28,12 +28,19 @@
   ## gRPC Maximum Message Size
   # max_msg_size = "4MB"
 
+  ## Subtree depth for depth extension (disables if < 1)
+  ## see https://github.com/openconfig/reference/blob/master/rpc/gnmi/gnmi-depth.md
+  # depth = 0
+
   ## Enable to get the canonical path as field-name
   # canonical_field_names = false
 
   ## Remove leading slashes and dots in field-name
   # trim_field_names = false
 
+  ## Only receive updates for the state, also suppresses receiving the initial state
+  # updates_only = false
+
   ## Guess the path-tag if an update does not contain a prefix-path
   ## Supported values are
   ##   none         -- do not add a 'path' tag
diff --git a/plugins/inputs/gnmi/sample.conf.in b/plugins/inputs/gnmi/sample.conf.in
index f4a9af6b9..8d5a3abdd 100644
--- a/plugins/inputs/gnmi/sample.conf.in
+++ b/plugins/inputs/gnmi/sample.conf.in
@@ -28,12 +28,19 @@
   ## gRPC Maximum Message Size
   # max_msg_size = "4MB"
 
+  ## Subtree depth for depth extension (disables if < 1)
+  ## see https://github.com/openconfig/reference/blob/master/rpc/gnmi/gnmi-depth.md
+  # depth = 0
+
   ## Enable to get the canonical path as field-name
   # canonical_field_names = false
 
   ## Remove leading slashes and dots in field-name
   # trim_field_names = false
 
+  ## Only receive updates for the state, also suppresses receiving the initial state
+  # updates_only = false
+
   ## Guess the path-tag if an update does not contain a prefix-path
   ## Supported values are
   ##   none         -- do not add a 'path' tag