wehub-resource-sync
349 lines
12 KiB
Go
349 lines
12 KiB
Go
// Copyright 2024 Dolthub, Inc.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
package pgcatalog
|
|
|
|
import (
|
|
"io"
|
|
"strings"
|
|
|
|
"github.com/dolthub/go-mysql-server/sql"
|
|
|
|
"github.com/dolthub/doltgresql/core/id"
|
|
"github.com/dolthub/doltgresql/server/functions"
|
|
"github.com/dolthub/doltgresql/server/tables"
|
|
pgtypes "github.com/dolthub/doltgresql/server/types"
|
|
)
|
|
|
|
// PgIndexName is a constant to the pg_index name.
|
|
const PgIndexName = "pg_index"
|
|
|
|
// InitPgIndex handles registration of the pg_index handler.
|
|
func InitPgIndex() {
|
|
tables.AddHandler(PgCatalogName, PgIndexName, PgIndexHandler{})
|
|
}
|
|
|
|
// PgIndexHandler is the handler for the pg_index table.
|
|
type PgIndexHandler struct{}
|
|
|
|
var _ tables.Handler = PgIndexHandler{}
|
|
var _ tables.IndexedTableHandler = PgIndexHandler{}
|
|
|
|
// Name implements the interface tables.Handler.
|
|
func (p PgIndexHandler) Name() string {
|
|
return PgIndexName
|
|
}
|
|
|
|
// RowIter implements the interface tables.Handler.
|
|
func (p PgIndexHandler) RowIter(ctx *sql.Context, partition sql.Partition) (sql.RowIter, error) {
|
|
// Use cached data from this session if it exists
|
|
pgCatalogCache, err := getPgCatalogCache(ctx)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
if pgCatalogCache.pgIndexes == nil {
|
|
err = cachePgIndexes(ctx, pgCatalogCache)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
|
|
if indexIdxPart, ok := partition.(inMemIndexPartition); ok {
|
|
return &inMemIndexScanIter[*pgIndex]{
|
|
lookup: indexIdxPart.lookup,
|
|
rangeConverter: p,
|
|
btreeAccess: pgCatalogCache.pgIndexes,
|
|
rowConverter: pgIndexToRow,
|
|
}, nil
|
|
}
|
|
|
|
return &pgIndexTableScanIter{
|
|
indexCache: pgCatalogCache.pgIndexes,
|
|
idx: 0,
|
|
}, nil
|
|
}
|
|
|
|
// PkSchema implements the interface tables.Handler.
|
|
func (p PgIndexHandler) PkSchema() sql.PrimaryKeySchema {
|
|
return sql.PrimaryKeySchema{
|
|
Schema: pgIndexSchema,
|
|
PkOrdinals: nil,
|
|
}
|
|
}
|
|
|
|
// Indexes implements tables.IndexedTableHandler.
|
|
func (p PgIndexHandler) Indexes() ([]sql.Index, error) {
|
|
return []sql.Index{
|
|
pgCatalogInMemIndex{
|
|
name: "pg_index_indexrelid_index",
|
|
tblName: "pg_index",
|
|
dbName: "pg_catalog",
|
|
uniq: true,
|
|
columnExprs: []sql.ColumnExpressionType{{Expression: "pg_index.indexrelid", Type: pgtypes.Oid}},
|
|
},
|
|
pgCatalogInMemIndex{
|
|
name: "pg_index_indrelid_index",
|
|
tblName: "pg_index",
|
|
dbName: "pg_catalog",
|
|
uniq: false,
|
|
columnExprs: []sql.ColumnExpressionType{{Expression: "pg_index.indrelid", Type: pgtypes.Oid}},
|
|
},
|
|
}, nil
|
|
}
|
|
|
|
// LookupPartitions implements tables.IndexedTableHandler.
|
|
func (p PgIndexHandler) LookupPartitions(context *sql.Context, lookup sql.IndexLookup) (sql.PartitionIter, error) {
|
|
return &inMemIndexPartIter{
|
|
part: inMemIndexPartition{
|
|
idxName: lookup.Index.(pgCatalogInMemIndex).name,
|
|
lookup: lookup,
|
|
},
|
|
}, nil
|
|
}
|
|
|
|
// getIndexScanRange implements the interface RangeConverter.
|
|
func (p PgIndexHandler) getIndexScanRange(rng sql.Range, index sql.Index) (*pgIndex, bool, *pgIndex, bool) {
|
|
var gte, lt *pgIndex
|
|
var hasLowerBound, hasUpperBound bool
|
|
|
|
switch index.(pgCatalogInMemIndex).name {
|
|
case "pg_index_indexrelid_index":
|
|
msrng := rng.(sql.MySQLRange)
|
|
oidRng := msrng[0]
|
|
if oidRng.HasLowerBound() {
|
|
lb := sql.GetMySQLRangeCutKey(oidRng.LowerBound)
|
|
if lb != nil {
|
|
lowerRangeCutKey := lb.(id.Id)
|
|
gte = &pgIndex{
|
|
indexOidNative: idToOid(lowerRangeCutKey),
|
|
}
|
|
hasLowerBound = true
|
|
}
|
|
}
|
|
if oidRng.HasUpperBound() {
|
|
ub := sql.GetMySQLRangeCutKey(oidRng.UpperBound)
|
|
if ub != nil {
|
|
upperRangeCutKey := ub.(id.Id)
|
|
lt = &pgIndex{
|
|
indexOidNative: idToOid(upperRangeCutKey) + 1,
|
|
}
|
|
hasUpperBound = true
|
|
}
|
|
}
|
|
|
|
case "pg_index_indrelid_index":
|
|
msrng := rng.(sql.MySQLRange)
|
|
oidRng := msrng[0]
|
|
if oidRng.HasLowerBound() {
|
|
lb := sql.GetMySQLRangeCutKey(oidRng.LowerBound)
|
|
if lb != nil {
|
|
lowerRangeCutKey := lb.(id.Id)
|
|
gte = &pgIndex{
|
|
tableOidNative: idToOid(lowerRangeCutKey),
|
|
}
|
|
hasLowerBound = true
|
|
}
|
|
}
|
|
if oidRng.HasUpperBound() {
|
|
ub := sql.GetMySQLRangeCutKey(oidRng.UpperBound)
|
|
if ub != nil {
|
|
upperRangeCutKey := ub.(id.Id)
|
|
lt = &pgIndex{
|
|
tableOidNative: idToOid(upperRangeCutKey) + 1,
|
|
}
|
|
hasUpperBound = true
|
|
}
|
|
}
|
|
default:
|
|
panic("unknown index name: " + index.(pgCatalogInMemIndex).name)
|
|
}
|
|
|
|
return gte, hasLowerBound, lt, hasUpperBound
|
|
}
|
|
|
|
// pgIndexSchema is the schema for pg_index.
|
|
var pgIndexSchema = sql.Schema{
|
|
{Name: "indexrelid", Type: pgtypes.Oid, Default: nil, Nullable: false, Source: PgIndexName},
|
|
{Name: "indrelid", Type: pgtypes.Oid, Default: nil, Nullable: false, Source: PgIndexName},
|
|
{Name: "indnatts", Type: pgtypes.Int16, Default: nil, Nullable: false, Source: PgIndexName},
|
|
{Name: "indnkeyatts", Type: pgtypes.Int16, Default: nil, Nullable: false, Source: PgIndexName},
|
|
{Name: "indisunique", Type: pgtypes.Bool, Default: nil, Nullable: false, Source: PgIndexName},
|
|
{Name: "indnullsnotdistinct", Type: pgtypes.Bool, Default: nil, Nullable: false, Source: PgIndexName},
|
|
{Name: "indisprimary", Type: pgtypes.Bool, Default: nil, Nullable: false, Source: PgIndexName},
|
|
{Name: "indisexclusion", Type: pgtypes.Bool, Default: nil, Nullable: false, Source: PgIndexName},
|
|
{Name: "indimmediate", Type: pgtypes.Bool, Default: nil, Nullable: false, Source: PgIndexName},
|
|
{Name: "indisclustered", Type: pgtypes.Bool, Default: nil, Nullable: false, Source: PgIndexName},
|
|
{Name: "indisvalid", Type: pgtypes.Bool, Default: nil, Nullable: false, Source: PgIndexName},
|
|
{Name: "indcheckxmin", Type: pgtypes.Bool, Default: nil, Nullable: false, Source: PgIndexName},
|
|
{Name: "indisready", Type: pgtypes.Bool, Default: nil, Nullable: false, Source: PgIndexName},
|
|
{Name: "indislive", Type: pgtypes.Bool, Default: nil, Nullable: false, Source: PgIndexName},
|
|
{Name: "indisreplident", Type: pgtypes.Bool, Default: nil, Nullable: false, Source: PgIndexName},
|
|
{Name: "indkey", Type: pgtypes.Int16vector, Default: nil, Nullable: false, Source: PgIndexName},
|
|
{Name: "indcollation", Type: pgtypes.Oidvector, Default: nil, Nullable: false, Source: PgIndexName},
|
|
{Name: "indclass", Type: pgtypes.Oidvector, Default: nil, Nullable: false, Source: PgIndexName},
|
|
{Name: "indoption", Type: pgtypes.Int16vector, Default: nil, Nullable: false, Source: PgIndexName},
|
|
{Name: "indexprs", Type: pgtypes.Text, Default: nil, Nullable: true, Source: PgIndexName}, // TODO: type pg_node_tree, collation C
|
|
{Name: "indpred", Type: pgtypes.Text, Default: nil, Nullable: true, Source: PgIndexName}, // TODO: type pg_node_tree, collation C
|
|
}
|
|
|
|
func extractColName(expr string) string {
|
|
// TODO: this breaks for column names that contain a `.`, but this is a problem that happens
|
|
// throughout index analysis in the engine
|
|
lastDot := strings.LastIndex(expr, ".")
|
|
return expr[lastDot+1:]
|
|
}
|
|
|
|
// pgIndex represents a row in the pg_index table.
|
|
// We store oids in their native format as well so that we can do range scans on them.
|
|
type pgIndex struct {
|
|
index sql.Index
|
|
schemaName string
|
|
indexOid id.Id
|
|
indexOidNative uint32
|
|
tableOid id.Id
|
|
tableOidNative uint32
|
|
indnatts int16
|
|
indisunique bool
|
|
indisprimary bool
|
|
indkey []any
|
|
}
|
|
|
|
// lessIndexOid is a sort function for pgIndex based on indexrelid.
|
|
func lessIndexOid(a, b *pgIndex) bool {
|
|
return a.indexOidNative < b.indexOidNative
|
|
}
|
|
|
|
// lessIndrelid is a sort function for pgIndex based on indrelid.
|
|
func lessIndrelid(a, b []*pgIndex) bool {
|
|
return a[0].tableOidNative < b[0].tableOidNative
|
|
}
|
|
|
|
// pgIndexTableScanIter is the sql.RowIter for the pg_index table.
|
|
type pgIndexTableScanIter struct {
|
|
indexCache *pgIndexCache
|
|
idx int
|
|
}
|
|
|
|
var _ sql.RowIter = (*pgIndexTableScanIter)(nil)
|
|
|
|
// Next implements the interface sql.RowIter.
|
|
func (iter *pgIndexTableScanIter) Next(ctx *sql.Context) (sql.Row, error) {
|
|
if iter.idx >= len(iter.indexCache.indexes) {
|
|
return nil, io.EOF
|
|
}
|
|
iter.idx++
|
|
index := iter.indexCache.indexes[iter.idx-1]
|
|
|
|
return pgIndexToRow(index), nil
|
|
}
|
|
|
|
// Close implements the interface sql.RowIter.
|
|
func (iter *pgIndexTableScanIter) Close(ctx *sql.Context) error {
|
|
return nil
|
|
}
|
|
|
|
// pgIndexToRow converts a pgIndex to a sql.Row.
|
|
func pgIndexToRow(index *pgIndex) sql.Row {
|
|
return sql.Row{
|
|
index.indexOid, // indexrelid
|
|
index.tableOid, // indrelid
|
|
index.indnatts, // indnatts
|
|
int16(0), // indnkeyatts
|
|
index.index.IsUnique(), // indisunique
|
|
false, // indnullsnotdistinct
|
|
index.indisprimary, // indisprimary
|
|
false, // indisexclusion
|
|
false, // indimmediate
|
|
false, // indisclustered
|
|
true, // indisvalid
|
|
false, // indcheckxmin
|
|
true, // indisready
|
|
true, // indislive
|
|
false, // indisreplident
|
|
index.indkey, // indkey
|
|
[]any{}, // indcollation
|
|
[]any{}, // indclass
|
|
[]any{int16(0)}, // indoption
|
|
nil, // indexprs
|
|
indexPred(index.index), // indpred
|
|
}
|
|
}
|
|
|
|
// indexPred returns the predicate expression string for partial indexes, or nil for full indexes.
|
|
func indexPred(idx sql.Index) interface{} {
|
|
if pi, ok := idx.(sql.PartialIndex); ok && pi.Predicate() != "" {
|
|
return pi.Predicate()
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// cachePgIndexes caches the pg_index data for the current database in the session.
|
|
func cachePgIndexes(ctx *sql.Context, pgCatalogCache *pgCatalogCache) error {
|
|
var indexes []*pgIndex
|
|
indexOidIdx := NewUniqueInMemIndexStorage[*pgIndex](lessIndexOid)
|
|
indrelidIdx := NewNonUniqueInMemIndexStorage[*pgIndex](lessIndrelid)
|
|
|
|
tableSchemas := make(map[id.Id]sql.Schema)
|
|
tableNames := make(map[id.Id]string)
|
|
|
|
err := functions.IterateCurrentDatabase(ctx, functions.Callbacks{
|
|
Index: func(ctx *sql.Context, schema functions.ItemSchema, table functions.ItemTable, index functions.ItemIndex) (cont bool, err error) {
|
|
if tableSchemas[table.OID.AsId()] == nil {
|
|
tableSchemas[table.OID.AsId()] = table.Item.Schema(ctx)
|
|
}
|
|
if tableNames[table.OID.AsId()] == "" {
|
|
tableNames[table.OID.AsId()] = table.Item.Name()
|
|
}
|
|
|
|
s := tableSchemas[table.OID.AsId()]
|
|
indKey := make([]any, len(index.Item.Expressions()))
|
|
for i, expr := range index.Item.Expressions() {
|
|
colName := extractColName(expr)
|
|
indKey[i] = int16(s.IndexOfColName(colName)) + 1
|
|
}
|
|
|
|
pgIdx := &pgIndex{
|
|
index: index.Item,
|
|
schemaName: schema.Item.SchemaName(),
|
|
indexOid: index.OID.AsId(),
|
|
indexOidNative: id.Cache().ToOID(index.OID.AsId()),
|
|
tableOid: table.OID.AsId(),
|
|
tableOidNative: id.Cache().ToOID(table.OID.AsId()),
|
|
indkey: indKey,
|
|
indnatts: int16(len(index.Item.Expressions())),
|
|
indisunique: index.Item.IsUnique(),
|
|
indisprimary: strings.ToLower(index.Item.ID()) == "primary",
|
|
}
|
|
|
|
indexOidIdx.Add(pgIdx)
|
|
indrelidIdx.Add(pgIdx)
|
|
indexes = append(indexes, pgIdx)
|
|
return true, nil
|
|
},
|
|
})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
pgCatalogCache.pgIndexes = &pgIndexCache{
|
|
indexes: indexes,
|
|
tableNames: tableNames,
|
|
indexOidIdx: indexOidIdx,
|
|
indrelidIdx: indrelidIdx,
|
|
}
|
|
|
|
return nil
|
|
}
|