fix: handle null getAspect responses, add UpdateColumnDescription (#58)

- getAspect() now returns ErrNotFound when DataHub returns 200 OK with
  a null or empty aspect value, preventing json.Unmarshal crashes on
  entities where the aspect has never been written
- Add UpdateColumnDescription() for read-modify-write on the
  editableSchemaMetadata aspect, enabling column-level description edits
- Add #nosec G704 annotations on httpClient.Do() calls where URLs are
  constructed from configured endpoints

Closes #57

Author: cjimti

pkg/client/client.go

@@ -175,7 +175,7 @@ func (c *Client) doRequest(ctx context.Context, jsonBody []byte, result any) err
 	req.Header.Set("Content-Type", "application/json")
 	req.Header.Set("Authorization", "Bearer "+c.token)
 
-	resp, err := c.httpClient.Do(req)
+	resp, err := c.httpClient.Do(req) //#nosec G704 -- URL is constructed from configured endpoint, not arbitrary user input
 	if err != nil {
 		return c.handleRequestError(ctx, err)
 	}

pkg/client/rest.go

@@ -58,7 +58,7 @@ func (c *Client) getAspect(ctx context.Context, entityURN, aspectName string) (j
 
 	c.setRESTHeaders(req)
 
-	resp, err := c.httpClient.Do(req)
+	resp, err := c.httpClient.Do(req) //#nosec G704 -- URL is constructed from configured endpoint, not arbitrary user input
 	if err != nil {
 		return nil, fmt.Errorf("REST GET failed: %w", err)
 	}
@@ -84,6 +84,13 @@ func (c *Client) getAspect(ctx context.Context, entityURN, aspectName string) (j
 		return nil, fmt.Errorf("failed to unmarshal aspect response: %w", err)
 	}
 
+	// DataHub may return 200 OK with a null or empty value when the entity
+	// exists but the requested aspect has never been written. Treat this the
+	// same as 404 so callers initialize a default struct.
+	if isNullOrEmptyJSON(aspectResp.Value) {
+		return nil, ErrNotFound
+	}
+
 	return aspectResp.Value, nil
 }
 
@@ -125,7 +132,7 @@ func (c *Client) postIngestProposal(ctx context.Context, proposal ingestProposal
 	c.setRESTHeaders(req)
 	req.Header.Set("Content-Type", "application/json")
 
-	resp, err := c.httpClient.Do(req)
+	resp, err := c.httpClient.Do(req) //#nosec G704 -- URL is constructed from configured endpoint, not arbitrary user input
 	if err != nil {
 		return fmt.Errorf("REST POST failed: %w", err)
 	}
@@ -151,6 +158,17 @@ func (c *Client) setRESTHeaders(req *http.Request) {
 	req.Header.Set("X-RestLi-Protocol-Version", "2.0.0")
 }
 
+// isNullOrEmptyJSON returns true if the raw JSON message is nil, empty,
+// or the JSON literal "null". This covers cases where DataHub returns
+// 200 OK but the aspect value is absent.
+func isNullOrEmptyJSON(raw json.RawMessage) bool {
+	if len(raw) == 0 {
+		return true
+	}
+	trimmed := bytes.TrimSpace(raw)
+	return len(trimmed) == 0 || string(trimmed) == "null"
+}
+
 // checkRESTStatus validates REST API response status codes.
 func (c *Client) checkRESTStatus(statusCode int, body []byte) error {
 	switch statusCode {

pkg/client/rest_test.go

@@ -158,6 +158,63 @@ func TestGetAspect_NotFound(t *testing.T) {
 	}
 }
 
+func TestGetAspect_NullValue(t *testing.T) {
+	tests := []struct {
+		name     string
+		response string
+	}{
+		{"null value", `{"value":null}`},
+		{"empty object no value", `{}`},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+				w.Header().Set("Content-Type", "application/json")
+				_, _ = w.Write([]byte(tt.response))
+			}))
+			defer server.Close()
+
+			c := &Client{
+				endpoint:   server.URL + "/api/graphql",
+				token:      "test-token",
+				httpClient: server.Client(),
+				logger:     NopLogger{},
+			}
+
+			_, err := c.getAspect(context.Background(), "urn:li:dataset:test", "editableDatasetProperties")
+			if err != ErrNotFound {
+				t.Errorf("expected ErrNotFound, got: %v", err)
+			}
+		})
+	}
+}
+
+func TestIsNullOrEmptyJSON(t *testing.T) {
+	tests := []struct {
+		name string
+		raw  json.RawMessage
+		want bool
+	}{
+		{"nil", nil, true},
+		{"empty bytes", json.RawMessage{}, true},
+		{"null literal", json.RawMessage(`null`), true},
+		{"null with spaces", json.RawMessage(`  null  `), true},
+		{"valid object", json.RawMessage(`{"key":"val"}`), false},
+		{"valid array", json.RawMessage(`[]`), false},
+		{"valid string", json.RawMessage(`"hello"`), false},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := isNullOrEmptyJSON(tt.raw)
+			if got != tt.want {
+				t.Errorf("isNullOrEmptyJSON(%q) = %v, want %v", string(tt.raw), got, tt.want)
+			}
+		})
+	}
+}
+
 func TestGetAspect_Unauthorized(t *testing.T) {
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
 		w.WriteHeader(http.StatusUnauthorized)

pkg/client/write.go

@@ -17,6 +17,20 @@ func entityTypeFromURN(urn string) (string, error) {
 	return parsed.EntityType, nil
 }
 
+// editableSchemaAspect is the REST API representation of editableSchemaMetadata.
+type editableSchemaAspect struct {
+	EditableSchemaFieldInfo []editableFieldInfo `json:"editableSchemaFieldInfo"`
+}
+
+// editableFieldInfo represents a field's editable metadata for REST API read-modify-write.
+// Uses json.RawMessage for tags and glossaryTerms to preserve existing data.
+type editableFieldInfo struct {
+	FieldPath     string          `json:"fieldPath"`
+	Description   string          `json:"description,omitempty"`
+	GlobalTags    json.RawMessage `json:"globalTags,omitempty"`
+	GlossaryTerms json.RawMessage `json:"glossaryTerms,omitempty"`
+}
+
 // editablePropertiesAspect represents the editableDatasetProperties aspect.
 type editablePropertiesAspect struct {
 	Description  string         `json:"description"`
@@ -336,3 +350,58 @@ func (c *Client) readInstitutionalMemory(ctx context.Context, urn string) (*inst
 	}
 	return &memory, nil
 }
+
+// UpdateColumnDescription sets the editable description for a specific column
+// using read-modify-write on the editableSchemaMetadata aspect.
+func (c *Client) UpdateColumnDescription(ctx context.Context, urn, fieldPath, description string) error {
+	entityType, err := entityTypeFromURN(urn)
+	if err != nil {
+		return fmt.Errorf("UpdateColumnDescription: %w", err)
+	}
+
+	schema, err := c.readEditableSchema(ctx, urn)
+	if err != nil {
+		return fmt.Errorf("UpdateColumnDescription: %w", err)
+	}
+
+	// Find or create the field entry
+	found := false
+	for i := range schema.EditableSchemaFieldInfo {
+		if schema.EditableSchemaFieldInfo[i].FieldPath == fieldPath {
+			schema.EditableSchemaFieldInfo[i].Description = description
+			found = true
+			break
+		}
+	}
+	if !found {
+		schema.EditableSchemaFieldInfo = append(schema.EditableSchemaFieldInfo, editableFieldInfo{
+			FieldPath:   fieldPath,
+			Description: description,
+		})
+	}
+
+	return c.postIngestProposal(ctx, ingestProposal{
+		EntityType: entityType,
+		EntityURN:  urn,
+		AspectName: "editableSchemaMetadata",
+		Aspect:     schema,
+	})
+}
+
+// readEditableSchema reads the current editableSchemaMetadata aspect.
+// Returns an empty aspect if none exists (not an error).
+func (c *Client) readEditableSchema(ctx context.Context, urn string) (*editableSchemaAspect, error) {
+	raw, err := c.getAspect(ctx, urn, "editableSchemaMetadata")
+	if err != nil {
+		if errors.Is(err, ErrNotFound) {
+			return &editableSchemaAspect{}, nil
+		}
+		return nil, fmt.Errorf("reading editableSchemaMetadata: %w", err)
+	}
+
+	var schema editableSchemaAspect
+	if err := json.Unmarshal(raw, &schema); err != nil {
+		return nil, fmt.Errorf("parsing editableSchemaMetadata: %w", err)
+	}
+	return &schema, nil
+}