Skip to content

Commit 7d354f1

Browse files
feat: instrument vscode, jupyter and 3p plugin usage (#925)
* feat: instrument vscode, jupyter and 3p plugin usage * fix the return value * add unit test coverage * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * use pathlib, reduce indentation, format --------- Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
1 parent d525eba commit 7d354f1

File tree

3 files changed

+180
-6
lines changed

3 files changed

+180
-6
lines changed

pandas_gbq/environment.py

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
1+
# Copyright (c) 2025 pandas-gbq Authors All rights reserved.
2+
# Use of this source code is governed by a BSD-style
3+
# license that can be found in the LICENSE file.
4+
5+
6+
import importlib
7+
import json
8+
import os
9+
import pathlib
10+
11+
Path = pathlib.Path
12+
13+
14+
# The identifier for GCP VS Code extension
15+
# https://cloud.google.com/code/docs/vscode/install
16+
GOOGLE_CLOUD_CODE_EXTENSION_NAME = "googlecloudtools.cloudcode"
17+
18+
19+
# The identifier for BigQuery Jupyter notebook plugin
20+
# https://cloud.google.com/bigquery/docs/jupyterlab-plugin
21+
BIGQUERY_JUPYTER_PLUGIN_NAME = "bigquery_jupyter_plugin"
22+
23+
24+
def _is_vscode_extension_installed(extension_id: str) -> bool:
25+
"""
26+
Checks if a given Visual Studio Code extension is installed.
27+
28+
Args:
29+
extension_id: The ID of the extension (e.g., "ms-python.python").
30+
31+
Returns:
32+
True if the extension is installed, False otherwise.
33+
"""
34+
try:
35+
# Determine the user's VS Code extensions directory.
36+
user_home = Path.home()
37+
vscode_extensions_dir = user_home / ".vscode" / "extensions"
38+
39+
# Check if the extensions directory exists.
40+
if not vscode_extensions_dir.exists():
41+
return False
42+
43+
# Iterate through the subdirectories in the extensions directory.
44+
for item in vscode_extensions_dir.iterdir():
45+
# Ignore non-directories.
46+
if not item.is_dir():
47+
continue
48+
49+
# Directory must start with the extension ID.
50+
if not item.name.startswith(extension_id + "-"):
51+
continue
52+
53+
# As a more robust check, the manifest file must exist.
54+
manifest_path = item / "package.json"
55+
if not manifest_path.exists() or not manifest_path.is_file():
56+
continue
57+
58+
# Finally, the manifest file must be a valid json
59+
with open(manifest_path, "r", encoding="utf-8") as f:
60+
json.load(f)
61+
62+
return True
63+
except Exception:
64+
pass
65+
66+
return False
67+
68+
69+
def _is_package_installed(package_name: str) -> bool:
70+
"""
71+
Checks if a Python package is installed.
72+
73+
Args:
74+
package_name: The name of the package to check (e.g., "requests", "numpy").
75+
76+
Returns:
77+
True if the package is installed, False otherwise.
78+
"""
79+
try:
80+
importlib.import_module(package_name)
81+
return True
82+
except Exception:
83+
return False
84+
85+
86+
def is_vscode() -> bool:
87+
return os.getenv("VSCODE_PID") is not None
88+
89+
90+
def is_jupyter() -> bool:
91+
return os.getenv("JPY_PARENT_PID") is not None
92+
93+
94+
def is_vscode_google_cloud_code_extension_installed() -> bool:
95+
return _is_vscode_extension_installed(GOOGLE_CLOUD_CODE_EXTENSION_NAME)
96+
97+
98+
def is_jupyter_bigquery_plugin_installed() -> bool:
99+
return _is_package_installed(BIGQUERY_JUPYTER_PLUGIN_NAME)

pandas_gbq/gbq_connector.py

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919

2020
import pandas_gbq.constants
2121
from pandas_gbq.contexts import context
22+
import pandas_gbq.environment as environment
2223
import pandas_gbq.exceptions
2324
from pandas_gbq.exceptions import (
2425
GenericGBQException,
@@ -517,11 +518,16 @@ def create_user_agent(
517518
)
518519
delimiter = "-"
519520

520-
identity = f"pandas{delimiter}{pd.__version__}"
521+
identities = [] if user_agent is None else [user_agent]
522+
identities.append(f"pandas{delimiter}{pd.__version__}")
521523

522-
if user_agent is None:
523-
user_agent = identity
524-
else:
525-
user_agent = f"{user_agent} {identity}"
524+
if environment.is_vscode():
525+
identities.append("vscode")
526+
if environment.is_vscode_google_cloud_code_extension_installed():
527+
identities.append(environment.GOOGLE_CLOUD_CODE_EXTENSION_NAME)
528+
elif environment.is_jupyter():
529+
identities.append("jupyter")
530+
if environment.is_jupyter_bigquery_plugin_installed():
531+
identities.append(environment.BIGQUERY_JUPYTER_PLUGIN_NAME)
526532

527-
return user_agent
533+
return " ".join(identities)

tests/unit/test_to_gbq.py

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,12 @@
22
# Use of this source code is governed by a BSD-style
33
# license that can be found in the LICENSE file.
44

5+
6+
import os
7+
import pathlib
8+
import tempfile
9+
import unittest.mock as mock
10+
511
import google.api_core.exceptions
612
import google.cloud.bigquery
713
import pandas as pd
@@ -10,6 +16,8 @@
1016

1117
from pandas_gbq import gbq
1218

19+
Path = pathlib.Path
20+
1321

1422
class FakeDataFrame:
1523
"""A fake bigframes DataFrame to avoid depending on bigframes."""
@@ -202,3 +210,64 @@ def test_create_user_agent(user_agent, rfc9110_delimiter, expected):
202210

203211
result = create_user_agent(user_agent, rfc9110_delimiter)
204212
assert result == expected
213+
214+
215+
@mock.patch.dict(os.environ, {"VSCODE_PID": "1234"}, clear=True)
216+
def test_create_user_agent_vscode():
217+
from pandas_gbq.gbq import create_user_agent
218+
219+
assert create_user_agent() == f"pandas-{pd.__version__} vscode"
220+
221+
222+
@mock.patch.dict(os.environ, {"VSCODE_PID": "1234"}, clear=True)
223+
def test_create_user_agent_vscode_plugin():
224+
from pandas_gbq.gbq import create_user_agent
225+
226+
with tempfile.TemporaryDirectory() as tmpdir:
227+
user_home = Path(tmpdir)
228+
plugin_dir = (
229+
user_home / ".vscode" / "extensions" / "googlecloudtools.cloudcode-0.12"
230+
)
231+
plugin_config = plugin_dir / "package.json"
232+
233+
# originally pluging config does not exist
234+
assert not plugin_config.exists()
235+
236+
# simulate plugin installation by creating plugin config on disk
237+
plugin_dir.mkdir(parents=True)
238+
with open(plugin_config, "w") as f:
239+
f.write("{}")
240+
241+
with mock.patch("pathlib.Path.home", return_value=user_home):
242+
assert (
243+
create_user_agent()
244+
== f"pandas-{pd.__version__} vscode googlecloudtools.cloudcode"
245+
)
246+
247+
248+
@mock.patch.dict(os.environ, {"JPY_PARENT_PID": "1234"}, clear=True)
249+
def test_create_user_agent_jupyter():
250+
from pandas_gbq.gbq import create_user_agent
251+
252+
assert create_user_agent() == f"pandas-{pd.__version__} jupyter"
253+
254+
255+
@mock.patch.dict(os.environ, {"JPY_PARENT_PID": "1234"}, clear=True)
256+
def test_create_user_agent_jupyter_extension():
257+
from pandas_gbq.gbq import create_user_agent
258+
259+
def custom_import_module_side_effect(name, package=None):
260+
if name == "bigquery_jupyter_plugin":
261+
return mock.MagicMock()
262+
else:
263+
import importlib
264+
265+
return importlib.import_module(name, package)
266+
267+
with mock.patch(
268+
"importlib.import_module", side_effect=custom_import_module_side_effect
269+
):
270+
assert (
271+
create_user_agent()
272+
== f"pandas-{pd.__version__} jupyter bigquery_jupyter_plugin"
273+
)

0 commit comments

Comments
 (0)