Fix for telemetry in sharaeable viz (#1551)

* fix telemetry for sharaeable viz Signed-off-by: ravi-kumar-pilla <[email protected]> * fix lint issue Signed-off-by: ravi-kumar-pilla <[email protected]> --------- Signed-off-by: ravi-kumar-pilla <[email protected]> Signed-off-by: Tynan DeBold <[email protected]> Co-authored-by: Tynan DeBold <[email protected]>
kedro-org · Oct 4, 2023 · ded96ff · ded96ff
1 parent 6ab25be
commit ded96ff
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 3 deletions.
diff --git a/package/kedro_viz/integrations/deployment/s3_deployer.py b/package/kedro_viz/integrations/deployment/s3_deployer.py
@@ -3,14 +3,17 @@
 
 import json
 import logging
+import tempfile
 from datetime import datetime
 from pathlib import Path
 
 import fsspec
+from jinja2 import Environment, FileSystemLoader
 from semver import VersionInfo
 
 from kedro_viz import __version__
 from kedro_viz.api.rest.responses import save_api_responses_to_fs
+from kedro_viz.integrations.kedro import telemetry as kedro_telemetry
 
 _HTML_DIR = Path(__file__).parent.parent.parent.absolute() / "html"
 _METADATA_PATH = "api/deploy-viz-metadata"
@@ -48,11 +51,41 @@ def _upload_api_responses(self):
         """Upload API responses to S3."""
         save_api_responses_to_fs(self._bucket_path)
 
+    def _ingest_heap_analytics(self):
+        """Ingest heap analytics to index file in the build folder."""
+        project_path = Path.cwd().absolute()
+        heap_app_id = kedro_telemetry.get_heap_app_id(project_path)
+        heap_user_identity = kedro_telemetry.get_heap_identity()
+        should_add_telemetry = bool(heap_app_id) and bool(heap_user_identity)
+        html_content = (_HTML_DIR / "index.html").read_text(encoding="utf-8")
+        injected_head_content = []
+
+        env = Environment(loader=FileSystemLoader(_HTML_DIR))
+
+        if should_add_telemetry:
+            logger.debug("Ingesting heap analytics.")
+            telemetry_content = env.get_template("telemetry.html").render(
+                heap_app_id=heap_app_id, heap_user_identity=heap_user_identity
+            )
+            injected_head_content.append(telemetry_content)
+
+        injected_head_content.append("</head>")
+        html_content = html_content.replace("</head>", "\n".join(injected_head_content))
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            temp_file_path = f"{temp_dir}/index.html"
+
+            with open(temp_file_path, "w", encoding="utf-8") as temp_index_file:
+                temp_index_file.write(html_content)
+
+            self._remote_fs.put(temp_file_path, f"{self._bucket_path}/")
+
     def _upload_static_files(self, html_dir: Path):
         """Upload static HTML files to S3."""
         logger.debug("Uploading static html files to %s.", self._bucket_path)
         try:
             self._remote_fs.put(f"{str(html_dir)}/*", self._bucket_path, recursive=True)
+            self._ingest_heap_analytics()
         except Exception as exc:  # pragma: no cover
             logger.exception("Upload failed: %s ", exc)
             raise exc

diff --git a/package/tests/test_integrations/test_s3_deployer.py b/package/tests/test_integrations/test_s3_deployer.py
@@ -28,11 +28,13 @@ def test_upload_api_responses(self, mocker, region, bucket_name):
 
     def test_upload_static_files(self, mocker, region, bucket_name):
         mocker.patch("fsspec.filesystem")
+        mocker.patch("kedro_viz.integrations.kedro.telemetry.get_heap_app_id")
+        mocker.patch("kedro_viz.integrations.kedro.telemetry.get_heap_identity")
+
         deployer = S3Deployer(region, bucket_name)
         deployer._upload_static_files(_HTML_DIR)
-        deployer._remote_fs.put.assert_called_once_with(
-            f"{str(_HTML_DIR)}/*", deployer._bucket_path, recursive=True
-        )
+
+        assert deployer._remote_fs.put.call_count == 2
 
     def test_upload_static_file_failed(self, mocker, region, bucket_name, caplog):
         mocker.patch("fsspec.filesystem")