Skip to content

Commit

Permalink
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add WithQueryOptions method to enable setting implicit limit and re…
Browse files Browse the repository at this point in the history
…adonly on queries
jaclarke committed Jan 17, 2025
1 parent 2daf9b2 commit 0b825b5
Showing 9 changed files with 149 additions and 12 deletions.
20 changes: 16 additions & 4 deletions internal/client/client.go
Original file line number Diff line number Diff line change
@@ -57,7 +57,8 @@ type Client struct {

cfg *connConfig
cacheCollection
state map[string]interface{}
state map[string]interface{}
queryOpts QueryOptions

warningHandler WarningHandler
}
@@ -332,6 +333,7 @@ func (p *Client) Execute(
args,
conn.capabilities1pX(),
copyState(p.state),
p.queryOpts,
nil,
true,
p.warningHandler,
@@ -357,7 +359,10 @@ func (p *Client) Query(
}

err = runQuery(
ctx, conn, "Query", cmd, out, args, p.state, p.warningHandler)
ctx, conn, "Query",
cmd, out, args,
p.state, p.queryOpts, p.warningHandler,
)
return firstError(err, p.release(conn, err))
}

@@ -384,6 +389,7 @@ func (p *Client) QuerySingle(
out,
args,
p.state,
p.queryOpts,
p.warningHandler,
)
return firstError(err, p.release(conn, err))
@@ -409,6 +415,7 @@ func (p *Client) QueryJSON(
out,
args,
p.state,
p.queryOpts,
p.warningHandler,
)
return firstError(err, p.release(conn, err))
@@ -436,6 +443,7 @@ func (p *Client) QuerySingleJSON(
out,
args,
p.state,
p.queryOpts,
p.warningHandler,
)
return firstError(err, p.release(conn, err))
@@ -454,7 +462,10 @@ func (p *Client) QuerySQL(
}

err = runQuery(
ctx, conn, "QuerySQL", cmd, out, args, p.state, p.warningHandler)
ctx, conn, "QuerySQL",
cmd, out, args,
p.state, p.queryOpts, p.warningHandler,
)
return firstError(err, p.release(conn, err))
}

@@ -475,6 +486,7 @@ func (p *Client) ExecuteSQL(
args,
conn.capabilities1pX(),
copyState(p.state),
p.queryOpts,
nil,
true,
p.warningHandler,
@@ -507,6 +519,6 @@ func (p *Client) Tx(ctx context.Context, action TxBlock) error {
return err
}

err = conn.tx(ctx, action, p.state, p.warningHandler)
err = conn.tx(ctx, action, p.state, p.queryOpts, p.warningHandler)
return firstError(err, p.release(conn, err))
}
1 change: 1 addition & 0 deletions internal/client/constants.go
Original file line number Diff line number Diff line change
@@ -37,6 +37,7 @@ var (
protocolVersion2p0 = internal.ProtocolVersion{Major: 2, Minor: 0}
protocolVersion3p0 = internal.ProtocolVersion{Major: 3, Minor: 0}

capabilitiesModifications uint64 = 0x1
capabilitiesSessionConfig uint64 = 0x2
capabilitiesTransaction uint64 = 0x4
capabilitiesDDL uint64 = 0x8
8 changes: 4 additions & 4 deletions internal/client/granularflow1pX.go
Original file line number Diff line number Diff line change
@@ -66,9 +66,9 @@ func (c *protocolConnection) parse1pX(
w := buff.NewWriter(c.writeMemory[:0])
w.BeginMessage(uint8(Parse))
w.PushUint16(0) // no headers
w.PushUint64(q.capabilities)
w.PushUint64(q.getCapabilities())
w.PushUint64(0) // no compilation_flags
w.PushUint64(0) // no implicit limit
w.PushUint64(q.queryOpts.ImplicitLimit)
w.PushUint8(uint8(q.fmt))
w.PushUint8(uint8(q.expCard))
w.PushString(q.cmd)
@@ -185,9 +185,9 @@ func (c *protocolConnection) execute1pX(
w := buff.NewWriter(c.writeMemory[:0])
w.BeginMessage(uint8(Execute))
w.PushUint16(0) // no headers
w.PushUint64(q.capabilities)
w.PushUint64(q.getCapabilities())
w.PushUint64(0) // no compilation_flags
w.PushUint64(0) // no implicit limit
w.PushUint64(q.queryOpts.ImplicitLimit)
w.PushUint8(uint8(q.fmt))
w.PushUint8(uint8(q.expCard))
w.PushString(q.cmd)
8 changes: 4 additions & 4 deletions internal/client/granularflow2pX.go
Original file line number Diff line number Diff line change
@@ -72,9 +72,9 @@ func (c *protocolConnection) parse2pX(
w := buff.NewWriter(c.writeMemory[:0])
w.BeginMessage(uint8(Parse))
w.PushUint16(0) // no headers
w.PushUint64(q.capabilities)
w.PushUint64(q.getCapabilities())
w.PushUint64(0) // no compilation_flags
w.PushUint64(0) // no implicit limit
w.PushUint64(q.queryOpts.ImplicitLimit)
if c.protocolVersion.GTE(protocolVersion3p0) {
w.PushUint8(uint8(q.lang))
}
@@ -194,9 +194,9 @@ func (c *protocolConnection) execute2pX(
w := buff.NewWriter(c.writeMemory[:0])
w.BeginMessage(uint8(Execute))
w.PushUint16(0) // no headers
w.PushUint64(q.capabilities)
w.PushUint64(q.getCapabilities())
w.PushUint64(0) // no compilation_flags
w.PushUint64(0) // no implicit limit
w.PushUint64(q.queryOpts.ImplicitLimit)
if c.protocolVersion.GTE(protocolVersion3p0) {
w.PushUint8(uint8(q.lang))
}
28 changes: 28 additions & 0 deletions internal/client/options.go
Original file line number Diff line number Diff line change
@@ -531,3 +531,31 @@ func (p Client) WithWarningHandler( // nolint:gocritic
p.warningHandler = warningHandler
return &p
}

// QueryOptions are config options that affect query behaviour.
type QueryOptions struct {
ReadOnly bool
ImplicitLimit uint64
}

// WithQueryOptions sets module name aliases for the returned client.
func (p Client) WithQueryOptions( // nolint:gocritic
readOnly *bool,
implicitLimit *uint64,
) *Client {
opts := QueryOptions{}

if readOnly != nil {
opts.ReadOnly = *readOnly
} else {
opts.ReadOnly = p.queryOpts.ReadOnly
}
if implicitLimit != nil {
opts.ImplicitLimit = *implicitLimit
} else {
opts.ImplicitLimit = p.queryOpts.ImplicitLimit
}

p.queryOpts = opts
return &p
}
14 changes: 14 additions & 0 deletions internal/client/query.go
Original file line number Diff line number Diff line change
@@ -44,10 +44,19 @@ type query struct {
args []interface{}
capabilities uint64
state map[string]interface{}
queryOpts QueryOptions
parse bool
warningHandler WarningHandler
}

func (q *query) getCapabilities() uint64 {
capabilities := q.capabilities
if q.queryOpts.ReadOnly {
capabilities &^= capabilitiesModifications
}
return capabilities
}

func (q *query) flat() bool {
if q.expCard != Many {
return true
@@ -73,6 +82,7 @@ func newQuery(
args []interface{},
capabilities uint64,
state map[string]interface{},
queryOpts QueryOptions,
out interface{},
parse bool,
warningHandler WarningHandler,
@@ -98,6 +108,7 @@ func newQuery(
args: args,
capabilities: capabilities,
state: state,
queryOpts: queryOpts,
parse: parse,
warningHandler: warningHandler,
}, nil
@@ -130,6 +141,7 @@ func newQuery(
args: args,
capabilities: capabilities,
state: state,
queryOpts: queryOpts,
parse: parse,
warningHandler: warningHandler,
}
@@ -173,6 +185,7 @@ func runQuery(
out interface{},
args []interface{},
state map[string]interface{},
queryOpts QueryOptions,
warningHandler WarningHandler,
) error {
if method == "QuerySingleJSON" {
@@ -191,6 +204,7 @@ func runQuery(
args,
c.capabilities1pX(),
state,
queryOpts,
out,
true,
warningHandler,
71 changes: 71 additions & 0 deletions internal/client/query_test.go
Original file line number Diff line number Diff line change
@@ -1179,3 +1179,74 @@ func TestWithWarningHandler(t *testing.T) {
require.NoError(t, err)
require.Greater(t, len(seen), 0)
}

func TestWithQueryOptionsReadonly(t *testing.T) {
ctx := context.Background()

err := client.Execute(ctx, `create type QueryOptsTest {
create property name -> str;
};`)
assert.NoError(t, err)
defer (func() {
err = client.Execute(ctx, `drop type QueryOptsTest;`)
assert.NoError(t, err)
})()

var res struct {
ID types.UUID `edgedb:"id"`
}
err = client.QuerySingle(ctx,
"insert QueryOptsTest {name := 'abc'}", &res)
assert.NoError(t, err)

readonly := true
readonlyClient := client.WithQueryOptions(&readonly, nil)

err = readonlyClient.QuerySingle(ctx,
"insert QueryOptsTest {name := 'def'}", &res)
assert.EqualError(t, err,
"edgedb.DisabledCapabilityError: "+
"cannot execute data modification queries: disabled by the client",
)

err = client.QuerySingle(ctx,
"select QueryOptsTest {id} limit 1", &res)
assert.NoError(t, err)

// check we didn't modify the original client
err = client.QuerySingle(ctx,
"insert QueryOptsTest {name := 'abc'}", &res)
assert.NoError(t, err)
}

func TestWithQueryOptionsImplicitLimit(t *testing.T) {
ctx := context.Background()

var res []struct {
Name string `edgedb:"name"`
}
err := client.Query(ctx,
"select schema::ObjectType {name}", &res)
assert.NoError(t, err)
assert.Greater(t, len(res), 10)

limit := uint64(10)
limitClient := client.WithQueryOptions(nil, &limit)

var limitRes []struct {
Name string `edgedb:"name"`
}
err = limitClient.Query(ctx,
"select schema::ObjectType {name}", &limitRes)
assert.NoError(t, err)
assert.Equal(t, len(limitRes), 10)

// check we didn't modify the original client
var res2 []struct {
Name string `edgedb:"name"`
}
err = client.Query(ctx,
"select schema::ObjectType {name}", &res2)
assert.NoError(t, err)
assert.Equal(t, len(res2), len(res))
}
2 changes: 2 additions & 0 deletions internal/client/transactable.go
Original file line number Diff line number Diff line change
@@ -77,6 +77,7 @@ func (c *transactableConn) tx(
ctx context.Context,
action TxBlock,
state map[string]interface{},
queryOpts QueryOptions,
warningHandler WarningHandler,
) (err error) {
conn, err := c.borrow("transaction")
@@ -102,6 +103,7 @@ func (c *transactableConn) tx(
txState: &txState{},
options: c.txOpts,
state: state,
queryOpts: queryOpts,
warningHandler: warningHandler,
}
err = tx.start(ctx)
9 changes: 9 additions & 0 deletions internal/client/transaction.go
Original file line number Diff line number Diff line change
@@ -78,6 +78,7 @@ type Tx struct {
*txState
options TxOptions
state map[string]interface{}
queryOpts QueryOptions
warningHandler WarningHandler
}

@@ -92,6 +93,7 @@ func (t *Tx) execute(
nil,
txCapabilities,
t.state,
t.queryOpts,
nil,
false,
t.warningHandler,
@@ -171,6 +173,7 @@ func (t *Tx) Execute(
args,
t.capabilities1pX(),
t.state,
t.queryOpts,
nil,
true,
t.warningHandler,
@@ -197,6 +200,7 @@ func (t *Tx) Query(
out,
args,
t.state,
t.queryOpts,
t.warningHandler,
)
}
@@ -219,6 +223,7 @@ func (t *Tx) QuerySingle(
out,
args,
t.state,
t.queryOpts,
t.warningHandler,
)
}
@@ -238,6 +243,7 @@ func (t *Tx) QueryJSON(
out,
args,
t.state,
t.queryOpts,
t.warningHandler,
)
}
@@ -259,6 +265,7 @@ func (t *Tx) QuerySingleJSON(
out,
args,
t.state,
t.queryOpts,
t.warningHandler,
)
}
@@ -275,6 +282,7 @@ func (t *Tx) ExecuteSQL(
args,
t.capabilities1pX(),
t.state,
t.queryOpts,
nil,
true,
t.warningHandler,
@@ -301,6 +309,7 @@ func (t *Tx) QuerySQL(
out,
args,
t.state,
t.queryOpts,
t.warningHandler,
)
}

0 comments on commit 0b825b5

Please sign in to comment.