WIP.

artie-labs · Feb 8, 2025 · c6fdac3 · c6fdac3
1 parent 740690b
commit c6fdac3
Showing 1 changed file with 20 additions and 16 deletions.
diff --git a/lib/apachelivy/types.go b/lib/apachelivy/types.go
@@ -48,30 +48,34 @@ type CreateStatementRequest struct {
 	Kind string `json:"kind"`
 }
 
-type CreateStatementResponse struct {
-	ID     int             `json:"id"`
-	State  string          `json:"state"`
-	Output StatementOutput `json:"output"`
+type CreateStatementResponse[T any] struct {
+	ID     int                `json:"id"`
+	State  string             `json:"state"`
+	Output StatementOutput[T] `json:"output"`
 }
 
-type StatementOutput struct {
-	Status    string         `json:"status"`
-	Data      map[string]any `json:"data,omitempty"`
-	EType     string         `json:"etype,omitempty"`
-	EValue    string         `json:"evalue,omitempty"`
-	TraceBack []string       `json:"traceback"`
+type JSONDataWrapper[T any] struct {
+	Data map[string]T `json:"application/json"`
 }
 
-type GetStatementResponse struct {
+type StatementOutput[T any] struct {
+	Status    string             `json:"status"`
+	Data      JSONDataWrapper[T] `json:"data,omitempty"`
+	EType     string             `json:"etype,omitempty"`
+	EValue    string             `json:"evalue,omitempty"`
+	TraceBack []string           `json:"traceback"`
+}
+
+type GetStatementResponse[T any] struct {
 	ID        int `json:"id"`
 	Code      string
-	State     string          `json:"state"`
-	Output    StatementOutput `json:"output"`
-	Started   int             `json:"started"`
-	Completed int             `json:"completed"`
+	State     string             `json:"state"`
+	Output    StatementOutput[T] `json:"output"`
+	Started   int                `json:"started"`
+	Completed int                `json:"completed"`
 }
 
-func (a GetStatementResponse) Error(sessionID int) error {
+func (a GetStatementResponse[T]) Error(sessionID int) error {
 	if a.Output.Status == "error" {
 		return fmt.Errorf("%s, stacktrace: %s for session %d, statement %d", a.Output.EValue, strings.Join(a.Output.TraceBack, "\n"), sessionID, a.ID)
 	}