Check if its pgconn.ConnectError{} instead of doing ping.

Author: 0xgouda

internal/sinks/postgres.go

@@ -438,8 +438,8 @@ func (pgw *PostgresWriter) flush(msgs []metrics.MeasurementEnvelope) {
 		rowsBatched += n
 		if err != nil {
 			logger.Error(err)
-			if err := pgw.sinkDb.Ping(pgw.ctx); err != nil {
-				logger.WithError(err).Error("Sink db is not reachable, dropping cached measurements")
+			if _, ok := err.(*pgconn.ConnectError); ok {
+				logger.Errorf("Sink DB not reachable, dropping %d cached measurements", len(msgs))
 				break
 			}
 			if PgError, ok := err.(*pgconn.PgError); ok {

internal/sinks/postgres_test.go

@@ -1083,19 +1083,25 @@ func TestEnsureMetricDbnameTime_IdempotentAcrossRestarts(t *testing.T) {
 }
 
 func TestFlush_CopyFromFailsAndPingSucceeds(t *testing.T) {
-	a := assert.New(t)
-	conn, err := pgxmock.NewPool()
-	a.NoError(err)
-	defer conn.Close()
-
-	pgw := &PostgresWriter{
-		ctx:                      ctx,
-		sinkDb:                   conn,
-		opts:                     &CmdOpts{PartitionInterval: "1 hour"},
-		partitionMapMetricDbname: make(map[string]map[string]ExistingPartitionInfo),
-		metricSchema:             DbStorageSchemaPostgres,
+	r := require.New(t)
+
+	pgContainer, pgTearDown, err := testutil.SetupPostgresContainer()
+	r.NoError(err)
+
+	connStr, err := pgContainer.ConnectionString(ctx, "sslmode=disable")
+	r.NoError(err)
+
+	opts := &CmdOpts{
+		PartitionInterval:   "1 day",
+		RetentionInterval:   "7 days",
+		MaintenanceInterval: "12 hours",
+		BatchingDelay:       time.Second,
 	}
 
+	pgw, err := NewPostgresWriter(ctx, connStr, opts)
+	r.NoError(err)
+	pgTearDown()
+
 	now := time.Now()
 	msgs := []metrics.MeasurementEnvelope{
 		{
@@ -1107,13 +1113,5 @@ func TestFlush_CopyFromFailsAndPingSucceeds(t *testing.T) {
 		},
 	}
 
-	conn.ExpectQuery("(?i)SELECT.*ensure_partition_metric_dbname_time").
-		WithArgs("test_metric", "test_db", pgxmock.AnyArg(), "1 hour").
-		WillReturnRows(pgxmock.NewRows([]string{"start_time", "end_time"}).AddRow(now.Add(-time.Hour), now.Add(time.Hour)))
-	conn.ExpectCopyFrom(pgx.Identifier{"test_metric"}, targetColumns[:]).WillReturnError(errors.New("copy failed"))
-	conn.ExpectPing().WillReturnError(nil)
-
 	pgw.flush(msgs)
-
-	a.NoError(conn.ExpectationsWereMet())
 }