Add support for connection key columns (#740)

Co-authored-by: ParthaI <[email protected]>
turbot · Apr 23, 2024 · 66f8840 · 66f8840
1 parent a330c2d
commit 66f8840
Showing 1 changed file with 15 additions and 0 deletions.
diff --git 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
+}