From 09b29affdaab9456b1d148d5e2e144819272a378 Mon Sep 17 00:00:00 2001
From: vladsaveliev <vladislav.savelyev@populationgenomics.org.au>
Date: Thu, 2 May 2024 18:11:38 +0200
Subject: [PATCH] Add is_notebook and interactive_function_name visit field

---
 app/db.py     |  4 +++-
 app/main.py   | 14 ++++++++++++++
 app/models.py |  4 ++++
 3 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/app/db.py b/app/db.py
index 42b6ff2..4d734dd 100644
--- a/app/db.py
+++ b/app/db.py
@@ -1,4 +1,5 @@
 """Functions to interact with the database."""
+
 from typing import Optional, Sequence
 
 import os
@@ -34,6 +35,8 @@ class VisitStats(SQLModel, table=True):
     is_singularity: bool
     is_conda: bool
     is_ci: bool
+    is_notebook: bool
+    interactive_function_name: Optional[str]
     count: int
 
 
@@ -129,4 +132,3 @@ def insert_download_stats(df: pd.DataFrame) -> pd.DataFrame:
                 session.add(new_entry)
         session.commit()
     return df
-
diff --git a/app/main.py b/app/main.py
index 2df1cab..bf3964b 100644
--- a/app/main.py
+++ b/app/main.py
@@ -89,6 +89,8 @@ def update_version():
     "is_singularity",
     "is_conda",
     "is_ci",
+    "is_notebook",
+    "interactive_function_name",
 ]
 
 # Thread-safe in-memory buffer to accumulate recent visits before writing to the CSV file
@@ -106,6 +108,8 @@ async def version(
     is_singularity: str = "",
     is_conda: str = "",
     is_ci: str = "",
+    is_notebook: str = "",
+    interactive_function_name: str = "",
 ):
     """
     Endpoint for MultiQC that returns the latest release, and logs
@@ -121,6 +125,8 @@ async def version(
         is_singularity=is_singularity,
         is_conda=is_conda,
         is_ci=is_ci,
+        is_notebook=is_notebook,
+        interactive_function_name=interactive_function_name,
     )
     return models.VersionResponse(latest_release=app.latest_release)
 
@@ -149,6 +155,8 @@ def _log_visit(
     is_singularity: str = "",
     is_conda: str = "",
     is_ci: str = "",
+    is_notebook: str = "",
+    interactive_function_name: str = "",
 ):
     global visit_buffer
     with visit_buffer_lock:
@@ -162,6 +170,8 @@ def _log_visit(
                 "is_singularity": is_singularity,
                 "is_conda": is_conda,
                 "is_ci": is_ci,
+                "is_notebook": is_notebook,
+                "interactive_function_name": interactive_function_name,
             }
         )
         logger.debug(f"Logged visit, total visits in buffer: {len(visit_buffer)}")
@@ -252,6 +262,8 @@ def strtobool(val) -> bool:
         df["is_singularity"] = df["is_singularity"].apply(strtobool)
         df["is_conda"] = df["is_conda"].apply(strtobool)
         df["is_ci"] = df["is_ci"].apply(strtobool)
+        df["is_notebook"] = df["is_notebook"].apply(strtobool)
+        df["interactive_function_name"] = df["interactive_function_name"].astype(str)
         df = df.drop(columns=["timestamp"])
 
         # Summarize visits per user per time interval
@@ -409,6 +421,8 @@ async def version_legacy(background_tasks: BackgroundTasks, v: str = ""):
         is_singularity="",
         is_conda="",
         is_ci="",
+        is_notebook="",
+        interactive_function_name="",
     )
     return app.latest_release.version
 
diff --git a/app/models.py b/app/models.py
index b9d828f..72ea4c7 100644
--- a/app/models.py
+++ b/app/models.py
@@ -66,6 +66,8 @@ class UsageCategory(str, Enum):
     is_singularity = "is_singularity"
     is_conda = "is_conda"
     is_ci = "is_ci"
+    is_notebook = "is_notebook"
+    interactive_function_name = "interactive_function_name"
 
 
 usage_category_nice_names = dict(
@@ -78,4 +80,6 @@ class UsageCategory(str, Enum):
     is_singularity="Singularity",
     is_conda="Conda",
     is_ci="CI environment",
+    is_notebook="Notebook",
+    interactive_function_name="Interactive function",
 )