attempting experiment with new clickhouse config

Author: entlein

k8s/vizier/bootstrap/adaptive_export_deployment.yaml

@@ -57,6 +57,29 @@ spec:
           value: "10"
         - name: DETECTION_LOOKBACK_SEC
           value: "30"
+        # EXPORT_MODE controls the reconcile behaviour:
+        #   auto    - detection drives on/off (default)
+        #   always  - plugin always enabled (bypass detection)
+        #   never   - plugin always disabled and ch-* scripts purged
+        - name: EXPORT_MODE
+          value: "auto"
+        # Number of consecutive empty detection ticks before auto-disable fires.
+        - name: EXPORT_QUIET_TICKS
+          value: "6"
+        # Optional overrides for the ClickHouse PxL scripts. When unset they are
+        # parsed from CLICKHOUSE_DSN. Individual fields win over the parsed DSN.
+        - name: KUBESCAPE_TABLE
+          value: "kubescape_logs"
+        # - name: CLICKHOUSE_HOST
+        #   value: "clickhouse-forensic-soc-db.clickhouse.svc.cluster.local"
+        # - name: CLICKHOUSE_PORT
+        #   value: "9000"
+        # - name: CLICKHOUSE_USER
+        #   value: "otelcollector"
+        # - name: CLICKHOUSE_PASSWORD
+        #   value: "changeme"
+        # - name: CLICKHOUSE_DATABASE
+        #   value: "default"
         securityContext:
           allowPrivilegeEscalation: false
           capabilities:

src/vizier/services/adaptive_export/cmd/main.go

@@ -42,21 +42,20 @@ const (
 )
 
 const (
-	// TODO(ddelnano): Clickhouse configuration should come from plugin config.
-	schemaCreationScript = `
+	schemaCreationScriptTmpl = `
 import px
 px.display(px.CreateClickHouseSchemas(
-  host="hyperdx-hdx-oss-v2-clickhouse.click.svc.cluster.local",
-  port=9000,
-  username="otelcollector",
-  password="otelcollectorpass",
-  database="default"
+  host="%s",
+  port=%s,
+  username="%s",
+  password="%s",
+  database="%s"
 ))
 `
-	detectionScript = `
+	detectionScriptTmpl = `
 import px
 
-df = px.DataFrame('kubescape_logs', clickhouse_dsn='otelcollector:otelcollectorpass@hyperdx-hdx-oss-v2-clickhouse.click.svc.cluster.local:9000/default', start_time='-%ds')
+df = px.DataFrame('%s', clickhouse_dsn='%s', start_time='-%ds')
 df.alert = df.message
 df.namespace = px.pluck(df.RuntimeK8sDetails, "podNamespace")
 df.podName = px.pluck(df.RuntimeK8sDetails, "podName")
@@ -66,6 +65,15 @@ px.display(df)
 `
 )
 
+func renderSchemaScript(cfg config.ClickHouse) string {
+	return fmt.Sprintf(schemaCreationScriptTmpl,
+		cfg.Host(), cfg.Port(), cfg.User(), cfg.Password(), cfg.Database())
+}
+
+func renderDetectionScript(cfg config.ClickHouse, lookback int64) string {
+	return fmt.Sprintf(detectionScriptTmpl, cfg.Table(), cfg.DSN(), lookback)
+}
+
 func main() {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
@@ -95,9 +103,9 @@ func main() {
 	}
 
 	// Start schema creation background task
-	go runSchemaCreationTask(ctx, pxClient, clusterId)
+	go runSchemaCreationTask(ctx, pxClient, clusterId, cfg.ClickHouse())
 
-	// Start detection script that monitors for when to enable persistence
+	// Start detection + reconcile loop that turns the retention plugin on/off
 	go runDetectionTask(ctx, pxClient, pluginClient, cfg, clusterId, clusterName)
 
 	// Wait for signal to shutdown
@@ -110,85 +118,148 @@ func main() {
 	time.Sleep(1 * time.Second)
 }
 
-func runSchemaCreationTask(ctx context.Context, client *pxapi.Client, clusterID string) {
+func runSchemaCreationTask(ctx context.Context, client *pxapi.Client, clusterID string, chCfg config.ClickHouse) {
 	ticker := time.NewTicker(schemaCreationInterval)
 	defer ticker.Stop()
 
-	// Run immediately on startup
-	log.Info("Running schema creation script")
-	execCtx, cancel := context.WithTimeout(ctx, scriptExecutionTimeout)
-	if _, err := pxl.ExecuteScript(execCtx, client, clusterID, schemaCreationScript); err != nil {
-		log.WithError(err).Error("failed to execute schema creation script")
-	} else {
+	runOnce := func() {
+		log.Info("Running schema creation script")
+		execCtx, cancel := context.WithTimeout(ctx, scriptExecutionTimeout)
+		defer cancel()
+		if _, err := pxl.ExecuteScript(execCtx, client, clusterID, renderSchemaScript(chCfg)); err != nil {
+			log.WithError(err).Error("failed to execute schema creation script")
+			return
+		}
 		log.Info("Schema creation script completed successfully")
 	}
-	cancel()
 
+	runOnce()
 	for {
 		select {
 		case <-ctx.Done():
 			log.Info("Schema creation task shutting down")
 			return
 		case <-ticker.C:
-			log.Info("Running schema creation script")
-			execCtx, cancel := context.WithTimeout(ctx, scriptExecutionTimeout)
-			if _, err := pxl.ExecuteScript(execCtx, client, clusterID, schemaCreationScript); err != nil {
-				log.WithError(err).Error("failed to execute schema creation script")
-			} else {
-				log.Info("Schema creation script completed successfully")
-			}
-			cancel()
+			runOnce()
 		}
 	}
 }
 
 func runDetectionTask(ctx context.Context, pxClient *pxapi.Client, pluginClient *pixie.Client, cfg config.Config, clusterID string, clusterName string) {
 	detectionInterval := time.Duration(cfg.Worker().DetectionInterval()) * time.Second
 	detectionLookback := cfg.Worker().DetectionLookback()
+	quietTicks := cfg.Worker().ExportQuietTicks()
+	mode := cfg.Worker().ExportMode()
 
 	ticker := time.NewTicker(detectionInterval)
 	defer ticker.Stop()
 
-	pluginEnabled := false
+	// pluginEnabled tracks our last-known retention-plugin state. A nil value means
+	// we haven't reconciled yet; we always query on the first tick.
+	var pluginEnabled *bool
+	quietStreak := int64(0)
+
+	reconcile := func(want bool) {
+		if pluginEnabled != nil && *pluginEnabled == want {
+			log.Debugf("export already in desired state (enabled=%v), no action taken", want)
+			return
+		}
+		pluginCtx, pluginCancel := context.WithTimeout(ctx, 2*time.Minute)
+		defer pluginCancel()
+		if want {
+			log.Info("Enabling forensic export")
+			if err := enableClickHousePlugin(pluginCtx, pluginClient, cfg, clusterID, clusterName); err != nil {
+				log.WithError(err).Error("failed to enable forensic export")
+				return
+			}
+			v := true
+			pluginEnabled = &v
+			log.Info("Forensic export enabled successfully")
+		} else {
+			log.Info("Disabling forensic export")
+			if err := disableClickHousePlugin(pluginCtx, pluginClient, cfg, clusterID, clusterName); err != nil {
+				log.WithError(err).Error("failed to disable forensic export")
+				return
+			}
+			v := false
+			pluginEnabled = &v
+			quietStreak = 0
+			log.Info("Forensic export disabled successfully")
+		}
+	}
+
+	log.Infof("Detection task starting (mode=%s, quietTicks=%d)", mode, quietTicks)
 
 	for {
 		select {
 		case <-ctx.Done():
 			log.Info("Detection task shutting down")
 			return
 		case <-ticker.C:
-			log.Info("Running detection script")
-			// Run detection script with lookback period
-			detectionPxl := fmt.Sprintf(detectionScript, detectionLookback)
+			switch mode {
+			case config.ExportModeAlways:
+				reconcile(true)
+				continue
+			case config.ExportModeNever:
+				reconcile(false)
+				continue
+			}
+
+			// auto mode: detection drives the state.
+			log.Debug("Running detection script")
 			execCtx, cancel := context.WithTimeout(ctx, scriptExecutionTimeout)
-			recordCount, err := pxl.ExecuteScript(execCtx, pxClient, clusterID, detectionPxl)
+			recordCount, err := pxl.ExecuteScript(execCtx, pxClient, clusterID, renderDetectionScript(cfg.ClickHouse(), detectionLookback))
 			cancel()
-
 			if err != nil {
 				log.WithError(err).Error("failed to execute detection script")
 				continue
 			}
-
 			log.Debugf("Detection script returned %d records", recordCount)
 
-			// If we have records and plugin is not enabled, enable it
-			if recordCount > 0 && !pluginEnabled {
-				log.Info("Detection script returned records - enabling forensic export")
-				pluginCtx, pluginCancel := context.WithTimeout(ctx, 2*time.Minute)
-				if err := enableClickHousePlugin(pluginCtx, pluginClient, cfg, clusterID, clusterName); err != nil {
-					log.WithError(err).Error("failed to enable forensic export")
-				} else {
-					pluginEnabled = true
-					log.Info("Forensic export enabled successfully")
+			if recordCount > 0 {
+				quietStreak = 0
+				reconcile(true)
+			} else {
+				quietStreak++
+				if quietStreak >= quietTicks {
+					reconcile(false)
 				}
-				pluginCancel()
-			} else if recordCount > 0 && pluginEnabled {
-				log.Info("Detection script returned records but forensic export already enabled, no action taken")
 			}
 		}
 	}
 }
 
+func disableClickHousePlugin(ctx context.Context, client *pixie.Client, cfg config.Config, clusterID string, clusterName string) error {
+	plugin, err := client.GetClickHousePlugin()
+	if err != nil {
+		return fmt.Errorf("getting data retention plugins failed: %w", err)
+	}
+	if !plugin.RetentionEnabled {
+		log.Info("ClickHouse plugin already disabled; removing any lingering ch-* scripts")
+	} else {
+		if err := client.DisableClickHousePlugin(plugin.LatestVersion); err != nil {
+			return fmt.Errorf("failed to disable ClickHouse plugin: %w", err)
+		}
+	}
+
+	// Tear down the per-cluster ch-* retention scripts so the demo can be re-run cleanly.
+	current, err := client.GetClusterScripts(clusterID, clusterName)
+	if err != nil {
+		return fmt.Errorf("failed to list retention scripts: %w", err)
+	}
+	var errs []error
+	for _, s := range current {
+		log.Infof("Deleting retention script %s", s.Name)
+		if err := client.DeleteDataRetentionScript(s.ScriptId); err != nil {
+			errs = append(errs, err)
+		}
+	}
+	if len(errs) > 0 {
+		return fmt.Errorf("errors while deleting retention scripts: %v", errs)
+	}
+	return nil
+}
+
 func enableClickHousePlugin(ctx context.Context, client *pixie.Client, cfg config.Config, clusterID string, clusterName string) error {
 	log.Info("Checking the current ClickHouse plugin configuration")
 	plugin, err := client.GetClickHousePlugin()