Add snapshot WebSocket API

Makefile

@@ -32,8 +32,13 @@ vendor:
 	modvendor -copy="**/*.c **/*.h **/*.proto" -v
 
 test: build
-	go test -coverprofile=coverage.out ./...
-	# go tool cover -html=coverage.out
+	go test -race -coverprofile=coverage.out ./...
+
+coverage: test
+	go tool cover -html=coverage.out
+
+run: build
+	go run -race .
 
 integration_test:
 	make -C integration_test

input/full.go

@@ -83,7 +83,7 @@ func CollectFull(ctx context.Context, server *state.Server, connection *sql.DB,
 	}
 
 	ps.StatementResetCounter = server.PrevState.StatementResetCounter + 1
-	if server.Grant.Config.Features.StatementResetFrequency != 0 && ps.StatementResetCounter >= server.Grant.Config.Features.StatementResetFrequency {
+	if server.Grant.Config.Features.StatementResetFrequency != 0 && ps.StatementResetCounter >= int(server.Grant.Config.Features.StatementResetFrequency) {
 		ps.StatementResetCounter = 0
 		err = postgres.ResetStatements(ctx, logger, connection, systemType)
 		if err != nil {

input/postgres/schema.go

@@ -81,7 +81,7 @@ func CollectAllSchemas(ctx context.Context, server *state.Server, collectionOpts
 		ts.DatabaseOidsWithLocalCatalog = append(ts.DatabaseOidsWithLocalCatalog, databaseOid)
 		server.SelfTest.MarkDbCollectionAspectOk(dbName, state.CollectionAspectSchema)
 	}
-	schemaTableLimit := server.Grant.Config.SchemaTableLimit
+	schemaTableLimit := int(server.Grant.Config.SchemaTableLimit)
 	if schemaTableLimit == 0 {
 		schemaTableLimit = defaultSchemaTableLimit
 	}

logs/querysample/tracing.go

@@ -25,7 +25,7 @@ func urlToSample(server *state.Server, grant state.GrantLogs, sample state.Postg
 
 	return fmt.Sprintf(
 		"%s/databases/%s/queries/%s/samples/%d?role=%s",
-		grant.Config.ServerURL,
+		grant.Config.ServerUrl,
 		sample.Database,
 		hex.EncodeToString(fpBin),
 		sample.OccurredAt.Unix(),

main.go

@@ -255,6 +255,8 @@ func run(ctx context.Context, wg *sync.WaitGroup, globalCollectionOpts state.Col
 		wg.Done()
 	}, logger, "high frequency query statistics of all servers", schedulerGroups["stats"])
 
+	runner.SetupWebsocketForAllServers(ctx, servers, globalCollectionOpts, logger)
+
 	keepRunning = true
 	return
 }

output/compact.go

@@ -58,6 +58,11 @@ func uploadAndSubmitCompactSnapshot(ctx context.Context, s pganalyze_collector.C
 		return nil
 	}
 
+	if server.WebSocket != nil {
+		server.SnapshotStream <- compressedData.Bytes()
+		return nil
+	}
+
 	s3Location, err := uploadSnapshot(ctx, server.Config.HTTPClientWithRetry, grant, logger, compressedData, snapshotUUID.String())
 	if err != nil {
 		logger.PrintError("Error uploading snapshot: %s", err)

output/full.go

@@ -82,6 +82,11 @@ func submitFull(ctx context.Context, s snapshot.FullSnapshot, server *state.Serv
 		return nil
 	}
 
+	if server.WebSocket != nil {
+		server.SnapshotStream <- compressedData.Bytes()
+		return nil
+	}
+
 	s3Location, err := uploadSnapshot(ctx, server.Config.HTTPClientWithRetry, server.Grant, logger, compressedData, snapshotUUID.String())
 	if err != nil {
 		logger.PrintError("Error uploading snapshot: %s", err)