Merge pull request #1517 from Permify/ufuk/tenantDeleteFix

fix: tenant delete request will delete all related data too

internal/storage/memory/tenantWriter.go

@@ -3,6 +3,7 @@ package memory
 import (
 	"context"
 	"errors"
+	"fmt"
 	"time"
 
 	"github.com/Permify/permify/internal/storage"
@@ -43,14 +44,47 @@ func (w *TenantWriter) CreateTenant(_ context.Context, id, name string) (result
 func (w *TenantWriter) DeleteTenant(_ context.Context, tenantID string) (result *base.Tenant, err error) {
 	txn := w.database.DB.Txn(true)
 	defer txn.Abort()
-	var raw interface{}
-	raw, err = txn.First(constants.TenantsTable, "id", tenantID)
-	if err != nil {
-		return nil, errors.New(base.ErrorCode_ERROR_CODE_EXECUTION.String())
+
+	// Define a slice of tables to delete associated records
+	tables := []string{
+		constants.AttributesTable,
+		constants.BundlesTable,
+		constants.RelationTuplesTable,
+		constants.SchemaDefinitionsTable,
 	}
-	if _, err = txn.DeleteAll(constants.TenantsTable, "id", tenantID); err != nil {
-		return nil, errors.New(base.ErrorCode_ERROR_CODE_EXECUTION.String())
+
+	// Iterate through each table and delete records associated with the tenant
+	totalDeleted := 0
+	for _, table := range tables {
+		if _, err := txn.First(table, "tenant_id", tenantID); err == nil {
+			numDeleted, deleteErr := txn.DeleteAll(table, "tenant_id", tenantID)
+			if deleteErr != nil {
+				return nil, errors.New(base.ErrorCode_ERROR_CODE_EXECUTION.String())
+			}
+			totalDeleted += numDeleted
+		}
 	}
+
+	// Retrieve the tenant first
+	raw, err := txn.First(constants.TenantsTable, "id", tenantID)
+	if err != nil {
+		if totalDeleted > 0 {
+			raw = storage.Tenant{
+				ID:        tenantID,
+				Name:      fmt.Sprintf("Affected rows: %d", totalDeleted),
+				CreatedAt: time.Now(),
+			}
+		} else {
+			return nil, errors.New(base.ErrorCode_ERROR_CODE_EXECUTION.String())
+		}
+	} else {
+		// Finally, delete the tenant record
+		if _, err = txn.DeleteAll(constants.TenantsTable, "id", tenantID); err != nil {
+			return nil, errors.New(base.ErrorCode_ERROR_CODE_EXECUTION.String())
+		}
+	}
+
 	txn.Commit()
 	return raw.(storage.Tenant).ToTenant(), nil
+
 }

internal/storage/postgres/tenantWriter.go

@@ -3,6 +3,7 @@ package postgres
 import (
 	"context"
 	"errors"
+	"fmt"
 	"log/slog"
 	"strings"
 	"time"
@@ -68,19 +69,70 @@ func (w *TenantWriter) DeleteTenant(ctx context.Context, tenantID string) (resul
 
 	slog.DebugContext(ctx, "deleting tenant", slog.Any("tenant_id", tenantID))
 
+	tx, err := w.database.WritePool.Begin(ctx)
+	if err != nil {
+		return nil, utils.HandleError(ctx, span, err, base.ErrorCode_ERROR_CODE_EXECUTION)
+	}
+	defer tx.Rollback(ctx)
+
+	// Prepare batch operations for deleting tenant-related records from multiple tables
+	tables := []string{"bundles", "relation_tuples", "attributes", "schema_definitions", "transactions"}
+	batch := &pgx.Batch{}
+	var totalDeleted int64
+	for _, table := range tables {
+		query := fmt.Sprintf(utils.DeleteAllByTenantTemplate, table)
+		batch.Queue(query, tenantID)
+	}
+	batch.Queue(utils.DeleteTenantTemplate, tenantID)
+
+	// Execute the batch of delete queries
+	br := tx.SendBatch(ctx, batch)
+
+	for i := 0; i < len(tables); i++ {
+		tag, err := br.Exec()
+		if err != nil {
+			err = br.Close()
+			if err != nil {
+				return nil, err
+			}
+			err = tx.Commit(ctx)
+			if err != nil {
+				return nil, utils.HandleError(ctx, span, err, base.ErrorCode_ERROR_CODE_EXECUTION)
+			}
+			return nil, utils.HandleError(ctx, span, err, base.ErrorCode_ERROR_CODE_EXECUTION)
+		} else {
+			totalDeleted += tag.RowsAffected()
+		}
+	}
+
+	// Retrieve the tenant details after deletion
 	var name string
 	var createdAt time.Time
+	err = br.QueryRow().Scan(&name, &createdAt)
 
-	err = w.database.WritePool.QueryRow(ctx, utils.DeleteTenantTemplate, tenantID).Scan(&name, &createdAt)
+	if err != nil {
+		if totalDeleted > 0 {
+			name = fmt.Sprintf("Affected rows: %d", totalDeleted)
+		} else {
+			return nil, utils.HandleError(ctx, span, err, base.ErrorCode_ERROR_CODE_EXECUTION)
+		}
+	}
+
+	err = br.Close()
 	if err != nil {
 		return nil, utils.HandleError(ctx, span, err, base.ErrorCode_ERROR_CODE_EXECUTION)
 	}
 
-	slog.DebugContext(ctx, "successfully deleted tenant")
+	err = tx.Commit(ctx)
+	if err != nil {
+		return nil, utils.HandleError(ctx, span, err, base.ErrorCode_ERROR_CODE_EXECUTION)
+	}
 
+	// Return the deleted tenant information
 	return &base.Tenant{
 		Id:        tenantID,
 		Name:      name,
 		CreatedAt: timestamppb.New(createdAt),
 	}, nil
+
 }

internal/storage/postgres/utils/common.go

@@ -37,9 +37,10 @@ SELECT entity_id
 FROM filtered_entities
 `
 
-	TransactionTemplate  = `INSERT INTO transactions (tenant_id) VALUES ($1) RETURNING id`
-	InsertTenantTemplate = `INSERT INTO tenants (id, name) VALUES ($1, $2) RETURNING created_at`
-	DeleteTenantTemplate = `DELETE FROM tenants WHERE id = $1 RETURNING name, created_at`
+	TransactionTemplate       = `INSERT INTO transactions (tenant_id) VALUES ($1) RETURNING id`
+	InsertTenantTemplate      = `INSERT INTO tenants (id, name) VALUES ($1, $2) RETURNING created_at`
+	DeleteTenantTemplate      = `DELETE FROM tenants WHERE id = $1 RETURNING name, created_at`
+	DeleteAllByTenantTemplate = `DELETE FROM %s WHERE tenant_id = $1`
 )
 
 // SnapshotQuery adds conditions to a SELECT query for checking transaction visibility based on created and expired transaction IDs.