featurebase/cli/queryer.go
Travis Turner 05ebdd15f0
Minor cleanup to some fbsql flags and meta-commands (#2299)
* Set fbsql prompt based on the connected database

This also changes the prompt to align with psql, where it begins with:
db=#
and the mid looks like:
db-#

* Require organizationID in on-prem, serverless queries

* Support meta-commands in `--file` command

I'm not sure why this was restricted before. Just an oversight.

* Change default history file name to fbsql_history

* Add port short flag: p

* Support meta command \list (for \l) and \out (for \o)

* cleanup while writing docs

* Have \cd with no arguments change to home directory

* Avoid shadowing `action`
2023-03-07 12:18:55 -06:00

108 lines
3.1 KiB
Go

package cli
import (
"fmt"
"io"
"net/http"
"time"
featurebase "github.com/featurebasedb/featurebase/v3"
"github.com/pkg/errors"
)
type Queryer interface {
Query(org string, db string, sql io.Reader) (*featurebase.WireQueryResponse, error)
}
// Ensure type implements interface.
var _ Queryer = (*nopQueryer)(nil)
type nopQueryer struct{}
func (qryr *nopQueryer) Query(org string, db string, sql io.Reader) (*featurebase.WireQueryResponse, error) {
return nil, errors.Errorf("no-op queryer")
}
// Ensure type implements interface.
var _ Queryer = (*standardQueryer)(nil)
// standardQueryer supports a standard featurebase deployment hitting the /sql
// endpoint with a payload containing only the sql statement.
type standardQueryer struct {
Host string
Port string
}
func (qryr *standardQueryer) Query(org string, db string, sql io.Reader) (*featurebase.WireQueryResponse, error) {
url := fmt.Sprintf("%s/sql", hostPort(qryr.Host, qryr.Port))
resp, err := http.Post(url, "application/json", sql)
if err != nil {
return nil, errors.Wrapf(err, "posting query")
}
fullbod, err := io.ReadAll(resp.Body)
if err != nil {
return nil, errors.Wrap(err, "reading response")
}
sqlResponse := &featurebase.WireQueryResponse{}
// TODO(tlt): switch this back once all responses are typed
// TODO(twg) 2023/03/01 using json.Number to decode large ints so care must be made
// if err := json.Unmarshal(fullbod, sqlResponse); err != nil {
if err := sqlResponse.UnmarshalJSONTyped(fullbod, true); err != nil {
return nil, errors.Wrapf(err, "unmarshaling query response, body:\n'%s'\n", fullbod)
}
return sqlResponse, nil
}
// Ensure type implements interface.
var _ Queryer = (*serverlessQueryer)(nil)
// serverlessQueryer is similar to the standardQueryer except that it hits a
// different endpoint, and its payload is database-aware.
type serverlessQueryer struct {
Host string
Port string
}
func (qryr *serverlessQueryer) Query(org string, db string, sql io.Reader) (*featurebase.WireQueryResponse, error) {
if org == "" {
return nil, NewErrOrganizationRequired()
}
url := fmt.Sprintf("%s/queryer/databases/%s/sql", hostPort(qryr.Host, qryr.Port), db)
if db == "" {
url = fmt.Sprintf("%s/queryer/sql", hostPort(qryr.Host, qryr.Port))
}
client := &http.Client{
Timeout: time.Second * 30,
}
req, err := http.NewRequest(http.MethodPost, url, sql)
if err != nil {
return nil, errors.Wrap(err, "creating new post request")
}
req.Header.Add("Content-Type", "text/plain")
req.Header.Add("OrganizationID", org)
var resp *http.Response
if resp, err = client.Do(req); err != nil {
return nil, errors.Wrap(err, "executing post request")
}
fullbod, err := io.ReadAll(resp.Body)
if err != nil {
return nil, errors.Wrap(err, "reading response")
}
sqlResponse := &featurebase.WireQueryResponse{}
// TODO(tlt): switch this back once all responses are typed
// if err := json.Unmarshal(fullbod, sqlResponse); err != nil {
if err := sqlResponse.UnmarshalJSONTyped(fullbod, true); err != nil {
return nil, errors.Wrapf(err, "unmarshaling query response, body:\n'%s'\n", fullbod)
}
return sqlResponse, nil
}