mirror of https://github.com/databricks/cli.git
recorder proxy for acceptance tests
This commit is contained in:
parent
8b4a61d123
commit
f0ad3d1846
|
@ -6,6 +6,7 @@ import (
|
||||||
"errors"
|
"errors"
|
||||||
"flag"
|
"flag"
|
||||||
"fmt"
|
"fmt"
|
||||||
|
"github.com/google/uuid"
|
||||||
"io"
|
"io"
|
||||||
"net/http"
|
"net/http"
|
||||||
"os"
|
"os"
|
||||||
|
@ -20,8 +21,6 @@ import (
|
||||||
"time"
|
"time"
|
||||||
"unicode/utf8"
|
"unicode/utf8"
|
||||||
|
|
||||||
"github.com/google/uuid"
|
|
||||||
|
|
||||||
"github.com/databricks/cli/internal/testutil"
|
"github.com/databricks/cli/internal/testutil"
|
||||||
"github.com/databricks/cli/libs/env"
|
"github.com/databricks/cli/libs/env"
|
||||||
"github.com/databricks/cli/libs/testdiff"
|
"github.com/databricks/cli/libs/testdiff"
|
||||||
|
@ -41,7 +40,7 @@ var (
|
||||||
// In order to debug CLI running under acceptance test, set this to full subtest name, e.g. "bundle/variables/empty"
|
// In order to debug CLI running under acceptance test, set this to full subtest name, e.g. "bundle/variables/empty"
|
||||||
// Then install your breakpoints and click "debug test" near TestAccept in VSCODE.
|
// Then install your breakpoints and click "debug test" near TestAccept in VSCODE.
|
||||||
// example: var SingleTest = "bundle/variables/empty"
|
// example: var SingleTest = "bundle/variables/empty"
|
||||||
var SingleTest = ""
|
var SingleTest = "bundle/deployment/bind/schema"
|
||||||
|
|
||||||
// If enabled, instead of compiling and running CLI externally, we'll start in-process server that accepts and runs
|
// If enabled, instead of compiling and running CLI externally, we'll start in-process server that accepts and runs
|
||||||
// CLI commands. The $CLI in test scripts is a helper that just forwards command-line arguments to this server (see bin/callserver.py).
|
// CLI commands. The $CLI in test scripts is a helper that just forwards command-line arguments to this server (see bin/callserver.py).
|
||||||
|
@ -58,7 +57,7 @@ const (
|
||||||
EntryPointScript = "script"
|
EntryPointScript = "script"
|
||||||
CleanupScript = "script.cleanup"
|
CleanupScript = "script.cleanup"
|
||||||
PrepareScript = "script.prepare"
|
PrepareScript = "script.prepare"
|
||||||
MaxFileSize = 100_000
|
MaxFileSize = 300_000
|
||||||
// Filename to save replacements to (used by diff.py)
|
// Filename to save replacements to (used by diff.py)
|
||||||
ReplsFile = "repls.json"
|
ReplsFile = "repls.json"
|
||||||
)
|
)
|
||||||
|
@ -77,6 +76,10 @@ func TestAccept(t *testing.T) {
|
||||||
testAccept(t, InprocessMode, SingleTest)
|
testAccept(t, InprocessMode, SingleTest)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestAcceptLocal(t *testing.T) {
|
||||||
|
testAccept(t, InprocessMode, SingleTest)
|
||||||
|
}
|
||||||
|
|
||||||
func TestInprocessMode(t *testing.T) {
|
func TestInprocessMode(t *testing.T) {
|
||||||
if InprocessMode {
|
if InprocessMode {
|
||||||
t.Skip("Already tested by TestAccept")
|
t.Skip("Already tested by TestAccept")
|
||||||
|
@ -223,10 +226,9 @@ func runTest(t *testing.T, dir, coverDir string, repls testdiff.ReplacementsCont
|
||||||
|
|
||||||
if !isTruePtr(config.Cloud) && cloudEnv != "" {
|
if !isTruePtr(config.Cloud) && cloudEnv != "" {
|
||||||
t.Skipf("Disabled via Cloud setting in %s (CLOUD_ENV=%s)", configPath, cloudEnv)
|
t.Skipf("Disabled via Cloud setting in %s (CLOUD_ENV=%s)", configPath, cloudEnv)
|
||||||
} else {
|
|
||||||
if isTruePtr(config.RequiresUnityCatalog) && os.Getenv("TEST_METASTORE_ID") == "" {
|
|
||||||
t.Skipf("Skipping on non-UC workspaces")
|
|
||||||
}
|
}
|
||||||
|
if cloudEnv != "" && isTruePtr(config.RequiresUnityCatalog) && os.Getenv("TEST_METASTORE_ID") == "" {
|
||||||
|
t.Skipf("Skipping on non-UC workspaces")
|
||||||
}
|
}
|
||||||
|
|
||||||
var tmpDir string
|
var tmpDir string
|
||||||
|
@ -334,6 +336,22 @@ func runTest(t *testing.T, dir, coverDir string, repls testdiff.ReplacementsCont
|
||||||
user = *pUser
|
user = *pUser
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if cloudEnv != "" && isTruePtr(config.RecordRequests) {
|
||||||
|
// Start a new recording proxy for this test
|
||||||
|
logPath := filepath.Join(tmpDir, "out.request-recordings.txt")
|
||||||
|
targetHost := os.Getenv("DATABRICKS_HOST")
|
||||||
|
proxyServer, err := testserver.NewProxyRecorder(targetHost, logPath)
|
||||||
|
t.Logf("Starting a recording proxy server; proxies to " + targetHost + "; recording requests to " + logPath)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Failed to create proxy server: %v", err)
|
||||||
|
}
|
||||||
|
t.Cleanup(proxyServer.Close)
|
||||||
|
|
||||||
|
t.Logf("Setting DATABRICKS_DEFAULT_HOST to " + proxyServer.URL())
|
||||||
|
t.Setenv("DATABRICKS_DEFAULT_HOST", proxyServer.URL())
|
||||||
|
cmd.Env = append(cmd.Env, "DATABRICKS_HOST="+proxyServer.URL())
|
||||||
|
}
|
||||||
|
|
||||||
testdiff.PrepareReplacementsUser(t, &repls, user)
|
testdiff.PrepareReplacementsUser(t, &repls, user)
|
||||||
testdiff.PrepareReplacementsWorkspaceClient(t, &repls, workspaceClient)
|
testdiff.PrepareReplacementsWorkspaceClient(t, &repls, workspaceClient)
|
||||||
|
|
||||||
|
|
|
@ -1,3 +1,32 @@
|
||||||
Local = false
|
Local = true
|
||||||
Cloud = true
|
Cloud = true
|
||||||
RequiresUnityCatalog = true
|
RequiresUnityCatalog = true
|
||||||
|
RecordRequests = true
|
||||||
|
|
||||||
|
[[Server]]
|
||||||
|
Pattern = "POST /api/2.1/unity-catalog/schemas"
|
||||||
|
Response.Body = '''
|
||||||
|
{
|
||||||
|
"name":"test-schema-6260d50f-e8ff-4905-8f28-812345678903",
|
||||||
|
"catalog_name":"main",
|
||||||
|
"enable_auto_maintenance":"INHERIT",
|
||||||
|
"enable_predictive_optimization":"INHERIT",
|
||||||
|
"full_name":"main.test-schema-6260d50f-e8ff-4905-8f28-812345678903",
|
||||||
|
"created_at":1741363224990,
|
||||||
|
"created_by":"[USERNAME]",
|
||||||
|
"updated_at":1741363224990,
|
||||||
|
"updated_by":"[USERNAME]",
|
||||||
|
"catalog_type":"MANAGED_CATALOG",
|
||||||
|
"schema_id":"6260d50f-e8ff-4905-8f28-812345678903",
|
||||||
|
"securable_type":"SCHEMA",
|
||||||
|
"securable_kind":"SCHEMA_STANDARD",
|
||||||
|
"browse_only":false,
|
||||||
|
"metastore_version":-1
|
||||||
|
}
|
||||||
|
'''
|
||||||
|
|
||||||
|
[[Server]]
|
||||||
|
Pattern = "GET /api/2.1/unity-catalog/schemas/{schema_fullname}"
|
||||||
|
|
||||||
|
[[Server]]
|
||||||
|
Pattern = "POST /api/2.1/unity-catalog/schemas"
|
|
@ -0,0 +1,287 @@
|
||||||
|
package testserver
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"compress/gzip"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"log"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"os"
|
||||||
|
"strings"
|
||||||
|
"sync"
|
||||||
|
"time"
|
||||||
|
"unicode/utf8"
|
||||||
|
)
|
||||||
|
|
||||||
|
// ProxyRecorder represents a proxy server that records HTTP traffic
|
||||||
|
type ProxyRecorder struct {
|
||||||
|
// Remote server to proxy to
|
||||||
|
RemoteURL string
|
||||||
|
// File to record traffic to
|
||||||
|
RecordFile *os.File
|
||||||
|
// Lock for concurrent writes to the file
|
||||||
|
mu sync.Mutex
|
||||||
|
// The test server
|
||||||
|
Server *httptest.Server
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewProxyRecorder creates a new proxy recorder
|
||||||
|
func NewProxyRecorder(remoteURL, recordFilePath string) (*ProxyRecorder, error) {
|
||||||
|
// Open or create the record file
|
||||||
|
file, err := os.OpenFile(recordFilePath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to open record file: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Create the proxy recorder
|
||||||
|
pr := &ProxyRecorder{
|
||||||
|
RemoteURL: remoteURL,
|
||||||
|
RecordFile: file,
|
||||||
|
}
|
||||||
|
|
||||||
|
// Create the test server
|
||||||
|
pr.Server = httptest.NewServer(http.HandlerFunc(pr.proxyHandler))
|
||||||
|
|
||||||
|
return pr, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Close stops the server and closes the record file
|
||||||
|
func (pr *ProxyRecorder) Close() {
|
||||||
|
pr.Server.Close()
|
||||||
|
err := pr.RecordFile.Close()
|
||||||
|
if err != nil {
|
||||||
|
log.Fatalf("Error closing record file")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// URL returns the URL of the test server
|
||||||
|
func (pr *ProxyRecorder) URL() string {
|
||||||
|
return pr.Server.URL
|
||||||
|
}
|
||||||
|
|
||||||
|
// proxyHandler handles HTTP requests, forwards them to the remote server,
|
||||||
|
// and records both the request and response
|
||||||
|
func (pr *ProxyRecorder) proxyHandler(w http.ResponseWriter, r *http.Request) {
|
||||||
|
// Record timestamp
|
||||||
|
timestamp := time.Now().Format(time.RFC3339)
|
||||||
|
|
||||||
|
// Dump the request to be recorded
|
||||||
|
var requestDump strings.Builder
|
||||||
|
requestDump.WriteString(fmt.Sprintf("=== REQUEST %s ===\n", timestamp))
|
||||||
|
requestDump.WriteString(fmt.Sprintf("%s %s %s\n", r.Method, r.URL.Path, r.Proto))
|
||||||
|
|
||||||
|
// Record headers
|
||||||
|
for name, values := range r.Header {
|
||||||
|
for _, value := range values {
|
||||||
|
requestDump.WriteString(fmt.Sprintf("%s: %s\n", name, value))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Read and record the request body, if any
|
||||||
|
var requestBody []byte
|
||||||
|
if r.Body != nil {
|
||||||
|
var err error
|
||||||
|
requestBody, err = io.ReadAll(r.Body)
|
||||||
|
if err != nil {
|
||||||
|
http.Error(w, "Failed to read request body", http.StatusInternalServerError)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
// Close the original body and replace it with a new reader
|
||||||
|
r.Body.Close()
|
||||||
|
r.Body = io.NopCloser(bytes.NewBuffer(requestBody))
|
||||||
|
|
||||||
|
if len(requestBody) > 0 {
|
||||||
|
requestDump.WriteString("\n")
|
||||||
|
|
||||||
|
// Check content type to handle binary data appropriately
|
||||||
|
contentType := r.Header.Get("Content-Type")
|
||||||
|
if isBinaryContent(contentType) {
|
||||||
|
// For binary content, log length and format instead of raw content
|
||||||
|
requestDump.WriteString(fmt.Sprintf("[Binary data, %d bytes, Content-Type: %s]\n",
|
||||||
|
len(requestBody), contentType))
|
||||||
|
} else {
|
||||||
|
// For text content, convert to UTF-8 if needed
|
||||||
|
bodyStr := safeStringConversion(requestBody)
|
||||||
|
requestDump.WriteString(bodyStr)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
requestDump.WriteString("\n\n")
|
||||||
|
|
||||||
|
// Create a new request to the remote server
|
||||||
|
remoteURL := pr.RemoteURL + r.URL.Path
|
||||||
|
if r.URL.RawQuery != "" {
|
||||||
|
remoteURL += "?" + r.URL.RawQuery
|
||||||
|
}
|
||||||
|
|
||||||
|
proxyReq, err := http.NewRequest(r.Method, remoteURL, bytes.NewBuffer(requestBody))
|
||||||
|
if err != nil {
|
||||||
|
http.Error(w, "Failed to create proxy request", http.StatusInternalServerError)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// Copy headers
|
||||||
|
for name, values := range r.Header {
|
||||||
|
for _, value := range values {
|
||||||
|
proxyReq.Header.Add(name, value)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Send the request to the remote server
|
||||||
|
client := &http.Client{}
|
||||||
|
resp, err := client.Do(proxyReq)
|
||||||
|
if err != nil {
|
||||||
|
http.Error(w, "Failed to proxy request", http.StatusBadGateway)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
|
||||||
|
// Record the response
|
||||||
|
var responseDump strings.Builder
|
||||||
|
responseDump.WriteString(fmt.Sprintf("=== RESPONSE %s ===\n", timestamp))
|
||||||
|
responseDump.WriteString(fmt.Sprintf("%s %d %s\n", resp.Proto, resp.StatusCode, resp.Status))
|
||||||
|
|
||||||
|
// Record response headers
|
||||||
|
for name, values := range resp.Header {
|
||||||
|
for _, value := range values {
|
||||||
|
responseDump.WriteString(fmt.Sprintf("%s: %s\n", name, value))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Read and record the response body
|
||||||
|
responseBody, err := io.ReadAll(resp.Body)
|
||||||
|
if err != nil {
|
||||||
|
http.Error(w, "Failed to read response body", http.StatusInternalServerError)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
var bodyForLogging []byte
|
||||||
|
isGzipped := false
|
||||||
|
|
||||||
|
// Check if response is gzipped
|
||||||
|
for _, encoding := range resp.Header.Values("Content-Encoding") {
|
||||||
|
if strings.Contains(strings.ToLower(encoding), "gzip") {
|
||||||
|
isGzipped = true
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if isGzipped {
|
||||||
|
// Decompress for logging
|
||||||
|
gzipReader, err := gzip.NewReader(bytes.NewReader(responseBody))
|
||||||
|
if err != nil {
|
||||||
|
// If decompression fails, just log that it's compressed
|
||||||
|
responseDump.WriteString(fmt.Sprintf("\n[Gzipped content, %d bytes]\n", len(responseBody)))
|
||||||
|
} else {
|
||||||
|
defer gzipReader.Close()
|
||||||
|
|
||||||
|
// Read the decompressed content
|
||||||
|
decompressed, err := io.ReadAll(gzipReader)
|
||||||
|
if err != nil {
|
||||||
|
responseDump.WriteString(fmt.Sprintf("\n[Error decompressing gzipped content: %v]\n", err))
|
||||||
|
} else {
|
||||||
|
bodyForLogging = decompressed
|
||||||
|
responseDump.WriteString("\n[Decompressed from gzip]\n")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
bodyForLogging = responseBody
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(bodyForLogging) > 0 {
|
||||||
|
// Check content type to handle binary data appropriately
|
||||||
|
contentType := resp.Header.Get("Content-Type")
|
||||||
|
if isBinaryContent(contentType) {
|
||||||
|
// For binary content, log length and format instead of raw content
|
||||||
|
responseDump.WriteString(fmt.Sprintf("[Binary data, %d bytes, Content-Type: %s]\n",
|
||||||
|
len(responseBody), contentType))
|
||||||
|
} else {
|
||||||
|
// For text content, convert to UTF-8 if needed
|
||||||
|
bodyStr := safeStringConversion(bodyForLogging)
|
||||||
|
responseDump.WriteString(bodyStr)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
responseDump.WriteString("\n\n")
|
||||||
|
|
||||||
|
// Write the request and response to the record file
|
||||||
|
pr.mu.Lock()
|
||||||
|
pr.RecordFile.WriteString(requestDump.String())
|
||||||
|
pr.RecordFile.WriteString(responseDump.String())
|
||||||
|
pr.mu.Unlock()
|
||||||
|
|
||||||
|
// Copy the response headers to the original response writer
|
||||||
|
for name, values := range resp.Header {
|
||||||
|
for _, value := range values {
|
||||||
|
w.Header().Add(name, value)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Set the status code
|
||||||
|
w.WriteHeader(resp.StatusCode)
|
||||||
|
|
||||||
|
// Write the response body to the original response writer
|
||||||
|
w.Write(responseBody)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ensureValidUTF8 converts a string to valid UTF-8, replacing invalid byte sequences
|
||||||
|
func ensureValidUTF8(s string) string {
|
||||||
|
if utf8.ValidString(s) {
|
||||||
|
return s
|
||||||
|
}
|
||||||
|
|
||||||
|
// Replace invalid UTF-8 sequences with the Unicode replacement character (U+FFFD)
|
||||||
|
v := make([]rune, 0, len(s))
|
||||||
|
for i, r := range s {
|
||||||
|
if r == utf8.RuneError {
|
||||||
|
_, size := utf8.DecodeRuneInString(s[i:])
|
||||||
|
if size == 1 {
|
||||||
|
// Invalid UTF-8 sequence, replace with Unicode replacement character
|
||||||
|
v = append(v, '\uFFFD')
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
}
|
||||||
|
v = append(v, r)
|
||||||
|
}
|
||||||
|
return string(v)
|
||||||
|
}
|
||||||
|
|
||||||
|
// isBinaryContent checks if the content type represents binary data
|
||||||
|
func isBinaryContent(contentType string) bool {
|
||||||
|
if contentType == "" {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
// List of common binary content types
|
||||||
|
binaryTypes := []string{
|
||||||
|
"image/", "audio/", "video/", "application/octet-stream",
|
||||||
|
"application/pdf", "application/zip", "application/gzip",
|
||||||
|
"application/x-tar", "application/x-rar-compressed",
|
||||||
|
"application/x-7z-compressed", "application/x-msdownload",
|
||||||
|
"application/vnd.ms-", "application/x-ms",
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, prefix := range binaryTypes {
|
||||||
|
if strings.HasPrefix(contentType, prefix) {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
// safeStringConversion attempts to convert bytes to a valid UTF-8 string
|
||||||
|
func safeStringConversion(data []byte) string {
|
||||||
|
// First try direct conversion
|
||||||
|
s := string(data)
|
||||||
|
|
||||||
|
// Check if it's valid UTF-8
|
||||||
|
if utf8.ValidString(s) {
|
||||||
|
return s
|
||||||
|
}
|
||||||
|
|
||||||
|
// Try to detect encoding and convert to UTF-8
|
||||||
|
// For simplicity, we'll just replace invalid characters
|
||||||
|
return ensureValidUTF8(s)
|
||||||
|
}
|
Loading…
Reference in New Issue