Add other common functions

connection/connection.go

@@ -163,6 +163,7 @@ func LogGRPC(ctx context.Context, method string, req, reply interface{}, cc *grp
 	return err
 }
 
+// GetDriverName returns name of CSI driver.
 func GetDriverName(ctx context.Context, conn *grpc.ClientConn) (string, error) {
 	client := csi.NewIdentityClient(conn)
 
@@ -177,3 +178,68 @@ func GetDriverName(ctx context.Context, conn *grpc.ClientConn) (string, error) {
 	}
 	return name, nil
 }
+
+// PluginCapabilitySet is set of CSI plugin capabilities. Only supported capabilities are in the map.
+type PluginCapabilitySet map[csi.PluginCapability_Service_Type]bool
+
+// GetPluginCapabilities returns set of supported capabilities of CSI driver.
+func GetPluginCapabilities(ctx context.Context, conn *grpc.ClientConn) (PluginCapabilitySet, error) {
+	client := csi.NewIdentityClient(conn)
+	req := csi.GetPluginCapabilitiesRequest{}
+	rsp, err := client.GetPluginCapabilities(ctx, &req)
+	if err != nil {
+		return nil, err
+	}
+	caps := PluginCapabilitySet{}
+	for _, cap := range rsp.GetCapabilities() {
+		if cap == nil {
+			continue
+		}
+		srv := cap.GetService()
+		if srv == nil {
+			continue
+		}
+		t := srv.GetType()
+		caps[t] = true
+	}
+	return caps, nil
+}
+
+// ControllerCapabilitySet is set of CSI controller capabilities. Only supported capabilities are in the map.
+type ControllerCapabilitySet map[csi.ControllerServiceCapability_RPC_Type]bool
+
+// GetControllerCapabilities returns set of supported controller capabilities of CSI driver.
+func GetControllerCapabilities(ctx context.Context, conn *grpc.ClientConn) (ControllerCapabilitySet, error) {
+	client := csi.NewControllerClient(conn)
+	req := csi.ControllerGetCapabilitiesRequest{}
+	rsp, err := client.ControllerGetCapabilities(ctx, &req)
+	if err != nil {
+		return nil, err
+	}
+
+	caps := ControllerCapabilitySet{}
+	for _, cap := range rsp.GetCapabilities() {
+		if cap == nil {
+			continue
+		}
+		rpc := cap.GetRpc()
+		if rpc == nil {
+			continue
+		}
+		t := rpc.GetType()
+		caps[t] = true
+	}
+	return caps, nil
+}
+
+// Probe calls Probe() of a CSI driver and waits until the driver becomes ready.
+// TODO: wait until the driver is ready? It may require timeout instead of context as a parameter.
+func Probe(ctx context.Context, conn *grpc.ClientConn) error {
+	client := csi.NewIdentityClient(conn)
+	req := csi.ProbeRequest{}
+	_, err := client.Probe(ctx, &req)
+	if err != nil {
+		return err
+	}
+	return nil
+}