From 8f437aa573ed9d8ad254b1bec840956d2d58c23b Mon Sep 17 00:00:00 2001
From: Rohan Yadav <rohany@alumni.cmu.edu>
Date: Tue, 27 Aug 2019 16:09:47 -0700
Subject: [PATCH] demo: Try to get a temporary license upon demo startup

Fixes #40222.

Release note (cli change): cockroach demo tries to get a temporary
enterprise license upon startup.
---
 pkg/cli/demo.go | 58 ++++++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 50 insertions(+), 8 deletions(-)

diff --git a/pkg/cli/demo.go b/pkg/cli/demo.go
index db9d08d6a2b2..5b19cfd4b817 100644
--- a/pkg/cli/demo.go
+++ b/pkg/cli/demo.go
@@ -14,7 +14,10 @@ import (
 	"context"
 	gosql "database/sql"
 	"fmt"
+	"io/ioutil"
+	"net/http"
 	"net/url"
+	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/base"
 	"github.com/cockroachdb/cockroach/pkg/cli/cliflags"
@@ -49,6 +52,11 @@ to avoid pre-loading a dataset.`,
 	}),
 }
 
+// TODO (rohany): change this once another endpoint is setup for getting licenses.
+// This URL grants a license that is valid for 1 hour.
+const licenseURL = "https://register.cockroachdb.com/api/prodtest"
+const demoOrg = "Cockroach Labs - Production Testing"
+
 const defaultGeneratorName = "movr"
 
 var defaultGenerator workload.Generator
@@ -95,6 +103,25 @@ func init() {
 	}
 }
 
+func getLicense() (string, error) {
+	client := &http.Client{
+		Timeout: time.Second,
+	}
+	resp, err := client.Get(licenseURL)
+	if err != nil {
+		return "", err
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode != http.StatusOK {
+		return "", errors.New("unable to connect to licensing endpoint")
+	}
+	bodyBytes, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		return "", err
+	}
+	return string(bodyBytes), nil
+}
+
 func setupTransientServers(
 	cmd *cobra.Command, gen workload.Generator,
 ) (connURL string, adminURL string, cleanup func(), err error) {
@@ -189,15 +216,30 @@ func setupTransientServers(
 	}
 	urlStr := url.String()
 
+	db, err := gosql.Open("postgres", urlStr)
+	if err != nil {
+		return ``, ``, cleanup, err
+	}
+	defer db.Close()
+
+	// Load a license.
+	if cliCtx.isInteractive {
+		license, err := getLicense()
+		if err == nil {
+			if _, err := db.Exec(`SET CLUSTER SETTING cluster.organization = $1`, demoOrg); err != nil {
+				return ``, ``, cleanup, err
+			}
+			if _, err := db.Exec(`SET CLUSTER SETTING enterprise.license = $1`, license); err != nil {
+				return ``, ``, cleanup, err
+			}
+		} else {
+			log.Warningf(ctx, "error when attempting to acquire demo license: %+v\n", err)
+		}
+	}
+
 	// If there is a load generator, create its database and load its
 	// fixture.
 	if gen != nil {
-		db, err := gosql.Open("postgres", urlStr)
-		if err != nil {
-			return ``, ``, cleanup, err
-		}
-		defer db.Close()
-
 		if _, err := db.Exec(`CREATE DATABASE ` + gen.Meta().Name); err != nil {
 			return ``, ``, cleanup, err
 		}
@@ -218,14 +260,14 @@ func runDemo(cmd *cobra.Command, gen workload.Generator) error {
 		gen = defaultGenerator
 	}
 
+	checkInteractive()
+
 	connURL, adminURL, cleanup, err := setupTransientServers(cmd, gen)
 	defer cleanup()
 	if err != nil {
 		return checkAndMaybeShout(err)
 	}
 
-	checkInteractive()
-
 	if cliCtx.isInteractive {
 		fmt.Printf(`#
 # Welcome to the CockroachDB demo database!