Skip to content
14 changes: 13 additions & 1 deletion google/cloud/documentai_toolbox/utilities/gcs_utilities.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
# limitations under the License.
#
"""Google Cloud Storage utilities."""
import importlib.metadata
import os
import re
from typing import Dict, List, Optional, Tuple
Expand Down Expand Up @@ -142,7 +143,18 @@ def get_blob(
if not re.match(constants.FILE_CHECK_REGEX, gcs_uri):
raise ValueError("gcs_uri must link to a single file.")

return storage.Blob.from_string(gcs_uri, _get_storage_client(module=module))
try:
version = importlib.metadata.version("google-cloud-storage")
except importlib.metadata.PackageNotFoundError:
raise ImportError("google-cloud-storage is not installed.")

client = _get_storage_client(module=module)

major, _, _ = map(int, version.split("."))
if major < 3:
return storage.Blob.from_string(gcs_uri, client)
else:
return storage.Blob.from_uri(gcs_uri, client)


def split_gcs_uri(gcs_uri: str) -> Tuple[str, str]:
Expand Down
Loading