From 66f8840401aebe48c72310e5325d47c77a9fc055 Mon Sep 17 00:00:00 2001
From: souravTurbot <78197905+bigdatasourav@users.noreply.github.com>
Date: Tue, 23 Apr 2024 19:14:43 +0530
Subject: [PATCH] Add support for connection key columns (#740)

Co-authored-by: ParthaI <partha@turbot.com>
---
 azure/plugin.go | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/azure/plugin.go b/azure/plugin.go
index c29b8a8c..821c7478 100644
--- a/azure/plugin.go
+++ b/azure/plugin.go
@@ -23,6 +23,12 @@ func Plugin(ctx context.Context) *plugin.Plugin {
 		DefaultIgnoreConfig: &plugin.IgnoreConfig{
 			ShouldIgnoreErrorFunc: shouldIgnoreErrorPluginDefault(),
 		},
+		ConnectionKeyColumns: []plugin.ConnectionKeyColumn{
+			{
+				Name:    "subscription_id",
+				Hydrate: getSubscriptionIdForConnection,
+			},
+		},
 		ConnectionConfigSchema: &plugin.ConnectionConfigSchema{
 			NewInstance: ConfigInstance,
 		},
@@ -189,3 +195,12 @@ func Plugin(ctx context.Context) *plugin.Plugin {
 
 	return p
 }
+
+func getSubscriptionIdForConnection(ctx context.Context, d *plugin.QueryData, h *plugin.HydrateData) (any, error) {
+	subscriptionID, err := getSubscriptionID(ctx, d, h)
+	if err != nil {
+		return nil, err
+	}
+
+	return subscriptionID, nil
+}