add table and column metadata for bigquery (#141)

Co-authored-by: poundifdef <jay@scratchdata.com>

models/models.go

@@ -13,3 +13,9 @@ type StorageServices struct {
 	Queue     queue.Queue
 	BlobStore blobstore.BlobStore
 }
+
+type Column struct {
+	Name     string `json:"name"`
+	Type     string `json:"type"`
+	JSONType string `json:"-"`
+}

pkg/api/router.go

@@ -24,6 +24,8 @@ var responseSize = promauto.NewHistogramVec(prometheus.HistogramOpts{
 type ScratchDataAPI interface {
 	Select(w http.ResponseWriter, r *http.Request)
 	Insert(w http.ResponseWriter, r *http.Request)
+	Tables(w http.ResponseWriter, r *http.Request)
+	Columns(w http.ResponseWriter, r *http.Request)
 
 	CreateQuery(w http.ResponseWriter, r *http.Request)
 	ShareData(w http.ResponseWriter, r *http.Request)
@@ -47,6 +49,8 @@ func CreateMux(apiFunctions ScratchDataAPI) *chi.Mux {
 	api.Post("/data/insert/{table}", apiFunctions.Insert)
 	api.Get("/data/query", apiFunctions.Select)
 	api.Post("/data/query", apiFunctions.Select)
+	api.Get("/tables", apiFunctions.Tables)
+	api.Get("/tables/{table}/columns", apiFunctions.Columns)
 
 	api.Get("/destinations", apiFunctions.GetDestinations)
 	api.Post("/destinations", apiFunctions.CreateDestination)

pkg/api/table.go

@@ -0,0 +1,44 @@
+package api
+
+import (
+	"net/http"
+
+	"github.com/go-chi/chi/v5"
+	"github.com/go-chi/render"
+)
+
+func (a *ScratchDataAPIStruct) Tables(w http.ResponseWriter, r *http.Request) {
+	databaseID := a.AuthGetDatabaseID(r.Context())
+	dest, err := a.destinationManager.Destination(r.Context(), databaseID)
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	tables, err := dest.Tables()
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	render.JSON(w, r, tables)
+}
+
+func (a *ScratchDataAPIStruct) Columns(w http.ResponseWriter, r *http.Request) {
+	table := chi.URLParam(r, "table")
+	databaseID := a.AuthGetDatabaseID(r.Context())
+
+	dest, err := a.destinationManager.Destination(r.Context(), databaseID)
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	columns, err := dest.Columns(table)
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	render.JSON(w, r, columns)
+}

pkg/destinations/bigquery/tables.go

@@ -0,0 +1,71 @@
+package bigquery
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"strings"
+
+	"github.com/rs/zerolog/log"
+	"github.com/scratchdata/scratchdata/models"
+	"google.golang.org/api/iterator"
+)
+
+func (b *BigQueryServer) Columns(table string) ([]models.Column, error) {
+	rc := []models.Column{}
+
+	tokens := strings.Split(table, ".")
+	if len(tokens) != 2 {
+		return nil, errors.New("Table should be in the format dataset.table")
+	}
+
+	datasetID := tokens[0]
+	tableID := tokens[1]
+
+	ctx := context.TODO()
+	tableInfo := b.conn.Dataset(datasetID).Table(tableID)
+	meta, err := tableInfo.Metadata(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, field := range meta.Schema {
+		rc = append(rc, models.Column{
+			Name: field.Name,
+			Type: string(field.Type),
+		})
+	}
+	return rc, nil
+}
+
+func (b *BigQueryServer) Tables() ([]string, error) {
+	rc := []string{}
+
+	ctx := context.TODO()
+	it := b.conn.Datasets(ctx)
+	for {
+		dataset, err := it.Next()
+		if err == iterator.Done {
+			break
+		}
+		if err != nil {
+			return nil, err
+		}
+
+		// For each dataset, list all tables.
+		tableIt := dataset.Tables(ctx)
+		for {
+			table, err := tableIt.Next()
+			if err == iterator.Done {
+				break
+			}
+			if err != nil {
+				log.Error().Err(err).Str("dataset_id", dataset.DatasetID).Msg("Failed to list tables")
+				continue
+			}
+			rc = append(rc, fmt.Sprintf("%s.%s", dataset.DatasetID, table.TableID))
+		}
+	}
+
+	return rc, nil
+}

pkg/destinations/clickhouse/tables.go

@@ -0,0 +1,15 @@
+package clickhouse
+
+import (
+	"errors"
+
+	"github.com/scratchdata/scratchdata/models"
+)
+
+func (b *ClickhouseServer) Columns(table string) ([]models.Column, error) {
+	return []models.Column{}, errors.New("not implemented")
+}
+
+func (b *ClickhouseServer) Tables() ([]string, error) {
+	return []string{}, errors.New("not implemented")
+}

pkg/destinations/destinations.go

@@ -25,6 +25,9 @@ type Destination interface {
 	QueryJSON(query string, writer io.Writer) error
 	QueryCSV(query string, writer io.Writer) error
 
+	Tables() ([]string, error)
+	Columns(table string) ([]models.Column, error)
+
 	CreateEmptyTable(name string) error
 	CreateColumns(table string, filePath string) error
 	InsertFromNDJsonFile(table string, filePath string) error

pkg/destinations/duckdb/tables.go

@@ -0,0 +1,15 @@
+package duckdb
+
+import (
+	"errors"
+
+	"github.com/scratchdata/scratchdata/models"
+)
+
+func (b *DuckDBServer) Columns(table string) ([]models.Column, error) {
+	return []models.Column{}, errors.New("not implemented")
+}
+
+func (b *DuckDBServer) Tables() ([]string, error) {
+	return []string{}, errors.New("not implemented")
+}

pkg/destinations/redshift/tables.go

@@ -0,0 +1,15 @@
+package redshift
+
+import (
+	"errors"
+
+	"github.com/scratchdata/scratchdata/models"
+)
+
+func (b *RedshiftServer) Columns(table string) ([]models.Column, error) {
+	return []models.Column{}, errors.New("not implemented")
+}
+
+func (b *RedshiftServer) Tables() ([]string, error) {
+	return []string{}, errors.New("not implemented")
+}