Skip to content

Commit

Permalink
Omit 'storage_options' in 'read/write_pandas' when using Dapla Lab
Browse files Browse the repository at this point in the history
  • Loading branch information
mallport committed Dec 19, 2024
1 parent dff00b0 commit cde6d92
Showing 1 changed file with 18 additions and 0 deletions.
18 changes: 18 additions & 0 deletions src/dapla/pandas.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from __future__ import annotations

import os
import typing as t
from enum import Enum
from typing import Any
Expand All @@ -18,6 +19,8 @@
from pandas import read_sas
from pandas import read_xml

from dapla.const import DaplaRegion

from .auth import AuthClient
from .files import FileClient

Expand Down Expand Up @@ -74,6 +77,12 @@ def read_pandas(
if isinstance(gcs_path, str):
gcs_path = FileClient._ensure_gcs_uri_prefix(gcs_path)

if (
os.getenv("DAPLA_REGION") == DaplaRegion.DAPLA_LAB.value
) and "storage_options" in kwargs:
## If using Dapla Lab, "storage_options" is not necessary because we derive authentication from the environment
kwargs.pop("storage_options", None)

if isinstance(gcs_path, list) and file_format != "parquet":
raise ValueError("Multiple paths are only supported for parquet format")
match SupportedFileFormat(file_format):
Expand Down Expand Up @@ -164,6 +173,12 @@ def write_pandas(
if isinstance(gcs_path, str):
gcs_path = FileClient._ensure_gcs_uri_prefix(gcs_path)

if (
os.getenv("DAPLA_REGION") == DaplaRegion.DAPLA_LAB.value
) and "storage_options" in kwargs:
## If using Dapla Lab, "storage_options" is not necessary because we derive authentication from the environment
kwargs.pop("storage_options", None)

match SupportedFileFormat(file_format):
case SupportedFileFormat.PARQUET:
# Transfom and write pandas dataframe
Expand Down Expand Up @@ -209,5 +224,8 @@ def _get_storage_options() -> Optional[dict[str, Optional[Credentials]]]:
An error will be raised by Pandas if providing this argument with a local path or a file-like buffer.
See the fsspec and backend storage implementation docs for the set of allowed keys and values
"""
if os.getenv("DAPLA_REGION") == DaplaRegion.DAPLA_LAB.value:
## If using Dapla Lab, "storage_options" is not necessary because we derive authentication from the environment
return None
credentials = AuthClient.fetch_google_credentials()
return {"token": credentials} if credentials is not None else None

0 comments on commit cde6d92

Please sign in to comment.