feat: instrument vscode, jupyter and 3p plugin usage

pandas_gbq/environment.py

@@ -0,0 +1,93 @@
+# Copyright (c) 2025 pandas-gbq Authors All rights reserved.
+# Use of this source code is governed by a BSD-style
+# license that can be found in the LICENSE file.
+
+
+import importlib
+import json
+import os
+
+# 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 = os.path.expanduser("~")
+        if os.name == "nt":  # Windows
+            vscode_extensions_dir = os.path.join(user_home, ".vscode", "extensions")
+        elif os.name == "posix":  # macOS and Linux
+            vscode_extensions_dir = os.path.join(user_home, ".vscode", "extensions")
+        else:
+            raise OSError("Unsupported operating system.")
+
+        # Check if the extensions directory exists.
+        if os.path.exists(vscode_extensions_dir):
+            # Iterate through the subdirectories in the extensions directory.
+            for item in os.listdir(vscode_extensions_dir):
+                item_path = os.path.join(vscode_extensions_dir, item)
+                if os.path.isdir(item_path) and item.startswith(extension_id + "-"):
+                    # Check if the folder starts with the extension ID.
+                    # Further check for manifest file, as a more robust check.
+                    manifest_path = os.path.join(item_path, "package.json")
+                    if os.path.exists(manifest_path):
+                        try:
+                            with open(manifest_path, "r", encoding="utf-8") as f:
+                                json.load(f)
+                            return True
+                        except (FileNotFoundError, json.JSONDecodeError):
+                            # Corrupted or incomplete extension, or manifest missing.
+                            pass
+    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)

pandas_gbq/gbq_connector.py

@@ -11,6 +11,8 @@
 
 import numpy as np
 
+from . import environment
+
 # Only import at module-level at type checking time to avoid circular
 # dependencies in the pandas package, which has an optional dependency on
 # pandas-gbq.
@@ -517,11 +519,16 @@ def create_user_agent(
         )
         delimiter = "-"
 
-    identity = f"pandas{delimiter}{pd.__version__}"
+    identities = [] if user_agent is None else [user_agent]
+    identities.append(f"pandas{delimiter}{pd.__version__}")
 
-    if user_agent is None:
-        user_agent = identity
-    else:
-        user_agent = f"{user_agent} {identity}"
+    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 user_agent
+    return " ".join(identities)

tests/unit/test_to_gbq.py

@@ -2,6 +2,10 @@
 # Use of this source code is governed by a BSD-style
 # license that can be found in the LICENSE file.
 
+
+import os
+import unittest.mock as mock
+
 import google.api_core.exceptions
 import google.cloud.bigquery
 import pandas as pd
@@ -202,3 +206,64 @@ def test_create_user_agent(user_agent, rfc9110_delimiter, expected):
 
     result = create_user_agent(user_agent, rfc9110_delimiter)
     assert result == expected
+
+
+@mock.patch.dict(os.environ, {"VSCODE_PID": "1234"}, clear=True)
+def test_create_user_agent_vscode():
+    from pandas_gbq.gbq import create_user_agent
+
+    assert create_user_agent() == f"pandas-{pd.__version__} vscode"
+
+
+@mock.patch.dict(os.environ, {"VSCODE_PID": "1234"}, clear=True)
+def test_create_user_agent_vscode_plugin():
+    from pandas_gbq.gbq import create_user_agent
+
+    # simulate plugin installation by creating plugin config on disk
+    user_home = os.path.expanduser("~")
+    plugin_dir = os.path.join(
+        user_home, ".vscode", "extensions", "googlecloudtools.cloudcode-0.12"
+    )
+    plugin_config = os.path.join(plugin_dir, "package.json")
+    assert not os.path.exists(plugin_config)  # initially does not exist
+    os.makedirs(plugin_dir, exist_ok=True)
+    with open(plugin_config, "w") as f:
+        f.write("{}")
+
+    # test
+    assert (
+        create_user_agent()
+        == f"pandas-{pd.__version__} vscode googlecloudtools.cloudcode"
+    )
+
+    # clean up disk
+    os.remove(plugin_config)
+
+
+@mock.patch.dict(os.environ, {"JPY_PARENT_PID": "1234"}, clear=True)
+def test_create_user_agent_jupyter():
+    from pandas_gbq.gbq import create_user_agent
+
+    assert create_user_agent() == f"pandas-{pd.__version__} jupyter"
+
+
+@mock.patch.dict(os.environ, {"JPY_PARENT_PID": "1234"}, clear=True)
+def test_create_user_agent_jupyter_extension():
+    from pandas_gbq.gbq import create_user_agent
+
+    def custom_import_module_side_effect(name, package=None):
+        print(f"  [Test Mock] Intercepted importlib.import_module for '{name}'")
+        if name == "bigquery_jupyter_plugin":
+            return mock.MagicMock()
+        else:
+            import importlib
+
+            return importlib.import_module(name, package)
+
+    with mock.patch(
+        "importlib.import_module", side_effect=custom_import_module_side_effect
+    ):
+        assert (
+            create_user_agent()
+            == f"pandas-{pd.__version__} jupyter bigquery_jupyter_plugin"
+        )