Merge pull request #71 from gregfurman/70-sql-connection-check

sql: Add optional verification ping on database client initialisation
warpstreamlabs · Jul 17, 2024 · 16ec8d7 · 16ec8d7
2 parents 30c2891 + cd05311
commit 16ec8d7
Show file tree

Hide file tree

Showing 17 changed files with 132 additions and 13 deletions.
diff --git a/internal/impl/sql/cache_sql.go b/internal/impl/sql/cache_sql.go
@@ -161,7 +161,7 @@ func newSQLCacheFromConfig(conf *service.ParsedConfig, mgr *service.Resources) (
 		return nil, err
 	}
 
-	if s.db, err = sqlOpenWithReworks(s.logger, s.driver, s.dsn); err != nil {
+	if s.db, err = sqlOpenWithReworks(context.Background(), s.logger, s.driver, s.dsn, connSettings.initVerifyConn); err != nil {
 		return nil, err
 	}
 	connSettings.apply(context.Background(), s.db, s.logger)

diff --git a/internal/impl/sql/conn_fields.go b/internal/impl/sql/conn_fields.go
@@ -78,6 +78,12 @@ CREATE TABLE IF NOT EXISTS some_table (
 			Optional().
 			Advanced().
 			Version("1.0.0"),
+		service.NewBoolField("init_verify_conn").
+			Description("Whether to verify the database connection on startup by performing a simple ping, by default this is disabled.").
+			Default(false).
+			Optional().
+			Advanced().
+			Version("1.2.0"),
 		service.NewDurationField("conn_max_idle_time").
 			Description("An optional maximum amount of time a connection may be idle. Expired connections may be closed lazily before reuse. If `value <= 0`, connections are not closed due to a connections idle time.").
 			Optional().
@@ -125,6 +131,7 @@ type connSettings struct {
 	initOnce           sync.Once
 	initFileStatements [][2]string // (path,statement)
 	initStatement      string
+	initVerifyConn     bool
 }
 
 func (c *connSettings) apply(ctx context.Context, db *sql.DB, log *service.Logger) {
@@ -206,14 +213,21 @@ func connSettingsFromParsed(
 			})
 		}
 	}
+
+	if conf.Contains("init_verify_conn") {
+		if c.initVerifyConn, err = conf.FieldBool("init_verify_conn"); err != nil {
+			return
+		}
+	}
+
 	return
 }
 
-func sqlOpenWithReworks(logger *service.Logger, driver, dsn string) (*sql.DB, error) {
+func reworkDSN(driver, dsn string) (string, error) {
 	if driver == "clickhouse" && strings.HasPrefix(dsn, "tcp") {
 		u, err := url.Parse(dsn)
 		if err != nil {
-			return nil, err
+			return "", err
 		}
 
 		u.Scheme = "clickhouse"
@@ -235,8 +249,33 @@ func sqlOpenWithReworks(logger *service.Logger, driver, dsn string) (*sql.DB, er
 		u.RawQuery = uq.Encode()
 		newDSN := u.String()
 
-		logger.Warnf("Detected old-style Clickhouse Data Source Name: '%v', replacing with new style: '%v'", dsn, newDSN)
-		dsn = newDSN
+		return newDSN, nil
 	}
-	return sql.Open(driver, dsn)
+
+	return dsn, nil
+}
+
+func sqlOpenWithReworks(ctx context.Context, logger *service.Logger, driver, dsn string, shouldPing bool) (*sql.DB, error) {
+	updatedDSN, err := reworkDSN(driver, dsn)
+	if err != nil {
+		return nil, err
+	}
+
+	if updatedDSN != dsn {
+		logger.Warnf("Detected old-style Clickhouse Data Source Name: '%v', replacing with new style: '%v'", dsn, updatedDSN)
+	}
+
+	db, err := sql.Open(driver, updatedDSN)
+	if err != nil {
+		return nil, err
+	}
+
+	if shouldPing {
+		if err := db.PingContext(ctx); err != nil {
+			_ = db.Close()
+			return nil, fmt.Errorf("could not establish connection to database: %w", err)
+		}
+	}
+
+	return db, nil
 }
diff --git a/internal/impl/sql/input_sql_raw.go b/internal/impl/sql/input_sql_raw.go
@@ -132,7 +132,7 @@ func (s *sqlRawInput) Connect(ctx context.Context) (err error) {
 	}
 
 	var db *sql.DB
-	if db, err = sqlOpenWithReworks(s.logger, s.driver, s.dsn); err != nil {
+	if db, err = sqlOpenWithReworks(ctx, s.logger, s.driver, s.dsn, s.connSettings.initVerifyConn); err != nil {
 		return err
 	}
 	defer func() {

diff --git a/internal/impl/sql/input_sql_select.go b/internal/impl/sql/input_sql_select.go
@@ -184,7 +184,7 @@ func (s *sqlSelectInput) Connect(ctx context.Context) (err error) {
 	}
 
 	var db *sql.DB
-	if db, err = sqlOpenWithReworks(s.logger, s.driver, s.dsn); err != nil {
+	if db, err = sqlOpenWithReworks(ctx, s.logger, s.driver, s.dsn, s.connSettings.initVerifyConn); err != nil {
 		return
 	}
 	defer func() {

diff --git a/internal/impl/sql/output_sql_insert.go b/internal/impl/sql/output_sql_insert.go
@@ -192,7 +192,7 @@ func (s *sqlInsertOutput) Connect(ctx context.Context) error {
 	}
 
 	var err error
-	if s.db, err = sqlOpenWithReworks(s.logger, s.driver, s.dsn); err != nil {
+	if s.db, err = sqlOpenWithReworks(ctx, s.logger, s.driver, s.dsn, s.connSettings.initVerifyConn); err != nil {
 		return err
 	}
 

diff --git a/internal/impl/sql/output_sql_raw.go b/internal/impl/sql/output_sql_raw.go
@@ -166,7 +166,7 @@ func (s *sqlRawOutput) Connect(ctx context.Context) error {
 	}
 
 	var err error
-	if s.db, err = sqlOpenWithReworks(s.logger, s.driver, s.dsn); err != nil {
+	if s.db, err = sqlOpenWithReworks(ctx, s.logger, s.driver, s.dsn, s.connSettings.initVerifyConn); err != nil {
 		return err
 	}
 

diff --git a/internal/impl/sql/processor_sql_insert.go b/internal/impl/sql/processor_sql_insert.go
@@ -171,7 +171,7 @@ func NewSQLInsertProcessorFromConfig(conf *service.ParsedConfig, mgr *service.Re
 		return nil, err
 	}
 
-	if s.db, err = sqlOpenWithReworks(mgr.Logger(), driverStr, dsnStr); err != nil {
+	if s.db, err = sqlOpenWithReworks(context.Background(), mgr.Logger(), driverStr, dsnStr, connSettings.initVerifyConn); err != nil {
 		return nil, err
 	}
 

diff --git a/internal/impl/sql/processor_sql_raw.go b/internal/impl/sql/processor_sql_raw.go
@@ -167,7 +167,7 @@ func newSQLRawProcessor(
 	}
 
 	var err error
-	if s.db, err = sqlOpenWithReworks(logger, driverStr, dsnStr); err != nil {
+	if s.db, err = sqlOpenWithReworks(context.Background(), logger, driverStr, dsnStr, connSettings.initVerifyConn); err != nil {
 		return nil, err
 	}
 	connSettings.apply(context.Background(), s.db, s.logger)

diff --git a/internal/impl/sql/processor_sql_select.go b/internal/impl/sql/processor_sql_select.go
@@ -171,7 +171,7 @@ func NewSQLSelectProcessorFromConfig(conf *service.ParsedConfig, mgr *service.Re
 		return nil, err
 	}
 
-	if s.db, err = sqlOpenWithReworks(mgr.Logger(), driverStr, dsnStr); err != nil {
+	if s.db, err = sqlOpenWithReworks(context.Background(), mgr.Logger(), driverStr, dsnStr, connSettings.initVerifyConn); err != nil {
 		return nil, err
 	}
 	connSettings.apply(context.Background(), s.db, s.logger)

diff --git a/website/docs/components/caches/sql.md b/website/docs/components/caches/sql.md
@@ -63,6 +63,7 @@ sql:
       baz varchar(50),
       primary key (foo)
     ) WITHOUT ROWID;
+  init_verify_conn: false
   conn_max_idle_time: "" # No default (optional)
   conn_max_life_time: "" # No default (optional)
   conn_max_idle: 2
@@ -249,6 +250,15 @@ init_statement: |2
   ) WITHOUT ROWID;
 ```
 
+### `init_verify_conn`
+
+Whether to verify the database connection on startup by performing a simple ping, by default this is disabled.
+
+
+Type: `bool`  
+Default: `false`  
+Requires version 1.2.0 or newer  
+
 ### `conn_max_idle_time`
 
 An optional maximum amount of time a connection may be idle. Expired connections may be closed lazily before reuse. If `value <= 0`, connections are not closed due to a connections idle time.

diff --git a/website/docs/components/inputs/sql_raw.md b/website/docs/components/inputs/sql_raw.md
@@ -63,6 +63,7 @@ input:
         baz varchar(50),
         primary key (foo)
       ) WITHOUT ROWID;
+    init_verify_conn: false
     conn_max_idle_time: "" # No default (optional)
     conn_max_life_time: "" # No default (optional)
     conn_max_idle: 2
@@ -246,6 +247,15 @@ init_statement: |2
   ) WITHOUT ROWID;
 ```
 
+### `init_verify_conn`
+
+Whether to verify the database connection on startup by performing a simple ping, by default this is disabled.
+
+
+Type: `bool`  
+Default: `false`  
+Requires version 1.2.0 or newer  
+
 ### `conn_max_idle_time`
 
 An optional maximum amount of time a connection may be idle. Expired connections may be closed lazily before reuse. If `value <= 0`, connections are not closed due to a connections idle time.

diff --git a/website/docs/components/inputs/sql_select.md b/website/docs/components/inputs/sql_select.md
@@ -69,6 +69,7 @@ input:
         baz varchar(50),
         primary key (foo)
       ) WITHOUT ROWID;
+    init_verify_conn: false
     conn_max_idle_time: "" # No default (optional)
     conn_max_life_time: "" # No default (optional)
     conn_max_idle: 2
@@ -288,6 +289,15 @@ init_statement: |2
   ) WITHOUT ROWID;
 ```
 
+### `init_verify_conn`
+
+Whether to verify the database connection on startup by performing a simple ping, by default this is disabled.
+
+
+Type: `bool`  
+Default: `false`  
+Requires version 1.2.0 or newer  
+
 ### `conn_max_idle_time`
 
 An optional maximum amount of time a connection may be idle. Expired connections may be closed lazily before reuse. If `value <= 0`, connections are not closed due to a connections idle time.

diff --git a/website/docs/components/outputs/sql_insert.md b/website/docs/components/outputs/sql_insert.md
@@ -69,6 +69,7 @@ output:
         baz varchar(50),
         primary key (foo)
       ) WITHOUT ROWID;
+    init_verify_conn: false
     conn_max_idle_time: "" # No default (optional)
     conn_max_life_time: "" # No default (optional)
     conn_max_idle: 2
@@ -283,6 +284,15 @@ init_statement: |2
   ) WITHOUT ROWID;
 ```
 
+### `init_verify_conn`
+
+Whether to verify the database connection on startup by performing a simple ping, by default this is disabled.
+
+
+Type: `bool`  
+Default: `false`  
+Requires version 1.2.0 or newer  
+
 ### `conn_max_idle_time`
 
 An optional maximum amount of time a connection may be idle. Expired connections may be closed lazily before reuse. If `value <= 0`, connections are not closed due to a connections idle time.

diff --git a/website/docs/components/outputs/sql_raw.md b/website/docs/components/outputs/sql_raw.md
@@ -66,6 +66,7 @@ output:
         baz varchar(50),
         primary key (foo)
       ) WITHOUT ROWID;
+    init_verify_conn: false
     conn_max_idle_time: "" # No default (optional)
     conn_max_life_time: "" # No default (optional)
     conn_max_idle: 2
@@ -263,6 +264,15 @@ init_statement: |2
   ) WITHOUT ROWID;
 ```
 
+### `init_verify_conn`
+
+Whether to verify the database connection on startup by performing a simple ping, by default this is disabled.
+
+
+Type: `bool`  
+Default: `false`  
+Requires version 1.2.0 or newer  
+
 ### `conn_max_idle_time`
 
 An optional maximum amount of time a connection may be idle. Expired connections may be closed lazily before reuse. If `value <= 0`, connections are not closed due to a connections idle time.

diff --git a/website/docs/components/processors/sql_insert.md b/website/docs/components/processors/sql_insert.md
@@ -60,6 +60,7 @@ sql_insert:
       baz varchar(50),
       primary key (foo)
     ) WITHOUT ROWID;
+  init_verify_conn: false
   conn_max_idle_time: "" # No default (optional)
   conn_max_life_time: "" # No default (optional)
   conn_max_idle: 2
@@ -263,6 +264,15 @@ init_statement: |2
   ) WITHOUT ROWID;
 ```
 
+### `init_verify_conn`
+
+Whether to verify the database connection on startup by performing a simple ping, by default this is disabled.
+
+
+Type: `bool`  
+Default: `false`  
+Requires version 1.2.0 or newer  
+
 ### `conn_max_idle_time`
 
 An optional maximum amount of time a connection may be idle. Expired connections may be closed lazily before reuse. If `value <= 0`, connections are not closed due to a connections idle time.

diff --git a/website/docs/components/processors/sql_raw.md b/website/docs/components/processors/sql_raw.md
@@ -59,6 +59,7 @@ sql_raw:
       baz varchar(50),
       primary key (foo)
     ) WITHOUT ROWID;
+  init_verify_conn: false
   conn_max_idle_time: "" # No default (optional)
   conn_max_life_time: "" # No default (optional)
   conn_max_idle: 2
@@ -269,6 +270,15 @@ init_statement: |2
   ) WITHOUT ROWID;
 ```
 
+### `init_verify_conn`
+
+Whether to verify the database connection on startup by performing a simple ping, by default this is disabled.
+
+
+Type: `bool`  
+Default: `false`  
+Requires version 1.2.0 or newer  
+
 ### `conn_max_idle_time`
 
 An optional maximum amount of time a connection may be idle. Expired connections may be closed lazily before reuse. If `value <= 0`, connections are not closed due to a connections idle time.

diff --git a/website/docs/components/processors/sql_select.md b/website/docs/components/processors/sql_select.md
@@ -62,6 +62,7 @@ sql_select:
       baz varchar(50),
       primary key (foo)
     ) WITHOUT ROWID;
+  init_verify_conn: false
   conn_max_idle_time: "" # No default (optional)
   conn_max_life_time: "" # No default (optional)
   conn_max_idle: 2
@@ -279,6 +280,15 @@ init_statement: |2
   ) WITHOUT ROWID;
 ```
 
+### `init_verify_conn`
+
+Whether to verify the database connection on startup by performing a simple ping, by default this is disabled.
+
+
+Type: `bool`  
+Default: `false`  
+Requires version 1.2.0 or newer  
+
 ### `conn_max_idle_time`
 
 An optional maximum amount of time a connection may be idle. Expired connections may be closed lazily before reuse. If `value <= 0`, connections are not closed due to a connections idle time.