Add WithQueryOptions method to enable setting implicit limit and re…

…adonly on queries

internal/client/client.go

@@ -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))
 }

internal/client/constants.go

@@ -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

internal/client/granularflow1pX.go

@@ -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)

internal/client/granularflow2pX.go

@@ -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))
 	}

internal/client/options.go

@@ -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
+}

internal/client/query.go

@@ -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,

internal/client/query_test.go

@@ -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))
+}

internal/client/transactable.go

@@ -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)

internal/client/transaction.go

@@ -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,
 	)
 }