Implement telemetry for commands (#84)

* lsClient has been remove from webview. clean up Signed-off-by: Nok <[email protected]> * add consent script Signed-off-by: Nok Lam Chan <[email protected]> * Install requirements from viz, telemetry separately Signed-off-by: Nok Lam Chan <[email protected]> * Refactor globalState into constants to have one place to keep track of all the globalstate we used Signed-off-by: Nok Lam Chan <[email protected]> * fixing telemetry schema Signed-off-by: Nok Lam Chan <[email protected]> * send telemetry Signed-off-by: Nok Lam Chan <[email protected]> * remove unused import Signed-off-by: Nok Lam Chan <[email protected]> * update version Signed-off-by: Nok Lam Chan <[email protected]> * remove true in logging, bug fix Signed-off-by: Nok Lam Chan <[email protected]> * bug fix - should use telemetry requirements file Signed-off-by: Nok Lam Chan <[email protected]> * bug fix -passing wrong parameters Signed-off-by: Nok Lam Chan <[email protected]> * bug fix Signed-off-by: Nok Lam Chan <[email protected]> * Properties Map to object Signed-off-by: Jitendra Gundaniya <[email protected]> * Update bundled/tool/check_consent.py Co-authored-by: Jitendra Gundaniya <[email protected]> * Update schema Signed-off-by: Nok Lam Chan <[email protected]> * fix telemetry schema Signed-off-by: Nok Lam Chan <[email protected]> * format & missing imports Signed-off-by: Nok Lam Chan <[email protected]> * fix consent check Signed-off-by: Nok Lam Chan <[email protected]> --------- Signed-off-by: Nok <[email protected]> Signed-off-by: Nok Lam Chan <[email protected]> Signed-off-by: Jitendra Gundaniya <[email protected]> Co-authored-by: Jitendra Gundaniya <[email protected]> Co-authored-by: Jitendra Gundaniya <[email protected]>
kedro-org · Sep 13, 2024 · a573818 · a573818
1 parent 4ed8552
commit a573818
Show file tree

Hide file tree

Showing 12 changed files with 272 additions and 27 deletions.
diff --git a/bundled/tool/check_consent.py b/bundled/tool/check_consent.py
@@ -0,0 +1,39 @@
+from common import update_sys_path
+from pathlib import Path
+import os
+import sys
+import json
+
+update_sys_path(
+    os.fspath(Path(__file__).parent.parent / "libs"),
+    os.getenv("LS_IMPORT_STRATEGY", "useBundled"),
+)
+
+# important to keep this after sys.path is updated
+from kedro_telemetry.plugin import _check_for_telemetry_consent
+from kedro_telemetry.plugin import (
+    _get_project_properties,
+    _get_or_create_uuid,
+)
+
+if __name__ == "__main__":
+    from pathlib import Path
+    import sys
+
+    if len(sys.argv) > 1:
+        project_path = Path(sys.argv[1])
+    else:
+        project_path = Path.cwd()
+    consent = _check_for_telemetry_consent(project_path)
+
+    # Project Metadata
+
+    user_uuid = _get_or_create_uuid()
+    properties = _get_project_properties(user_uuid, project_path)
+    # Extension will parse this message
+    properties["consent"] = consent
+    print("telemetry consent: ", end="")
+    # It is important to use json.dump, if the message is printed together Python
+    # convert it to single quote and the result is no longer valid JSON. The message
+    # will be parsed by the extension client.
+    result = json.dump(properties, sys.stdout)
diff --git a/bundled/tool/install_telemetry_dependencies.py b/bundled/tool/install_telemetry_dependencies.py
@@ -0,0 +1,47 @@
+import subprocess
+import sys
+from pathlib import Path
+
+def install_dependencies(extension_root_dir):
+    """
+    Install dependencies required for the Kedro extension.
+
+    Args:
+        extension_root_dir (str): The root directory of the extension.
+    Raises:
+        ImportError: If the required dependencies are not found.
+    """
+    ...
+    libs_path = Path(extension_root_dir) / "bundled" / "libs"
+    requirements_path = Path(extension_root_dir) / "kedro-telemetry-requirements.txt"
+
+    try:
+        import kedro_telemetry
+        from packaging.version import parse
+
+        version = parse(kedro_telemetry.__version__)
+        if version.major<1 and version.minor<6: # at least >0.6.0
+            raise ImportError("kedro-telemetry version must be >=0.6.0")
+    except ImportError:
+        subprocess.check_call(
+            [
+                sys.executable,
+                "-m",
+                "pip",
+                "install",
+                "-r",
+                Path(requirements_path),
+                "-t",
+                Path(libs_path),
+                "--no-cache-dir",
+                "--no-deps"
+            ]
+        )
+
+
+if __name__ == "__main__":
+    if len(sys.argv) > 1:
+        extension_root_dir = sys.argv[1]
+    else:
+        extension_root_dir = None
+    install_dependencies(extension_root_dir)
diff --git a/bundled/tool/install_dependencies.py → bundled/tool/install_viz_dependencies.py b/bundled/tool/install_dependencies.py → bundled/tool/install_viz_dependencies.py
@@ -18,6 +18,7 @@ def install_dependencies(extension_root_dir):
     try:
         import fastapi
         import orjson
+
     except ImportError:
         subprocess.check_call(
             [

diff --git a/bundled/tool/lsp_server.py b/bundled/tool/lsp_server.py
@@ -556,7 +556,7 @@ def definition_from_flowchart(ls, word):
     return result
 
 @LSP_SERVER.command("kedro.getProjectData")
-def get_porject_data_from_viz(lsClient):
+def get_project_data_from_viz(lsClient):
     """Get project data from kedro viz
     """
     data = None

diff --git a/kedro-telemetry-requirements.txt b/kedro-telemetry-requirements.txt
@@ -0,0 +1,2 @@
+packaging
+kedro-telemetry>=0.6.0 # First version that does not prompt for telemetry
diff --git a/kedro-viz-requirements.txt b/kedro-viz-requirements.txt
@@ -1,3 +1,3 @@
 fastapi>=0.100.0,<0.200.0
 pydantic>=2.0.0 # In case of FastAPI installs pydantic==1
-orjson>=3.9, <4.0
+orjson>=3.9, <4.0
diff --git a/package-lock.json b/package-lock.json
diff --git a/package.json b/package.json
@@ -168,11 +168,11 @@
                 "category": "kedro",
                 "title": "Run Kedro Viz"
             }
-
         ]
     },
     "dependencies": {
         "@vscode/python-extension": "^1.0.5",
+        "axios": "^1.7.7",
         "fs-extra": "^11.2.0",
         "vscode-languageclient": "^8.1.0"
     },

diff --git a/src/common/constants.ts b/src/common/constants.ts
@@ -9,3 +9,9 @@ export const EXTENSION_ROOT_DIR =
 export const BUNDLED_PYTHON_SCRIPTS_DIR = path.join(EXTENSION_ROOT_DIR, 'bundled');
 export const SERVER_SCRIPT_PATH = path.join(BUNDLED_PYTHON_SCRIPTS_DIR, 'tool', `lsp_server.py`);
 export const DEBUG_SERVER_SCRIPT_PATH = path.join(BUNDLED_PYTHON_SCRIPTS_DIR, 'tool', `_debug_server.py`);
+
+
+// Global state
+export const DEPENDENCIES_INSTALLED = 'kedro.dependenciesInstalled'
+export const TELEMETRY_CONSENT = 'kedro.telemetryConsent';
+export const PROJECT_METADATA = 'kedro.projectMetadata';
diff --git a/src/common/telemetry.ts b/src/common/telemetry.ts
@@ -0,0 +1,48 @@
+import axios from 'axios';
+
+// from kedro_telemetry.plugin
+interface HeapData {
+    app_id: string;
+    event: string;
+    timestamp: string;
+    properties: any;
+    identity?: string;
+}
+
+const HEAP_APPID_PROD = '4039408868'; // todo: Dev server, change it back to prod
+const HEAP_ENDPOINT = 'https://heapanalytics.com/api/track';
+const HEAP_HEADERS = { 'Content-Type': 'application/json' };
+
+export async function sendHeapEvent(commandName: string, properties?: any, identity?: string): Promise<void> {
+    var telemetryData;
+    const eventName = "Kedro VSCode Command";
+
+    if (properties) {
+        telemetryData = { ...properties };
+        telemetryData["command_name"] = commandName;
+    }
+
+
+    const data: HeapData = {
+        app_id: HEAP_APPID_PROD,
+        event: eventName,
+        timestamp: new Date().toISOString(),
+        properties: telemetryData || {},
+    };
+
+    if (identity) {
+        data.identity = identity;
+    }
+
+    try {
+        const response = await axios.post(HEAP_ENDPOINT, data, {
+            headers: HEAP_HEADERS,
+            timeout: 10000, // 10 seconds
+        });
+
+        // Handle the response if needed
+        console.log('Heap event sent successfully:', response.status);
+    } catch (error) {
+        console.error('Error sending Heap event:', error);
+    }
+}
diff --git a/src/common/utilities.ts b/src/common/utilities.ts
@@ -9,7 +9,7 @@ import { LogLevel, Uri, WorkspaceFolder } from 'vscode';
 import { Trace } from 'vscode-jsonrpc/node';
 import { getWorkspaceFolders } from './vscodeapi';
 import { callPythonScript } from './callPythonScript';
-import { EXTENSION_ROOT_DIR } from './constants';
+import { DEPENDENCIES_INSTALLED, EXTENSION_ROOT_DIR, PROJECT_METADATA, TELEMETRY_CONSENT } from './constants';
 import { traceError, traceLog } from './log/logging';
 import { executeGetProjectDataCommand } from './commands';
 import KedroVizPanel from '../webview/vizWebView';
@@ -74,27 +74,77 @@ export async function getProjectRoot(): Promise<WorkspaceFolder> {
 }
 
 export async function installDependenciesIfNeeded(context: vscode.ExtensionContext): Promise<void> {
-    const alreadyInstalled = context.globalState.get('dependenciesInstalled', false);
+    // Install necessary dependencies for the flowcharts and telemetry
+    const alreadyInstalled = context.globalState.get(DEPENDENCIES_INSTALLED, false);
 
     if (!alreadyInstalled) {
-        const pathToScript = 'bundled/tool/install_dependencies.py';
+        const vizPathToScript = 'bundled/tool/install_viz_dependencies.py';
+        const telemetryPathToScript = 'bundled/tool/install_telemetry_dependencies.py';
         try {
-            const stdout = await callPythonScript(pathToScript, EXTENSION_ROOT_DIR, context);
-
+            const stdoutViz = await callPythonScript(vizPathToScript, EXTENSION_ROOT_DIR, context);
+            const stdoutTelemetry = await callPythonScript(telemetryPathToScript, EXTENSION_ROOT_DIR, context);
             // Check if the script output contains the success message
-            if (stdout.includes('Successfully installed')) {
-                context.globalState.update('dependenciesInstalled', true);
-                traceLog(`Python dependencies installed!`);
-                console.log('Python dependencies installed!');
+            if (stdoutViz.includes('Successfully installed')) {
+                traceLog(`Kedro-viz dependencies installed!`);
+                console.log('Kedro-viz dependencies installed!');
+            }
+            if (stdoutTelemetry.includes('Successfully installed')) {
+                traceLog(`kedro-telemetry dependencies installed!`);
+                console.log('kedro-telemetry dependencies installed!');
             }
+            context.globalState.update(DEPENDENCIES_INSTALLED, true);
         } catch (error) {
             traceError(`Failed to install Python dependencies:: ${error}`);
             console.error(`Failed to install Python dependencies:: ${error}`);
         }
     }
 }
 
+
+export async function checkKedroProjectConsent(context: vscode.ExtensionContext): Promise<Boolean> {
+    const pathToScript = 'bundled/tool/check_consent.py';
+    try {
+        const stdout = await callPythonScript(pathToScript, EXTENSION_ROOT_DIR, context);
+        const telemetryResult = parseTelemetryConsent(stdout);
+
+        // Check if the script output contains the success message
+        if (telemetryResult) {
+            const consent = telemetryResult['consent'];
+            context.globalState.update(PROJECT_METADATA, telemetryResult);
+            delete telemetryResult['consent'];
+
+            context.globalState.update(TELEMETRY_CONSENT, consent);
+            console.log(`Consent from Kedro Project: ${consent}`);
+            return consent;
+        }
+        return false;
+    } catch (error) {
+        traceError(`Failed to check for telemetry consent:: ${error}`);
+    }
+    return false;
+}
+
+function parseTelemetryConsent(logMessage: string): Record<string, any> | null {
+    // Step 1: Define a regular expression to match the telemetry consent data
+    const telemetryRegex = /telemetry consent: ({.*})/;
+    const match = logMessage.match(telemetryRegex);
+
+    if (match && match[1]) {
+        try {
+            const telemetryData = JSON.parse(match[1]);
+            return telemetryData;
+        } catch (error) {
+            console.error('Failed to parse telemetry consent data:', error);
+            return null;
+        }
+    } else {
+        console.log('Telemetry consent data not found in log message.');
+        return null;
+    }
+}
+
 export async function updateKedroVizPanel(lsClient: LanguageClient | undefined): Promise<void> {
     const projectData = await executeGetProjectDataCommand(lsClient);
     KedroVizPanel.currentPanel?.updateData(projectData);
+
 }