From 58d1ee874b778683f684f0a4ce8bf7341b69f0ad Mon Sep 17 00:00:00 2001 From: Sabri Karagonen Date: Sat, 21 Feb 2026 17:30:53 +0100 Subject: [PATCH] feat: add read_sheet() and write_sheet() with GoogleSheetsConnection support Add zero-boilerplate Google Sheets read/write that works with both google_cloud_platform (reuse GCP creds) and google_sheets (standalone) connection types. Fix GCPConnection.sheets() to scope credentials for Sheets+Drive. Co-Authored-By: Claude Opus 4.6 --- src/bruin/__init__.py | 3 +- src/bruin/_connection.py | 142 +++++++++- src/bruin/_sheets.py | 129 +++++++++ tests/conftest.py | 16 ++ tests/test_sheets.py | 556 +++++++++++++++++++++++++++++++++++++++ 5 files changed, 835 insertions(+), 11 deletions(-) create mode 100644 src/bruin/_sheets.py create mode 100644 tests/test_sheets.py diff --git a/src/bruin/__init__.py b/src/bruin/__init__.py index 0fbb831..0e512f4 100644 --- a/src/bruin/__init__.py +++ b/src/bruin/__init__.py @@ -3,6 +3,7 @@ from bruin._connection import get_connection from bruin._context import context from bruin._query import query +from bruin._sheets import read_sheet, write_sheet __version__ = "0.3.1" -__all__ = ["__version__", "context", "get_connection", "query"] +__all__ = ["__version__", "context", "get_connection", "query", "read_sheet", "write_sheet"] diff --git a/src/bruin/_connection.py b/src/bruin/_connection.py index 8f50fc4..bbab79f 100644 --- a/src/bruin/_connection.py +++ b/src/bruin/_connection.py @@ -70,6 +70,7 @@ def __init__(self, name: str, raw: dict): super().__init__(name, "google_cloud_platform", raw) self._credentials = None self._bigquery_client = None + self._sheets_client = None def _uses_adc(self) -> bool: """Return True when the connection is configured for Application Default Credentials.""" @@ -145,15 +146,32 @@ def bigquery(self): return self._bigquery_client def sheets(self): - """Return an authorized pygsheets client.""" - try: - import pygsheets - except ImportError: - raise ImportError( - "Install bruin-sdk[sheets] to use Google Sheets connections: " - "pip install 'bruin-sdk[sheets]'" + """Return a cached, authorized pygsheets client with Sheets + Drive scopes.""" + if self._sheets_client is None: + try: + import pygsheets + except ImportError: + raise ImportError( + "Install bruin-sdk[sheets] to use Google Sheets connections: " + "pip install 'bruin-sdk[sheets]'" + ) + _SHEETS_SCOPES = ( + "https://www.googleapis.com/auth/spreadsheets", + "https://www.googleapis.com/auth/drive", ) - return pygsheets.authorize(custom_credentials=self.credentials) + scoped = self.credentials.with_scopes(_SHEETS_SCOPES) + self._sheets_client = pygsheets.authorize(custom_credentials=scoped) + return self._sheets_client + + def read_sheet(self, spreadsheet, worksheet="Sheet1"): + """Read a Google Sheets worksheet into a pandas DataFrame.""" + from bruin._sheets import _read_sheet_impl + return _read_sheet_impl(self, spreadsheet, worksheet) + + def write_sheet(self, df, spreadsheet, worksheet="Sheet1", fit=True): + """Write a pandas DataFrame to a Google Sheets worksheet.""" + from bruin._sheets import _write_sheet_impl + _write_sheet_impl(self, df, spreadsheet, worksheet, fit) def storage(self): """Return a google.cloud.storage.Client.""" @@ -175,15 +193,116 @@ def client(self): return self.bigquery() def close(self): - """Close the BigQuery client if it was initialized.""" + """Close the BigQuery and Sheets clients if initialized.""" if self._bigquery_client is not None: logger.debug("Closing BigQuery client for connection '%s'", self.name) close = getattr(self._bigquery_client, "close", None) if callable(close): close() self._bigquery_client = None + self._sheets_client = None + self._credentials = None + + +class GoogleSheetsConnection(Connection): + """Standalone Google Sheets connection (mirrors Go CLI's ``google_sheets`` type).""" + + def __init__(self, name: str, raw: dict): + super().__init__(name, "google_sheets", raw) + self._credentials = None + self._sheets_client = None + + def _parse_sa_info(self): + """Parse the service account JSON from the connection payload.""" + sa_json = self.raw.get("service_account_json", "") + if not sa_json: + sa_file = self.raw.get("service_account_file", "") + if sa_file: + return sa_file # sentinel — handled in credentials property + return None + try: + return json.loads(sa_json) + except (json.JSONDecodeError, TypeError) as exc: + raise ConnectionParseError( + f"Failed to parse service_account_json for '{self.name}': {exc}" + ) from exc + + @property + def credentials(self): + """Return scoped google credentials for Sheets + Drive.""" + if self._credentials is not None: + return self._credentials + + _SCOPES = [ + "https://www.googleapis.com/auth/spreadsheets", + "https://www.googleapis.com/auth/drive", + ] + + sa_info = self._parse_sa_info() + + if sa_info is None: + raise ConnectionParseError( + f"Google Sheets connection '{self.name}' has no " + f"'service_account_json' or 'service_account_file'." + ) + + try: + from google.oauth2 import service_account + except ImportError: + raise ImportError( + "Install bruin-sdk[sheets] to use Google Sheets connections: " + "pip install 'bruin-sdk[sheets]'" + ) + + if isinstance(sa_info, str): + # sa_info is a file path (from service_account_file) + self._credentials = service_account.Credentials.from_service_account_file( + sa_info, scopes=_SCOPES, + ) + else: + self._credentials = service_account.Credentials.from_service_account_info( + sa_info, scopes=_SCOPES, + ) + + logger.debug("Using service account credentials for '%s'", self.name) + return self._credentials + + def sheets(self): + """Return a cached, authorized pygsheets client.""" + if self._sheets_client is None: + try: + import pygsheets + except ImportError: + raise ImportError( + "Install bruin-sdk[sheets] to use Google Sheets connections: " + "pip install 'bruin-sdk[sheets]'" + ) + self._sheets_client = pygsheets.authorize(custom_credentials=self.credentials) + return self._sheets_client + + def read_sheet(self, spreadsheet, worksheet="Sheet1"): + """Read a Google Sheets worksheet into a pandas DataFrame.""" + from bruin._sheets import _read_sheet_impl + return _read_sheet_impl(self, spreadsheet, worksheet) + + def write_sheet(self, df, spreadsheet, worksheet="Sheet1", fit=True): + """Write a pandas DataFrame to a Google Sheets worksheet.""" + from bruin._sheets import _write_sheet_impl + _write_sheet_impl(self, df, spreadsheet, worksheet, fit) + + @property + def client(self): + """Alias for sheets() — the primary use case for this connection type.""" + return self.sheets() + + def close(self): + """Clear cached client and credentials.""" + self._sheets_client = None self._credentials = None + def __repr__(self): + return f"GoogleSheetsConnection(name={self.name!r})" + def _create_client(conn_type: str, raw): """Create a database client based on connection type.""" @@ -212,7 +331,7 @@ def _create_client(conn_type: str, raw): if factory is None: raise ConnectionTypeError( f"Unsupported connection type '{conn_type}'. " - f"Supported types: google_cloud_platform, {', '.join(sorted(factories))}." + f"Supported types: google_cloud_platform, google_sheets, {', '.join(sorted(factories))}." ) return factory(raw) @@ -584,4 +703,7 @@ def get_connection(name: str) -> "Connection | GCPConnection": if conn_type == "google_cloud_platform": return GCPConnection(name, raw) + if conn_type == "google_sheets": + return GoogleSheetsConnection(name, raw) + return Connection(name, conn_type, raw) diff --git a/src/bruin/_sheets.py b/src/bruin/_sheets.py new file mode 100644 index 0000000..7ae3a69 --- /dev/null +++ b/src/bruin/_sheets.py @@ -0,0 +1,129 @@ +import logging + +from bruin._connection import get_connection +from bruin.exceptions import ConnectionNotFoundError, ConnectionTypeError + +logger = logging.getLogger("bruin") + + +def _get_sheets_client(connection): + """Resolve *connection* to a ``(Connection, pygsheets.Client)`` pair. + + Works with both ``google_cloud_platform`` and ``google_sheets`` connection types. + """ + if connection is None: + from bruin._context import context + + connection = context.connection + if connection is None: + raise ConnectionNotFoundError( + "No connection specified and no default connection set " + "(BRUIN_CONNECTION env var is missing). " + "Pass a connection name explicitly: read_sheet(spreadsheet, connection='my_gcp')" + ) + + conn = get_connection(connection) + + if not hasattr(conn, "sheets"): + raise ConnectionTypeError( + f"Connection '{conn.name}' ({conn.type}) does not support Google Sheets." + ) + return conn, conn.sheets() + + +def _read_sheet_impl(conn, spreadsheet, worksheet="Sheet1"): + """Core read logic that operates on an already-resolved connection object.""" + gc = conn.sheets() + logger.debug("Reading '%s'.'%s' via '%s'", spreadsheet, worksheet, conn.name) + + sh = gc.open_by_key(spreadsheet) + wks = sh.worksheet_by_title(worksheet) + df = wks.get_as_df( + numerize=True, + empty_value="", + include_tailing_empty=True, + include_tailing_empty_rows=False, + ) + + logger.debug("Read %d rows x %d cols", len(df), len(df.columns)) + return df + + +def _write_sheet_impl(conn, df, spreadsheet, worksheet="Sheet1", fit=True): + """Core write logic that operates on an already-resolved connection object.""" + gc = conn.sheets() + logger.debug( + "Writing %d rows to '%s'.'%s' via '%s'", + len(df), spreadsheet, worksheet, conn.name, + ) + + sh = gc.open_by_key(spreadsheet) + wks = sh.worksheet_by_title(worksheet) + wks.clear() + wks.set_dataframe(df, start="A1", fit=fit, nan="", escape_formulae=True) + + logger.debug("Write complete") + + +def read_sheet(spreadsheet, worksheet="Sheet1", connection=None): + """Read a Google Sheets worksheet into a pandas DataFrame. + + Parameters + ---------- + spreadsheet : str + The spreadsheet ID (the long string in the Google Sheets URL). + worksheet : str + The worksheet tab title. Defaults to ``"Sheet1"``. + connection : str, optional + Connection name. When *None*, falls back to the asset's default + connection (``BRUIN_CONNECTION`` env var). + + Returns + ------- + pandas.DataFrame + """ + conn, gc = _get_sheets_client(connection) + logger.debug("Reading '%s'.'%s' via '%s'", spreadsheet, worksheet, conn.name) + + sh = gc.open_by_key(spreadsheet) + wks = sh.worksheet_by_title(worksheet) + df = wks.get_as_df( + numerize=True, + empty_value="", + include_tailing_empty=True, + include_tailing_empty_rows=False, + ) + + logger.debug("Read %d rows x %d cols", len(df), len(df.columns)) + return df + + +def write_sheet(df, spreadsheet, worksheet="Sheet1", connection=None, fit=True): + """Write a pandas DataFrame to a Google Sheets worksheet. + + Parameters + ---------- + df : pandas.DataFrame + The data to write. + spreadsheet : str + The spreadsheet ID (the long string in the Google Sheets URL). + worksheet : str + The worksheet tab title. Defaults to ``"Sheet1"``. + connection : str, optional + Connection name. When *None*, falls back to the asset's default + connection (``BRUIN_CONNECTION`` env var). + fit : bool + If *True* (default), resize the sheet to match the DataFrame dimensions. + """ + conn, gc = _get_sheets_client(connection) + logger.debug( + "Writing %d rows to '%s'.'%s' via '%s'", + len(df), spreadsheet, worksheet, conn.name, + ) + + sh = gc.open_by_key(spreadsheet) + wks = sh.worksheet_by_title(worksheet) + wks.clear() + wks.set_dataframe(df, start="A1", fit=fit, nan="", escape_formulae=True) + + logger.debug("Write complete") diff --git a/tests/conftest.py b/tests/conftest.py index 2ac211c..f18e99c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -224,3 +224,19 @@ def vertica_connection_json(): "password": "s3cret", "database": "analytics", } + + +@pytest.fixture +def google_sheets_connection_json(): + """Standalone Google Sheets connection payload (service_account_json is double-serialized).""" + sa_info = { + "type": "service_account", + "project_id": "my-gcp-project", + "private_key_id": "key-id-123", + "private_key": "-----BEGIN RSA PRIVATE KEY-----\nfake\n-----END RSA PRIVATE KEY-----\n", + "client_email": "sa@my-gcp-project.iam.gserviceaccount.com", + "client_id": "123456789", + "auth_uri": "https://accounts.google.com/o/oauth2/auth", + "token_uri": "https://oauth2.googleapis.com/token", + } + return {"service_account_json": json.dumps(sa_info)} diff --git a/tests/test_sheets.py b/tests/test_sheets.py new file mode 100644 index 0000000..02a72b4 --- /dev/null +++ b/tests/test_sheets.py @@ -0,0 +1,556 @@ +import json +import sys +from unittest.mock import MagicMock, patch, call + +import pandas as pd +import pytest + +from bruin import get_connection, read_sheet, write_sheet +from bruin._connection import GCPConnection, GoogleSheetsConnection +from bruin._sheets import _get_sheets_client, _read_sheet_impl, _write_sheet_impl +from bruin.exceptions import ( + ConnectionNotFoundError, + ConnectionParseError, + ConnectionTypeError, +) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _mock_google_modules(): + """Create mock google.oauth2.service_account module hierarchy.""" + mock_sa_module = MagicMock() + mock_creds = MagicMock() + mock_sa_module.Credentials.from_service_account_info.return_value = mock_creds + mock_sa_module.Credentials.from_service_account_file.return_value = mock_creds + + mock_google = MagicMock() + mock_google.oauth2.service_account = mock_sa_module + + modules = { + "google": mock_google, + "google.oauth2": mock_google.oauth2, + "google.oauth2.service_account": mock_sa_module, + } + return modules, mock_sa_module, mock_creds + + +def _mock_pygsheets_module(mock_gc=None): + """Create mock pygsheets module.""" + mock_pygsheets = MagicMock() + if mock_gc is None: + mock_gc = MagicMock() + mock_pygsheets.authorize.return_value = mock_gc + return mock_pygsheets, mock_gc + + +def _mock_sheets_chain(df=None): + """Create full mock chain: pygsheets.authorize → open_by_key → worksheet_by_title → wks.""" + if df is None: + df = pd.DataFrame({"col1": [1, 2], "col2": ["a", "b"]}) + + mock_wks = MagicMock() + mock_wks.get_as_df.return_value = df + + mock_sh = MagicMock() + mock_sh.worksheet_by_title.return_value = mock_wks + + mock_gc = MagicMock() + mock_gc.open_by_key.return_value = mock_sh + + return mock_gc, mock_sh, mock_wks + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + +@pytest.fixture +def gcp_sheets_env(monkeypatch, bq_connection_json): + """Set up env vars for a GCP connection used for Sheets.""" + monkeypatch.setenv( + "BRUIN_CONNECTION_TYPES", + json.dumps({"my_gcp": "google_cloud_platform"}), + ) + monkeypatch.setenv("my_gcp", json.dumps(bq_connection_json)) + + +@pytest.fixture +def google_sheets_env(monkeypatch, google_sheets_connection_json): + """Set up env vars for a standalone google_sheets connection.""" + monkeypatch.setenv( + "BRUIN_CONNECTION_TYPES", + json.dumps({"my_sheets": "google_sheets"}), + ) + monkeypatch.setenv("my_sheets", json.dumps(google_sheets_connection_json)) + + +@pytest.fixture +def postgres_env(monkeypatch, postgres_connection_json): + monkeypatch.setenv( + "BRUIN_CONNECTION_TYPES", + json.dumps({"my_pg": "postgres"}), + ) + monkeypatch.setenv("my_pg", json.dumps(postgres_connection_json)) + + +@pytest.fixture +def default_conn_env(monkeypatch, bq_connection_json): + """Set up env with BRUIN_CONNECTION pointing to a GCP connection.""" + monkeypatch.setenv( + "BRUIN_CONNECTION_TYPES", + json.dumps({"default_gcp": "google_cloud_platform"}), + ) + monkeypatch.setenv("default_gcp", json.dumps(bq_connection_json)) + monkeypatch.setenv("BRUIN_CONNECTION", "default_gcp") + + +# --------------------------------------------------------------------------- +# GoogleSheetsConnection +# --------------------------------------------------------------------------- + +class TestGoogleSheetsConnection: + def test_get_connection_returns_google_sheets(self, google_sheets_env): + conn = get_connection("my_sheets") + assert isinstance(conn, GoogleSheetsConnection) + assert conn.type == "google_sheets" + assert conn.name == "my_sheets" + + def test_repr(self, google_sheets_env): + conn = get_connection("my_sheets") + assert repr(conn) == "GoogleSheetsConnection(name='my_sheets')" + + def test_credentials_parsed_from_sa_json(self, google_sheets_env): + google_mods, mock_sa_mod, mock_creds = _mock_google_modules() + + conn = get_connection("my_sheets") + with patch.dict(sys.modules, google_mods): + creds = conn.credentials + + mock_sa_mod.Credentials.from_service_account_info.assert_called_once() + sa_info = mock_sa_mod.Credentials.from_service_account_info.call_args[0][0] + assert sa_info["type"] == "service_account" + assert sa_info["project_id"] == "my-gcp-project" + + # Scopes should be passed + assert mock_sa_mod.Credentials.from_service_account_info.call_args[1]["scopes"] == [ + "https://www.googleapis.com/auth/spreadsheets", + "https://www.googleapis.com/auth/drive", + ] + + def test_credentials_are_cached(self, google_sheets_env): + google_mods, mock_sa_mod, _ = _mock_google_modules() + + conn = get_connection("my_sheets") + with patch.dict(sys.modules, google_mods): + _ = conn.credentials + _ = conn.credentials + mock_sa_mod.Credentials.from_service_account_info.assert_called_once() + + def test_missing_sa_json_raises(self, monkeypatch): + monkeypatch.setenv( + "BRUIN_CONNECTION_TYPES", + json.dumps({"bad_sheets": "google_sheets"}), + ) + monkeypatch.setenv("bad_sheets", json.dumps({})) + + conn = get_connection("bad_sheets") + with pytest.raises(ConnectionParseError, match="no.*service_account_json"): + _ = conn.credentials + + def test_invalid_sa_json_raises(self, monkeypatch): + monkeypatch.setenv( + "BRUIN_CONNECTION_TYPES", + json.dumps({"bad_sheets": "google_sheets"}), + ) + monkeypatch.setenv("bad_sheets", json.dumps({"service_account_json": "not-json"})) + + conn = get_connection("bad_sheets") + with pytest.raises(ConnectionParseError, match="Failed to parse"): + _ = conn.credentials + + def test_credentials_from_service_account_file(self, monkeypatch): + monkeypatch.setenv( + "BRUIN_CONNECTION_TYPES", + json.dumps({"file_sheets": "google_sheets"}), + ) + monkeypatch.setenv("file_sheets", json.dumps({"service_account_file": "/path/to/sa.json"})) + + google_mods, mock_sa_mod, _ = _mock_google_modules() + + conn = get_connection("file_sheets") + with patch.dict(sys.modules, google_mods): + _ = conn.credentials + + mock_sa_mod.Credentials.from_service_account_file.assert_called_once_with( + "/path/to/sa.json", + scopes=[ + "https://www.googleapis.com/auth/spreadsheets", + "https://www.googleapis.com/auth/drive", + ], + ) + + def test_sheets_returns_pygsheets_client(self, google_sheets_env): + google_mods, _, mock_creds = _mock_google_modules() + mock_gc, _, _ = _mock_sheets_chain() + mock_pyg, _ = _mock_pygsheets_module(mock_gc) + + conn = get_connection("my_sheets") + with patch.dict(sys.modules, {**google_mods, "pygsheets": mock_pyg}): + result = conn.sheets() + + assert result is mock_gc + mock_pyg.authorize.assert_called_once_with(custom_credentials=mock_creds) + + def test_sheets_client_is_cached(self, google_sheets_env): + google_mods, _, _ = _mock_google_modules() + mock_pyg, mock_gc = _mock_pygsheets_module() + + conn = get_connection("my_sheets") + with patch.dict(sys.modules, {**google_mods, "pygsheets": mock_pyg}): + c1 = conn.sheets() + c2 = conn.sheets() + assert c1 is c2 + mock_pyg.authorize.assert_called_once() + + def test_close_clears_cached_state(self, google_sheets_env): + google_mods, _, _ = _mock_google_modules() + conn = get_connection("my_sheets") + with patch.dict(sys.modules, google_mods): + _ = conn.credentials + + conn.close() + assert conn._sheets_client is None + assert conn._credentials is None + + def test_context_manager(self, google_sheets_env): + with get_connection("my_sheets") as conn: + assert isinstance(conn, GoogleSheetsConnection) + + def test_client_property_aliases_sheets(self, google_sheets_env): + conn = get_connection("my_sheets") + with patch.object(conn, "sheets", return_value=MagicMock()) as mock_sheets: + result = conn.client + mock_sheets.assert_called_once() + assert result is mock_sheets.return_value + + +# --------------------------------------------------------------------------- +# GCPConnection.sheets() — scoped credentials +# --------------------------------------------------------------------------- + +class TestGCPConnectionSheets: + def test_sheets_adds_scopes(self, gcp_sheets_env): + google_mods, mock_sa_mod, mock_creds = _mock_google_modules() + mock_scoped = MagicMock() + mock_creds.with_scopes.return_value = mock_scoped + + mock_pyg, mock_gc = _mock_pygsheets_module() + + conn = get_connection("my_gcp") + with patch.dict(sys.modules, {**google_mods, "pygsheets": mock_pyg}): + result = conn.sheets() + + assert result is mock_gc + mock_creds.with_scopes.assert_called_once_with(( + "https://www.googleapis.com/auth/spreadsheets", + "https://www.googleapis.com/auth/drive", + )) + mock_pyg.authorize.assert_called_once_with(custom_credentials=mock_scoped) + + def test_sheets_client_is_cached(self, gcp_sheets_env): + google_mods, _, mock_creds = _mock_google_modules() + mock_creds.with_scopes.return_value = MagicMock() + mock_pyg, _ = _mock_pygsheets_module() + + conn = get_connection("my_gcp") + with patch.dict(sys.modules, {**google_mods, "pygsheets": mock_pyg}): + c1 = conn.sheets() + c2 = conn.sheets() + assert c1 is c2 + mock_pyg.authorize.assert_called_once() + + def test_close_clears_sheets_client(self, gcp_sheets_env): + conn = get_connection("my_gcp") + conn._sheets_client = MagicMock() + + conn.close() + assert conn._sheets_client is None + + +# --------------------------------------------------------------------------- +# _get_sheets_client — connection resolution +# --------------------------------------------------------------------------- + +class TestGetSheetsClient: + def test_gcp_connection_works(self, gcp_sheets_env): + google_mods, _, mock_creds = _mock_google_modules() + mock_creds.with_scopes.return_value = MagicMock() + mock_pyg, _ = _mock_pygsheets_module() + + with patch.dict(sys.modules, {**google_mods, "pygsheets": mock_pyg}): + conn, gc = _get_sheets_client("my_gcp") + assert isinstance(conn, GCPConnection) + assert gc is not None + + def test_google_sheets_connection_works(self, google_sheets_env): + google_mods, _, _ = _mock_google_modules() + mock_pyg, _ = _mock_pygsheets_module() + + with patch.dict(sys.modules, {**google_mods, "pygsheets": mock_pyg}): + conn, gc = _get_sheets_client("my_sheets") + assert isinstance(conn, GoogleSheetsConnection) + assert gc is not None + + def test_postgres_connection_raises(self, postgres_env): + with pytest.raises(ConnectionTypeError, match="does not support Google Sheets"): + _get_sheets_client("my_pg") + + def test_no_connection_no_default_raises(self, monkeypatch): + monkeypatch.delenv("BRUIN_CONNECTION", raising=False) + with pytest.raises(ConnectionNotFoundError, match="No connection specified"): + _get_sheets_client(None) + + def test_default_connection_fallback(self, default_conn_env): + google_mods, _, mock_creds = _mock_google_modules() + mock_creds.with_scopes.return_value = MagicMock() + mock_pyg, _ = _mock_pygsheets_module() + + with patch.dict(sys.modules, {**google_mods, "pygsheets": mock_pyg}): + conn, gc = _get_sheets_client(None) + assert conn.name == "default_gcp" + + +# --------------------------------------------------------------------------- +# read_sheet() +# --------------------------------------------------------------------------- + +class TestReadSheet: + def test_read_returns_dataframe(self, gcp_sheets_env): + expected_df = pd.DataFrame({"col1": [1, 2], "col2": ["a", "b"]}) + google_mods, _, mock_creds = _mock_google_modules() + mock_creds.with_scopes.return_value = MagicMock() + mock_gc, mock_sh, mock_wks = _mock_sheets_chain(expected_df) + mock_pyg, _ = _mock_pygsheets_module(mock_gc) + + with patch.dict(sys.modules, {**google_mods, "pygsheets": mock_pyg}): + result = read_sheet("spreadsheet_id_123", connection="my_gcp") + + pd.testing.assert_frame_equal(result, expected_df) + mock_gc.open_by_key.assert_called_once_with("spreadsheet_id_123") + mock_sh.worksheet_by_title.assert_called_once_with("Sheet1") + mock_wks.get_as_df.assert_called_once_with( + numerize=True, + empty_value="", + include_tailing_empty=True, + include_tailing_empty_rows=False, + ) + + def test_read_custom_worksheet(self, gcp_sheets_env): + google_mods, _, mock_creds = _mock_google_modules() + mock_creds.with_scopes.return_value = MagicMock() + mock_gc, mock_sh, _ = _mock_sheets_chain(pd.DataFrame()) + mock_pyg, _ = _mock_pygsheets_module(mock_gc) + + with patch.dict(sys.modules, {**google_mods, "pygsheets": mock_pyg}): + read_sheet("id123", worksheet="Revenue", connection="my_gcp") + mock_sh.worksheet_by_title.assert_called_once_with("Revenue") + + def test_read_uses_default_connection(self, default_conn_env): + google_mods, _, mock_creds = _mock_google_modules() + mock_creds.with_scopes.return_value = MagicMock() + mock_gc, _, _ = _mock_sheets_chain(pd.DataFrame()) + mock_pyg, _ = _mock_pygsheets_module(mock_gc) + + with patch.dict(sys.modules, {**google_mods, "pygsheets": mock_pyg}): + # No connection= arg — should fall back to BRUIN_CONNECTION + read_sheet("id123") + + +# --------------------------------------------------------------------------- +# write_sheet() +# --------------------------------------------------------------------------- + +class TestWriteSheet: + def test_write_clears_then_sets_dataframe(self, gcp_sheets_env): + google_mods, _, mock_creds = _mock_google_modules() + mock_creds.with_scopes.return_value = MagicMock() + mock_gc, mock_sh, mock_wks = _mock_sheets_chain() + mock_pyg, _ = _mock_pygsheets_module(mock_gc) + + df = pd.DataFrame({"x": [1, 2, 3]}) + with patch.dict(sys.modules, {**google_mods, "pygsheets": mock_pyg}): + write_sheet(df, "spreadsheet_id", connection="my_gcp") + + mock_gc.open_by_key.assert_called_once_with("spreadsheet_id") + mock_sh.worksheet_by_title.assert_called_once_with("Sheet1") + + # clear() must be called before set_dataframe() + assert mock_wks.clear.call_count == 1 + assert mock_wks.set_dataframe.call_count == 1 + + # Verify call order: clear before set_dataframe + method_names = [c[0] for c in mock_wks.method_calls] + clear_idx = method_names.index("clear") + set_idx = method_names.index("set_dataframe") + assert clear_idx < set_idx + + mock_wks.set_dataframe.assert_called_once_with( + df, start="A1", fit=True, nan="", escape_formulae=True, + ) + + def test_write_fit_false(self, gcp_sheets_env): + google_mods, _, mock_creds = _mock_google_modules() + mock_creds.with_scopes.return_value = MagicMock() + mock_gc, _, mock_wks = _mock_sheets_chain() + mock_pyg, _ = _mock_pygsheets_module(mock_gc) + + df = pd.DataFrame({"x": [1]}) + with patch.dict(sys.modules, {**google_mods, "pygsheets": mock_pyg}): + write_sheet(df, "id123", connection="my_gcp", fit=False) + + mock_wks.set_dataframe.assert_called_once_with( + df, start="A1", fit=False, nan="", escape_formulae=True, + ) + + def test_write_custom_worksheet(self, gcp_sheets_env): + google_mods, _, mock_creds = _mock_google_modules() + mock_creds.with_scopes.return_value = MagicMock() + mock_gc, mock_sh, _ = _mock_sheets_chain() + mock_pyg, _ = _mock_pygsheets_module(mock_gc) + + df = pd.DataFrame({"x": [1]}) + with patch.dict(sys.modules, {**google_mods, "pygsheets": mock_pyg}): + write_sheet(df, "id123", worksheet="Output", connection="my_gcp") + mock_sh.worksheet_by_title.assert_called_once_with("Output") + + +# --------------------------------------------------------------------------- +# Connection methods — conn.read_sheet() / conn.write_sheet() +# --------------------------------------------------------------------------- + +class TestConnectionMethods: + def test_gcp_conn_read_sheet(self, gcp_sheets_env): + expected_df = pd.DataFrame({"a": [1]}) + google_mods, _, mock_creds = _mock_google_modules() + mock_creds.with_scopes.return_value = MagicMock() + mock_gc, mock_sh, mock_wks = _mock_sheets_chain(expected_df) + mock_pyg, _ = _mock_pygsheets_module(mock_gc) + + conn = get_connection("my_gcp") + with patch.dict(sys.modules, {**google_mods, "pygsheets": mock_pyg}): + result = conn.read_sheet("id123", "Tab1") + + pd.testing.assert_frame_equal(result, expected_df) + mock_gc.open_by_key.assert_called_once_with("id123") + mock_sh.worksheet_by_title.assert_called_once_with("Tab1") + + def test_gcp_conn_write_sheet(self, gcp_sheets_env): + google_mods, _, mock_creds = _mock_google_modules() + mock_creds.with_scopes.return_value = MagicMock() + mock_gc, _, mock_wks = _mock_sheets_chain() + mock_pyg, _ = _mock_pygsheets_module(mock_gc) + + conn = get_connection("my_gcp") + df = pd.DataFrame({"x": [10]}) + with patch.dict(sys.modules, {**google_mods, "pygsheets": mock_pyg}): + conn.write_sheet(df, "id456", "Out") + + mock_wks.clear.assert_called_once() + mock_wks.set_dataframe.assert_called_once_with( + df, start="A1", fit=True, nan="", escape_formulae=True, + ) + + def test_google_sheets_conn_read_sheet(self, google_sheets_env): + expected_df = pd.DataFrame({"b": [2]}) + google_mods, _, _ = _mock_google_modules() + mock_gc, _, mock_wks = _mock_sheets_chain(expected_df) + mock_pyg, _ = _mock_pygsheets_module(mock_gc) + + conn = get_connection("my_sheets") + with patch.dict(sys.modules, {**google_mods, "pygsheets": mock_pyg}): + result = conn.read_sheet("id789") + + pd.testing.assert_frame_equal(result, expected_df) + + def test_google_sheets_conn_write_sheet(self, google_sheets_env): + google_mods, _, _ = _mock_google_modules() + mock_gc, _, mock_wks = _mock_sheets_chain() + mock_pyg, _ = _mock_pygsheets_module(mock_gc) + + conn = get_connection("my_sheets") + df = pd.DataFrame({"z": [99]}) + with patch.dict(sys.modules, {**google_mods, "pygsheets": mock_pyg}): + conn.write_sheet(df, "id789", fit=False) + + mock_wks.set_dataframe.assert_called_once_with( + df, start="A1", fit=False, nan="", escape_formulae=True, + ) + + +# --------------------------------------------------------------------------- +# Import errors +# --------------------------------------------------------------------------- + +class TestImportErrors: + def test_gcp_sheets_missing_pygsheets(self, gcp_sheets_env): + conn = get_connection("my_gcp") + # Force credentials so we get past that step + conn._credentials = MagicMock() + + with patch.dict("sys.modules", {"pygsheets": None}): + with pytest.raises(ImportError, match="bruin-sdk\\[sheets\\]"): + conn.sheets() + + def test_google_sheets_missing_google_auth(self, google_sheets_env): + conn = get_connection("my_sheets") + + with patch.dict("sys.modules", {"google.oauth2": None, "google.oauth2.service_account": None}): + with pytest.raises(ImportError, match="bruin-sdk\\[sheets\\]"): + _ = conn.credentials + + def test_google_sheets_missing_pygsheets(self, google_sheets_env): + conn = get_connection("my_sheets") + conn._credentials = MagicMock() + + with patch.dict("sys.modules", {"pygsheets": None}): + with pytest.raises(ImportError, match="bruin-sdk\\[sheets\\]"): + conn.sheets() + + +# --------------------------------------------------------------------------- +# _read_sheet_impl / _write_sheet_impl (internal helpers) +# --------------------------------------------------------------------------- + +class TestImplHelpers: + def test_read_sheet_impl(self, gcp_sheets_env): + expected_df = pd.DataFrame({"v": [42]}) + google_mods, _, mock_creds = _mock_google_modules() + mock_creds.with_scopes.return_value = MagicMock() + mock_gc, mock_sh, mock_wks = _mock_sheets_chain(expected_df) + mock_pyg, _ = _mock_pygsheets_module(mock_gc) + + conn = get_connection("my_gcp") + with patch.dict(sys.modules, {**google_mods, "pygsheets": mock_pyg}): + result = _read_sheet_impl(conn, "sid", "Tab") + + pd.testing.assert_frame_equal(result, expected_df) + mock_sh.worksheet_by_title.assert_called_once_with("Tab") + + def test_write_sheet_impl(self, gcp_sheets_env): + google_mods, _, mock_creds = _mock_google_modules() + mock_creds.with_scopes.return_value = MagicMock() + mock_gc, _, mock_wks = _mock_sheets_chain() + mock_pyg, _ = _mock_pygsheets_module(mock_gc) + + conn = get_connection("my_gcp") + df = pd.DataFrame({"w": [7]}) + with patch.dict(sys.modules, {**google_mods, "pygsheets": mock_pyg}): + _write_sheet_impl(conn, df, "sid", "Tab", fit=False) + + mock_wks.clear.assert_called_once() + mock_wks.set_dataframe.assert_called_once_with( + df, start="A1", fit=False, nan="", escape_formulae=True, + )