mi: make timeout configurable for build functions (#2377)

internal/mi/session.go

@@ -22,6 +22,7 @@ import (
 	"fmt"
 	"reflect"
 	"syscall"
+	"time"
 	"unsafe"
 
 	"golang.org/x/sys/windows"
@@ -229,7 +230,7 @@ func (s *Session) QueryUnmarshal(dst any,
 	)
 
 	if result := ResultError(r0); !errors.Is(result, MI_RESULT_OK) {
-		return result
+		return fmt.Errorf("failed to query instances: %w", result)
 	}
 
 	defer func() {
@@ -308,11 +309,28 @@ func (s *Session) QueryUnmarshal(dst any,
 }
 
 // Query queries for a set of instances based on a query expression.
-func (s *Session) Query(dst any, namespaceName Namespace, queryExpression Query) error {
-	err := s.QueryUnmarshal(dst, OperationFlagsStandardRTTI, nil, namespaceName, QueryDialectWQL, queryExpression)
-	if err != nil {
-		return fmt.Errorf("WMI query failed: %w", err)
+//
+//nolint:nestif
+func (s *Session) Query(dst any, namespaceName Namespace, queryExpression Query, queryTimeout time.Duration) error {
+	var operationOptions *OperationOptions
+
+	if queryTimeout >= 0 {
+		app, err := s.GetApplication()
+		if err != nil {
+			return fmt.Errorf("failed to get application: %w", err)
+		}
+
+		operationOptions, err = app.NewOperationOptions()
+		if err != nil {
+			return fmt.Errorf("failed to create operation options: %w", err)
+		}
+
+		if queryTimeout > 0 {
+			if err = operationOptions.SetTimeout(queryTimeout); err != nil {
+				return fmt.Errorf("failed to set timeout: %w", err)
+			}
+		}
 	}
 
-	return nil
+	return s.QueryUnmarshal(dst, OperationFlagsStandardRTTI, operationOptions, namespaceName, QueryDialectWQL, queryExpression)
 }