Better peer checks (#572)

Fixes #482 Fixes #484 Also checks for dataset existence for BQ Peer. (This will also be applicable to UI)
PeerDB-io · Dec 1, 2023 · 27fe4dc · 27fe4dc
1 parent 40117c5
commit 27fe4dc
Show file tree

Hide file tree

Showing 20 changed files with 596 additions and 509 deletions.
diff --git a/flow/cmd/handler.go b/flow/cmd/handler.go
@@ -503,12 +503,12 @@ func (h *FlowRequestHandler) ValidatePeer(
 		}, nil
 	}
 
-	status := conn.ConnectionActive()
-	if !status {
+	connErr := conn.ConnectionActive()
+	if connErr != nil {
 		return &protos.ValidatePeerResponse{
 			Status: protos.ValidatePeerStatus_INVALID,
-			Message: fmt.Sprintf("failed to establish active connection to %s peer %s.",
-				req.Peer.Type, req.Peer.Name),
+			Message: fmt.Sprintf("failed to establish active connection to %s peer %s: %v",
+				req.Peer.Type, req.Peer.Name, connErr),
 		}, nil
 	}
 

diff --git a/flow/connectors/bigquery/bigquery.go b/flow/connectors/bigquery/bigquery.go
@@ -204,8 +204,16 @@ func (c *BigQueryConnector) Close() error {
 }
 
 // ConnectionActive returns true if the connection is active.
-func (c *BigQueryConnector) ConnectionActive() bool {
-	return c.client != nil
+func (c *BigQueryConnector) ConnectionActive() error {
+	_, err := c.client.Dataset(c.datasetID).Metadata(c.ctx)
+	if err != nil {
+		return fmt.Errorf("failed to get dataset metadata: %v", err)
+	}
+
+	if c.client == nil {
+		return fmt.Errorf("BigQuery client is nil")
+	}
+	return nil
 }
 
 // NeedsSetupMetadataTables returns true if the metadata tables need to be set up.

diff --git a/flow/connectors/core.go b/flow/connectors/core.go
@@ -21,7 +21,7 @@ var ErrUnsupportedFunctionality = errors.New("requested connector does not suppo
 
 type Connector interface {
 	Close() error
-	ConnectionActive() bool
+	ConnectionActive() error
 }
 
 type CDCPullConnector interface {
@@ -223,7 +223,6 @@ func GetConnector(ctx context.Context, peer *protos.Peer) (Connector, error) {
 			return nil, fmt.Errorf("missing snowflake config for %s peer %s", peer.Type.String(), peer.Name)
 		}
 		return connsnowflake.NewSnowflakeConnector(ctx, sfConfig)
-
 	case protos.DBType_SQLSERVER:
 		sqlServerConfig := peer.GetSqlserverConfig()
 		if sqlServerConfig == nil {

diff --git a/flow/connectors/eventhub/eventhub.go b/flow/connectors/eventhub/eventhub.go
@@ -65,8 +65,8 @@ func (c *EventHubConnector) Close() error {
 	return allErrors
 }
 
-func (c *EventHubConnector) ConnectionActive() bool {
-	return true
+func (c *EventHubConnector) ConnectionActive() error {
+	return nil
 }
 
 func (c *EventHubConnector) InitializeTableSchema(req map[string]*protos.TableSchema) error {

diff --git a/flow/connectors/postgres/postgres.go b/flow/connectors/postgres/postgres.go
@@ -104,11 +104,12 @@ func (c *PostgresConnector) Close() error {
 }
 
 // ConnectionActive returns true if the connection is active.
-func (c *PostgresConnector) ConnectionActive() bool {
+func (c *PostgresConnector) ConnectionActive() error {
 	if c.pool == nil {
-		return false
+		return fmt.Errorf("connection pool is nil")
 	}
-	return c.pool.Ping(c.ctx) == nil
+	pingErr := c.pool.Ping(c.ctx)
+	return pingErr
 }
 
 // NeedsSetupMetadataTables returns true if the metadata tables need to be set up.

diff --git a/flow/connectors/postgres/postgres_repl_test.go b/flow/connectors/postgres/postgres_repl_test.go
@@ -94,12 +94,12 @@ func (suite *PostgresReplicationSnapshotTestSuite) TearDownSuite() {
 	err = teardownTx.Commit(context.Background())
 	require.NoError(suite.T(), err)
 
-	suite.True(suite.connector.ConnectionActive())
+	suite.True(suite.connector.ConnectionActive() == nil)
 
 	err = suite.connector.Close()
 	require.NoError(suite.T(), err)
 
-	suite.False(suite.connector.ConnectionActive())
+	suite.False(suite.connector.ConnectionActive() == nil)
 }
 
 func (suite *PostgresReplicationSnapshotTestSuite) TestSimpleSlotCreation() {

diff --git a/flow/connectors/postgres/postgres_schema_delta_test.go b/flow/connectors/postgres/postgres_schema_delta_test.go
@@ -67,10 +67,10 @@ func (suite *PostgresSchemaDeltaTestSuite) TearDownSuite() {
 	err = teardownTx.Commit(context.Background())
 	suite.failTestError(err)
 
-	suite.True(suite.connector.ConnectionActive())
+	suite.True(suite.connector.ConnectionActive() == nil)
 	err = suite.connector.Close()
 	suite.failTestError(err)
-	suite.False(suite.connector.ConnectionActive())
+	suite.False(suite.connector.ConnectionActive() == nil)
 }
 
 func (suite *PostgresSchemaDeltaTestSuite) TestSimpleAddColumn() {

diff --git a/flow/connectors/s3/s3.go b/flow/connectors/s3/s3.go
@@ -68,12 +68,6 @@ func NewS3Connector(ctx context.Context,
 		return nil, err
 	}
 
-	validErr := ValidCheck(s3Client, config.Url, pgMetadata)
-	if validErr != nil {
-		log.Errorf("failed to validate s3 connector: %v", validErr)
-		return nil, validErr
-	}
-
 	return &S3Connector{
 		ctx:        ctx,
 		url:        config.Url,
@@ -140,9 +134,19 @@ func ValidCheck(s3Client *s3.S3, bucketURL string, metadataDB *metadataStore.Pos
 	return nil
 }
 
-func (c *S3Connector) ConnectionActive() bool {
-	_, err := c.client.ListBuckets(nil)
-	return err == nil
+func (c *S3Connector) ConnectionActive() error {
+	_, listErr := c.client.ListBuckets(nil)
+	if listErr != nil {
+		return listErr
+	}
+
+	validErr := ValidCheck(&c.client, c.url, c.pgMetadata)
+	if validErr != nil {
+		log.Errorf("failed to validate s3 connector: %v", validErr)
+		return validErr
+	}
+
+	return nil
 }
 
 func (c *S3Connector) NeedsSetupMetadataTables() bool {

diff --git a/flow/connectors/snowflake/snowflake.go b/flow/connectors/snowflake/snowflake.go
@@ -168,11 +168,14 @@ func (c *SnowflakeConnector) Close() error {
 	return nil
 }
 
-func (c *SnowflakeConnector) ConnectionActive() bool {
+func (c *SnowflakeConnector) ConnectionActive() error {
 	if c == nil || c.database == nil {
-		return false
+		return fmt.Errorf("SnowflakeConnector is nil")
 	}
-	return c.database.PingContext(c.ctx) == nil
+
+	// This also checks if database exists
+	err := c.database.PingContext(c.ctx)
+	return err
 }
 
 func (c *SnowflakeConnector) NeedsSetupMetadataTables() bool {

diff --git a/flow/connectors/sql/query_executor.go b/flow/connectors/sql/query_executor.go
@@ -17,7 +17,7 @@ import (
 )
 
 type SQLQueryExecutor interface {
-	ConnectionActive() bool
+	ConnectionActive() error
 	Close() error
 
 	CreateSchema(schemaName string) error

diff --git a/flow/connectors/sqlserver/sqlserver.go b/flow/connectors/sqlserver/sqlserver.go
@@ -53,9 +53,9 @@ func (c *SQLServerConnector) Close() error {
 }
 
 // ConnectionActive checks if the connection is still active
-func (c *SQLServerConnector) ConnectionActive() bool {
+func (c *SQLServerConnector) ConnectionActive() error {
 	if err := c.db.Ping(); err != nil {
-		return false
+		return err
 	}
-	return true
+	return nil
 }
diff --git a/flow/e2e/sqlserver/sqlserver_helper.go b/flow/e2e/sqlserver/sqlserver_helper.go
@@ -41,9 +41,9 @@ func NewSQLServerHelper(name string) (*SQLServerHelper, error) {
 		return nil, err
 	}
 
-	testConn := connector.ConnectionActive()
-	if !testConn {
-		return nil, fmt.Errorf("invalid connection configs")
+	connErr := connector.ConnectionActive()
+	if connErr != nil {
+		return nil, fmt.Errorf("invalid connection configs: %v", connErr)
 	}
 
 	rndNum, err := util.RandomUInt64()