13719691815

Committed 07 Mar 2025 11:25AM UTC coverage: 81.488% (+0.02%) from 81.468%

Build # 13719691815

Build Type

push

github

Committed by

web-flow

Commit Message

Merge pull request #2103 from Permify/dependabot/go_modules/go.opentelemetry.io/contrib/instrumentation/host-0.60.0

build(deps): bump go.opentelemetry.io/contrib/instrumentation/host from 0.59.0 to 0.60.0

Run Details

8280 of 10161 relevant lines covered (81.49%)

156.8 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

71.19

/internal/storage/postgres/dataWriter.go

package postgres

import (
        "context"
        "errors"
        "log/slog"

        "github.com/jackc/pgx/v5"

        "github.com/Masterminds/squirrel"
        "github.com/jackc/pgx/v5/pgconn"
        "google.golang.org/protobuf/encoding/protojson"

        "github.com/Permify/permify/internal"
        "github.com/Permify/permify/internal/storage/postgres/snapshot"
        "github.com/Permify/permify/internal/storage/postgres/types"
        "github.com/Permify/permify/internal/storage/postgres/utils"
        "github.com/Permify/permify/internal/validation"
        "github.com/Permify/permify/pkg/bundle"
        "github.com/Permify/permify/pkg/database"
        db "github.com/Permify/permify/pkg/database/postgres"
        base "github.com/Permify/permify/pkg/pb/base/v1"
        "github.com/Permify/permify/pkg/token"
        "github.com/Permify/permify/pkg/tuple"
)

// DataWriter - Structure for Data Writer
type DataWriter struct {
        database *db.Postgres
        // options
        txOptions pgx.TxOptions
}

func NewDataWriter(database *db.Postgres) *DataWriter {
        return &DataWriter{
                database:  database,
                txOptions: pgx.TxOptions{IsoLevel: pgx.Serializable, AccessMode: pgx.ReadWrite},
        }
}

// Write method writes a collection of tuples and attributes to the database for a specific tenant.
// It returns an EncodedSnapToken upon successful write or an error if the write fails.
func (w *DataWriter) Write(
        ctx context.Context,
        tenantID string,
        tupleCollection *database.TupleCollection,
        attributeCollection *database.AttributeCollection,
) (token token.EncodedSnapToken, err error) {
        // Start a new tracing span for this operation.
        ctx, span := internal.Tracer.Start(ctx, "data-writer.write")
        defer span.End() // Ensure that the span is ended when the function returns.

        // Log the start of a data write operation.
        slog.DebugContext(ctx, "writing data for tenant_id", slog.String("tenant_id", tenantID), "max retries", slog.Any("max_retries", w.database.GetMaxRetries()))

        // Check if the total number of tuples and attributes exceeds the maximum allowed per write.
        if len(tupleCollection.GetTuples())+len(attributeCollection.GetAttributes()) > w.database.GetMaxDataPerWrite() {
                return nil, errors.New(base.ErrorCode_ERROR_CODE_MAX_DATA_PER_WRITE_EXCEEDED.String())
        }

        // Retry loop for handling transient errors like serialization issues.
        for i := 0; i <= w.database.GetMaxRetries(); i++ {
                // Attempt to write the data to the database.
                tkn, err := w.write(ctx, tenantID, tupleCollection, attributeCollection)
                if err != nil {
                        // Check if the error is due to serialization, and if so, retry.
                        if utils.IsSerializationRelatedError(err) || pgconn.SafeToRetry(err) {
                                slog.WarnContext(ctx, "serialization error occurred", slog.String("tenant_id", tenantID), slog.Int("retry", i))
                                utils.WaitWithBackoff(ctx, tenantID, i)
                                continue // Retry the operation.
                        }
                        // If the error is not serialization-related, handle it and return.
                        return nil, utils.HandleError(ctx, span, err, base.ErrorCode_ERROR_CODE_DATASTORE)
                }
                // If to write is successful, return the token.
                return tkn, nil
        }

        // Log an error if the operation failed after reaching the maximum number of retries.
        slog.ErrorContext(ctx, "max retries reached", slog.Any("error", errors.New(base.ErrorCode_ERROR_CODE_ERROR_MAX_RETRIES.String())))

        // Return an error indicating that the maximum number of retries has been reached.
        return nil, errors.New(base.ErrorCode_ERROR_CODE_ERROR_MAX_RETRIES.String())
}

// Delete method removes data from the database based on the provided tuple and attribute filters.
// It returns an EncodedSnapToken upon successful deletion or an error if the deletion fails.
func (w *DataWriter) Delete(
        ctx context.Context,
        tenantID string,
        tupleFilter *base.TupleFilter,
        attributeFilter *base.AttributeFilter,
) (token.EncodedSnapToken, error) {
        // Start a new tracing span for this delete operation.
        ctx, span := internal.Tracer.Start(ctx, "data-writer.delete")
        defer span.End() // Ensure that the span is ended when the function returns.

        // Log the start of a data deletion operation.
        slog.DebugContext(ctx, "deleting data for tenant_id", slog.String("tenant_id", tenantID), "max retries", slog.Any("max_retries", w.database.GetMaxRetries()))

        // Retry loop for handling transient errors like serialization issues.
        for i := 0; i <= w.database.GetMaxRetries(); i++ {
                // Attempt to delete the data from the database.
                tkn, err := w.delete(ctx, tenantID, tupleFilter, attributeFilter)
                if err != nil {
                        // Check if the error is due to serialization, and if so, retry.
                        if utils.IsSerializationRelatedError(err) || pgconn.SafeToRetry(err) {
                                slog.WarnContext(ctx, "serialization error occurred", slog.String("tenant_id", tenantID), slog.Int("retry", i))
                                utils.WaitWithBackoff(ctx, tenantID, i)
                                continue // Retry the operation.
                        }
                        // If the error is not serialization-related, handle it and return.
                        return nil, utils.HandleError(ctx, span, err, base.ErrorCode_ERROR_CODE_DATASTORE)
                }
                // If the delete operation is successful, return the token.
                return tkn, nil
        }

        // Log an error if the operation failed after reaching the maximum number of retries.
        slog.DebugContext(ctx, "max retries reached", slog.Any("error", errors.New(base.ErrorCode_ERROR_CODE_ERROR_MAX_RETRIES.String())))

        // Return an error indicating that the maximum number of retries has been reached.
        return nil, errors.New(base.ErrorCode_ERROR_CODE_ERROR_MAX_RETRIES.String())
}

// RunBundle executes a bundle of operations in the context of a given tenant.
// It returns an EncodedSnapToken upon successful completion or an error if the operation fails.
func (w *DataWriter) RunBundle(
        ctx context.Context,
        tenantID string,
        arguments map[string]string,
        b *base.DataBundle,
) (token.EncodedSnapToken, error) {
        // Start a new tracing span for this operation.
        ctx, span := internal.Tracer.Start(ctx, "data-writer.run-bundle")
        defer span.End() // Ensure that the span is ended when the function returns.

        // Log the start of running a bundle operation.
        slog.DebugContext(ctx, "running bundle for tenant_id", slog.String("tenant_id", tenantID), "max retries", slog.Any("max_retries", w.database.GetMaxRetries()))

        // Retry loop for handling transient errors like serialization issues.
        for i := 0; i <= w.database.GetMaxRetries(); i++ {
                // Attempt to run the bundle operation.
                tkn, err := w.runBundle(ctx, tenantID, arguments, b)
                if err != nil {
                        // Check if the error is due to serialization, and if so, retry.
                        if utils.IsSerializationRelatedError(err) || pgconn.SafeToRetry(err) {
                                slog.WarnContext(ctx, "serialization error occurred", slog.String("tenant_id", tenantID), slog.Int("retry", i))
                                utils.WaitWithBackoff(ctx, tenantID, i)
                                continue // Retry the operation.
                        }
                        // If the error is not serialization-related, handle it and return.
                        return nil, utils.HandleError(ctx, span, err, base.ErrorCode_ERROR_CODE_DATASTORE)
                }
                // If the operation is successful, return the token.
                return tkn, nil
        }

        // Log an error if the operation failed after reaching the maximum number of retries.
        slog.ErrorContext(ctx, "max retries reached", slog.Any("error", errors.New(base.ErrorCode_ERROR_CODE_ERROR_MAX_RETRIES.String())))

        // Return an error indicating that the maximum number of retries has been reached.
        return nil, errors.New(base.ErrorCode_ERROR_CODE_ERROR_MAX_RETRIES.String())
}

// write handles the database writing of tuple and attribute collections for a given tenant.
// It returns an EncodedSnapToken upon successful write or an error if the write fails.
func (w *DataWriter) write(
        ctx context.Context,
        tenantID string,
        tupleCollection *database.TupleCollection,
        attributeCollection *database.AttributeCollection,
) (token token.EncodedSnapToken, err error) {
        var tx pgx.Tx
        tx, err = w.database.WritePool.BeginTx(ctx, w.txOptions)
        if err != nil {
                return nil, err
        }

        defer func() {
                _ = tx.Rollback(ctx)
        }()

        var xid types.XID8
        err = tx.QueryRow(ctx, utils.TransactionTemplate, tenantID).Scan(&xid)
        if err != nil {
                return nil, err
        }

        slog.DebugContext(ctx, "retrieved transaction", slog.Any("xid", xid), "for tenant", slog.Any("tenant_id", tenantID))

        slog.DebugContext(ctx, "processing tuples and executing insert query")

        batch := &pgx.Batch{}

        if len(tupleCollection.GetTuples()) > 0 {
                err = w.batchInsertRelationships(batch, xid, tenantID, tupleCollection)
                if err != nil {
                        return nil, err
                }
        }

        if len(attributeCollection.GetAttributes()) > 0 {
                err = w.batchUpdateAttributes(batch, xid, tenantID, buildDeleteClausesForAttributes(attributeCollection))
                if err != nil {
                        return nil, err
                }
                err = w.batchInsertAttributes(batch, xid, tenantID, attributeCollection)
                if err != nil {
                        return nil, err
                }
        }

        batchResult := tx.SendBatch(ctx, batch)
        for i := 0; i < batch.Len(); i++ {
                _, err = batchResult.Exec()
                if err != nil {
                        err = batchResult.Close()
                        if err != nil {
                                return nil, err
                        }
                        return nil, err
                }
        }

        err = batchResult.Close()
        if err != nil {
                return nil, err
        }

        if err = tx.Commit(ctx); err != nil {
                return nil, err
        }

        slog.DebugContext(ctx, "data successfully written to the database")

        return snapshot.NewToken(xid).Encode(), nil
}

// delete handles the deletion of tuples and attributes from the database based on provided filters.
// It returns an EncodedSnapToken upon successful deletion or an error if the deletion fails.
func (w *DataWriter) delete(
        ctx context.Context,
        tenantID string,
        tupleFilter *base.TupleFilter,
        attributeFilter *base.AttributeFilter,
) (token token.EncodedSnapToken, err error) {
        var tx pgx.Tx
        tx, err = w.database.WritePool.BeginTx(ctx, w.txOptions)
        if err != nil {
                return nil, err
        }

        defer func() {
                _ = tx.Rollback(ctx)
        }()

        var xid types.XID8
        err = tx.QueryRow(ctx, utils.TransactionTemplate, tenantID).Scan(&xid)
        if err != nil {
                return nil, err
        }

        slog.DebugContext(ctx, "retrieved transaction", slog.Any("xid", xid), "for tenant", slog.Any("tenant_id", tenantID))

        slog.DebugContext(ctx, "processing tuple and executing update query")

        if !validation.IsTupleFilterEmpty(tupleFilter) {
                tbuilder := w.database.Builder.Update(RelationTuplesTable).Set("expired_tx_id", xid).Where(squirrel.Eq{"expired_tx_id": "0", "tenant_id": tenantID})
                tbuilder = utils.TuplesFilterQueryForUpdateBuilder(tbuilder, tupleFilter)

                var tquery string
                var targs []interface{}

                tquery, targs, err = tbuilder.ToSql()
                if err != nil {
                        return nil, err
                }

                _, err = tx.Exec(ctx, tquery, targs...)
                if err != nil {
                        return nil, err
                }
        }

        slog.DebugContext(ctx, "processing attribute and executing update query")

        if !validation.IsAttributeFilterEmpty(attributeFilter) {
                abuilder := w.database.Builder.Update(AttributesTable).Set("expired_tx_id", xid).Where(squirrel.Eq{"expired_tx_id": "0", "tenant_id": tenantID})
                abuilder = utils.AttributesFilterQueryForUpdateBuilder(abuilder, attributeFilter)

                var aquery string
                var aargs []interface{}

                aquery, aargs, err = abuilder.ToSql()
                if err != nil {
                        return nil, err
                }

                _, err = tx.Exec(ctx, aquery, aargs...)
                if err != nil {
                        return nil, err
                }
        }

        if err = tx.Commit(ctx); err != nil {
                return nil, err
        }

        slog.DebugContext(ctx, "data successfully deleted from the database")

        return snapshot.NewToken(xid).Encode(), nil
}

// runBundle executes a series of operations defined in a DataBundle within a single database transaction.
// It returns an EncodedSnapToken upon successful execution of all operations or an error if any operation fails.
func (w *DataWriter) runBundle(
        ctx context.Context,
        tenantID string,
        arguments map[string]string,
        b *base.DataBundle,
) (token token.EncodedSnapToken, err error) {
        var tx pgx.Tx
        tx, err = w.database.WritePool.BeginTx(ctx, w.txOptions)
        if err != nil {
                return nil, err
        }

        defer func() {
                _ = tx.Rollback(ctx)
        }()

        var xid types.XID8
        err = tx.QueryRow(ctx, utils.TransactionTemplate, tenantID).Scan(&xid)
        if err != nil {
                return nil, err
        }

        slog.DebugContext(ctx, "retrieved transaction", slog.Any("xid", xid), "for tenant", slog.Any("tenant_id", tenantID))

        batch := &pgx.Batch{}

        for _, op := range b.GetOperations() {
                tb, ab, err := bundle.Operation(arguments, op)
                if err != nil {
                        return nil, err
                }

                err = w.runOperation(batch, xid, tenantID, tb, ab)
                if err != nil {
                        return nil, err
                }
        }

        batchResult := tx.SendBatch(ctx, batch)
        for i := 0; i < batch.Len(); i++ {
                _, err = batchResult.Exec()
                if err != nil {
                        err = batchResult.Close()
                        if err != nil {
                                return nil, err
                        }
                        return nil, err
                }
        }

        err = batchResult.Close()
        if err != nil {
                return nil, err
        }

        if err = tx.Commit(ctx); err != nil {
                return nil, err
        }

        return snapshot.NewToken(xid).Encode(), nil
}

// runOperation processes and executes database operations defined in TupleBundle and AttributeBundle within a given transaction.
func (w *DataWriter) runOperation(
        batch *pgx.Batch,
        xid types.XID8,
        tenantID string,
        tb database.TupleBundle,
        ab database.AttributeBundle,
) (err error) {
        slog.Debug("processing bundles queries")
        if len(tb.Write.GetTuples()) > 0 {
                err = w.batchInsertRelationships(batch, xid, tenantID, &tb.Write)
                if err != nil {
                        return err
                }
        }

        if len(ab.Write.GetAttributes()) > 0 {
                deleteClauses := buildDeleteClausesForAttributes(&ab.Write)
                err = w.batchUpdateAttributes(batch, xid, tenantID, deleteClauses)
                if err != nil {
                        return err
                }

                err = w.batchInsertAttributes(batch, xid, tenantID, &ab.Write)
                if err != nil {
                        return err
                }
        }

        if len(tb.Delete.GetTuples()) > 0 {
                deleteClauses := buildDeleteClausesForRelationships(&tb.Delete)
                err = w.batchUpdateRelationships(batch, xid, tenantID, deleteClauses)
                if err != nil {
                        return err
                }
        }

        if len(ab.Delete.GetAttributes()) > 0 {
                deleteClauses := buildDeleteClausesForAttributes(&ab.Delete)
                err = w.batchUpdateAttributes(batch, xid, tenantID, deleteClauses)
                if err != nil {
                        return err
                }
        }

        return nil
}

// batchInsertTuples function for batch inserting tuples
func (w *DataWriter) batchInsertRelationships(batch *pgx.Batch, xid types.XID8, tenantID string, tupleCollection *database.TupleCollection) error {
        titer := tupleCollection.CreateTupleIterator()
        for titer.HasNext() {
                t := titer.GetNext()
                srelation := t.GetSubject().GetRelation()
                if srelation == tuple.ELLIPSIS {
                        srelation = ""
                }
                batch.Queue(
                        "INSERT INTO relation_tuples (entity_type, entity_id, relation, subject_type, subject_id, subject_relation, created_tx_id, tenant_id) VALUES ($1, $2, $3, $4, $5, $6, $7, $8) ON CONFLICT ON CONSTRAINT uq_relation_tuple_not_expired DO NOTHING",
                        t.GetEntity().GetType(), t.GetEntity().GetId(), t.GetRelation(), t.GetSubject().GetType(), t.GetSubject().GetId(), srelation, xid, tenantID,
                )
        }
        return nil
}

// batchUpdateTuples function for batch updating tuples
func (w *DataWriter) batchUpdateRelationships(batch *pgx.Batch, xid types.XID8, tenantID string, deleteClauses []squirrel.Eq) error {
        for _, condition := range deleteClauses {
                query, args, err := w.database.Builder.Update(RelationTuplesTable).
                        Set("expired_tx_id", xid).
                        Where(squirrel.Eq{"expired_tx_id": "0", "tenant_id": tenantID}).
                        Where(condition).
                        ToSql()
                if err != nil {
                        return err
                }
                batch.Queue(query, args...)
        }
        return nil
}

// buildDeleteClauses function to build delete clauses for tuples
func buildDeleteClausesForRelationships(tupleCollection *database.TupleCollection) []squirrel.Eq {
        deleteClauses := make([]squirrel.Eq, 0)

        titer := tupleCollection.CreateTupleIterator()
        for titer.HasNext() {
                t := titer.GetNext()
                srelation := t.GetSubject().GetRelation()
                if srelation == tuple.ELLIPSIS {
                        srelation = ""
                }

                condition := squirrel.Eq{
                        "entity_type":      t.GetEntity().GetType(),
                        "entity_id":        t.GetEntity().GetId(),
                        "relation":         t.GetRelation(),
                        "subject_type":     t.GetSubject().GetType(),
                        "subject_id":       t.GetSubject().GetId(),
                        "subject_relation": srelation,
                }

                deleteClauses = append(deleteClauses, condition)
        }

        return deleteClauses
}

// batchInsertAttributes function for batch inserting attributes
func (w *DataWriter) batchInsertAttributes(batch *pgx.Batch, xid types.XID8, tenantID string, attributeCollection *database.AttributeCollection) error {
        aiter := attributeCollection.CreateAttributeIterator()
        for aiter.HasNext() {
                a := aiter.GetNext()

                jsonBytes, err := protojson.Marshal(a.GetValue())
                if err != nil {
                        return err
                }

                jsonStr := string(jsonBytes)

                batch.Queue(
                        "INSERT INTO attributes (entity_type, entity_id, attribute, value, created_tx_id, tenant_id) VALUES ($1, $2, $3, $4, $5, $6)",
                        a.GetEntity().GetType(), a.GetEntity().GetId(), a.GetAttribute(), jsonStr, xid, tenantID,
                )
        }
        return nil
}

// batchUpdateAttributes function for batch updating attributes
func (w *DataWriter) batchUpdateAttributes(batch *pgx.Batch, xid types.XID8, tenantID string, deleteClauses []squirrel.Eq) error {
        for _, condition := range deleteClauses {
                query, args, err := w.database.Builder.Update(AttributesTable).
                        Set("expired_tx_id", xid).
                        Where(squirrel.Eq{"expired_tx_id": "0", "tenant_id": tenantID}).
                        Where(condition).
                        ToSql()
                if err != nil {
                        return err
                }

                batch.Queue(query, args...)
        }
        return nil
}

// buildDeleteClausesForAttributes function to build delete clauses for attributes
func buildDeleteClausesForAttributes(attributeCollection *database.AttributeCollection) []squirrel.Eq {
        deleteClauses := make([]squirrel.Eq, 0)

        aiter := attributeCollection.CreateAttributeIterator()
        for aiter.HasNext() {
                a := aiter.GetNext()

                condition := squirrel.Eq{
                        "entity_type": a.GetEntity().GetType(),
                        "entity_id":   a.GetEntity().GetId(),
                        "attribute":   a.GetAttribute(),
                }

                deleteClauses = append(deleteClauses, condition)
        }

        return deleteClauses
}

1	package postgres
2
3	import (
4	"context"
5	"errors"
6	"log/slog"
7
8	"github.com/jackc/pgx/v5"
9
10	"github.com/Masterminds/squirrel"
11	"github.com/jackc/pgx/v5/pgconn"
12	"google.golang.org/protobuf/encoding/protojson"
13
14	"github.com/Permify/permify/internal"
15	"github.com/Permify/permify/internal/storage/postgres/snapshot"
16	"github.com/Permify/permify/internal/storage/postgres/types"
17	"github.com/Permify/permify/internal/storage/postgres/utils"
18	"github.com/Permify/permify/internal/validation"
19	"github.com/Permify/permify/pkg/bundle"
20	"github.com/Permify/permify/pkg/database"
21	db "github.com/Permify/permify/pkg/database/postgres"
22	base "github.com/Permify/permify/pkg/pb/base/v1"
23	"github.com/Permify/permify/pkg/token"
24	"github.com/Permify/permify/pkg/tuple"
25	)
26
27	// DataWriter - Structure for Data Writer
28	type DataWriter struct {
29	database *db.Postgres
30	// options
31	txOptions pgx.TxOptions
32	}
33
34	func NewDataWriter(database db.Postgres) DataWriter {	14✔
35	return &DataWriter{	14✔
36	database: database,	14✔
37	txOptions: pgx.TxOptions{IsoLevel: pgx.Serializable, AccessMode: pgx.ReadWrite},	14✔
38	}	14✔
39	}	14✔
40
41	// Write method writes a collection of tuples and attributes to the database for a specific tenant.
42	// It returns an EncodedSnapToken upon successful write or an error if the write fails.
43	func (w *DataWriter) Write(
44	ctx context.Context,
45	tenantID string,
46	tupleCollection *database.TupleCollection,
47	attributeCollection *database.AttributeCollection,
48	) (token token.EncodedSnapToken, err error) {	20✔
49	// Start a new tracing span for this operation.	20✔
50	ctx, span := internal.Tracer.Start(ctx, "data-writer.write")	20✔
51	defer span.End() // Ensure that the span is ended when the function returns.	20✔
52		20✔
53	// Log the start of a data write operation.	20✔
54	slog.DebugContext(ctx, "writing data for tenant_id", slog.String("tenant_id", tenantID), "max retries", slog.Any("max_retries", w.database.GetMaxRetries()))	20✔
55		20✔
56	// Check if the total number of tuples and attributes exceeds the maximum allowed per write.	20✔
57	if len(tupleCollection.GetTuples())+len(attributeCollection.GetAttributes()) > w.database.GetMaxDataPerWrite() {	20✔
58	return nil, errors.New(base.ErrorCode_ERROR_CODE_MAX_DATA_PER_WRITE_EXCEEDED.String())	×
59	}	×
60
61	// Retry loop for handling transient errors like serialization issues.
62	for i := 0; i <= w.database.GetMaxRetries(); i++ {	40✔
63	// Attempt to write the data to the database.	20✔
64	tkn, err := w.write(ctx, tenantID, tupleCollection, attributeCollection)	20✔
65	if err != nil {	20✔
66	// Check if the error is due to serialization, and if so, retry.	×
67	if utils.IsSerializationRelatedError(err) \|\| pgconn.SafeToRetry(err) {	×
68	slog.WarnContext(ctx, "serialization error occurred", slog.String("tenant_id", tenantID), slog.Int("retry", i))	×
69	utils.WaitWithBackoff(ctx, tenantID, i)	×
70	continue // Retry the operation.	×
71	}
72	// If the error is not serialization-related, handle it and return.
73	return nil, utils.HandleError(ctx, span, err, base.ErrorCode_ERROR_CODE_DATASTORE)	×
74	}
75	// If to write is successful, return the token.
76	return tkn, nil	20✔
77	}
78
79	// Log an error if the operation failed after reaching the maximum number of retries.
80	slog.ErrorContext(ctx, "max retries reached", slog.Any("error", errors.New(base.ErrorCode_ERROR_CODE_ERROR_MAX_RETRIES.String())))	×
81		×
82	// Return an error indicating that the maximum number of retries has been reached.	×
83	return nil, errors.New(base.ErrorCode_ERROR_CODE_ERROR_MAX_RETRIES.String())	×
84	}
85
86	// Delete method removes data from the database based on the provided tuple and attribute filters.
87	// It returns an EncodedSnapToken upon successful deletion or an error if the deletion fails.
88	func (w *DataWriter) Delete(
89	ctx context.Context,
90	tenantID string,
91	tupleFilter *base.TupleFilter,
92	attributeFilter *base.AttributeFilter,
93	) (token.EncodedSnapToken, error) {	5✔
94	// Start a new tracing span for this delete operation.	5✔
95	ctx, span := internal.Tracer.Start(ctx, "data-writer.delete")	5✔
96	defer span.End() // Ensure that the span is ended when the function returns.	5✔
97		5✔
98	// Log the start of a data deletion operation.	5✔
99	slog.DebugContext(ctx, "deleting data for tenant_id", slog.String("tenant_id", tenantID), "max retries", slog.Any("max_retries", w.database.GetMaxRetries()))	5✔
100		5✔
101	// Retry loop for handling transient errors like serialization issues.	5✔
102	for i := 0; i <= w.database.GetMaxRetries(); i++ {	10✔
103	// Attempt to delete the data from the database.	5✔
104	tkn, err := w.delete(ctx, tenantID, tupleFilter, attributeFilter)	5✔
105	if err != nil {	5✔
106	// Check if the error is due to serialization, and if so, retry.	×
107	if utils.IsSerializationRelatedError(err) \|\| pgconn.SafeToRetry(err) {	×
108	slog.WarnContext(ctx, "serialization error occurred", slog.String("tenant_id", tenantID), slog.Int("retry", i))	×
109	utils.WaitWithBackoff(ctx, tenantID, i)	×
110	continue // Retry the operation.	×
111	}
112	// If the error is not serialization-related, handle it and return.
113	return nil, utils.HandleError(ctx, span, err, base.ErrorCode_ERROR_CODE_DATASTORE)	×
114	}
115	// If the delete operation is successful, return the token.
116	return tkn, nil	5✔
117	}
118
119	// Log an error if the operation failed after reaching the maximum number of retries.
120	slog.DebugContext(ctx, "max retries reached", slog.Any("error", errors.New(base.ErrorCode_ERROR_CODE_ERROR_MAX_RETRIES.String())))	×
121		×
122	// Return an error indicating that the maximum number of retries has been reached.	×
123	return nil, errors.New(base.ErrorCode_ERROR_CODE_ERROR_MAX_RETRIES.String())	×
124	}
125
126	// RunBundle executes a bundle of operations in the context of a given tenant.
127	// It returns an EncodedSnapToken upon successful completion or an error if the operation fails.
128	func (w *DataWriter) RunBundle(
129	ctx context.Context,
130	tenantID string,
131	arguments map[string]string,
132	b *base.DataBundle,
133	) (token.EncodedSnapToken, error) {	1✔
134	// Start a new tracing span for this operation.	1✔
135	ctx, span := internal.Tracer.Start(ctx, "data-writer.run-bundle")	1✔
136	defer span.End() // Ensure that the span is ended when the function returns.	1✔
137		1✔
138	// Log the start of running a bundle operation.	1✔
139	slog.DebugContext(ctx, "running bundle for tenant_id", slog.String("tenant_id", tenantID), "max retries", slog.Any("max_retries", w.database.GetMaxRetries()))	1✔
140		1✔
141	// Retry loop for handling transient errors like serialization issues.	1✔
142	for i := 0; i <= w.database.GetMaxRetries(); i++ {	2✔
143	// Attempt to run the bundle operation.	1✔
144	tkn, err := w.runBundle(ctx, tenantID, arguments, b)	1✔
145	if err != nil {	1✔
146	// Check if the error is due to serialization, and if so, retry.	×
147	if utils.IsSerializationRelatedError(err) \|\| pgconn.SafeToRetry(err) {	×
148	slog.WarnContext(ctx, "serialization error occurred", slog.String("tenant_id", tenantID), slog.Int("retry", i))	×
149	utils.WaitWithBackoff(ctx, tenantID, i)	×
150	continue // Retry the operation.	×
151	}
152	// If the error is not serialization-related, handle it and return.
153	return nil, utils.HandleError(ctx, span, err, base.ErrorCode_ERROR_CODE_DATASTORE)	×
154	}
155	// If the operation is successful, return the token.
156	return tkn, nil	1✔
157	}
158
159	// Log an error if the operation failed after reaching the maximum number of retries.
160	slog.ErrorContext(ctx, "max retries reached", slog.Any("error", errors.New(base.ErrorCode_ERROR_CODE_ERROR_MAX_RETRIES.String())))	×
161		×
162	// Return an error indicating that the maximum number of retries has been reached.	×
163	return nil, errors.New(base.ErrorCode_ERROR_CODE_ERROR_MAX_RETRIES.String())	×
164	}
165
166	// write handles the database writing of tuple and attribute collections for a given tenant.
167	// It returns an EncodedSnapToken upon successful write or an error if the write fails.
168	func (w *DataWriter) write(
169	ctx context.Context,
170	tenantID string,
171	tupleCollection *database.TupleCollection,
172	attributeCollection *database.AttributeCollection,
173	) (token token.EncodedSnapToken, err error) {	20✔
174	var tx pgx.Tx	20✔
175	tx, err = w.database.WritePool.BeginTx(ctx, w.txOptions)	20✔
176	if err != nil {	20✔
177	return nil, err	×
178	}	×
179
180	defer func() {	40✔
181	_ = tx.Rollback(ctx)	20✔
182	}()	20✔
183
184	var xid types.XID8	20✔
185	err = tx.QueryRow(ctx, utils.TransactionTemplate, tenantID).Scan(&xid)	20✔
186	if err != nil {	20✔
187	return nil, err	×
188	}	×
189
190	slog.DebugContext(ctx, "retrieved transaction", slog.Any("xid", xid), "for tenant", slog.Any("tenant_id", tenantID))	20✔
191		20✔
192	slog.DebugContext(ctx, "processing tuples and executing insert query")	20✔
193		20✔
194	batch := &pgx.Batch{}	20✔
195		20✔
196	if len(tupleCollection.GetTuples()) > 0 {	34✔
197	err = w.batchInsertRelationships(batch, xid, tenantID, tupleCollection)	14✔
198	if err != nil {	14✔
199	return nil, err	×
200	}	×
201	}
202
203	if len(attributeCollection.GetAttributes()) > 0 {	35✔
204	err = w.batchUpdateAttributes(batch, xid, tenantID, buildDeleteClausesForAttributes(attributeCollection))	15✔
205	if err != nil {	15✔
206	return nil, err	×
207	}	×
208	err = w.batchInsertAttributes(batch, xid, tenantID, attributeCollection)	15✔
209	if err != nil {	15✔
210	return nil, err	×
211	}	×
212	}
213
214	batchResult := tx.SendBatch(ctx, batch)	20✔
215	for i := 0; i < batch.Len(); i++ {	141✔
216	_, err = batchResult.Exec()	121✔
217	if err != nil {	121✔
218	err = batchResult.Close()	×
219	if err != nil {	×
220	return nil, err	×
221	}	×
222	return nil, err	×
223	}
224	}
225
226	err = batchResult.Close()	20✔
227	if err != nil {	20✔
228	return nil, err	×
229	}	×
230
231	if err = tx.Commit(ctx); err != nil {	20✔
232	return nil, err	×
233	}	×
234
235	slog.DebugContext(ctx, "data successfully written to the database")	20✔
236		20✔
237	return snapshot.NewToken(xid).Encode(), nil	20✔
238	}
239
240	// delete handles the deletion of tuples and attributes from the database based on provided filters.
241	// It returns an EncodedSnapToken upon successful deletion or an error if the deletion fails.
242	func (w *DataWriter) delete(
243	ctx context.Context,
244	tenantID string,
245	tupleFilter *base.TupleFilter,
246	attributeFilter *base.AttributeFilter,
247	) (token token.EncodedSnapToken, err error) {	5✔
248	var tx pgx.Tx	5✔
249	tx, err = w.database.WritePool.BeginTx(ctx, w.txOptions)	5✔
250	if err != nil {	5✔
251	return nil, err	×
252	}	×
253
254	defer func() {	10✔
255	_ = tx.Rollback(ctx)	5✔
256	}()	5✔
257
258	var xid types.XID8	5✔
259	err = tx.QueryRow(ctx, utils.TransactionTemplate, tenantID).Scan(&xid)	5✔
260	if err != nil {	5✔
261	return nil, err	×
262	}	×
263
264	slog.DebugContext(ctx, "retrieved transaction", slog.Any("xid", xid), "for tenant", slog.Any("tenant_id", tenantID))	5✔
265		5✔
266	slog.DebugContext(ctx, "processing tuple and executing update query")	5✔
267		5✔
268	if !validation.IsTupleFilterEmpty(tupleFilter) {	8✔
269	tbuilder := w.database.Builder.Update(RelationTuplesTable).Set("expired_tx_id", xid).Where(squirrel.Eq{"expired_tx_id": "0", "tenant_id": tenantID})	3✔
270	tbuilder = utils.TuplesFilterQueryForUpdateBuilder(tbuilder, tupleFilter)	3✔
271		3✔
272	var tquery string	3✔
273	var targs []interface{}	3✔
274		3✔
275	tquery, targs, err = tbuilder.ToSql()	3✔
276	if err != nil {	3✔
277	return nil, err	×
278	}	×
279
280	_, err = tx.Exec(ctx, tquery, targs...)	3✔
281	if err != nil {	3✔
282	return nil, err	×
283	}	×
284	}
285
286	slog.DebugContext(ctx, "processing attribute and executing update query")	5✔
287		5✔
288	if !validation.IsAttributeFilterEmpty(attributeFilter) {	9✔
289	abuilder := w.database.Builder.Update(AttributesTable).Set("expired_tx_id", xid).Where(squirrel.Eq{"expired_tx_id": "0", "tenant_id": tenantID})	4✔
290	abuilder = utils.AttributesFilterQueryForUpdateBuilder(abuilder, attributeFilter)	4✔
291		4✔
292	var aquery string	4✔
293	var aargs []interface{}	4✔
294		4✔
295	aquery, aargs, err = abuilder.ToSql()	4✔
296	if err != nil {	4✔
297	return nil, err	×
298	}	×
299
300	_, err = tx.Exec(ctx, aquery, aargs...)	4✔
301	if err != nil {	4✔
302	return nil, err	×
303	}	×
304	}
305
306	if err = tx.Commit(ctx); err != nil {	5✔
307	return nil, err	×
308	}	×
309
310	slog.DebugContext(ctx, "data successfully deleted from the database")	5✔
311		5✔
312	return snapshot.NewToken(xid).Encode(), nil	5✔
313	}
314
315	// runBundle executes a series of operations defined in a DataBundle within a single database transaction.
316	// It returns an EncodedSnapToken upon successful execution of all operations or an error if any operation fails.
317	func (w *DataWriter) runBundle(
318	ctx context.Context,
319	tenantID string,
320	arguments map[string]string,
321	b *base.DataBundle,
322	) (token token.EncodedSnapToken, err error) {	1✔
323	var tx pgx.Tx	1✔
324	tx, err = w.database.WritePool.BeginTx(ctx, w.txOptions)	1✔
325	if err != nil {	1✔
326	return nil, err	×
327	}	×
328
329	defer func() {	2✔
330	_ = tx.Rollback(ctx)	1✔
331	}()	1✔
332
333	var xid types.XID8	1✔
334	err = tx.QueryRow(ctx, utils.TransactionTemplate, tenantID).Scan(&xid)	1✔
335	if err != nil {	1✔
336	return nil, err	×
337	}	×
338
339	slog.DebugContext(ctx, "retrieved transaction", slog.Any("xid", xid), "for tenant", slog.Any("tenant_id", tenantID))	1✔
340		1✔
341	batch := &pgx.Batch{}	1✔
342		1✔
343	for _, op := range b.GetOperations() {	2✔
344	tb, ab, err := bundle.Operation(arguments, op)	1✔
345	if err != nil {	1✔
346	return nil, err	×
347	}	×
348
349	err = w.runOperation(batch, xid, tenantID, tb, ab)	1✔
350	if err != nil {	1✔
351	return nil, err	×
352	}	×
353	}
354
355	batchResult := tx.SendBatch(ctx, batch)	1✔
356	for i := 0; i < batch.Len(); i++ {	10✔
357	_, err = batchResult.Exec()	9✔
358	if err != nil {	9✔
359	err = batchResult.Close()	×
360	if err != nil {	×
361	return nil, err	×
362	}	×
363	return nil, err	×
364	}
365	}
366
367	err = batchResult.Close()	1✔
368	if err != nil {	1✔
369	return nil, err	×
370	}	×
371
372	if err = tx.Commit(ctx); err != nil {	1✔
373	return nil, err	×
374	}	×
375
376	return snapshot.NewToken(xid).Encode(), nil	1✔
377	}
378
379	// runOperation processes and executes database operations defined in TupleBundle and AttributeBundle within a given transaction.
380	func (w *DataWriter) runOperation(
381	batch *pgx.Batch,
382	xid types.XID8,
383	tenantID string,
384	tb database.TupleBundle,
385	ab database.AttributeBundle,
386	) (err error) {	1✔
387	slog.Debug("processing bundles queries")	1✔
388	if len(tb.Write.GetTuples()) > 0 {	2✔
389	err = w.batchInsertRelationships(batch, xid, tenantID, &tb.Write)	1✔
390	if err != nil {	1✔
391	return err	×
392	}	×
393	}
394
395	if len(ab.Write.GetAttributes()) > 0 {	2✔
396	deleteClauses := buildDeleteClausesForAttributes(&ab.Write)	1✔
397	err = w.batchUpdateAttributes(batch, xid, tenantID, deleteClauses)	1✔
398	if err != nil {	1✔
399	return err	×
400	}	×
401
402	err = w.batchInsertAttributes(batch, xid, tenantID, &ab.Write)	1✔
403	if err != nil {	1✔
404	return err	×
405	}	×
406	}
407
408	if len(tb.Delete.GetTuples()) > 0 {	2✔
409	deleteClauses := buildDeleteClausesForRelationships(&tb.Delete)	1✔
410	err = w.batchUpdateRelationships(batch, xid, tenantID, deleteClauses)	1✔
411	if err != nil {	1✔
412	return err	×
413	}	×
414	}
415
416	if len(ab.Delete.GetAttributes()) > 0 {	2✔
417	deleteClauses := buildDeleteClausesForAttributes(&ab.Delete)	1✔
418	err = w.batchUpdateAttributes(batch, xid, tenantID, deleteClauses)	1✔
419	if err != nil {	1✔
420	return err	×
421	}	×
422	}
423
424	return nil	1✔
425	}
426
427	// batchInsertTuples function for batch inserting tuples
428	func (w DataWriter) batchInsertRelationships(batch pgx.Batch, xid types.XID8, tenantID string, tupleCollection *database.TupleCollection) error {	15✔
429	titer := tupleCollection.CreateTupleIterator()	15✔
430	for titer.HasNext() {	59✔
431	t := titer.GetNext()	44✔
432	srelation := t.GetSubject().GetRelation()	44✔
433	if srelation == tuple.ELLIPSIS {	44✔
434	srelation = ""	×
435	}	×
436	batch.Queue(	44✔
437	"INSERT INTO relation_tuples (entity_type, entity_id, relation, subject_type, subject_id, subject_relation, created_tx_id, tenant_id) VALUES ($1, $2, $3, $4, $5, $6, $7, $8) ON CONFLICT ON CONSTRAINT uq_relation_tuple_not_expired DO NOTHING",	44✔
438	t.GetEntity().GetType(), t.GetEntity().GetId(), t.GetRelation(), t.GetSubject().GetType(), t.GetSubject().GetId(), srelation, xid, tenantID,	44✔
439	)	44✔
440	}
441	return nil	15✔
442	}
443
444	// batchUpdateTuples function for batch updating tuples
445	func (w DataWriter) batchUpdateRelationships(batch pgx.Batch, xid types.XID8, tenantID string, deleteClauses []squirrel.Eq) error {	1✔
446	for _, condition := range deleteClauses {	2✔
447	query, args, err := w.database.Builder.Update(RelationTuplesTable).	1✔
448	Set("expired_tx_id", xid).	1✔
449	Where(squirrel.Eq{"expired_tx_id": "0", "tenant_id": tenantID}).	1✔
450	Where(condition).	1✔
451	ToSql()	1✔
452	if err != nil {	1✔
453	return err	×
454	}	×
455	batch.Queue(query, args...)	1✔
456	}
457	return nil	1✔
458	}
459
460	// buildDeleteClauses function to build delete clauses for tuples
461	func buildDeleteClausesForRelationships(tupleCollection *database.TupleCollection) []squirrel.Eq {	1✔
462	deleteClauses := make([]squirrel.Eq, 0)	1✔
463		1✔
464	titer := tupleCollection.CreateTupleIterator()	1✔
465	for titer.HasNext() {	2✔
466	t := titer.GetNext()	1✔
467	srelation := t.GetSubject().GetRelation()	1✔
468	if srelation == tuple.ELLIPSIS {	1✔
469	srelation = ""	×
470	}	×
471
472	condition := squirrel.Eq{	1✔
473	"entity_type": t.GetEntity().GetType(),	1✔
474	"entity_id": t.GetEntity().GetId(),	1✔
475	"relation": t.GetRelation(),	1✔
476	"subject_type": t.GetSubject().GetType(),	1✔
477	"subject_id": t.GetSubject().GetId(),	1✔
478	"subject_relation": srelation,	1✔
479	}	1✔
480		1✔
481	deleteClauses = append(deleteClauses, condition)	1✔
482	}
483
484	return deleteClauses	1✔
485	}
486
487	// batchInsertAttributes function for batch inserting attributes
488	func (w DataWriter) batchInsertAttributes(batch pgx.Batch, xid types.XID8, tenantID string, attributeCollection *database.AttributeCollection) error {	16✔
489	aiter := attributeCollection.CreateAttributeIterator()	16✔
490	for aiter.HasNext() {	58✔
491	a := aiter.GetNext()	42✔
492		42✔
493	jsonBytes, err := protojson.Marshal(a.GetValue())	42✔
494	if err != nil {	42✔
495	return err	×
496	}	×
497
498	jsonStr := string(jsonBytes)	42✔
499		42✔
500	batch.Queue(	42✔
501	"INSERT INTO attributes (entity_type, entity_id, attribute, value, created_tx_id, tenant_id) VALUES ($1, $2, $3, $4, $5, $6)",	42✔
502	a.GetEntity().GetType(), a.GetEntity().GetId(), a.GetAttribute(), jsonStr, xid, tenantID,	42✔
503	)	42✔
504	}
505	return nil	16✔
506	}
507
508	// batchUpdateAttributes function for batch updating attributes
509	func (w DataWriter) batchUpdateAttributes(batch pgx.Batch, xid types.XID8, tenantID string, deleteClauses []squirrel.Eq) error {	17✔
510	for _, condition := range deleteClauses {	60✔
511	query, args, err := w.database.Builder.Update(AttributesTable).	43✔
512	Set("expired_tx_id", xid).	43✔
513	Where(squirrel.Eq{"expired_tx_id": "0", "tenant_id": tenantID}).	43✔
514	Where(condition).	43✔
515	ToSql()	43✔
516	if err != nil {	43✔
517	return err	×
518	}	×
519
520	batch.Queue(query, args...)	43✔
521	}
522	return nil	17✔
523	}
524
525	// buildDeleteClausesForAttributes function to build delete clauses for attributes
526	func buildDeleteClausesForAttributes(attributeCollection *database.AttributeCollection) []squirrel.Eq {	17✔
527	deleteClauses := make([]squirrel.Eq, 0)	17✔
528		17✔
529	aiter := attributeCollection.CreateAttributeIterator()	17✔
530	for aiter.HasNext() {	60✔
531	a := aiter.GetNext()	43✔
532		43✔
533	condition := squirrel.Eq{	43✔
534	"entity_type": a.GetEntity().GetType(),	43✔
535	"entity_id": a.GetEntity().GetId(),	43✔
536	"attribute": a.GetAttribute(),	43✔
537	}	43✔
538		43✔
539	deleteClauses = append(deleteClauses, condition)	43✔
540	}	43✔
541
542	return deleteClauses	17✔
543	}

Permify / permify / 13719691815

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous