feat(gsheets): file upload (#24921)

tests/unit_tests/db_engine_specs/test_gsheets.py

@@ -19,9 +19,13 @@
 
 import json
 
+import pandas as pd
+import pytest
 from pytest_mock import MockFixture
 
 from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
+from superset.exceptions import SupersetException
+from superset.sql_parse import Table
 
 
 class ProgrammingError(Exception):
@@ -307,3 +311,91 @@ def test_unmask_encrypted_extra_when_new_is_none() -> None:
     new = None
 
     assert GSheetsEngineSpec.unmask_encrypted_extra(old, new) is None
+
+
+def test_upload_new(mocker: MockFixture) -> None:
+    """
+    Test file upload when the table does not exist.
+    """
+    from superset.db_engine_specs.gsheets import GSheetsEngineSpec
+
+    mocker.patch("superset.db_engine_specs.gsheets.db")
+    get_adapter_for_table_name = mocker.patch(
+        "shillelagh.backends.apsw.dialects.base.get_adapter_for_table_name"
+    )
+    session = get_adapter_for_table_name()._get_session()
+    session.post().json.return_value = {
+        "spreadsheetId": 1,
+        "spreadsheetUrl": "https://docs.example.org",
+        "sheets": [{"properties": {"title": "sample_data"}}],
+    }
+
+    database = mocker.MagicMock()
+    database.get_extra.return_value = {}
+
+    df = pd.DataFrame([1, "foo", 3.0])
+    table = Table("sample_data")
+
+    GSheetsEngineSpec.df_to_sql(database, table, df, {})
+    assert database.extra == json.dumps(
+        {"engine_params": {"catalog": {"sample_data": "https://docs.example.org"}}}
+    )
+
+
+def test_upload_existing(mocker: MockFixture) -> None:
+    """
+    Test file upload when the table does exist.
+    """
+    from superset.db_engine_specs.gsheets import GSheetsEngineSpec
+
+    mocker.patch("superset.db_engine_specs.gsheets.db")
+    get_adapter_for_table_name = mocker.patch(
+        "shillelagh.backends.apsw.dialects.base.get_adapter_for_table_name"
+    )
+    adapter = get_adapter_for_table_name()
+    adapter._spreadsheet_id = 1
+    adapter._sheet_name = "sheet0"
+    session = adapter._get_session()
+    session.post().json.return_value = {
+        "spreadsheetId": 1,
+        "spreadsheetUrl": "https://docs.example.org",
+        "sheets": [{"properties": {"title": "sample_data"}}],
+    }
+
+    database = mocker.MagicMock()
+    database.get_extra.return_value = {
+        "engine_params": {"catalog": {"sample_data": "https://docs.example.org"}}
+    }
+
+    df = pd.DataFrame([1, "foo", 3.0])
+    table = Table("sample_data")
+
+    with pytest.raises(SupersetException) as excinfo:
+        GSheetsEngineSpec.df_to_sql(database, table, df, {"if_exists": "append"})
+    assert str(excinfo.value) == "Append operation not currently supported"
+
+    with pytest.raises(SupersetException) as excinfo:
+        GSheetsEngineSpec.df_to_sql(database, table, df, {"if_exists": "fail"})
+    assert str(excinfo.value) == "Table already exists"
+
+    GSheetsEngineSpec.df_to_sql(database, table, df, {"if_exists": "replace"})
+    session.post.assert_has_calls(
+        [
+            mocker.call(),
+            mocker.call(
+                "https://sheets.googleapis.com/v4/spreadsheets/1/values/sheet0:clear",
+                json={},
+            ),
+            mocker.call().json(),
+            mocker.call(
+                "https://sheets.googleapis.com/v4/spreadsheets/1/values/sheet0:append",
+                json={
+                    "range": "sheet0",
+                    "majorDimension": "ROWS",
+                    "values": [[1], ["foo"], [3.0]],
+                },
+                params={"valueInputOption": "USER_ENTERED"},
+            ),
+            mocker.call().json(),
+        ]
+    )