From 9b7b84b183f9959e8b8db23d0d9d932c0f3f4eea Mon Sep 17 00:00:00 2001
From: bwebs <bfweber@gmail.com>
Date: Tue, 20 May 2025 00:40:20 -0700
Subject: [PATCH] Add embed observability feature with API and HTML container

- Introduced `main.py` to define a Typer command group for embed observability.
- Created `constants.py` for session length and permissions.
- Implemented `create_sso_embed_url.py` for generating SSO embed URLs.
- Added `embed_container.html` for the frontend iframe container.
- Developed `embed_server.py` to handle API requests for SSO URL creation and event logging.
---
 lkr/embed/main.py                             | 31 +++++++
 lkr/embed/observability/constants.py          |  8 ++
 .../observability/create_sso_embed_url.py     | 53 +++++++++++
 lkr/embed/observability/embed_container.html  | 88 +++++++++++++++++++
 lkr/embed/observability/embed_server.py       | 86 ++++++++++++++++++
 5 files changed, 266 insertions(+)
 create mode 100644 lkr/embed/main.py
 create mode 100644 lkr/embed/observability/constants.py
 create mode 100644 lkr/embed/observability/create_sso_embed_url.py
 create mode 100644 lkr/embed/observability/embed_container.html
 create mode 100644 lkr/embed/observability/embed_server.py

diff --git a/lkr/embed/main.py b/lkr/embed/main.py
new file mode 100644
index 0000000..bc69d05
--- /dev/null
+++ b/lkr/embed/main.py
@@ -0,0 +1,31 @@
+from typing import Annotated
+
+import typer
+
+from lkr.auth_service import get_auth
+from lkr.embed.observability.embed_server import run_server
+
+__all__ = ["group"]
+
+group = typer.Typer(name="embed", help="Embed commands for LookML Repository")
+
+
+@group.command()
+def observability(
+    ctx: typer.Context,
+    port: Annotated[
+        int, typer.Option("--port", help="Port to run the API on", envvar="PORT")
+    ] = 8080,
+    log_event_prefix: Annotated[
+        str, typer.Option("--log-event-prefix", help="Prefix to use for the log events")
+    ] = "lkr:embed:observability",
+):
+    """
+    Spin up an API to do embed observability. **Important:** requires a user with the `admin` role. This will create
+    an endpoint where you can send embed user properties to it, it will create the sso embed url, then load it up into a selenium
+    browser. The browser will log the user in using the sso embed url, then it will navigate to the dashboard and start the observability tests. It will log structured payloads
+    with metadata in each step. You can then ingest these using tools your of choice
+    """
+    sdk = get_auth(ctx).get_current_sdk()
+
+    run_server(sdk=sdk, port=port, log_event_prefix=log_event_prefix)
diff --git a/lkr/embed/observability/constants.py b/lkr/embed/observability/constants.py
new file mode 100644
index 0000000..3e50dbc
--- /dev/null
+++ b/lkr/embed/observability/constants.py
@@ -0,0 +1,8 @@
+
+MAX_SESSION_LENGTH = 2592000
+
+PERMISSIONS = [
+    "access_data",
+    "see_user_dashboards",
+    "see_lookml_dashboards"
+]
diff --git a/lkr/embed/observability/create_sso_embed_url.py b/lkr/embed/observability/create_sso_embed_url.py
new file mode 100644
index 0000000..b145f6b
--- /dev/null
+++ b/lkr/embed/observability/create_sso_embed_url.py
@@ -0,0 +1,53 @@
+from typing import Any, Dict, List, Optional, TypedDict
+from uuid import uuid4
+
+from looker_sdk.sdk.api40.methods import Looker40SDK
+from looker_sdk.sdk.api40.models import EmbedSsoParams
+from pydantic import BaseModel, Field
+
+from lkr.embed.observability.constants import MAX_SESSION_LENGTH, PERMISSIONS
+
+
+class CreateSSOEmbedUrlParams(BaseModel):
+    external_user_id: Optional[str] = Field(description="The external user id to create the sso embed url for", default_factory=lambda: f"embed-observability-{str(uuid4())}")
+    external_group_id: Optional[str] = None
+    models: Optional[List[str]] = Field(description="The models to create the sso embed url for")
+    permissions: Optional[List[str]] = Field(description="The permissions to create the sso embed url for")
+    dashboard: Optional[str] = Field(description="The dashboard to create the sso embed url for")
+    user_attribute: Optional[List[str]] = Field(description="The user attributes to create the sso embed url for")
+    user_timezone: Optional[str] = Field(description="The user timezone to create the sso embed url for")
+    group_ids: Optional[List[str]] = Field(description="The group ids to create the sso embed url for")
+    embed_domain: Optional[str] = Field(description="The embed domain to create the sso embed url for")
+
+    def to_embed_sso_params(self) -> EmbedSsoParams:
+        all_permissions = list(set(PERMISSIONS + (self.permissions or [])))
+        return EmbedSsoParams(
+            external_user_id=self.external_user_id,
+            external_group_id=self.external_group_id,
+            models=self.models,
+            permissions=all_permissions,
+            target_url=f"/embed/dashboards/{self.dashboard}",
+        session_length=MAX_SESSION_LENGTH,
+        force_logout_login=True,
+        first_name=None,
+        last_name=None,
+        user_timezone=None,
+        group_ids=None,
+        user_attributes=None,
+        embed_domain=None,
+
+        )
+class URLResponse(TypedDict):
+    url: str
+    external_user_id: str
+    
+def create_sso_embed_url(sdk: Looker40SDK, *, data: Dict[str, Any]) -> URLResponse:
+    params = CreateSSOEmbedUrlParams(
+        external_user_id=data.get("external_user_id"),
+        external_group_id=data.get("external_group_id"),
+        models=data.get("models"),
+        permissions=data.get("permissions"),
+        dashboard=data.get("dashboard"),
+    )
+    sso_url = sdk.create_sso_embed_url(body=params.to_embed_sso_params())
+    return dict(url=sso_url.url, external_user_id=sso_url.external_user_id)
\ No newline at end of file
diff --git a/lkr/embed/observability/embed_container.html b/lkr/embed/observability/embed_container.html
new file mode 100644
index 0000000..4678e7a
--- /dev/null
+++ b/lkr/embed/observability/embed_container.html
@@ -0,0 +1,88 @@
+<!DOCTYPE html>
+<html>
+<head>
+    <title>Looker Embed Container</title>
+    <style>
+        body, html {
+            margin: 0;
+            padding: 0;
+            width: 100%;
+            height: 100%;
+            overflow: hidden;
+        }
+        iframe {
+            width: 100%;
+            height: 100%;
+            border: none;
+        }
+        #completion-indicator {
+            display: none;
+            position: absolute;
+            top: 0;
+            left: 0;
+            width: 1px;
+            height: 1px;
+        }
+    </style>
+</head>
+<body>
+    <iframe id="looker-iframe" src=""></iframe>
+    <script>
+        // Get the iframe URL from query parameters
+        const urlParams = new URLSearchParams(window.location.search);
+        const iframeUrl = urlParams.get('iframe_url');
+        const origin = new URL(iframeUrl).origin;
+        const dashboard = urlParams.get('dashboard_id');
+        const user_id = urlParams.get('user_id');
+        const task_id = urlParams.get('task_id');
+        
+        // Set the iframe source
+        document.getElementById('looker-iframe').src = iframeUrl;
+
+        // Track which events we've received
+        
+        const trackedEvents = new Set(['dashboard:loaded', 'dashboard:run:complete', 'dashboard:tile:complete', 'dashboard:run:start', 'dashboard:tile:start']);
+
+        // Listen for Looker embed events
+        window.addEventListener('message', function(event) {
+            if (event.origin !== origin) {
+                return;
+            }
+            const {type, ...data} = JSON.parse(event.data)
+            if (!trackedEvents.has(type)) {
+                return;
+            }
+            if (type) {
+                const now = new Date();
+                const eventData = {
+                    event_type: type,
+                    event_data: data,
+                    timestamp: now.toISOString(),
+                    duration_ms: now.getTime() - window.taskStart.getTime(),
+                    dashboard_id: dashboard,
+                    user_id: user_id,
+                    task_id: task_id
+                };
+                // Send event data to the server
+                fetch('/log_event', {
+                    method: 'POST',
+                    headers: {
+                        'Content-Type': 'application/json',
+                    },
+                    body: JSON.stringify(eventData)
+                });
+
+                // Track received events
+                if (type === 'dashboard:run:complete') {
+                    complete = document.createElement("div");
+                    complete.id = "completion-indicator";
+                    document.body.appendChild(complete);
+                }
+            }
+        });
+
+        // Store task start time
+        window.taskStart = new Date();
+    </script>
+</body>
+</html> 
\ No newline at end of file
diff --git a/lkr/embed/observability/embed_server.py b/lkr/embed/observability/embed_server.py
new file mode 100644
index 0000000..981dcba
--- /dev/null
+++ b/lkr/embed/observability/embed_server.py
@@ -0,0 +1,86 @@
+import json
+from http.server import BaseHTTPRequestHandler, HTTPServer
+from pathlib import Path
+
+from looker_sdk.sdk.api40.methods import Looker40SDK
+
+from lkr.embed.observability.create_sso_embed_url import create_sso_embed_url
+from lkr.logging import logger, structured_logger
+
+
+def log_event(prefix: str, event: str, **kwargs):
+    logger.info(f"{prefix}:{event}", **kwargs)
+
+class EmbedHandler(BaseHTTPRequestHandler):
+    def __init__(self, *args, **kwargs):
+        self.sdk: Looker40SDK | None = None
+        super().__init__(*args, **kwargs)
+
+    def log_message(self, format, *args):
+        # Override to disable default server logging
+        pass
+
+    def do_GET(self):
+        path, *rest = self.path.split('?')
+        if path =='/':
+        # Serve the HTML file
+            self.send_response(200)
+            self.send_header('Content-type', 'text/html')
+            self.end_headers()
+            
+            html_path = Path(__file__).parent / 'embed_container.html'
+            with open(html_path, 'rb') as f:
+                self.wfile.write(f.read())
+        else:
+            self.send_response(204)
+            self.end_headers()
+
+
+    def do_POST(self):
+        path, *rest = self.path.split('?')
+        if path == '/create_sso_embed_url':
+            if not self.sdk:
+                self.send_response(500)
+                self.end_headers()
+                return
+            else:
+                content_length = int(self.headers['Content-Length'])
+                post_data = self.rfile.read(content_length)
+                data = json.loads(post_data)
+
+                sso_url = create_sso_embed_url(self.sdk, data=data)
+                self.send_response(200)
+                self.send_header('Content-type', 'application/json')
+                self.end_headers()
+                self.wfile.write(json.dumps(sso_url).encode('utf-8'))
+        if path == '/log_event':
+            content_length = int(self.headers['Content-Length'])
+            post_data = self.rfile.read(content_length)
+            event_data = json.loads(post_data)
+            
+            # # Log the event using structlog
+            # structured_logger.info(
+            #     event_data['event_type'],
+            #     timestamp=event_data['timestamp'],
+            #     duration_ms=event_data['duration_ms'],
+            #     dashboard_id=event_data['dashboard_id'],
+            #     user_id=event_data['user_id'],
+            #     **event_data['event_data'],
+            # )
+            
+            self.send_response(200)
+            self.end_headers()
+        else:
+            self.send_response(404)
+            self.end_headers()
+
+def run_server(*, sdk: Looker40SDK, port:int =3000, log_event_prefix="looker_embed_observability"):
+    class EmbedHandlerWithPrefix(EmbedHandler):
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, **kwargs)
+            self.log_event_prefix = log_event_prefix
+            self.sdk = sdk
+    server_address = ('', port)
+    httpd = HTTPServer(server_address, EmbedHandlerWithPrefix)
+    structured_logger.info(f"{log_event_prefix}:embed_server_started", port=port, embed_domain=f"http://localhost:{port}")
+    httpd.serve_forever()
\ No newline at end of file