diff --git a/sqlitestore.go b/sqlitestore.go index dfa2caa..c89a60e 100644 --- a/sqlitestore.go +++ b/sqlitestore.go @@ -486,3 +486,20 @@ func (s *SQLiteStore) ReadTransaction(ctx context.Context, fn func(q *store.Quer return fn(st) } + +func (s *SQLiteStore) GetNumberOfEntities(ctx context.Context) (numberOfEntities uint64, err error) { + err = s.ReadTransaction(ctx, func(q *store.Queries) error { + ni, err := q.GetNumberOfEntities(ctx) + if err != nil { + return fmt.Errorf("failed to get number of entities: %w", err) + } + numberOfEntities = uint64(ni) + return nil + }) + + if err != nil { + return 0, err + } + + return numberOfEntities, nil +} diff --git a/store/db.go b/store/db.go index 69376ca..d164fbe 100644 --- a/store/db.go +++ b/store/db.go @@ -93,6 +93,9 @@ func Prepare(ctx context.Context, db DBTX) (*Queries, error) { if q.getLastBlockStmt, err = db.PrepareContext(ctx, getLastBlock); err != nil { return nil, fmt.Errorf("error preparing query GetLastBlock: %w", err) } + if q.getNumberOfEntitiesStmt, err = db.PrepareContext(ctx, getNumberOfEntities); err != nil { + return nil, fmt.Errorf("error preparing query GetNumberOfEntities: %w", err) + } if q.getNumericAttributeValueBitmapStmt, err = db.PrepareContext(ctx, getNumericAttributeValueBitmap); err != nil { return nil, fmt.Errorf("error preparing query GetNumericAttributeValueBitmap: %w", err) } @@ -237,6 +240,11 @@ func (q *Queries) Close() error { err = fmt.Errorf("error closing getLastBlockStmt: %w", cerr) } } + if q.getNumberOfEntitiesStmt != nil { + if cerr := q.getNumberOfEntitiesStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing getNumberOfEntitiesStmt: %w", cerr) + } + } if q.getNumericAttributeValueBitmapStmt != nil { if cerr := q.getNumericAttributeValueBitmapStmt.Close(); cerr != nil { err = fmt.Errorf("error closing getNumericAttributeValueBitmapStmt: %w", cerr) @@ -339,6 +347,7 @@ type Queries struct { evaluateStringAttributeValueNotGlobStmt *sql.Stmt evaluateStringAttributeValueNotInclusionStmt *sql.Stmt getLastBlockStmt *sql.Stmt + getNumberOfEntitiesStmt *sql.Stmt getNumericAttributeValueBitmapStmt *sql.Stmt getPayloadForEntityKeyStmt *sql.Stmt getStringAttributeValueBitmapStmt *sql.Stmt @@ -376,6 +385,7 @@ func (q *Queries) WithTx(tx *sql.Tx) *Queries { evaluateStringAttributeValueNotGlobStmt: q.evaluateStringAttributeValueNotGlobStmt, evaluateStringAttributeValueNotInclusionStmt: q.evaluateStringAttributeValueNotInclusionStmt, getLastBlockStmt: q.getLastBlockStmt, + getNumberOfEntitiesStmt: q.getNumberOfEntitiesStmt, getNumericAttributeValueBitmapStmt: q.getNumericAttributeValueBitmapStmt, getPayloadForEntityKeyStmt: q.getPayloadForEntityKeyStmt, getStringAttributeValueBitmapStmt: q.getStringAttributeValueBitmapStmt, diff --git a/store/querier.go b/store/querier.go index 5cbcd20..a8e27ee 100644 --- a/store/querier.go +++ b/store/querier.go @@ -32,6 +32,7 @@ type Querier interface { EvaluateStringAttributeValueNotGlob(ctx context.Context, arg EvaluateStringAttributeValueNotGlobParams) ([]*Bitmap, error) EvaluateStringAttributeValueNotInclusion(ctx context.Context, arg EvaluateStringAttributeValueNotInclusionParams) ([]*Bitmap, error) GetLastBlock(ctx context.Context) (uint64, error) + GetNumberOfEntities(ctx context.Context) (int64, error) GetNumericAttributeValueBitmap(ctx context.Context, arg GetNumericAttributeValueBitmapParams) (*Bitmap, error) GetPayloadForEntityKey(ctx context.Context, entityKey []byte) (GetPayloadForEntityKeyRow, error) GetStringAttributeValueBitmap(ctx context.Context, arg GetStringAttributeValueBitmapParams) (*Bitmap, error) diff --git a/store/queries/retrieve.sql b/store/queries/retrieve.sql index 8a77c4c..98bf897 100644 --- a/store/queries/retrieve.sql +++ b/store/queries/retrieve.sql @@ -3,3 +3,6 @@ SELECT entity_key, id, payload, content_type, string_attributes, numeric_attribu FROM payloads WHERE id IN (sqlc.slice(ids)) ORDER BY id DESC; + +-- name: GetNumberOfEntities :one +SELECT COUNT(*) FROM payloads; \ No newline at end of file diff --git a/store/retrieve.sql.go b/store/retrieve.sql.go index 8e82730..108badb 100644 --- a/store/retrieve.sql.go +++ b/store/retrieve.sql.go @@ -10,6 +10,17 @@ import ( "strings" ) +const getNumberOfEntities = `-- name: GetNumberOfEntities :one +SELECT COUNT(*) FROM payloads +` + +func (q *Queries) GetNumberOfEntities(ctx context.Context) (int64, error) { + row := q.queryRow(ctx, q.getNumberOfEntitiesStmt, getNumberOfEntities) + var count int64 + err := row.Scan(&count) + return count, err +} + const retrievePayloads = `-- name: RetrievePayloads :many SELECT entity_key, id, payload, content_type, string_attributes, numeric_attributes FROM payloads