From da0b1911de6e865d9d419b769cffe00575a545d9 Mon Sep 17 00:00:00 2001
From: abhishuraina <45932588+abhishuraina@users.noreply.github.com>
Date: Wed, 20 Nov 2024 17:17:38 +0530
Subject: [PATCH] add argument to disable disk-space-metrics (#186)

* add argument to disable disk-space-metrics
---
 CHANGELOG.md                               |  2 ++
 mssql-config.yml.sample                    |  1 +
 src/args/argument_list.go                  |  1 +
 src/metrics/instance_metric_definitions.go | 29 ++++++++++++----------
 src/metrics/metrics.go                     |  3 +++
 src/metrics/metrics_test.go                |  3 ++-
 6 files changed, 25 insertions(+), 14 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 46fa0d3b..cdb1d589 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,6 +8,8 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 Unreleased section should follow [Release Toolkit](https://github.com/newrelic/release-toolkit#render-markdown-and-update-markdown)
 
 ## Unreleased
+### Enhancement
+- Add `enable_disk_metrics_in_bytes` argument (default true), allows disabling disk-space-metrics
 
 ## v2.13.0 - 2024-10-08
 
diff --git a/mssql-config.yml.sample b/mssql-config.yml.sample
index f6105d5f..52434ce1 100644
--- a/mssql-config.yml.sample
+++ b/mssql-config.yml.sample
@@ -19,6 +19,7 @@ integrations:
 
     # ENABLE_BUFFER_METRICS: true
     # ENABLE_DATABASE_RESERVE_METRICS: true 
+    # ENABLE_DISK_METRICS_IN_BYTES: true
 
     # YAML configuration with one or more SQL queries to collect custom metrics
     # CUSTOM_METRICS_CONFIG: ""
diff --git a/src/args/argument_list.go b/src/args/argument_list.go
index 1e8ca8b4..b8b242b8 100644
--- a/src/args/argument_list.go
+++ b/src/args/argument_list.go
@@ -27,6 +27,7 @@ type ArgumentList struct {
 	CustomMetricsConfig          string `default:"" help:"YAML configuration with one or more SQL queries to collect custom metrics"`
 	ShowVersion                  bool   `default:"false" help:"Print build information and exit"`
 	ExtraConnectionURLArgs       string `default:"" help:"Appends additional parameters to connection url. Ex. 'applicationintent=readonly&foo=bar'"`
+	EnableDiskMetricsInBytes     bool   `default:"true" help:"Enable collection of instance.diskInBytes."`
 }
 
 // Validate validates SQL specific arguments
diff --git a/src/metrics/instance_metric_definitions.go b/src/metrics/instance_metric_definitions.go
index 5a811ca8..325d96ba 100644
--- a/src/metrics/instance_metric_definitions.go
+++ b/src/metrics/instance_metric_definitions.go
@@ -108,19 +108,6 @@ var instanceDefinitions = []*QueryDefinition{
 			Sleeping   *int64 `db:"sleeping" metric_name:"instance.sleepingProcessesCount" source_type:"gauge"`
 		}{},
 	},
-	{
-		query: `SELECT Sum(total_bytes) AS total_disk_space FROM (
-			SELECT DISTINCT
-			dovs.volume_mount_point,
-			dovs.available_bytes available_bytes,
-			dovs.total_bytes total_bytes
-			FROM sys.master_files mf WITH (nolock)
-			CROSS apply sys.Dm_os_volume_stats(mf.database_id, mf.file_id) dovs
-			) drives`,
-		dataModels: &[]struct {
-			TotalDiskSpace *int64 `db:"total_disk_space" metric_name:"instance.diskInBytes" source_type:"gauge"`
-		}{},
-	},
 	{
 		query: `SELECT Sum(runnable_tasks_count) AS runnable_tasks_count
 		FROM sys.dm_os_schedulers
@@ -172,3 +159,19 @@ type waitTimeModel struct {
 	WaitTime  *int64  `db:"wait_time"`
 	WaitCount *int64  `db:"waiting_tasks_count"`
 }
+
+var diskMetricInBytesDefination = []*QueryDefinition{
+	{
+		query: `SELECT Sum(total_bytes) AS total_disk_space FROM (
+			SELECT DISTINCT
+			dovs.volume_mount_point,
+			dovs.available_bytes available_bytes,
+			dovs.total_bytes total_bytes
+			FROM sys.master_files mf WITH (nolock)
+			CROSS apply sys.Dm_os_volume_stats(mf.database_id, mf.file_id) dovs
+			) drives`,
+		dataModels: &[]struct {
+			TotalDiskSpace *int64 `db:"total_disk_space" metric_name:"instance.diskInBytes" source_type:"gauge"`
+		}{},
+	},
+}
diff --git a/src/metrics/metrics.go b/src/metrics/metrics.go
index 0f9d2134..bfcd0190 100644
--- a/src/metrics/metrics.go
+++ b/src/metrics/metrics.go
@@ -48,6 +48,9 @@ func PopulateInstanceMetrics(instanceEntity *integration.Entity, connection *con
 	if arguments.EnableBufferMetrics {
 		collectionList = append(collectionList, instanceBufferDefinitions...)
 	}
+	if arguments.EnableDiskMetricsInBytes {
+		collectionList = append(collectionList, diskMetricInBytesDefination...)
+	}
 
 	for _, queryDef := range collectionList {
 		models := queryDef.GetDataModels()
diff --git a/src/metrics/metrics_test.go b/src/metrics/metrics_test.go
index 83138102..099af355 100644
--- a/src/metrics/metrics_test.go
+++ b/src/metrics/metrics_test.go
@@ -159,7 +159,8 @@ func Test_populateInstanceMetrics(t *testing.T) {
 	mock.ExpectClose()
 
 	args := args.ArgumentList{
-		EnableBufferMetrics: true,
+		EnableBufferMetrics:      true,
+		EnableDiskMetricsInBytes: true,
 	}
 	PopulateInstanceMetrics(e, conn, args)