mongo: move schema logic to shared logic (#3357)

Moved ListDatabases and ListSchema to shared logic for reuse.

Test: validated logic locally in peerdb, system dbs/tables and views are not included.

Author: jgao54

flow/connectors/mongo/schema.go

@@ -3,24 +3,22 @@ package connmongo
 import (
 	"context"
 	"fmt"
-	"slices"
-
-	"go.mongodb.org/mongo-driver/v2/bson"
 
 	"github.com/PeerDB-io/peerdb/flow/generated/protos"
+	shared_mongo "github.com/PeerDB-io/peerdb/flow/shared/mongo"
 )
 
 func (c *MongoConnector) GetAllTables(ctx context.Context) (*protos.AllTablesResponse, error) {
 	tableNames := make([]string, 0)
 
-	dbNames, err := c.getAllDatabaseNames(ctx)
+	dbNames, err := shared_mongo.GetDatabaseNames(ctx, c.client)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get databases: %w", err)
 	}
 	for _, dbName := range dbNames {
-		collNames, err := c.getCollectionNames(ctx, dbName)
+		collNames, err := shared_mongo.GetCollectionNames(ctx, c.client, dbName)
 		if err != nil {
-			return nil, err
+			return nil, fmt.Errorf("failed to get collections: %w", err)
 		}
 		for _, collName := range collNames {
 			tableNames = append(tableNames, fmt.Sprintf("%s.%s", dbName, collName))
@@ -32,7 +30,7 @@ func (c *MongoConnector) GetAllTables(ctx context.Context) (*protos.AllTablesRes
 }
 
 func (c *MongoConnector) GetSchemas(ctx context.Context) (*protos.PeerSchemasResponse, error) {
-	dbNames, err := c.getAllDatabaseNames(ctx)
+	dbNames, err := shared_mongo.GetDatabaseNames(ctx, c.client)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get databases: %w", err)
 	}
@@ -42,9 +40,9 @@ func (c *MongoConnector) GetSchemas(ctx context.Context) (*protos.PeerSchemasRes
 }
 
 func (c *MongoConnector) GetTablesInSchema(ctx context.Context, schema string, cdcEnabled bool) (*protos.SchemaTablesResponse, error) {
-	collectionNames, err := c.getCollectionNames(ctx, schema)
+	collectionNames, err := shared_mongo.GetCollectionNames(ctx, c.client, schema)
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("failed to get collections: %w", err)
 	}
 
 	response := protos.SchemaTablesResponse{
@@ -69,39 +67,3 @@ func (c *MongoConnector) GetColumns(ctx context.Context, version uint32, schema
 		Columns: []*protos.ColumnsItem{},
 	}, nil
 }
-
-func (c *MongoConnector) getCollectionNames(ctx context.Context, databaseName string) ([]string, error) {
-	collectionNames, err := c.client.Database(databaseName).ListCollectionNames(ctx, bson.M{
-		"name": bson.M{
-			"$not": bson.Regex{
-				Pattern: "^system\\.",
-			},
-		},
-		"type": bson.M{
-			"$ne": "view",
-		},
-	})
-	if err != nil {
-		return nil, fmt.Errorf("failed to get collections: %w", err)
-	}
-	slices.Sort(collectionNames)
-	return collectionNames, nil
-}
-
-// Get all database names, but excluding MongoDB's default databases
-func (c *MongoConnector) getAllDatabaseNames(ctx context.Context) ([]string, error) {
-	dbs, err := c.client.ListDatabaseNames(ctx, bson.M{
-		"name": bson.M{
-			"$not": bson.Regex{
-				Pattern: "^(admin|local|config)$",
-			},
-		},
-	})
-	if err != nil {
-		return nil, err
-	}
-	filteredDbNames := make([]string, 0, len(dbs))
-	filteredDbNames = append(filteredDbNames, dbs...)
-	slices.Sort(filteredDbNames)
-	return filteredDbNames, nil
-}

flow/shared/mongo/schema.go

@@ -0,0 +1,67 @@
+package mongo
+
+import (
+	"context"
+	"slices"
+	"strings"
+
+	"go.mongodb.org/mongo-driver/v2/bson"
+	"go.mongodb.org/mongo-driver/v2/mongo"
+)
+
+// Filter out system databases after listing all databases.
+// Some MongoDB Atlas tiers (e.g., free tier) don't support regex filters in ListDatabases.
+func GetDatabaseNames(ctx context.Context, client *mongo.Client) ([]string, error) {
+	dbs, err := client.ListDatabases(ctx, bson.M{})
+	if err != nil {
+		return nil, err
+	}
+
+	filteredDbNames := make([]string, 0, len(dbs.Databases))
+	for _, db := range dbs.Databases {
+		if db.Name == "admin" || db.Name == "local" || db.Name == "config" {
+			continue
+		}
+		filteredDbNames = append(filteredDbNames, db.Name)
+	}
+	slices.Sort(filteredDbNames)
+	return filteredDbNames, nil
+}
+
+// Filter out views and system collections after listing all collections.
+// Some MongoDB Atlas tiers (e.g., free tier) don't support regex filters in ListCollections.
+func GetCollectionNames(ctx context.Context, client *mongo.Client, databaseName string) ([]string, error) {
+	db := client.Database(databaseName)
+	cur, err := db.ListCollections(ctx, bson.D{})
+	if err != nil {
+		return nil, err
+	}
+	defer cur.Close(ctx)
+
+	type CollectionSpec struct {
+		Name string `bson:"name"`
+		Type string `bson:"type"` // "collection" | "view" | etc
+	}
+
+	filteredCollNames := make([]string, 0, 100)
+	for cur.Next(ctx) {
+		if err := cur.Err(); err != nil {
+			return nil, err
+		}
+
+		var coll CollectionSpec
+		if err := cur.Decode(&coll); err != nil {
+			return nil, err
+		}
+
+		if strings.HasPrefix(coll.Name, "system.") {
+			continue
+		}
+		if strings.EqualFold(coll.Type, "view") {
+			continue
+		}
+		filteredCollNames = append(filteredCollNames, coll.Name)
+	}
+	slices.Sort(filteredCollNames)
+	return filteredCollNames, nil
+}