add an option to open a postgres transaction in lookup so it isnt actually a lookup

This commit is contained in:
Nia Weiss 2021-05-26 13:55:15 -04:00
parent a94dc50b1c
commit a72f6425af
No known key found for this signature in database
GPG key ID: 895E83409BFDA1BB
3 changed files with 30 additions and 1 deletions

View file

@ -3912,6 +3912,11 @@ func (e *executor) executeExternalLookup(ctx context.Context, qcx *Qcx, index st
return ExtractedTable{}, errors.New("too many inputs to lookup query")
}
write, _, err := c.BoolArg("write")
if err != nil {
return ExtractedTable{}, errors.Wrap(err, "parsing write argument")
}
rawArg, err := e.executeCall(ctx, qcx, index, c.Children[0], shards, opt)
if err != nil {
return ExtractedTable{}, errors.Wrapf(err, "evaluating SQL argument call %q", c.String())
@ -3937,6 +3942,21 @@ func (e *executor) executeExternalLookup(ctx context.Context, qcx *Qcx, index st
arg = argRow.Columns()
}
if write {
tx, err := e.Holder.lookupDB.BeginTx(ctx, nil)
if err != nil {
return ExtractedTable{}, errors.Wrap(err, "creating postgres transaction")
}
defer tx.Rollback() //nolint:errcheck
_, err = tx.ExecContext(ctx, query, pq.Array(arg))
if err != nil {
return ExtractedTable{}, errors.Wrap(err, "executing postgres write")
}
return ExtractedTable{}, errors.Wrap(tx.Commit(), "committing postgres transaction")
}
result, err := e.Holder.lookupDB.QueryContext(ctx, query, pq.Array(arg))
if err != nil {
return ExtractedTable{}, errors.Wrapf(err, "SQL query failed")

View file

@ -7913,7 +7913,7 @@ func TestExternalLookup(t *testing.T) {
defer c.Close()
// Populate a field with some data that can be used in queries.
c.CreateField(t, "i", pilosa.IndexOptions{}, "f")
c.CreateField(t, "i", pilosa.IndexOptions{TrackExistence: true}, "f")
c.ImportBits(t, "i", "f", [][2]uint64{
{1, 1},
{1, 3},
@ -8053,4 +8053,12 @@ func TestExternalLookup(t *testing.T) {
})
}
})
t.Run("Delete", func(t *testing.T) {
c.Query(t, "i", `ExternalLookup(All(), query="delete from lookup where id = ANY($1)", write=true)`)
res := c.Query(t, "i", `ExternalLookup(All(), query="select id from lookup where id = ANY($1)")`)
tbl := res.Results[0].(pilosa.ExtractedTable)
if len(tbl.Columns) != 0 {
t.Errorf("unexpected remaining records: %v", tbl)
}
})
}

View file

@ -437,6 +437,7 @@ var callInfoByFunc = map[string]callInfo{
allowUnknown: false,
prototypes: map[string]interface{}{
"query": "",
"write": true,
},
},
"Limit": {