<commit_before>package ds

import (
	"context"
	"reflect"

	"cloud.google.com/go/datastore"
)

// Tx is the datastore transaction wrapper
type Tx struct {
	*datastore.Transaction
}

// RunInTx is the RunInTransaction wrapper
func (client *Client) RunInTx(ctx context.Context, f func(tx *Tx) error, opts ...datastore.TransactionOption) (*datastore.Commit, error) {
	return client.RunInTransaction(ctx, func(t *datastore.Transaction) error {
		return f(&Tx{t})
	})
}

// GetByKey retrieves model from datastore by key
func (tx *Tx) GetByKey(key *datastore.Key, dst interface{}) error {
	err := tx.Get(key, dst)
	if err != nil {
		return err
	}
	return nil
}

// GetByKeys retrieves models from datastore by keys
func (tx *Tx) GetByKeys(keys []*datastore.Key, dst interface{}) error {
	// prepare slice if dst is pointer to 0 len slice
	if rf := reflect.ValueOf(dst); rf.Kind() == reflect.Ptr {
		rs := rf.Elem()
		if rs.Kind() == reflect.Slice && rs.Len() == 0 {
			l := len(keys)
			rs.Set(reflect.MakeSlice(rs.Type(), l, l))
		}
		dst = rs.Interface()
	}

	err := tx.GetMulti(keys, dst)
	if err != nil {
		return err
	}
	return nil
}

// GetByModel retrieves model from datastore by key from model
func (tx *Tx) GetByModel(dst interface{}) error {
	key := ExtractKey(dst)
	return tx.GetByKey(key, dst)
}

// GetByModels retrieves models from datastore by keys from models
func (tx *Tx) GetByModels(dst interface{}) error {
	keys := ExtractKeys(dst)<commit_msg>tx set keys after get<commit_after>package ds

import (
	"context"
	"reflect"

	"cloud.google.com/go/datastore"
)

// Tx is the datastore transaction wrapper
type Tx struct {
	*datastore.Transaction
}

// RunInTx is the RunInTransaction wrapper
func (client *Client) RunInTx(ctx context.Context, f func(tx *Tx) error, opts ...datastore.TransactionOption) (*datastore.Commit, error) {
	return client.RunInTransaction(ctx, func(t *datastore.Transaction) error {
		return f(&Tx{t})
	})
}

// GetByKey retrieves model from datastore by key
func (tx *Tx) GetByKey(key *datastore.Key, dst interface{}) error {
	err := tx.Get(key, dst)
	SetKey(key, dst)
	if err != nil {
		return err
	}
	return nil
}

// GetByKeys retrieves models from datastore by keys
func (tx *Tx) GetByKeys(keys []*datastore.Key, dst interface{}) error {
	// prepare slice if dst is pointer to 0 len slice
	if rf := reflect.ValueOf(dst); rf.Kind() == reflect.Ptr {
		rs := rf.Elem()
		if rs.Kind() == reflect.Slice && rs.Len() == 0 {
			l := len(keys)
			rs.Set(reflect.MakeSlice(rs.Type(), l, l))
		}
		dst = rs.Interface()
	}

	err := tx.GetMulti(keys, dst)
	SetKeys(keys, dst)
	if err != nil {
		return err
	}
	return nil
}

// GetByModel retrieves model from datastore by key from model
func (tx *Tx) GetByModel(dst interface{}) error {
	key := ExtractKey(dst)
	return tx.GetByKey(key, dst)
}

// GetByModels retrieves models from datastore by keys from models
func (tx *Tx) GetByModels(dst interface{}) error {
	keys := ExtractKeys(dst)
{'fraction_non_alphanumeric': 0.13553607552258934, 'fraction_numerical': 0.0013486176668914363, 'mean_word_length': 3.9865546218487395, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1278607', 'n_tokens_mistral': 1077, 'n_tokens_neox': 1043, 'n_words': 321}