Add Associative support

command/encoding/json.go

@@ -27,6 +27,15 @@ type ResultRows struct {
 	Time         float64         `json:"time,omitempty"`
 }
 
+type AssociativeResultRows struct {
+	LastInsertID int64                    `json:"last_insert_id,omitempty"`
+	RowsAffected int64                    `json:"rows_affected,omitempty"`
+	Types        map[string]string        `json:"types,omitempty"`
+	Rows         []map[string]interface{} `json:"rows,omitempty"`
+	Error        string                   `json:"error,omitempty"`
+	Time         float64                  `json:"time,omitempty"`
+}
+
 // Result represents the outcome of an operation that changes rows.
 type Result struct {
 	LastInsertID int64   `json:"last_insert_id,omitempty"`
@@ -84,6 +93,51 @@ func NewResultRowsFromExecuteQueryResponse(e *command.ExecuteQueryResponse) (*Re
 	return nil, errors.New("no ExecuteResult or QueryRows")
 }
 
+func NewAssociativeResultRowsFromExecuteQueryResponse(e *command.ExecuteQueryResponse) (*AssociativeResultRows, error) {
+	er := e.GetE()
+	qr := e.GetQ()
+
+	if er != nil {
+		return &AssociativeResultRows{
+			LastInsertID: er.LastInsertId,
+			RowsAffected: er.RowsAffected,
+			Error:        er.Error,
+			Time:         er.Time,
+		}, nil
+	} else if qr != nil {
+		if len(qr.Columns) != len(qr.Types) {
+			return nil, ErrTypesColumnsLengthViolation
+		}
+		values := make([][]interface{}, len(qr.Values))
+		if err := NewValuesFromQueryValues(values, qr.Values); err != nil {
+			return nil, err
+		}
+
+		rows := make([]map[string]interface{}, len(values))
+		for i := range rows {
+			m := make(map[string]interface{})
+			for ii, c := range qr.Columns {
+				m[c] = values[i][ii]
+			}
+			rows[i] = m
+		}
+
+		types := make(map[string]string)
+		for i := range qr.Types {
+			types[qr.Columns[i]] = qr.Types[i]
+		}
+
+		return &AssociativeResultRows{
+			Types: types,
+			Rows:  rows,
+			Error: qr.Error,
+			Time:  qr.Time,
+		}, nil
+	}
+
+	return nil, errors.New("no ExecuteResult or QueryRows")
+}
+
 // NewResultFromExecuteResult returns an API Result object from an ExecuteResult.
 func NewResultFromExecuteResult(e *command.ExecuteResult) (*Result, error) {
 	return &Result{
@@ -284,15 +338,27 @@ func jsonMarshal(i interface{}, f marshalFunc, assoc bool) ([]byte, error) {
 			return f(rows)
 		}
 	case []*command.ExecuteQueryResponse:
-		res := make([]*ResultRows, len(v))
-		for j := range v {
-			r, err := NewResultRowsFromExecuteQueryResponse(v[j])
-			if err != nil {
-				return nil, err
+		if assoc {
+			res := make([]*AssociativeResultRows, len(v))
+			for j := range v {
+				r, err := NewAssociativeResultRowsFromExecuteQueryResponse(v[j])
+				if err != nil {
+					return nil, err
+				}
+				res[j] = r
+			}
+			return f(res)
+		} else {
+			res := make([]*ResultRows, len(v))
+			for j := range v {
+				r, err := NewResultRowsFromExecuteQueryResponse(v[j])
+				if err != nil {
+					return nil, err
+				}
+				res[j] = r
 			}
-			res[j] = r
+			return f(res)
 		}
-		return f(res)
 	case []*command.Values:
 		values := make([][]interface{}, len(v))
 		if err := NewValuesFromQueryValues(values, v); err != nil {

command/encoding/json_test.go

@@ -458,3 +458,118 @@ func Test_MarshalExecuteQueryResponse(t *testing.T) {
 		})
 	}
 }
+
+func Test_MarshalExecuteQueryAssociativeResponse(t *testing.T) {
+	enc := Encoder{
+		Associative: true,
+	}
+
+	tests := []struct {
+		name      string
+		responses []*command.ExecuteQueryResponse
+		expected  string
+	}{
+		{
+			name: "Test with ExecuteResult",
+			responses: []*command.ExecuteQueryResponse{
+				{
+					Result: &command.ExecuteQueryResponse_E{
+						E: &command.ExecuteResult{
+							LastInsertId: 123,
+							RowsAffected: 456,
+						},
+					},
+				},
+			},
+			expected: `[{"last_insert_id":123,"rows_affected":456}]`,
+		},
+		{
+			name: "Test with QueryRows",
+			responses: []*command.ExecuteQueryResponse{
+				{
+					Result: &command.ExecuteQueryResponse_Q{
+						Q: &command.QueryRows{
+							Columns: []string{"column1", "column2"},
+							Types:   []string{"type1", "type2"},
+							Values: []*command.Values{
+								{
+									Parameters: []*command.Parameter{
+										{
+											Value: &command.Parameter_I{
+												I: 123,
+											},
+										},
+										{
+											Value: &command.Parameter_S{
+												S: "fiona",
+											},
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+			expected: `[{"types":{"column1":"type1","column2":"type2"},"rows":[{"column1":123,"column2":"fiona"}]}]`,
+		},
+		{
+			name: "Test with ExecuteResult and QueryRows",
+			responses: []*command.ExecuteQueryResponse{
+				{
+					Result: &command.ExecuteQueryResponse_E{
+						E: &command.ExecuteResult{
+							LastInsertId: 123,
+							RowsAffected: 456,
+						},
+					},
+				},
+				{
+					Result: &command.ExecuteQueryResponse_E{
+						E: &command.ExecuteResult{
+							Error: "unique constraint failed",
+						},
+					},
+				},
+				{
+					Result: &command.ExecuteQueryResponse_Q{
+						Q: &command.QueryRows{
+							Columns: []string{"column1", "column2"},
+							Types:   []string{"int", "text"},
+							Values: []*command.Values{
+								{
+									Parameters: []*command.Parameter{
+										{
+											Value: &command.Parameter_I{
+												I: 456,
+											},
+										},
+										{
+											Value: &command.Parameter_S{
+												S: "declan",
+											},
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+			expected: `[{"last_insert_id":123,"rows_affected":456},{"error":"unique constraint failed"},{"types":{"column1":"int","column2":"text"},"rows":[{"column1":456,"column2":"declan"}]}]`,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			data, err := enc.JSONMarshal(tt.responses)
+			if err != nil {
+				t.Errorf("failed to marshal ExecuteQueryResponse: %v", err)
+			}
+
+			if string(data) != tt.expected {
+				t.Errorf("unexpected JSON output: got %v want %v", string(data), tt.expected)
+			}
+		})
+	}
+}