feat: instrument vscode, jupyter and 3p plugin usage (#115)

shobsi · web-flow · commit 0dd163f3e5a2 · 2025-05-16T06:20:37.000Z
* feat: instrument vscode, jupyter and 3p plugin usage

* add more test coverage

* improve code coverage
diff --git a/bigquery_magics/bigquery.py b/bigquery_magics/bigquery.py
@@ -126,6 +126,7 @@
 from google.cloud.bigquery.job import QueryJobConfig
 import pandas
 
+from bigquery_magics import environment
 from bigquery_magics import line_arg_parser as lap
 import bigquery_magics._versions_helpers
 import bigquery_magics.config
@@ -143,10 +144,27 @@
 except ImportError:
     bpd = None
 
-USER_AGENT = f"ipython-{IPython.__version__} bigquery-magics/{bigquery_magics.version.__version__}"
 context = bigquery_magics.config.context
 
 
+def _get_user_agent():
+    identities = [
+        f"ipython-{IPython.__version__}",
+        f"bigquery-magics/{bigquery_magics.version.__version__}",
+    ]
+
+    if environment.is_vscode():
+        identities.append("vscode")
+        if environment.is_vscode_google_cloud_code_extension_installed():
+            identities.append(environment.GOOGLE_CLOUD_CODE_EXTENSION_NAME)
+    elif environment.is_jupyter():
+        identities.append("jupyter")
+        if environment.is_jupyter_bigquery_plugin_installed():
+            identities.append(environment.BIGQUERY_JUPYTER_PLUGIN_NAME)
+
+    return " ".join(identities)
+
+
 def _handle_error(error, destination_var=None):
     """Process a query execution error.
 
@@ -558,7 +576,7 @@ def _create_clients(args: Any) -> Tuple[bigquery.Client, Any]:
         project=args.project or context.project,
         credentials=context.credentials,
         default_query_job_config=context.default_query_job_config,
-        client_info=client_info.ClientInfo(user_agent=USER_AGENT),
+        client_info=client_info.ClientInfo(user_agent=_get_user_agent()),
         client_options=bigquery_client_options,
         location=args.location,
     )
@@ -885,7 +903,7 @@ def _make_bqstorage_client(client, client_options):
 
     return client._ensure_bqstorage_client(
         client_options=client_options,
-        client_info=gapic_client_info.ClientInfo(user_agent=USER_AGENT),
+        client_info=gapic_client_info.ClientInfo(user_agent=_get_user_agent()),
     )
 
 
diff --git a/bigquery_magics/environment.py b/bigquery_magics/environment.py
@@ -0,0 +1,105 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import importlib
+import json
+import os
+import pathlib
+
+Path = pathlib.Path
+
+
+# The identifier for GCP VS Code extension
+# https://cloud.google.com/code/docs/vscode/install
+GOOGLE_CLOUD_CODE_EXTENSION_NAME = "googlecloudtools.cloudcode"
+
+
+# The identifier for BigQuery Jupyter notebook plugin
+# https://cloud.google.com/bigquery/docs/jupyterlab-plugin
+BIGQUERY_JUPYTER_PLUGIN_NAME = "bigquery_jupyter_plugin"
+
+
+def _is_vscode_extension_installed(extension_id: str) -> bool:
+    """
+    Checks if a given Visual Studio Code extension is installed.
+
+    Args:
+        extension_id: The ID of the extension (e.g., "ms-python.python").
+
+    Returns:
+        True if the extension is installed, False otherwise.
+    """
+    try:
+        # Determine the user's VS Code extensions directory.
+        user_home = Path.home()
+        vscode_extensions_dir = user_home / ".vscode" / "extensions"
+
+        # Check if the extensions directory exists.
+        if not vscode_extensions_dir.exists():
+            return False
+
+        # Iterate through the subdirectories in the extensions directory.
+        extension_dirs = filter(
+            lambda p: p.is_dir() and p.name.startswith(extension_id + "-"),
+            vscode_extensions_dir.iterdir(),
+        )
+        for extension_dir in extension_dirs:
+            # As a more robust check, the manifest file must exist.
+            manifest_path = extension_dir / "package.json"
+            if not manifest_path.exists() or not manifest_path.is_file():
+                continue
+
+            # Finally, the manifest file must be a valid json
+            with open(manifest_path, "r", encoding="utf-8") as f:
+                json.load(f)
+
+            return True
+    except Exception:
+        pass
+
+    return False
+
+
+def _is_package_installed(package_name: str) -> bool:
+    """
+    Checks if a Python package is installed.
+
+    Args:
+        package_name: The name of the package to check (e.g., "requests", "numpy").
+
+    Returns:
+        True if the package is installed, False otherwise.
+    """
+    try:
+        importlib.import_module(package_name)
+        return True
+    except Exception:
+        return False
+
+
+def is_vscode() -> bool:
+    return os.getenv("VSCODE_PID") is not None
+
+
+def is_jupyter() -> bool:
+    return os.getenv("JPY_PARENT_PID") is not None
+
+
+def is_vscode_google_cloud_code_extension_installed() -> bool:
+    return _is_vscode_extension_installed(GOOGLE_CLOUD_CODE_EXTENSION_NAME)
+
+
+def is_jupyter_bigquery_plugin_installed() -> bool:
+    return _is_package_installed(BIGQUERY_JUPYTER_PLUGIN_NAME)
diff --git a/tests/unit/bigquery/test_bigquery.py b/tests/unit/bigquery/test_bigquery.py
@@ -16,8 +16,11 @@
 import contextlib
 import copy
 import json
+import os
+import pathlib
 import re
 import sys
+import tempfile
 from unittest import mock
 import warnings
 
@@ -59,6 +62,8 @@
 except ImportError:
     gpd = None
 
+Path = pathlib.Path
+
 
 def make_connection(*args):
     # TODO(tswast): Remove this in favor of a mock google.cloud.bigquery.Client
@@ -929,6 +934,180 @@ def test_bigquery_magic_default_connection_user_agent():
     )
 
 
+def test_bigquery_magic_default_connection_user_agent_vscode():
+    globalipapp.start_ipython()
+    ip = globalipapp.get_ipython()
+    ip.extension_manager.load_extension("bigquery_magics")
+    bigquery_magics.context._connection = None
+
+    credentials_mock = mock.create_autospec(
+        google.auth.credentials.Credentials, instance=True
+    )
+    default_patch = mock.patch(
+        "google.auth.default", return_value=(credentials_mock, "general-project")
+    )
+    run_query_patch = mock.patch("bigquery_magics.bigquery._run_query", autospec=True)
+    conn_patch = mock.patch("google.cloud.bigquery.client.Connection", autospec=True)
+    env_patch = mock.patch.dict(os.environ, {"VSCODE_PID": "1234"}, clear=True)
+
+    with conn_patch as conn, run_query_patch, default_patch, env_patch:
+        ip.run_cell_magic("bigquery", "", "SELECT 17 as num")
+
+    client_info_arg = conn.call_args[1].get("client_info")
+    assert client_info_arg is not None
+    assert (
+        client_info_arg.user_agent
+        == f"ipython-{IPython.__version__} bigquery-magics/{bigquery_magics.__version__} vscode"
+    )
+
+
+@pytest.mark.parametrize(
+    (
+        "install_dir_exists",
+        "manifest_exists",
+        "manifest_valid",
+        "expect_extension_user_agent",
+    ),
+    [
+        pytest.param(False, False, False, False, id="no-install"),
+        pytest.param(True, False, False, False, id="no-manifest"),
+        pytest.param(True, True, False, False, id="invalid-manifest"),
+        pytest.param(True, True, True, True, id="good-install"),
+    ],
+)
+def test_bigquery_magic_default_connection_user_agent_vscode_extension(
+    install_dir_exists, manifest_exists, manifest_valid, expect_extension_user_agent
+):
+    globalipapp.start_ipython()
+    ip = globalipapp.get_ipython()
+    ip.extension_manager.load_extension("bigquery_magics")
+    bigquery_magics.context._connection = None
+
+    credentials_mock = mock.create_autospec(
+        google.auth.credentials.Credentials, instance=True
+    )
+    default_patch = mock.patch(
+        "google.auth.default", return_value=(credentials_mock, "general-project")
+    )
+    run_query_patch = mock.patch("bigquery_magics.bigquery._run_query", autospec=True)
+    conn_patch = mock.patch("google.cloud.bigquery.client.Connection", autospec=True)
+    env_patch = mock.patch.dict(os.environ, {"VSCODE_PID": "1234"}, clear=True)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        user_home = Path(tmpdir)
+        extension_dir = (
+            user_home / ".vscode" / "extensions" / "googlecloudtools.cloudcode-0.12"
+        )
+        extension_config = extension_dir / "package.json"
+
+        # originally extension config does not exist
+        assert not extension_config.exists()
+
+        # simulate extension installation by creating extension config on disk
+        if install_dir_exists:
+            extension_dir.mkdir(parents=True)
+
+        if manifest_exists:
+            if manifest_valid:
+                with open(extension_config, "w") as f:
+                    f.write("{}")
+            else:
+                extension_config.touch()
+
+        home_dir_patch = mock.patch("pathlib.Path.home", return_value=user_home)
+
+        with conn_patch as conn, (
+            run_query_patch
+        ), default_patch, env_patch, home_dir_patch:
+            ip.run_cell_magic("bigquery", "", "SELECT 17 as num")
+
+    expected_user_agents = [
+        f"ipython-{IPython.__version__}",
+        f"bigquery-magics/{bigquery_magics.__version__}",
+        "vscode",
+    ]
+    if expect_extension_user_agent:
+        expected_user_agents.append("googlecloudtools.cloudcode")
+    expected_user_agent = " ".join(expected_user_agents)
+
+    client_info_arg = conn.call_args[1].get("client_info")
+    assert client_info_arg is not None
+    assert client_info_arg.user_agent == expected_user_agent
+
+
+def test_bigquery_magic_default_connection_user_agent_jupyter():
+    globalipapp.start_ipython()
+    ip = globalipapp.get_ipython()
+    ip.extension_manager.load_extension("bigquery_magics")
+    bigquery_magics.context._connection = None
+
+    credentials_mock = mock.create_autospec(
+        google.auth.credentials.Credentials, instance=True
+    )
+    default_patch = mock.patch(
+        "google.auth.default", return_value=(credentials_mock, "general-project")
+    )
+    run_query_patch = mock.patch("bigquery_magics.bigquery._run_query", autospec=True)
+    conn_patch = mock.patch("google.cloud.bigquery.client.Connection", autospec=True)
+    env_patch = mock.patch.dict(os.environ, {"JPY_PARENT_PID": "1234"}, clear=True)
+
+    with conn_patch as conn, run_query_patch, default_patch, env_patch:
+        ip.run_cell_magic("bigquery", "", "SELECT 17 as num")
+
+    client_info_arg = conn.call_args[1].get("client_info")
+    assert client_info_arg is not None
+    assert (
+        client_info_arg.user_agent
+        == f"ipython-{IPython.__version__} bigquery-magics/{bigquery_magics.__version__} jupyter"
+    )
+
+
+def test_bigquery_magic_default_connection_user_agent_jupyter_plugin():
+    globalipapp.start_ipython()
+    ip = globalipapp.get_ipython()
+    ip.extension_manager.load_extension("bigquery_magics")
+    bigquery_magics.context._connection = None
+
+    credentials_mock = mock.create_autospec(
+        google.auth.credentials.Credentials, instance=True
+    )
+    default_patch = mock.patch(
+        "google.auth.default", return_value=(credentials_mock, "general-project")
+    )
+    run_query_patch = mock.patch("bigquery_magics.bigquery._run_query", autospec=True)
+    conn_patch = mock.patch("google.cloud.bigquery.client.Connection", autospec=True)
+    env_patch = mock.patch.dict(os.environ, {"JPY_PARENT_PID": "1234"}, clear=True)
+
+    def custom_import_module_side_effect(name, package=None):
+        if name == "bigquery_jupyter_plugin":
+            return mock.MagicMock()
+        else:
+            import importlib
+
+            return importlib.import_module(name, package)
+
+    assert isinstance(
+        custom_import_module_side_effect("bigquery_jupyter_plugin"), mock.MagicMock
+    )
+    assert custom_import_module_side_effect("bigquery_magics") is bigquery_magics
+
+    extension_import_patch = mock.patch(
+        "importlib.import_module", side_effect=custom_import_module_side_effect
+    )
+
+    with conn_patch as conn, (
+        run_query_patch
+    ), default_patch, env_patch, extension_import_patch:
+        ip.run_cell_magic("bigquery", "", "SELECT 17 as num")
+
+    client_info_arg = conn.call_args[1].get("client_info")
+    assert client_info_arg is not None
+    assert (
+        client_info_arg.user_agent
+        == f"ipython-{IPython.__version__} bigquery-magics/{bigquery_magics.__version__} jupyter bigquery_jupyter_plugin"
+    )
+
+
 def test_bigquery_magic_with_legacy_sql():
     globalipapp.start_ipython()
     ip = globalipapp.get_ipython()
