Fix pgdb helpers

FerretDB · Apr 1, 2022 · 335707d · 335707d
1 parent 0b17030
commit 335707d
Show file tree

Hide file tree

Showing 4 changed files with 67 additions and 12 deletions.
diff --git a/internal/handlers/handler_test.go b/internal/handlers/handler_test.go
@@ -2365,7 +2365,7 @@ func TestCreateListDropCollection(t *testing.T) {
 	})
 
 	t.Run("existing", func(t *testing.T) {
-		collection := testutil.CreateTable(ctx, t, pool, db)
+		collection := testutil.Table(ctx, t, pool, db)
 
 		actual := handle(ctx, t, handler, types.MustNewDocument(
 			"create", collection,

diff --git a/internal/handlers/pg/pgdb/pool.go b/internal/handlers/pg/pgdb/pool.go
@@ -276,16 +276,22 @@ func (pgPool *Pool) DropSchema(ctx context.Context, schema string) error {
 		return nil
 	}
 
-	if pgErr, ok := err.(*pgconn.PgError); ok && pgErr.Code == pgerrcode.InvalidSchemaName {
-		return ErrNotExist
+	pgErr, ok := err.(*pgconn.PgError)
+	if !ok {
+		return lazyerrors.Errorf("pg.DropSchema: %w", err)
 	}
 
-	return lazyerrors.Errorf("pg.DropSchema: %w", err)
+	switch pgErr.Code {
+	case pgerrcode.InvalidSchemaName:
+		return ErrNotExist
+	default:
+		return lazyerrors.Errorf("pg.DropSchema: %w", err)
+	}
 }
 
-// CreateTable creates a new FerretDB collection / PostgreSQL jsonb1 table.
+// CreateTable creates a new FerretDB collection / PostgreSQL jsonb1 table in existing schema.
 //
-// It returns ErrAlreadyExist if table already exist.
+// It returns ErrAlreadyExist if table already exist, ErrNotExist is schema does not exist.
 func (pgPool *Pool) CreateTable(ctx context.Context, schema, table string) error {
 	sql := `CREATE TABLE ` + pgx.Identifier{schema, table}.Sanitize() + ` (_jsonb jsonb)`
 	_, err := pgPool.Exec(ctx, sql)
@@ -299,6 +305,8 @@ func (pgPool *Pool) CreateTable(ctx context.Context, schema, table string) error
 	}
 
 	switch pgErr.Code {
+	case pgerrcode.InvalidSchemaName:
+		return ErrNotExist
 	case pgerrcode.DuplicateTable:
 		return ErrAlreadyExist
 	case pgerrcode.UniqueViolation, pgerrcode.DuplicateObject:
@@ -312,7 +320,7 @@ func (pgPool *Pool) CreateTable(ctx context.Context, schema, table string) error
 
 // DropTable drops FerretDB collection / PostgreSQL table.
 //
-// It returns ErrNotExist is table does not exist.
+// It returns ErrNotExist if schema or table does not exist.
 func (pgPool *Pool) DropTable(ctx context.Context, schema, table string) error {
 	// TODO probably not CASCADE
 	sql := `DROP TABLE ` + pgx.Identifier{schema, table}.Sanitize() + `CASCADE`
@@ -321,11 +329,17 @@ func (pgPool *Pool) DropTable(ctx context.Context, schema, table string) error {
 		return nil
 	}
 
-	if pgErr, ok := err.(*pgconn.PgError); ok && pgErr.Code == pgerrcode.UndefinedTable {
-		return ErrNotExist
+	pgErr, ok := err.(*pgconn.PgError)
+	if !ok {
+		return lazyerrors.Errorf("pg.DropTable: %w", err)
 	}
 
-	return lazyerrors.Errorf("pg.DropTable: %w", err)
+	switch pgErr.Code {
+	case pgerrcode.InvalidSchemaName, pgerrcode.UndefinedTable:
+		return ErrNotExist
+	default:
+		return lazyerrors.Errorf("pg.DropTable: %w", err)
+	}
 }
 
 // EnsureTableExist ensures that given FerretDB database / PostgreSQL schema and FerretDB collection / PostgreSQL table exist.

diff --git a/internal/handlers/pg/pgdb/pool_test.go b/internal/handlers/pg/pgdb/pool_test.go
@@ -88,6 +88,47 @@ func TestTables(t *testing.T) {
 	assert.Empty(t, tables)
 }
 
+func TestCreateDrop(t *testing.T) {
+	t.Parallel()
+
+	ctx := testutil.Ctx(t)
+	pool := testutil.Pool(ctx, t, nil, zaptest.NewLogger(t))
+
+	schemaName := testutil.SchemaName(t)
+	tableName := testutil.TableName(t)
+
+	t.Cleanup(func() {
+		pool.DropSchema(ctx, schemaName)
+	})
+
+	err := pool.CreateTable(ctx, schemaName, tableName)
+	require.Equal(t, pgdb.ErrNotExist, err)
+
+	err = pool.CreateSchema(ctx, schemaName)
+	require.NoError(t, err)
+
+	err = pool.CreateSchema(ctx, schemaName)
+	require.Equal(t, pgdb.ErrAlreadyExist, err)
+
+	err = pool.CreateTable(ctx, schemaName, tableName)
+	require.NoError(t, err)
+
+	err = pool.CreateTable(ctx, schemaName, tableName)
+	require.Equal(t, pgdb.ErrAlreadyExist, err)
+
+	err = pool.DropTable(ctx, schemaName, tableName)
+	require.NoError(t, err)
+
+	err = pool.DropTable(ctx, schemaName, tableName)
+	require.Equal(t, pgdb.ErrNotExist, err)
+
+	err = pool.DropSchema(ctx, schemaName)
+	require.NoError(t, err)
+
+	err = pool.DropSchema(ctx, schemaName)
+	require.Equal(t, pgdb.ErrNotExist, err)
+}
+
 func TestConcurrentCreate(t *testing.T) {
 	t.Parallel()
 

diff --git a/internal/util/testutil/pg.go b/internal/util/testutil/pg.go
@@ -103,10 +103,10 @@ func TableName(tb testing.TB) string {
 	return strings.ReplaceAll(strings.ToLower(tb.Name()), "/", "_")
 }
 
-// CreateTable creates FerretDB collection / PostgreSQL table for testing.
+// Table creates FerretDB collection / PostgreSQL table for testing.
 //
 // Name is stable for that test.
-func CreateTable(ctx context.Context, tb testing.TB, pool *pgdb.Pool, db string) string {
+func Table(ctx context.Context, tb testing.TB, pool *pgdb.Pool, db string) string {
 	tb.Helper()
 
 	table := TableName(tb)