Add async logger for telemetry

acceptance/bin/wait_pid.py

@@ -0,0 +1,49 @@
+#!/usr/bin/env python3
+import sys
+import os
+import time
+import platform
+import subprocess
+
+
+def wait_pid(pid):
+    max_attempts = 600  # 600 * 0.1 seconds = 1 minute
+    sleep_time = 0.1
+
+    for i in range(max_attempts):
+        # Check if we are on Windows or a Unix system.
+        if platform.system().lower() == "windows":
+            # Windows approach: use tasklist to check for the existence of the process.
+            try:
+                # Get the output of 'tasklist'
+                output = subprocess.check_output(["tasklist"], text=True)
+            except subprocess.CalledProcessError:
+                print("[wait_pid] Error retrieving tasklist. Assuming process has ended.")
+                return 0
+
+            # if the PID is not found in the list then assume the process ended.
+            if str(pid) not in output:
+                print("[wait_pid] process has ended")
+                return 0
+        else:
+            # Linux/macOS approach: using os.kill with signal 0 to check if the process exists.
+            try:
+                os.kill(pid, 0)
+            except OSError:
+                print("[wait_pid] process has ended")
+                return 0
+
+        time.sleep(sleep_time)
+
+    print(f"Timeout: Process {pid} did not end within 1 minute")
+    return 1
+
+
+try:
+    pid = int(sys.argv[1])
+except ValueError:
+    print("Error: PID must be an integer.")
+    sys.exit(1)
+
+exit_code = wait_pid(pid)
+sys.exit(exit_code)

acceptance/telemetry/test.toml

@@ -0,0 +1,15 @@
+IncludeRequestHeaders = ["Authorization"]
+Local = true
+Cloud = false
+
+[[Repls]]
+Old = '17\d{11}'
+New = '"UNIX_TIME_MILLIS"'
+
+[[Repls]]
+Old = 'darwin|linux|windows'
+New = 'OS'
+
+[[Repls]]
+Old = 'execution_time_ms\\\":\d{1,5},'
+New = 'execution_time_ms\":\"SMALL_INT\",'

acceptance/telemetry/upload-command/out.requests.txt

@@ -0,0 +1,17 @@
+{
+  "headers": {
+    "Authorization": [
+      "Bearer [DATABRICKS_TOKEN]"
+    ]
+  },
+  "method": "POST",
+  "path": "/telemetry-ext",
+  "body": {
+    "uploadTime": "UNIX_TIME_MILLIS",
+    "items": [],
+    "protoLogs": [
+      "{\"frontend_log_event_id\":\"[UUID]\",\"entry\":{\"databricks_cli_log\":{\"cli_test_event\":{\"name\":\"VALUE1\"}}}}",
+      "{\"frontend_log_event_id\":\"[UUID]\",\"entry\":{\"databricks_cli_log\":{\"cli_test_event\":{\"name\":\"VALUE2\"}}}}"
+    ]
+  }
+}

acceptance/telemetry/upload-command/output.txt

@@ -0,0 +1,5 @@
+
+>>> [CLI] telemetry upload
+Telemetry logs uploaded successfully
+Response:
+{"errors":[],"numProtoSuccess":2}

acceptance/telemetry/upload-command/script

@@ -0,0 +1,5 @@
+export DATABRICKS_CLI_TELEMETRY_UPLOAD_LOGS_FILE=./out.upload.txt
+
+# This command / test cannot be run in inprocess / debug mode. This is because
+# it does not go through the [root.Execute] function.
+trace $CLI telemetry upload < stdin

acceptance/telemetry/upload-command/stdin

@@ -0,0 +1,24 @@
+{
+    "logs": [
+        {
+            "frontend_log_event_id": "AAAAAAAA-AAAA-AAAA-AAAA-AAAAAAAAAAAA",
+            "entry": {
+                "databricks_cli_log": {
+                    "cli_test_event": {
+                        "name": "VALUE1"
+                    }
+                }
+            }
+        },
+        {
+            "frontend_log_event_id": "BBBBBBBB-BBBB-BBBB-BBBB-BBBBBBBBBBBB",
+            "entry": {
+                "databricks_cli_log": {
+                    "cli_test_event": {
+                        "name": "VALUE2"
+                    }
+                }
+            }
+        }
+    ]
+}

acceptance/telemetry/upload-command/test.toml

@@ -0,0 +1,10 @@
+RecordRequests = true
+
+[[Server]]
+Pattern = "POST /telemetry-ext"
+Response.Body = '''
+{
+    "errors": [],
+    "numProtoSuccess": 2
+}
+'''

acceptance/telemetry/upload-fails/out.requests.txt

@@ -0,0 +1,17 @@
+{
+  "headers": {
+    "Authorization": [
+      "Bearer [DATABRICKS_TOKEN]"
+    ]
+  },
+  "method": "POST",
+  "path": "/telemetry-ext",
+  "body": {
+    "uploadTime": "UNIX_TIME_MILLIS",
+    "items": [],
+    "protoLogs": [
+      "{\"frontend_log_event_id\":\"[UUID]\",\"entry\":{\"databricks_cli_log\":{\"execution_context\":{\"cmd_exec_id\":\"[UUID]\",\"version\":\"[DEV_VERSION]\",\"command\":\"selftest_send-telemetry\",\"operating_system\":\"OS\",\"execution_time_ms\":\"SMALL_INT\",\"exit_code\":0},\"cli_test_event\":{\"name\":\"VALUE1\"}}}}",
+      "{\"frontend_log_event_id\":\"[UUID]\",\"entry\":{\"databricks_cli_log\":{\"execution_context\":{\"cmd_exec_id\":\"[UUID]\",\"version\":\"[DEV_VERSION]\",\"command\":\"selftest_send-telemetry\",\"operating_system\":\"OS\",\"execution_time_ms\":\"SMALL_INT\",\"exit_code\":0},\"cli_test_event\":{\"name\":\"VALUE2\"}}}}"
+    ]
+  }
+}

acceptance/telemetry/upload-fails/out.upload_process.txt

@@ -0,0 +1,2 @@
+error: Failed to upload telemetry logs: Endpoint not implemented.
+

acceptance/telemetry/upload-fails/output.txt

@@ -0,0 +1,3 @@
+
+>>> [CLI] selftest send-telemetry
+[wait_pid] process has ended

acceptance/telemetry/upload-fails/script

@@ -0,0 +1,12 @@
+export DATABRICKS_CLI_TELEMETRY_PID_FILE=./telemetry.pid
+export DATABRICKS_CLI_TELEMETRY_UPLOAD_LOGS_FILE=./out.upload_process.txt
+
+# This test ensures that the main CLI command does not error even if
+# telemetry upload fails.
+trace $CLI selftest send-telemetry
+
+# Wait for the child telemetry process to finish
+wait_pid.py $(cat ./telemetry.pid)
+
+# cleanup the pid file
+rm -f ./telemetry.pid

acceptance/telemetry/upload-fails/test.toml

@@ -0,0 +1,11 @@
+RecordRequests = true
+
+[[Server]]
+Pattern = "POST /telemetry-ext"
+Response.Body = '''
+{
+    "error_code": "ERROR_CODE",
+    "message": "Endpoint not implemented."
+}
+'''
+Response.StatusCode = 501

acceptance/telemetry/upload-skipped/output.txt

@@ -0,0 +1,3 @@
+
+>>> [CLI] selftest send-telemetry
+telemetry process not started. No pid file found

acceptance/telemetry/upload-skipped/script

@@ -0,0 +1,12 @@
+export DATABRICKS_CLI_TELEMETRY_PID_FILE=./telemetry.pid
+export DATABRICKS_CLI_TELEMETRY_UPLOAD_LOGS_FILE=./out.upload_process.txt
+export DATABRICKS_CLI_DISABLE_TELEMETRY="true"
+
+trace $CLI selftest send-telemetry
+
+if [ -f ./telemetry.pid ]; then
+  echo "telemetry process was started"
+  exit 1
+else
+  echo "telemetry process not started. No pid file found"
+fi

acceptance/telemetry/upload-skipped/test.toml

@@ -0,0 +1,10 @@
+RecordRequests = true
+
+[[Server]]
+Pattern = "POST /telemetry-ext"
+Response.Body = '''
+{
+    "errors": [],
+    "numProtoSuccess": 2
+}
+'''

acceptance/telemetry/upload-succeeds-integration/out.upload_process.txt

@@ -0,0 +1,3 @@
+Telemetry logs uploaded successfully
+Response:
+{"errors":[],"numProtoSuccess":2}

acceptance/telemetry/upload-succeeds-integration/output.txt

@@ -0,0 +1,4 @@
+
+>>> [CLI] selftest send-telemetry
+waiting for telemetry process to finish
+[wait_pid] process has ended

acceptance/telemetry/upload-succeeds-integration/script

@@ -0,0 +1,12 @@
+export DATABRICKS_CLI_TELEMETRY_PID_FILE=./telemetry.pid
+export DATABRICKS_CLI_TELEMETRY_UPLOAD_LOGS_FILE=./out.upload_process.txt
+
+trace $CLI selftest send-telemetry
+
+echo "waiting for telemetry process to finish"
+
+# Wait for the child telemetry process to finish
+wait_pid.py $(cat ./telemetry.pid)
+
+# cleanup the pid file
+rm -f ./telemetry.pid

acceptance/telemetry/upload-succeeds-integration/test.toml

@@ -0,0 +1,14 @@
+# RecordRequests is not supported for integration tests yet. Once we have support
+# for running integration tests with record requests we can remove this test since
+# it's a copy of acceptance/telemetry/upload-succeeds.
+Local = true
+Cloud = true
+
+[[Server]]
+Pattern = "POST /telemetry-ext"
+Response.Body = '''
+{
+    "errors": [],
+    "numProtoSuccess": 2
+}
+'''

acceptance/telemetry/upload-succeeds-oauth/out.requests.txt

@@ -0,0 +1,45 @@
+{
+  "method": "GET",
+  "path": "/oidc/.well-known/oauth-authorization-server"
+}
+{
+  "headers": {
+    "Authorization": [
+      "Basic [ENCODED_AUTH]"
+    ]
+  },
+  "method": "POST",
+  "path": "/oidc/v1/token",
+  "raw_body": "grant_type=client_credentials\u0026scope=all-apis"
+}
+{
+  "method": "GET",
+  "path": "/oidc/.well-known/oauth-authorization-server"
+}
+{
+  "headers": {
+    "Authorization": [
+      "Basic [ENCODED_AUTH]"
+    ]
+  },
+  "method": "POST",
+  "path": "/oidc/v1/token",
+  "raw_body": "grant_type=client_credentials\u0026scope=all-apis"
+}
+{
+  "headers": {
+    "Authorization": [
+      "Bearer oauth-token"
+    ]
+  },
+  "method": "POST",
+  "path": "/telemetry-ext",
+  "body": {
+    "uploadTime": "UNIX_TIME_MILLIS",
+    "items": [],
+    "protoLogs": [
+      "{\"frontend_log_event_id\":\"[UUID]\",\"entry\":{\"databricks_cli_log\":{\"execution_context\":{\"cmd_exec_id\":\"[UUID]\",\"version\":\"[DEV_VERSION]\",\"command\":\"selftest_send-telemetry\",\"operating_system\":\"OS\",\"execution_time_ms\":\"SMALL_INT\",\"exit_code\":0},\"cli_test_event\":{\"name\":\"VALUE1\"}}}}",
+      "{\"frontend_log_event_id\":\"[UUID]\",\"entry\":{\"databricks_cli_log\":{\"execution_context\":{\"cmd_exec_id\":\"[UUID]\",\"version\":\"[DEV_VERSION]\",\"command\":\"selftest_send-telemetry\",\"operating_system\":\"OS\",\"execution_time_ms\":\"SMALL_INT\",\"exit_code\":0},\"cli_test_event\":{\"name\":\"VALUE2\"}}}}"
+    ]
+  }
+}

acceptance/telemetry/upload-succeeds-oauth/out.upload_process.txt

@@ -0,0 +1,3 @@
+Telemetry logs uploaded successfully
+Response:
+{"errors":[],"numProtoSuccess":2}

acceptance/telemetry/upload-succeeds-oauth/output.txt

@@ -0,0 +1,4 @@
+
+>>> [CLI] selftest send-telemetry
+waiting for telemetry process to finish
+[wait_pid] process has ended

acceptance/telemetry/upload-succeeds-oauth/script

@@ -0,0 +1,19 @@
+export DATABRICKS_CLI_TELEMETRY_PID_FILE=./telemetry.pid
+export DATABRICKS_CLI_TELEMETRY_UPLOAD_LOGS_FILE=./out.upload_process.txt
+
+# Unset the token which is configured by default
+# in acceptance tests
+export DATABRICKS_TOKEN=""
+
+export DATABRICKS_CLIENT_ID=client_id
+export DATABRICKS_CLIENT_SECRET=client_secret
+
+trace $CLI selftest send-telemetry
+
+echo "waiting for telemetry process to finish"
+
+# Wait for the child telemetry process to finish
+wait_pid.py $(cat ./telemetry.pid)
+
+# cleanup the pid file
+rm -f ./telemetry.pid

acceptance/telemetry/upload-succeeds-oauth/test.toml

@@ -0,0 +1,17 @@
+RecordRequests = true
+
+[[Server]]
+Pattern = "POST /telemetry-ext"
+Response.Body = '''
+{
+    "errors": [],
+    "numProtoSuccess": 2
+}
+'''
+
+
+# "client_id:client_secret" in base64 is Y2xpZW50X2lkOmNsaWVudF9zZWNyZXQ=, expect to
+# see this in Authorization header
+[[Repls]]
+Old = "Y2xpZW50X2lkOmNsaWVudF9zZWNyZXQ="
+New = "[ENCODED_AUTH]"

acceptance/telemetry/upload-succeeds/out.requests.txt

@@ -0,0 +1,17 @@
+{
+  "headers": {
+    "Authorization": [
+      "Bearer [DATABRICKS_TOKEN]"
+    ]
+  },
+  "method": "POST",
+  "path": "/telemetry-ext",
+  "body": {
+    "uploadTime": "UNIX_TIME_MILLIS",
+    "items": [],
+    "protoLogs": [
+      "{\"frontend_log_event_id\":\"[UUID]\",\"entry\":{\"databricks_cli_log\":{\"execution_context\":{\"cmd_exec_id\":\"[UUID]\",\"version\":\"[DEV_VERSION]\",\"command\":\"selftest_send-telemetry\",\"operating_system\":\"OS\",\"execution_time_ms\":\"SMALL_INT\",\"exit_code\":0},\"cli_test_event\":{\"name\":\"VALUE1\"}}}}",
+      "{\"frontend_log_event_id\":\"[UUID]\",\"entry\":{\"databricks_cli_log\":{\"execution_context\":{\"cmd_exec_id\":\"[UUID]\",\"version\":\"[DEV_VERSION]\",\"command\":\"selftest_send-telemetry\",\"operating_system\":\"OS\",\"execution_time_ms\":\"SMALL_INT\",\"exit_code\":0},\"cli_test_event\":{\"name\":\"VALUE2\"}}}}"
+    ]
+  }
+}

acceptance/telemetry/upload-succeeds/out.upload_process.txt

@@ -0,0 +1,3 @@
+Telemetry logs uploaded successfully
+Response:
+{"errors":[],"numProtoSuccess":2}

acceptance/telemetry/upload-succeeds/output.txt

@@ -0,0 +1,4 @@
+
+>>> [CLI] selftest send-telemetry
+waiting for telemetry process to finish
+[wait_pid] process has ended

acceptance/telemetry/upload-succeeds/script

@@ -0,0 +1,12 @@
+export DATABRICKS_CLI_TELEMETRY_PID_FILE=./telemetry.pid
+export DATABRICKS_CLI_TELEMETRY_UPLOAD_LOGS_FILE=./out.upload_process.txt
+
+trace $CLI selftest send-telemetry
+
+echo "waiting for telemetry process to finish"
+
+# Wait for the child telemetry process to finish
+wait_pid.py $(cat ./telemetry.pid)
+
+# cleanup the pid file
+rm -f ./telemetry.pid

acceptance/telemetry/upload-succeeds/test.toml

@@ -0,0 +1,10 @@
+RecordRequests = true
+
+[[Server]]
+Pattern = "POST /telemetry-ext"
+Response.Body = '''
+{
+    "errors": [],
+    "numProtoSuccess": 2
+}
+'''

acceptance/telemetry/upload-timeout/out.upload_process.txt

@@ -0,0 +1 @@
+error: Failed to flush telemetry log due to timeout

acceptance/telemetry/upload-timeout/output.txt

@@ -0,0 +1,4 @@
+
+>>> [CLI] selftest send-telemetry
+waiting for telemetry process to finish
+[wait_pid] process has ended