diff --git a/docs/generated/sql/functions.md b/docs/generated/sql/functions.md
index 6cf96ce88c52..08d0973dbeb4 100644
--- a/docs/generated/sql/functions.md
+++ b/docs/generated/sql/functions.md
@@ -991,6 +991,9 @@ SELECT * FROM crdb_internal.check_consistency(true, ‘\x02’, ‘\x04’)
current_user() → string | Returns the current user. This function is provided for compatibility with PostgreSQL.
|
+locality() → tuple | Returns the hierarchical location of the current node as a tuple of labeled values, ordered from most inclusive to least inclusive.
+For example: region=east,datacenter=us-east-1.
+ |
version() → string | Returns the node’s version of CockroachDB.
|
diff --git a/pkg/sql/conn_executor.go b/pkg/sql/conn_executor.go
index dbcbb59e6509..30630a44b2a6 100644
--- a/pkg/sql/conn_executor.go
+++ b/pkg/sql/conn_executor.go
@@ -1877,6 +1877,7 @@ func (ex *connExecutor) resetPlanner(
p.semaCtx.Location = &ex.sessionData.DataConversion.Location
p.semaCtx.SearchPath = ex.sessionData.SearchPath
p.semaCtx.AsOfTimestamp = nil
+ p.semaCtx.Locality = ex.server.cfg.Locality
ex.resetEvalCtx(&p.extendedEvalCtx, txn, stmtTS)
diff --git a/pkg/sql/logictest/testdata/logic_test/locality b/pkg/sql/logictest/testdata/logic_test/locality
new file mode 100644
index 000000000000..49854a96701d
--- /dev/null
+++ b/pkg/sql/logictest/testdata/logic_test/locality
@@ -0,0 +1,19 @@
+# LogicTest: 5node-dist 5node-dist-opt
+
+query T
+SELECT locality()
+----
+(test,dc1)
+
+query T
+SELECT (locality()).region
+----
+test
+
+query T
+SELECT (locality()).dc
+----
+dc1
+
+statement error could not identify column "unk" in tuple{string AS region, string AS dc}
+SELECT (locality()).unk
diff --git a/pkg/sql/opt/memo/typing.go b/pkg/sql/opt/memo/typing.go
index a1643980564f..4d60ff4d1d56 100644
--- a/pkg/sql/opt/memo/typing.go
+++ b/pkg/sql/opt/memo/typing.go
@@ -296,7 +296,7 @@ func typeAsAggregate(e opt.ScalarExpr) *types.T {
// types (i.e. pass nil to the ReturnTyper). Aggregates with return types
// that depend on argument types are handled separately.
_, overload := FindAggregateOverload(e)
- t := overload.ReturnType(nil)
+ t := overload.ReturnType(nil /* ctx */, nil /* args */)
if t == tree.UnknownReturnType {
panic(pgerror.AssertionFailedf("unknown aggregate return type. e:\n%s", e))
}
@@ -307,7 +307,7 @@ func typeAsAggregate(e opt.ScalarExpr) *types.T {
// typeAsAggregate.
func typeAsWindow(e opt.ScalarExpr) *types.T {
_, overload := FindWindowOverload(e)
- t := overload.ReturnType(nil)
+ t := overload.ReturnType(nil /* ctx */, nil /* args */)
if t == tree.UnknownReturnType {
panic(pgerror.AssertionFailedf("unknown window return type. e:\n%s", e))
}
diff --git a/pkg/sql/opt/memo/typing_test.go b/pkg/sql/opt/memo/typing_test.go
index b36d4a46da92..b8ad88e894c9 100644
--- a/pkg/sql/opt/memo/typing_test.go
+++ b/pkg/sql/opt/memo/typing_test.go
@@ -205,7 +205,7 @@ func TestTypingAggregateAssumptions(t *testing.T) {
}
// Check for fixed return types.
- retType := overload.ReturnType(nil)
+ retType := overload.ReturnType(nil /* ctx */, nil /* args */)
if retType == tree.UnknownReturnType {
t.Errorf("return type is not fixed for %s: %+v", name, overload.Types.Types())
}
diff --git a/pkg/sql/opt/norm/custom_funcs.go b/pkg/sql/opt/norm/custom_funcs.go
index 12c65bd4174d..cdc7bdd75239 100644
--- a/pkg/sql/opt/norm/custom_funcs.go
+++ b/pkg/sql/opt/norm/custom_funcs.go
@@ -1837,4 +1837,5 @@ var FoldFunctionWhitelist = map[string]struct{}{
"jsonb_strip_nulls": {},
"json_array_length": {},
"jsonb_array_length": {},
+ "locality": {},
}
diff --git a/pkg/sql/opt/optbuilder/testdata/scalar b/pkg/sql/opt/optbuilder/testdata/scalar
index 5b5347f25950..66e605ad213a 100644
--- a/pkg/sql/opt/optbuilder/testdata/scalar
+++ b/pkg/sql/opt/optbuilder/testdata/scalar
@@ -1175,3 +1175,29 @@ if-err [type=decimal]
│ └── variable: @1 [type=decimal]
└── err-code
└── const: '10000' [type=string]
+
+build-scalar
+locality()
+----
+function: locality [type=tuple]
+
+build locality=(region=east,dc=east1-b)
+SELECT (locality()).dc
+----
+project
+ ├── columns: dc:1(string)
+ ├── values
+ │ └── tuple [type=tuple]
+ └── projections
+ └── column-access: 1 [type=string]
+ └── function: locality [type=tuple{string AS region, string AS dc}]
+
+build
+SELECT (locality()).unk
+----
+error (42804): could not identify column "unk" in tuple
+
+build locality=(region=east,dc=east1-b)
+SELECT (locality()).unk
+----
+error (42804): could not identify column "unk" in tuple{string AS region, string AS dc}
diff --git a/pkg/sql/opt/testutils/opttester/opt_tester.go b/pkg/sql/opt/testutils/opttester/opt_tester.go
index 4586186d3d6a..b140047f1d4d 100644
--- a/pkg/sql/opt/testutils/opttester/opt_tester.go
+++ b/pkg/sql/opt/testutils/opttester/opt_tester.go
@@ -257,6 +257,7 @@ func (ot *OptTester) RunCommand(tb testing.TB, d *datadriven.TestData) string {
ot.Flags.Verbose = testing.Verbose()
ot.evalCtx.TestingKnobs.OptimizerCostPerturbation = ot.Flags.PerturbCost
ot.evalCtx.Locality = ot.Flags.Locality
+ ot.semaCtx.Locality = ot.Flags.Locality
switch d.Cmd {
case "exec-ddl":
diff --git a/pkg/sql/pg_catalog.go b/pkg/sql/pg_catalog.go
index 4653abfc84fc..e083effb6713 100644
--- a/pkg/sql/pg_catalog.go
+++ b/pkg/sql/pg_catalog.go
@@ -1510,7 +1510,7 @@ CREATE TABLE pg_catalog.pg_operator (
panic(fmt.Sprintf("Unexpected operator %s with %d params",
opName, params.Length()))
}
- returnType := tree.NewDOid(tree.DInt(returnTyper(nil).Oid()))
+ returnType := tree.NewDOid(tree.DInt(returnTyper(nil /* ctx */, nil /* args */).Oid()))
err := addRow(
h.OperatorOid(opName, leftType, rightType, returnType), // oid
diff --git a/pkg/sql/planner.go b/pkg/sql/planner.go
index 3ad6fb35b7f0..eae705e76226 100644
--- a/pkg/sql/planner.go
+++ b/pkg/sql/planner.go
@@ -252,6 +252,7 @@ func newInternalPlanner(
p.semaCtx = tree.MakeSemaContext()
p.semaCtx.Location = &sd.DataConversion.Location
p.semaCtx.SearchPath = sd.SearchPath
+ p.semaCtx.Locality = execCfg.Locality
plannerMon := mon.MakeUnlimitedMonitor(ctx,
fmt.Sprintf("internal-planner.%s.%s", user, opName),
diff --git a/pkg/sql/sem/builtins/aggregate_builtins.go b/pkg/sql/sem/builtins/aggregate_builtins.go
index e92b190dcd69..d48d7785cf12 100644
--- a/pkg/sql/sem/builtins/aggregate_builtins.go
+++ b/pkg/sql/sem/builtins/aggregate_builtins.go
@@ -107,7 +107,7 @@ var aggregates = map[string]builtinDefinition{
arrayBuiltin(func(t *types.T) tree.Overload {
return makeAggOverloadWithReturnType(
[]*types.T{t},
- func(args []tree.TypedExpr) *types.T {
+ func(ctx *tree.SemaContext, args []tree.TypedExpr) *types.T {
if len(args) == 0 {
return types.MakeArray(t)
}
diff --git a/pkg/sql/sem/builtins/builtins.go b/pkg/sql/sem/builtins/builtins.go
index 7f71a818b805..f05e556fdd92 100644
--- a/pkg/sql/sem/builtins/builtins.go
+++ b/pkg/sql/sem/builtins/builtins.go
@@ -2756,6 +2756,31 @@ may increase either contention or retry errors, or both.`,
},
),
+ "locality": makeBuiltin(
+ tree.FunctionProperties{Category: categorySystemInfo},
+ tree.Overload{
+ Types: tree.ArgTypes{},
+ ReturnType: func(ctx *tree.SemaContext, _ []tree.TypedExpr) *types.T {
+ if ctx == nil {
+ // Use simplified tuple type for signature (e.g. for docs).
+ return types.AnyTuple
+ }
+ return localityReturnType(&ctx.Locality)
+ },
+ Fn: func(ctx *tree.EvalContext, args tree.Datums) (tree.Datum, error) {
+ datums := make(tree.Datums, len(ctx.Locality.Tiers))
+ for i := range ctx.Locality.Tiers {
+ datums[i] = tree.NewDString(ctx.Locality.Tiers[i].Value)
+ }
+ typ := localityReturnType(&ctx.Locality)
+ return tree.NewDTuple(typ, datums...), nil
+ },
+ Info: "Returns the hierarchical location of the current node as a tuple " +
+ "of labeled values, ordered from most inclusive to least inclusive. \n\n" +
+ "For example: `region=east,datacenter=us-east-1`.",
+ },
+ ),
+
"crdb_internal.node_executable_version": makeBuiltin(
tree.FunctionProperties{Category: categorySystemInfo},
tree.Overload{
@@ -4628,3 +4653,27 @@ func recentTimestamp(ctx *tree.EvalContext) (time.Time, error) {
}
return ctx.StmtTimestamp.Add(offset), nil
}
+
+// localityReturnType returns the type that the locality() function will return.
+// This is a tuple type with labeled string fields, one for each locality tier:
+//
+// tuple{string AS , string AS , ...}
+//
+// For example, consider this locality:
+//
+// region=east,datacenter=us-east-1
+//
+// The tuple type will be:
+//
+// tuple{string AS region, string AS datacenter}
+//
+func localityReturnType(locality *roachpb.Locality) *types.T {
+ contents := make([]types.T, len(locality.Tiers))
+ labels := make([]string, len(locality.Tiers))
+ for i := range locality.Tiers {
+ tier := &locality.Tiers[i]
+ contents[i] = *types.String
+ labels[i] = tier.Key
+ }
+ return types.MakeLabeledTuple(contents, labels)
+}
diff --git a/pkg/sql/sem/builtins/generator_builtins.go b/pkg/sql/sem/builtins/generator_builtins.go
index 79f0297908f1..0947cb9081f8 100644
--- a/pkg/sql/sem/builtins/generator_builtins.go
+++ b/pkg/sql/sem/builtins/generator_builtins.go
@@ -109,7 +109,7 @@ var generators = map[string]builtinDefinition{
// See https://www.postgresql.org/docs/current/static/functions-array.html
makeGeneratorOverloadWithReturnType(
tree.ArgTypes{{"input", types.AnyArray}},
- func(args []tree.TypedExpr) *types.T {
+ func(ctx *tree.SemaContext, args []tree.TypedExpr) *types.T {
if len(args) == 0 || args[0].ResolvedType().Family() == types.UnknownFamily {
return tree.UnknownReturnType
}
@@ -123,7 +123,7 @@ var generators = map[string]builtinDefinition{
"information_schema._pg_expandarray": makeBuiltin(genProps(expandArrayValueGeneratorLabels),
makeGeneratorOverloadWithReturnType(
tree.ArgTypes{{"input", types.AnyArray}},
- func(args []tree.TypedExpr) *types.T {
+ func(ctx *tree.SemaContext, args []tree.TypedExpr) *types.T {
if len(args) == 0 || args[0].ResolvedType().Family() == types.UnknownFamily {
return tree.UnknownReturnType
}
diff --git a/pkg/sql/sem/tree/overload.go b/pkg/sql/sem/tree/overload.go
index b9faeebc17ce..0fc431d119cf 100644
--- a/pkg/sql/sem/tree/overload.go
+++ b/pkg/sql/sem/tree/overload.go
@@ -327,17 +327,17 @@ var UnknownReturnType *types.T
// ReturnTyper defines the type-level function in which a builtin function's return type
// is determined. ReturnTypers should make sure to return unknownReturnType when necessary.
-type ReturnTyper func(args []TypedExpr) *types.T
+type ReturnTyper func(ctx *SemaContext, args []TypedExpr) *types.T
// FixedReturnType functions simply return a fixed type, independent of argument types.
func FixedReturnType(typ *types.T) ReturnTyper {
- return func(args []TypedExpr) *types.T { return typ }
+ return func(_ *SemaContext, args []TypedExpr) *types.T { return typ }
}
// IdentityReturnType creates a returnType that is a projection of the idx'th
// argument type.
func IdentityReturnType(idx int) ReturnTyper {
- return func(args []TypedExpr) *types.T {
+ return func(_ *SemaContext, args []TypedExpr) *types.T {
if len(args) == 0 {
return UnknownReturnType
}
@@ -351,7 +351,7 @@ func IdentityReturnType(idx int) ReturnTyper {
// with HomogeneousType functions, in which all arguments have been checked to
// have the same type (or be null).
func FirstNonNullReturnType() ReturnTyper {
- return func(args []TypedExpr) *types.T {
+ return func(_ *SemaContext, args []TypedExpr) *types.T {
if len(args) == 0 {
return UnknownReturnType
}
@@ -365,7 +365,7 @@ func FirstNonNullReturnType() ReturnTyper {
}
func returnTypeToFixedType(s ReturnTyper) *types.T {
- if t := s(nil); t != UnknownReturnType {
+ if t := s(nil, nil); t != UnknownReturnType {
return t
}
return types.Any
@@ -498,7 +498,7 @@ func typeCheckOverloadedExprs(
// fixed return types. This could be improved, but is not currently
// critical because we have no cases of functions with multiple
// overloads that do not all expose FixedReturnTypes.
- if t := o.returnType()(nil); t != UnknownReturnType {
+ if t := o.returnType()(ctx, nil); t != UnknownReturnType {
return t.Equivalent(desired)
}
return true
diff --git a/pkg/sql/sem/tree/type_check.go b/pkg/sql/sem/tree/type_check.go
index b18ab6977931..21391aee20fb 100644
--- a/pkg/sql/sem/tree/type_check.go
+++ b/pkg/sql/sem/tree/type_check.go
@@ -19,6 +19,7 @@ import (
"strings"
"time"
+ "github.com/cockroachdb/cockroach/pkg/roachpb"
"github.com/cockroachdb/cockroach/pkg/server/telemetry"
"github.com/cockroachdb/cockroach/pkg/sql/lex"
"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
@@ -58,6 +59,14 @@ type SemaContext struct {
// globally for the entire txn and this field would not be needed.
AsOfTimestamp *hlc.Timestamp
+ // Locality contains the location of the current node as a set of user-defined
+ // key/value pairs, ordered from most inclusive to least inclusive. If there
+ // are no tiers, then the node's location is not known. Example:
+ //
+ // [region=us,dc=east]
+ //
+ Locality roachpb.Locality
+
Properties SemaProperties
}
@@ -357,7 +366,7 @@ func (expr *BinaryExpr) TypeCheck(ctx *SemaContext, desired *types.T) (TypedExpr
expr.Left, expr.Right = leftTyped, rightTyped
expr.fn = binOp
- expr.typ = binOp.returnType()(typedSubExprs)
+ expr.typ = binOp.returnType()(ctx, typedSubExprs)
return expr, nil
}
@@ -559,7 +568,7 @@ func (expr *TupleStar) TypeCheck(ctx *SemaContext, desired *types.T) (TypedExpr,
// Alghough we're going to elide the tuple star, we need to ensure
// the expression is indeed a labeled tuple first.
- if resolvedType.Family() != types.TupleFamily || len(resolvedType.TupleLabels()) == 0 {
+ if resolvedType.Family() != types.TupleFamily || resolvedType.TupleLabels() == nil {
return nil, NewTypeIsNotCompositeError(resolvedType)
}
@@ -585,7 +594,7 @@ func (expr *ColumnAccessExpr) TypeCheck(ctx *SemaContext, desired *types.T) (Typ
expr.Expr = subExpr
resolvedType := subExpr.ResolvedType()
- if resolvedType.Family() != types.TupleFamily || len(resolvedType.TupleLabels()) == 0 {
+ if resolvedType.Family() != types.TupleFamily || resolvedType.TupleLabels() == nil {
return nil, NewTypeIsNotCompositeError(resolvedType)
}
@@ -922,7 +931,7 @@ func (expr *FuncExpr) TypeCheck(ctx *SemaContext, desired *types.T) (TypedExpr,
}
expr.fn = overloadImpl
expr.fnProps = &def.FunctionProperties
- expr.typ = overloadImpl.returnType()(typedSubExprs)
+ expr.typ = overloadImpl.returnType()(ctx, typedSubExprs)
if expr.typ == UnknownReturnType {
typeNames := make([]string, 0, len(expr.Exprs))
for _, expr := range typedSubExprs {
@@ -1213,7 +1222,7 @@ func (expr *UnaryExpr) TypeCheck(ctx *SemaContext, desired *types.T) (TypedExpr,
expr.Expr = exprTyped
expr.fn = unaryOp
- expr.typ = unaryOp.returnType()(typedSubExprs)
+ expr.typ = unaryOp.returnType()(ctx, typedSubExprs)
return expr, nil
}