Method write

internal/storage/postgres/data_writer.go:167–238

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.

(
	ctx context.Context,
	tenantID string,
	tupleCollection *database.TupleCollection,
	attributeCollection *database.AttributeCollection,
)

// 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 rollback
	defer func() {
		_ = tx.Rollback(ctx)
	}()
	// Get transaction ID and snapshot
	var xid db.XID8
	var snapshotValue string
	err = tx.QueryRow(ctx, utils.TransactionTemplate, tenantID).Scan(&xid, &snapshotValue)
	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
		}
	}