gitbase: implement sql.PartitionCounter in all tables

Signed-off-by: Miguel Molina <[email protected]>

blobs.go

@@ -26,6 +26,7 @@ var (
 )
 
 type blobsTable struct {
+	partitioned
 	filters    []sql.Expression
 	projection []string
 	index      sql.IndexLookup
@@ -83,10 +84,6 @@ func (r *blobsTable) IndexLookup() sql.IndexLookup { return r.index }
 func (r *blobsTable) Filters() []sql.Expression    { return r.filters }
 func (r *blobsTable) Projection() []string         { return r.projection }
 
-func (r *blobsTable) Partitions(ctx *sql.Context) (sql.PartitionIter, error) {
-	return newRepositoryPartitionIter(ctx)
-}
-
 func (r *blobsTable) PartitionRows(
 	ctx *sql.Context,
 	p sql.Partition,

commit_blobs.go

@@ -10,6 +10,7 @@ import (
 )
 
 type commitBlobsTable struct {
+	partitioned
 	filters []sql.Expression
 	index   sql.IndexLookup
 }
@@ -55,10 +56,6 @@ func (t *commitBlobsTable) WithIndexLookup(idx sql.IndexLookup) sql.Table {
 func (t *commitBlobsTable) IndexLookup() sql.IndexLookup { return t.index }
 func (t *commitBlobsTable) Filters() []sql.Expression    { return t.filters }
 
-func (t *commitBlobsTable) Partitions(ctx *sql.Context) (sql.PartitionIter, error) {
-	return newRepositoryPartitionIter(ctx)
-}
-
 func (t *commitBlobsTable) PartitionRows(
 	ctx *sql.Context,
 	p sql.Partition,

commit_files.go

@@ -14,6 +14,7 @@ import (
 )
 
 type commitFilesTable struct {
+	partitioned
 	filters []sql.Expression
 	index   sql.IndexLookup
 }
@@ -60,10 +61,6 @@ func (t *commitFilesTable) WithIndexLookup(idx sql.IndexLookup) sql.Table {
 func (t *commitFilesTable) IndexLookup() sql.IndexLookup { return t.index }
 func (t *commitFilesTable) Filters() []sql.Expression    { return t.filters }
 
-func (t *commitFilesTable) Partitions(ctx *sql.Context) (sql.PartitionIter, error) {
-	return newRepositoryPartitionIter(ctx)
-}
-
 func (t *commitFilesTable) PartitionRows(
 	ctx *sql.Context,
 	p sql.Partition,

commit_trees.go

@@ -12,6 +12,7 @@ import (
 )
 
 type commitTreesTable struct {
+	partitioned
 	filters []sql.Expression
 	index   sql.IndexLookup
 }
@@ -56,10 +57,6 @@ func (t *commitTreesTable) WithIndexLookup(idx sql.IndexLookup) sql.Table {
 func (t *commitTreesTable) IndexLookup() sql.IndexLookup { return t.index }
 func (t *commitTreesTable) Filters() []sql.Expression    { return t.filters }
 
-func (t *commitTreesTable) Partitions(ctx *sql.Context) (sql.PartitionIter, error) {
-	return newRepositoryPartitionIter(ctx)
-}
-
 func (t *commitTreesTable) PartitionRows(
 	ctx *sql.Context,
 	p sql.Partition,

commits.go

@@ -10,6 +10,7 @@ import (
 )
 
 type commitsTable struct {
+	partitioned
 	filters []sql.Expression
 	index   sql.IndexLookup
 }
@@ -66,10 +67,6 @@ func (r *commitsTable) WithIndexLookup(idx sql.IndexLookup) sql.Table {
 func (r *commitsTable) IndexLookup() sql.IndexLookup { return r.index }
 func (r *commitsTable) Filters() []sql.Expression    { return r.filters }
 
-func (r *commitsTable) Partitions(ctx *sql.Context) (sql.PartitionIter, error) {
-	return newRepositoryPartitionIter(ctx)
-}
-
 func (r *commitsTable) PartitionRows(
 	ctx *sql.Context,
 	p sql.Partition,

files.go

@@ -11,6 +11,7 @@ import (
 )
 
 type filesTable struct {
+	partitioned
 	filters    []sql.Expression
 	projection []string
 	index      sql.IndexLookup
@@ -61,10 +62,6 @@ func (r *filesTable) IndexLookup() sql.IndexLookup { return r.index }
 func (r *filesTable) Filters() []sql.Expression    { return r.filters }
 func (r *filesTable) Projection() []string         { return r.projection }
 
-func (r *filesTable) Partitions(ctx *sql.Context) (sql.PartitionIter, error) {
-	return newRepositoryPartitionIter(ctx)
-}
-
 func (r *filesTable) PartitionRows(
 	ctx *sql.Context,
 	p sql.Partition,

partition.go

@@ -7,8 +7,27 @@ import (
 	"gopkg.in/src-d/go-mysql-server.v0/sql"
 )
 
+// partitioned is an embeddable helper that contains the methods for a table
+// that is partitioned by repository.
+type partitioned struct{}
+
+func (partitioned) Partitions(ctx *sql.Context) (sql.PartitionIter, error) {
+	return newRepositoryPartitionIter(ctx)
+}
+
+func (partitioned) PartitionCount(ctx *sql.Context) (int64, error) {
+	s, err := getSession(ctx)
+	if err != nil {
+		return 0, err
+	}
+
+	return int64(len(s.Pool.repositories)), nil
+}
+
+// RepositoryPartition represents a partition which is a repository id.
 type RepositoryPartition string
 
+// Key implements the sql.Partition interface.
 func (p RepositoryPartition) Key() []byte {
 	return []byte(p)
 }
@@ -41,6 +60,8 @@ func (i *repositoryPartitionIter) Close() error {
 	return nil
 }
 
+// ErrNoRepositoryPartition is returned when the partition is not a valid
+// repository partition.
 var ErrNoRepositoryPartition = errors.NewKind("%T not a valid repository partition")
 
 func getPartitionRepo(ctx *sql.Context, p sql.Partition) (*Repository, error) {

ref_commits.go

@@ -13,6 +13,7 @@ import (
 )
 
 type refCommitsTable struct {
+	partitioned
 	filters []sql.Expression
 	index   sql.IndexLookup
 }
@@ -59,10 +60,6 @@ func (t *refCommitsTable) WithIndexLookup(idx sql.IndexLookup) sql.Table {
 func (t *refCommitsTable) IndexLookup() sql.IndexLookup { return t.index }
 func (t *refCommitsTable) Filters() []sql.Expression    { return t.filters }
 
-func (t *refCommitsTable) Partitions(ctx *sql.Context) (sql.PartitionIter, error) {
-	return newRepositoryPartitionIter(ctx)
-}
-
 func (t *refCommitsTable) PartitionRows(
 	ctx *sql.Context,
 	p sql.Partition,

references.go

@@ -14,6 +14,7 @@ import (
 )
 
 type referencesTable struct {
+	partitioned
 	filters []sql.Expression
 	index   sql.IndexLookup
 }
@@ -62,10 +63,6 @@ func (r *referencesTable) WithIndexLookup(idx sql.IndexLookup) sql.Table {
 func (r *referencesTable) IndexLookup() sql.IndexLookup { return r.index }
 func (r *referencesTable) Filters() []sql.Expression    { return r.filters }
 
-func (r *referencesTable) Partitions(ctx *sql.Context) (sql.PartitionIter, error) {
-	return newRepositoryPartitionIter(ctx)
-}
-
 func (r *referencesTable) PartitionRows(
 	ctx *sql.Context,
 	p sql.Partition,

remotes.go

@@ -10,6 +10,7 @@ import (
 )
 
 type remotesTable struct {
+	partitioned
 	filters []sql.Expression
 	index   sql.IndexLookup
 }
@@ -61,10 +62,6 @@ func (r *remotesTable) WithIndexLookup(idx sql.IndexLookup) sql.Table {
 func (r *remotesTable) IndexLookup() sql.IndexLookup { return r.index }
 func (r *remotesTable) Filters() []sql.Expression    { return r.filters }
 
-func (r *remotesTable) Partitions(ctx *sql.Context) (sql.PartitionIter, error) {
-	return newRepositoryPartitionIter(ctx)
-}
-
 func (r *remotesTable) PartitionRows(
 	ctx *sql.Context,
 	p sql.Partition,

repositories.go

@@ -7,6 +7,7 @@ import (
 )
 
 type repositoriesTable struct {
+	partitioned
 	filters []sql.Expression
 	index   sql.IndexLookup
 }
@@ -54,10 +55,6 @@ func (r *repositoriesTable) WithIndexLookup(idx sql.IndexLookup) sql.Table {
 	return &nt
 }
 
-func (r *repositoriesTable) Partitions(ctx *sql.Context) (sql.PartitionIter, error) {
-	return newRepositoryPartitionIter(ctx)
-}
-
 func (r *repositoriesTable) PartitionRows(
 	ctx *sql.Context,
 	p sql.Partition,

squash.go

@@ -10,6 +10,7 @@ import (
 // SquashedTable is a table that combines the output of some tables as the
 // inputs of others with chaining so it's less expensive to compute.
 type SquashedTable struct {
+	partitioned
 	iter           ChainableIter
 	tables         []string
 	schemaMappings []int
@@ -26,10 +27,17 @@ func NewSquashedTable(
 	indexedTables []string,
 	tables ...string,
 ) *SquashedTable {
-	return &SquashedTable{iter, tables, mapping, filters, indexedTables, nil}
+	return &SquashedTable{
+		iter:           iter,
+		tables:         tables,
+		schemaMappings: mapping,
+		filters:        filters,
+		indexedTables:  indexedTables,
+	}
 }
 
 var _ sql.Table = (*SquashedTable)(nil)
+var _ sql.PartitionCounter = (*SquashedTable)(nil)
 
 // Name implements the sql.Table interface.
 func (t *SquashedTable) Name() string {
@@ -53,11 +61,6 @@ func (t *SquashedTable) Schema() sql.Schema {
 	return t.schema
 }
 
-// Partitions implements the sql.Table interface.
-func (t *SquashedTable) Partitions(ctx *sql.Context) (sql.PartitionIter, error) {
-	return newRepositoryPartitionIter(ctx)
-}
-
 // PartitionRows implements the sql.Table interface.
 func (t *SquashedTable) PartitionRows(ctx *sql.Context, p sql.Partition) (sql.RowIter, error) {
 	span, ctx := ctx.Span("gitbase.SquashedTable")

table.go

@@ -9,6 +9,7 @@ import (
 // Table represents a gitbase table.
 type Table interface {
 	sql.FilteredTable
+	sql.PartitionCounter
 	gitBase
 }
 

tree_entries.go

@@ -12,6 +12,7 @@ import (
 )
 
 type treeEntriesTable struct {
+	partitioned
 	filters []sql.Expression
 	index   sql.IndexLookup
 }
@@ -58,10 +59,6 @@ func (r *treeEntriesTable) WithIndexLookup(idx sql.IndexLookup) sql.Table {
 func (r *treeEntriesTable) IndexLookup() sql.IndexLookup { return r.index }
 func (r *treeEntriesTable) Filters() []sql.Expression    { return r.filters }
 
-func (r *treeEntriesTable) Partitions(ctx *sql.Context) (sql.PartitionIter, error) {
-	return newRepositoryPartitionIter(ctx)
-}
-
 func (r *treeEntriesTable) PartitionRows(
 	ctx *sql.Context,
 	p sql.Partition,