fix: align Query mutations with DataHub v1.3+ GraphQL schema (#64)

The CreateQuery and UpdateQuery mutations used a fabricated
QuerySubjects wrapper type that doesn't exist in DataHub's schema.
This caused GraphQL validation errors at runtime:

  Field 'datasets' in type 'QuerySubject' is undefined

Three issues fixed:

1. Response selection: subjects { datasets { dataset { urn } } }
   changed to subjects { dataset { urn } } — QueryEntity.subjects
   returns [QuerySubject!] directly, not a wrapper.

2. Input format: {"subjects": {"datasets": [{datasetUrn: ...}]}}
   changed to {"subjects": [{datasetUrn: ...}]} — CreateQueryInput
   expects [CreateQuerySubjectInput!]!, a flat array.

3. subjects is now always sent (required field in schema) as an empty
   array when no DatasetURNs are provided, instead of being omitted.

Author: cjimti

pkg/client/queries.go

@@ -600,10 +600,8 @@ mutation createQuery($input: CreateQueryInput!) {
       }
     }
     subjects {
-      datasets {
-        dataset {
-          urn
-        }
+      dataset {
+        urn
       }
     }
   }

pkg/client/write_queries.go

@@ -63,9 +63,9 @@ type deleteQueryResponse struct {
 
 // queryEntityResponse is the nested query entity returned by mutations.
 type queryEntityResponse struct {
-	URN        string                 `json:"urn"`
-	Properties *queryPropertiesRaw    `json:"properties"`
-	Subjects   *querySubjectsRawOuter `json:"subjects"`
+	URN        string              `json:"urn"`
+	Properties *queryPropertiesRaw `json:"properties"`
+	Subjects   []querySubjectRaw   `json:"subjects"`
 }
 
 // queryPropertiesRaw maps the GraphQL QueryProperties type.
@@ -89,13 +89,8 @@ type auditStampGQL struct {
 	Actor string `json:"actor"`
 }
 
-// querySubjectsRawOuter maps the QuerySubjects type.
-type querySubjectsRawOuter struct {
-	Datasets []queryDatasetRef `json:"datasets"`
-}
-
-// queryDatasetRef maps QuerySubjects.datasets[].dataset.
-type queryDatasetRef struct {
+// querySubjectRaw maps the QuerySubject type (one per subject).
+type querySubjectRaw struct {
 	Dataset struct {
 		URN string `json:"urn"`
 	} `json:"dataset"`
@@ -123,15 +118,11 @@ func (c *Client) CreateQuery(ctx context.Context, input CreateQueryInput) (*type
 		},
 	}
 
-	if len(input.DatasetURNs) > 0 {
-		datasets := make([]map[string]string, len(input.DatasetURNs))
-		for i, urn := range input.DatasetURNs {
-			datasets[i] = map[string]string{"datasetUrn": urn}
-		}
-		gqlInput["subjects"] = map[string]any{
-			"datasets": datasets,
-		}
+	subjects := make([]map[string]string, len(input.DatasetURNs))
+	for i, urn := range input.DatasetURNs {
+		subjects[i] = map[string]string{"datasetUrn": urn}
 	}
+	gqlInput["subjects"] = subjects
 
 	variables := map[string]any{"input": gqlInput}
 
@@ -178,13 +169,11 @@ func (c *Client) UpdateQuery(ctx context.Context, input UpdateQueryInput) (*type
 	}
 
 	if input.DatasetURNs != nil {
-		datasets := make([]map[string]string, len(input.DatasetURNs))
+		subjects := make([]map[string]string, len(input.DatasetURNs))
 		for i, urn := range input.DatasetURNs {
-			datasets[i] = map[string]string{"datasetUrn": urn}
-		}
-		gqlInput["subjects"] = map[string]any{
-			"datasets": datasets,
+			subjects[i] = map[string]string{"datasetUrn": urn}
 		}
+		gqlInput["subjects"] = subjects
 	}
 
 	variables := map[string]any{

pkg/client/write_queries_test.go

@@ -81,21 +81,27 @@ func TestCreateQuery_WithDatasetURNs(t *testing.T) {
 			t.Fatalf("failed to decode request: %v", err)
 		}
 
-		// Verify subjects are passed through
+		// Verify subjects are passed as a flat array
 		inputRaw, ok := req.Variables["input"].(map[string]any)
 		if !ok {
 			t.Fatal("expected input variable")
 		}
-		subjects, ok := inputRaw["subjects"].(map[string]any)
+		subjects, ok := inputRaw["subjects"].([]any)
 		if !ok {
-			t.Fatal("expected subjects in input")
+			t.Fatalf("expected subjects as array in input, got %T", inputRaw["subjects"])
 		}
-		datasets, ok := subjects["datasets"].([]any)
-		if !ok {
-			t.Fatal("expected datasets in subjects")
+		if len(subjects) != 2 {
+			t.Errorf("expected 2 subjects, got %d", len(subjects))
 		}
-		if len(datasets) != 2 {
-			t.Errorf("expected 2 datasets, got %d", len(datasets))
+		// Verify each subject has datasetUrn
+		for i, s := range subjects {
+			subj, ok := s.(map[string]any)
+			if !ok {
+				t.Fatalf("subject[%d]: expected map, got %T", i, s)
+			}
+			if _, ok := subj["datasetUrn"]; !ok {
+				t.Errorf("subject[%d]: missing datasetUrn field", i)
+			}
 		}
 
 		resp := `{
@@ -270,16 +276,12 @@ func TestUpdateQuery_WithDatasetURNs(t *testing.T) {
 		if !ok {
 			t.Fatal("expected input variable")
 		}
-		subjects, ok := inputRaw["subjects"].(map[string]any)
-		if !ok {
-			t.Fatal("expected subjects in input")
-		}
-		datasets, ok := subjects["datasets"].([]any)
+		subjects, ok := inputRaw["subjects"].([]any)
 		if !ok {
-			t.Fatal("expected datasets in subjects")
+			t.Fatalf("expected subjects as array in input, got %T", inputRaw["subjects"])
 		}
-		if len(datasets) != 1 {
-			t.Errorf("expected 1 dataset, got %d", len(datasets))
+		if len(subjects) != 1 {
+			t.Errorf("expected 1 subject, got %d", len(subjects))
 		}
 
 		resp := `{"data":{"updateQuery":{"urn":"urn:li:query:abc123",` +