<commit_before>package object

import (
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"github.com/jmoiron/sqlx"

	"github.com/tapglue/snaas/platform/flake"
	"github.com/tapglue/snaas/platform/pg"
)

const (
	orderNone ordering = iota
	orderCreatedAt
)

const (
	pgInsertObject = `INSERT INTO %s.objects(json_data) VALUES($1)`
	pgUpdateObject = `UPDATE %s.objects SET json_data = $1
		WHERE (json_data->>'id')::BIGINT = $2::BIGINT`
	pgDeleteObject = `DELETE FROM %s.objects
		WHERE (json_data->>'id')::BIGINT = $1::BIGINT`

	pgCountObjects = `SELECT count(json_data) FROM %s.objects
		%s`
	pgListObjects = `SELECT json_data FROM %s.objects
		%s`

	pgClauseAfter      = `(json_data->>'created_at') > ?`
	pgClauseBefore     = `(json_data->>'created_at') < ?`
	pgClauseDeleted    = `(json_data->>'deleted')::BOOL = ?::BOOL`
	pgClauseExternalID = `(json_data->>'external_id')::TEXT IN (?)`
	pgClauseID         = `(json_data->>'id')::BIGINT = ?::BIGINT`
	pgClauseObjectID   = `(json_data->>'object_id')::BIGINT IN (?)`
	pgClauseOwnerID    = `(json_data->>'owner_id')::BIGINT IN (?)`
	pgClauseOwned      = `(json_data->>'owned')::BOOL = ?::BOOL`
	pgClauseTags       = `(json_data->'tags')::JSONB @> '[%s]'`
	pgClauseType       = `(json_data->>'type')::TEXT IN (?)`
	pgClauseVisibility = `(json_data->>'visibility')::INT IN (?)`
	pgOrderCreatedAt   = `ORDER BY json_data->>'created_at' DESC`

	pgCreateSchema = `CREATE SCHEMA IF NOT EXISTS %s`
	pgCreateTable  = `CREATE TABLE IF NOT EXISTS %s.objects
		(json_data JSONB NOT NULL)`

	pgCreateIndexCreatedAt = `CREATE INDEX %s ON %s.objects
		USING btree ((json_data->>'created_at') DESC)`
	pgCreateIndexExternalID = `CREATE INDEX %s ON %s.objects
		USING btree (((json_data->>'external_id')::TEXT))`
	pgCreateIndexID = `CREATE INDEX %s ON %s.objects
		USING btree (((json_data->>'id')::BIGINT))`
	pgCreateIndexObjectID = `CREATE INDEX %s ON %s.objects
		USING btree (((json_data->>'object_id')::BIGINT))`
	pgCreateIndexOwnerID = `CREATE INDEX %s ON %s.objects
		USING btree (((json_data->>'owner_id')::BIGINT))`
	pgCreateIndexOwned = `CREATE INDEX %s ON %s.objects
		USING btree (((json_data->>'owned')::BOOL))`
	pgCreateIndexTags = `CREATE INDEX %s ON %s.objects
		USING gin ((json_data->'tags'))`
	pgCreateIndexType = `CREATE INDEX %s ON %s.objects
		USING btree (((json_data->>'type')::TEXT))`
	pgCreateIndexVisibility = `CREATE INDEX %s ON %s.objects
		USING btree (((json_data->>'visibility')::INT))`
	pgCreateIndexPostAll = `
		CREATE INDEX
			%s
		ON
			%s.objects ((json_data->>'created_at') DESC)
		WHERE
			(json_data->>'deleted')::BOOL = false
		    AND (json_data->>'owned')::BOOL = true
		    AND (json_data->>'type')::TEXT IN ('tg_post')
		    AND (json_data->>'visibility')::INT IN (30, 40)`

	pgDropTable = `DROP TABLE IF EXISTS %s.objects`
)

type ordering int

type pgService struct {
	db *sqlx.DB
}

// PostgresService returns a Postgres based Service implementation.
func PostgresService(db *sqlx.DB) Service {
	return &pgService{
		db: db,
	}
}

func (s *pgService) Count(ns string, opts QueryOptions) (int, error) {
	where, params, err := convertOpts(opts, orderNone)
	if err != nil {
		return 0, err
	}

	return s.countObjects(ns, where, params...)
}

func (s *pgService) CountMulti(
	ns string,
	objectIDs ...uint64,
) (m CountsMap, err error) {
	tx, err := s.db.Beginx()
	if err != nil {
		return nil, err
	}

	defer func(tx *sqlx.Tx) {
		if err != nil {
			_ = tx.Rollback()
		}
	}(tx)

	var (
		countsMap = CountsMap{}
		owned     = true
	)

	for _, oid := range objectIDs {
		where, params, err := convertOpts(QueryOptions{
			Deleted: false,
			ObjectIDs: []uint64{
				oid,
			},
			Owned: &owned,
			Types: []string{
				TypeComment,
			},
		}, orderNone)
		if err != nil {
			return nil, err
		}

		var (
			query = fmt.Sprintf(pgCountObjects, ns, where)

			count uint64
		)

		err = tx.Get(&count, query, params...)
		if err != nil {
			return nil, err
		}

		countsMap[oid] = Counts{
			Comments: uint64(count),
		}
	}

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

	return countsMap, nil
}

func (s *pgService) Put(ns string, object *Object) (*Object, error) {<commit_msg>Be smarter comment multiple count<commit_after>package object

import (
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"github.com/jmoiron/sqlx"

	"github.com/tapglue/snaas/platform/flake"
	"github.com/tapglue/snaas/platform/pg"
)

const (
	orderNone ordering = iota
	orderCreatedAt
)

const (
	pgInsertObject = `INSERT INTO %s.objects(json_data) VALUES($1)`
	pgUpdateObject = `UPDATE %s.objects SET json_data = $1
		WHERE (json_data->>'id')::BIGINT = $2::BIGINT`
	pgDeleteObject = `DELETE FROM %s.objects
		WHERE (json_data->>'id')::BIGINT = $1::BIGINT`

	pgCountObjects = `SELECT count(json_data) FROM %s.objects
		%s`
	pgCountObjectsMulti = `
		SELECT
			json_data->>'object_id',
			count(*)
		FROM
			%s.objects
		WHERE
			(json_data->>'deleted')::BOOL = false
			AND (json_data->>'object_id')::BIGINT IN (?)
			AND (json_data->>'owned')::BOOL = true
			AND (json_data->>'type')::TEXT = 'tg_comment'
		GROUP BY
			json_data->>'object_id'`
	pgListObjects = `SELECT json_data FROM %s.objects
		%s`

	pgClauseAfter      = `(json_data->>'created_at') > ?`
	pgClauseBefore     = `(json_data->>'created_at') < ?`
	pgClauseDeleted    = `(json_data->>'deleted')::BOOL = ?::BOOL`
	pgClauseExternalID = `(json_data->>'external_id')::TEXT IN (?)`
	pgClauseID         = `(json_data->>'id')::BIGINT = ?::BIGINT`
	pgClauseObjectID   = `(json_data->>'object_id')::BIGINT IN (?)`
	pgClauseOwnerID    = `(json_data->>'owner_id')::BIGINT IN (?)`
	pgClauseOwned      = `(json_data->>'owned')::BOOL = ?::BOOL`
	pgClauseTags       = `(json_data->'tags')::JSONB @> '[%s]'`
	pgClauseType       = `(json_data->>'type')::TEXT IN (?)`
	pgClauseVisibility = `(json_data->>'visibility')::INT IN (?)`
	pgOrderCreatedAt   = `ORDER BY json_data->>'created_at' DESC`

	pgCreateSchema = `CREATE SCHEMA IF NOT EXISTS %s`
	pgCreateTable  = `CREATE TABLE IF NOT EXISTS %s.objects
		(json_data JSONB NOT NULL)`

	pgCreateIndexCreatedAt = `CREATE INDEX %s ON %s.objects
		USING btree ((json_data->>'created_at') DESC)`
	pgCreateIndexExternalID = `CREATE INDEX %s ON %s.objects
		USING btree (((json_data->>'external_id')::TEXT))`
	pgCreateIndexID = `CREATE INDEX %s ON %s.objects
		USING btree (((json_data->>'id')::BIGINT))`
	pgCreateIndexObjectID = `CREATE INDEX %s ON %s.objects
		USING btree (((json_data->>'object_id')::BIGINT))`
	pgCreateIndexOwnerID = `CREATE INDEX %s ON %s.objects
		USING btree (((json_data->>'owner_id')::BIGINT))`
	pgCreateIndexOwned = `CREATE INDEX %s ON %s.objects
		USING btree (((json_data->>'owned')::BOOL))`
	pgCreateIndexTags = `CREATE INDEX %s ON %s.objects
		USING gin ((json_data->'tags'))`
	pgCreateIndexType = `CREATE INDEX %s ON %s.objects
		USING btree (((json_data->>'type')::TEXT))`
	pgCreateIndexVisibility = `CREATE INDEX %s ON %s.objects
		USING btree (((json_data->>'visibility')::INT))`
	pgCreateIndexPostAll = `
		CREATE INDEX
			%s
		ON
			%s.objects ((json_data->>'created_at') DESC)
		WHERE
			(json_data->>'deleted')::BOOL = false
		    AND (json_data->>'owned')::BOOL = true
		    AND (json_data->>'type')::TEXT IN ('tg_post')
		    AND (json_data->>'visibility')::INT IN (30, 40)`

	pgDropTable = `DROP TABLE IF EXISTS %s.objects`
)

type ordering int

type pgService struct {
	db *sqlx.DB
}

// PostgresService returns a Postgres based Service implementation.
func PostgresService(db *sqlx.DB) Service {
	return &pgService{
		db: db,
	}
}

func (s *pgService) Count(ns string, opts QueryOptions) (int, error) {
	where, params, err := convertOpts(opts, orderNone)
	if err != nil {
		return 0, err
	}

	return s.countObjects(ns, where, params...)
}

func (s *pgService) CountMulti(
	ns string,
	objectIDs ...uint64,
) (m CountsMap, err error) {
	var (
		countsMap = CountsMap{}
		ps        = []interface{}{}
	)

	if len(objectIDs) == 0 {
		return countsMap, nil
	}

	for _, id := range objectIDs {
		ps = append(ps, id)
	}

	query, _, err := sqlx.In(pgCountObjectsMulti, ps)
	if err != nil {
		return nil, err
	}

	query = sqlx.Rebind(sqlx.DOLLAR, query)
	query = fmt.Sprintf(query, ns)

	rows, err := s.db.Query(query, ps...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	for rows.Next() {
		var (
			objectID uint64
			count    uint64
		)

		err := rows.Scan(&objectID, &count)
		if err != nil {
			return nil, err
		}

		countsMap[objectID] = Counts{
			Comments: count,
		}
	}

	if err := rows.Err(); err != nil {
		return nil, err
	}

	return countsMap, nil
}

func (s *pgService) Put(ns string, object *Object) (*Object, error) {
{'fraction_non_alphanumeric': 0.16933240611961056, 'fraction_numerical': 0.0038247566063977748, 'mean_word_length': 4.052107728337236, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 127, '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': '1176766', 'n_tokens_mistral': 3484, 'n_tokens_neox': 3274, 'n_words': 745}