feat: instrument vscode, jupyter and 3p plugin usage

pandas_gbq/environment.py

@@ -0,0 +1,99 @@
+# 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
+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.
+        for item in vscode_extensions_dir.iterdir():
+            # Ignore non-directories.
+            if not item.is_dir():
+                continue
+
+            # Directory must start with the extension ID.
+            if not item.name.startswith(extension_id + "-"):
+                continue
+
+            # As a more robust check, the manifest file must exist.
+            manifest_path = item / "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)

pandas_gbq/gbq_connector.py

@@ -19,6 +19,7 @@
 
 import pandas_gbq.constants
 from pandas_gbq.contexts import context
+import pandas_gbq.environment as environment
 import pandas_gbq.exceptions
 from pandas_gbq.exceptions import (
     GenericGBQException,
@@ -517,11 +518,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,12 @@
 # Use of this source code is governed by a BSD-style
 # license that can be found in the LICENSE file.
 
+
+import os
+import pathlib
+import tempfile
+import unittest.mock as mock
+
 import google.api_core.exceptions
 import google.cloud.bigquery
 import pandas as pd
@@ -10,6 +16,8 @@
 
 from pandas_gbq import gbq
 
+Path = pathlib.Path
+
 
 class FakeDataFrame:
     """A fake bigframes DataFrame to avoid depending on bigframes."""
@@ -202,3 +210,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
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        user_home = Path(tmpdir)
+        plugin_dir = (
+            user_home / ".vscode" / "extensions" / "googlecloudtools.cloudcode-0.12"
+        )
+        plugin_config = plugin_dir / "package.json"
+
+        # originally pluging config does not exist
+        assert not plugin_config.exists()
+
+        # simulate plugin installation by creating plugin config on disk
+        plugin_dir.mkdir(parents=True)
+        with open(plugin_config, "w") as f:
+            f.write("{}")
+
+        with mock.patch("pathlib.Path.home", return_value=user_home):
+            assert (
+                create_user_agent()
+                == f"pandas-{pd.__version__} vscode googlecloudtools.cloudcode"
+            )
+
+
+@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):
+        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"
+        )