finally working

cmd/query.go

@@ -11,8 +11,10 @@ import (
 	"github.com/spf13/viper"
 	"github.com/thediveo/enumflag/v2"
 	"github.com/turbot/go-kit/helpers"
+	"github.com/turbot/pipe-fittings/cloud"
 	pconstants "github.com/turbot/pipe-fittings/constants"
 	"github.com/turbot/pipe-fittings/contexthelpers"
+	"github.com/turbot/pipe-fittings/steampipeconfig"
 	"github.com/turbot/pipe-fittings/utils"
 	"github.com/turbot/steampipe-plugin-sdk/v5/sperr"
 	"github.com/turbot/steampipe/pkg/cmdconfig"
@@ -201,3 +203,29 @@ func getPipedStdinData() string {
 	}
 	return stdinData
 }
+
+func publishSnapshotIfNeeded(ctx context.Context, snapshot *steampipeconfig.SteampipeSnapshot) error {
+	shouldShare := viper.GetBool(pconstants.ArgShare)
+	shouldUpload := viper.GetBool(pconstants.ArgSnapshot)
+
+	if !(shouldShare || shouldUpload) {
+		return nil
+	}
+
+	message, err := cloud.PublishSnapshot(ctx, snapshot, shouldShare)
+	if err != nil {
+		// reword "402 Payment Required" error
+		return handlePublishSnapshotError(err)
+	}
+	if viper.GetBool(pconstants.ArgProgress) {
+		fmt.Println(message)
+	}
+	return nil
+}
+
+func handlePublishSnapshotError(err error) error {
+	if err.Error() == "402 Payment Required" {
+		return fmt.Errorf("maximum number of snapshots reached")
+	}
+	return err
+}

pkg/export/snapshot_exporter.go

@@ -5,8 +5,8 @@ import (
 	"fmt"
 	"strings"
 
+	"github.com/turbot/pipe-fittings/steampipeconfig"
 	"github.com/turbot/steampipe/pkg/constants"
-	"github.com/turbot/steampipe/pkg/snapshot"
 )
 
 type SnapshotExporter struct {
@@ -14,7 +14,7 @@ type SnapshotExporter struct {
 }
 
 func (e *SnapshotExporter) Export(_ context.Context, input ExportSourceData, filePath string) error {
-	snapshot, ok := input.(*snapshot.SteampipeSnapshot)
+	snapshot, ok := input.(*steampipeconfig.SteampipeSnapshot)
 	if !ok {
 		return fmt.Errorf("SnapshotExporter input must be *dashboardtypes.SteampipeSnapshot")
 	}

pkg/query/queryexecute/execute.go

@@ -13,6 +13,7 @@ import (
 	"github.com/turbot/pipe-fittings/contexthelpers"
 	"github.com/turbot/pipe-fittings/modconfig"
 	"github.com/turbot/pipe-fittings/querydisplay"
+	"github.com/turbot/pipe-fittings/steampipeconfig"
 	"github.com/turbot/pipe-fittings/utils"
 	"github.com/turbot/steampipe/pkg/cmdconfig"
 	"github.com/turbot/steampipe/pkg/connection_sync"
@@ -108,7 +109,7 @@ func executeQuery(ctx context.Context, initData *query.InitData, resolvedQuery *
 	utils.LogTime("query.execute.executeQuery start")
 	defer utils.LogTime("query.execute.executeQuery end")
 
-	var snap *snapshot.SteampipeSnapshot
+	var snap *steampipeconfig.SteampipeSnapshot
 
 	// the db executor sends result data over resultsStreamer
 	resultsStreamer, err := db_common.ExecuteQuery(ctx, initData.Client, resolvedQuery.ExecuteSQL, resolvedQuery.Args...)

pkg/snapshot/snapshot.go

@@ -2,7 +2,6 @@ package snapshot
 
 import (
 	"context"
-	"encoding/json"
 	"fmt"
 	"strings"
 	"time"
@@ -12,75 +11,12 @@ import (
 	"github.com/turbot/pipe-fittings/querydisplay"
 	"github.com/turbot/pipe-fittings/queryresult"
 	pqueryresult "github.com/turbot/pipe-fittings/queryresult"
+	"github.com/turbot/pipe-fittings/steampipeconfig"
 	"github.com/turbot/pipe-fittings/utils"
-	steampipecloud "github.com/turbot/pipes-sdk-go"
 )
 
 const schemaVersion = "20221222"
 
-// SteampipeSnapshot struct definition
-type SteampipeSnapshot struct {
-	SchemaVersion string                 `json:"schema_version"`
-	Panels        map[string]PanelData   `json:"panels"`
-	Inputs        map[string]interface{} `json:"inputs"`
-	Variables     map[string]interface{} `json:"variables"`
-	SearchPath    []string               `json:"search_path"`
-	StartTime     string                 `json:"start_time"`
-	EndTime       string                 `json:"end_time"`
-	Layout        LayoutData             `json:"layout"`
-}
-
-func (s *SteampipeSnapshot) IsExportSourceData() {}
-
-func (s *SteampipeSnapshot) AsCloudSnapshot() (*steampipecloud.WorkspaceSnapshotData, error) {
-	jsonbytes, err := json.Marshal(s)
-	if err != nil {
-		return nil, err
-	}
-
-	res := &steampipecloud.WorkspaceSnapshotData{}
-	if err := json.Unmarshal(jsonbytes, res); err != nil {
-		return nil, err
-	}
-
-	return res, nil
-}
-
-func (s *SteampipeSnapshot) AsStrippedJson(indent bool) ([]byte, error) {
-	res, err := s.AsCloudSnapshot()
-	if err != nil {
-		return nil, err
-	}
-	if err = StripSnapshot(res); err != nil {
-		return nil, err
-	}
-	if indent {
-		return json.MarshalIndent(res, "", "  ")
-	}
-	return json.Marshal(res)
-}
-
-func StripSnapshot(snapshot *steampipecloud.WorkspaceSnapshotData) error {
-	propertiesToStrip := []string{
-		"sql",
-		"source_definition",
-		"documentation",
-		"search_path",
-		"search_path_prefix"}
-	for _, p := range snapshot.Panels {
-		panel := p.(map[string]any)
-		properties, _ := panel["properties"].(map[string]any)
-		for _, property := range propertiesToStrip {
-			// look both at top level and under properties
-			delete(panel, property)
-			if properties != nil {
-				delete(properties, property)
-			}
-		}
-	}
-	return nil
-}
-
 type PanelData struct {
 	Dashboard        string                 `json:"dashboard"`
 	Name             string                 `json:"name"`
@@ -93,6 +29,9 @@ type PanelData struct {
 	Data             map[string]interface{} `json:"data,omitempty"`
 }
 
+// IsSnapshotPanel implements SnapshotPanel
+func (*PanelData) IsSnapshotPanel() {}
+
 type LayoutData struct {
 	Name      string        `json:"name"`
 	Children  []LayoutChild `json:"children"` // Slice of LayoutChild structs
@@ -105,52 +44,67 @@ type LayoutChild struct {
 }
 
 // QueryResultToSnapshot function to generate a snapshot from a query result
-func QueryResultToSnapshot[T any](ctx context.Context, result *queryresult.Result[T], resolvedQuery *modconfig.ResolvedQuery, searchPath []string, startTime time.Time) (*SteampipeSnapshot, error) {
+func QueryResultToSnapshot[T any](ctx context.Context, result *queryresult.Result[T], resolvedQuery *modconfig.ResolvedQuery, searchPath []string, startTime time.Time) (*steampipeconfig.SteampipeSnapshot, error) {
 
 	endTime := time.Now()
+	hash, err := utils.Base36Hash(resolvedQuery.RawSQL, 8)
+	if err != nil {
+		return nil, err
+	}
+	dashboardName := fmt.Sprintf("custom.dashboard.sql_%s", hash)
 	// Build the snapshot data (use the new getData function to retrieve data)
-	snapshotData := &SteampipeSnapshot{
+	snapshotData := &steampipeconfig.SteampipeSnapshot{
 		SchemaVersion: schemaVersion,
-		Panels:        getPanels[T](ctx, result, resolvedQuery),
-		Inputs:        map[string]interface{}{},
-		Variables:     map[string]interface{}{},
-		SearchPath:    searchPath,
-		StartTime:     startTime.Format(time.RFC3339),
-		EndTime:       endTime.Format(time.RFC3339),
-		Layout:        getLayout[T](result, resolvedQuery),
+		Panels: map[string]steampipeconfig.SnapshotPanel{
+			dashboardName:          getPanelDashboard[T](ctx, result, resolvedQuery),
+			"custom.table.results": getPanelTable[T](ctx, result, resolvedQuery),
+		},
+		Inputs:     map[string]interface{}{},
+		Variables:  map[string]string{},
+		SearchPath: searchPath,
+		StartTime:  startTime,
+		EndTime:    endTime,
+		Layout:     getLayout[T](result, resolvedQuery),
 	}
 	// Return the snapshot data
 	return snapshotData, nil
 }
 
-func getPanels[T any](ctx context.Context, result *queryresult.Result[T], resolvedQuery *modconfig.ResolvedQuery) map[string]PanelData {
+func getPanelDashboard[T any](ctx context.Context, result *queryresult.Result[T], resolvedQuery *modconfig.ResolvedQuery) *PanelData {
 	hash, err := utils.Base36Hash(resolvedQuery.RawSQL, 8)
 	if err != nil {
-		return nil
+		return &PanelData{}
 	}
 	dashboardName := fmt.Sprintf("custom.dashboard.sql_%s", hash)
 	// Build panel data with proper fields
-	return map[string]PanelData{
-		dashboardName: {
-			Dashboard:        dashboardName,
-			Name:             dashboardName,
-			PanelType:        "dashboard",
-			SourceDefinition: "",
-			Status:           "complete",
-			Title:            fmt.Sprintf("Custom query [%s]", hash),
-		},
-		"custom.table.results": {
-			Dashboard:        dashboardName,
-			Name:             "custom.table.results",
-			PanelType:        "table",
-			SourceDefinition: "",
-			Status:           "complete",
-			SQL:              resolvedQuery.RawSQL,
-			Properties: map[string]string{
-				"name": "results",
-			},
-			Data: getData(ctx, result),
+	return &PanelData{
+		Dashboard:        dashboardName,
+		Name:             dashboardName,
+		PanelType:        "dashboard",
+		SourceDefinition: "",
+		Status:           "complete",
+		Title:            fmt.Sprintf("Custom query [%s]", hash),
+	}
+}
+
+func getPanelTable[T any](ctx context.Context, result *queryresult.Result[T], resolvedQuery *modconfig.ResolvedQuery) *PanelData {
+	hash, err := utils.Base36Hash(resolvedQuery.RawSQL, 8)
+	if err != nil {
+		return &PanelData{}
+	}
+	dashboardName := fmt.Sprintf("custom.dashboard.sql_%s", hash)
+	// Build panel data with proper fields
+	return &PanelData{
+		Dashboard:        dashboardName,
+		Name:             "custom.table.results",
+		PanelType:        "table",
+		SourceDefinition: "",
+		Status:           "complete",
+		SQL:              resolvedQuery.RawSQL,
+		Properties: map[string]string{
+			"name": "results",
 		},
+		Data: getData(ctx, result),
 	}
 }
 
@@ -190,21 +144,21 @@ func getData[T any](ctx context.Context, result *queryresult.Result[T]) map[stri
 	}
 }
 
-func getLayout[T any](result *queryresult.Result[T], resolvedQuery *modconfig.ResolvedQuery) LayoutData {
+func getLayout[T any](result *queryresult.Result[T], resolvedQuery *modconfig.ResolvedQuery) *steampipeconfig.SnapshotTreeNode {
 	hash, err := utils.Base36Hash(resolvedQuery.RawSQL, 8)
 	if err != nil {
-		return LayoutData{}
+		return nil
 	}
 	dashboardName := fmt.Sprintf("custom.dashboard.sql_%s", hash)
 	// Define layout structure
-	return LayoutData{
+	return &steampipeconfig.SnapshotTreeNode{
 		Name: dashboardName,
-		Children: []LayoutChild{
+		Children: []*steampipeconfig.SnapshotTreeNode{
 			{
-				Name:      "custom.table.results",
-				PanelType: "table",
+				Name:     "custom.table.results",
+				NodeType: "table",
 			},
 		},
-		PanelType: "dashboard",
+		NodeType: "dashboard",
 	}
 }