From dd462aa2716e9987d880c8272faff89c34429de7 Mon Sep 17 00:00:00 2001
From: Javier Goizueta <jgoizueta@gmail.com>
Date: Fri, 22 Mar 2024 13:06:48 +0100
Subject: [PATCH] fix: error when bigquery dependencies not installed (#133)

---
 raster_loader/errors.py      | 49 +++++-------------------------------
 raster_loader/io/bigquery.py | 27 ++++++++++++++------
 raster_loader/io/common.py   | 14 +----------
 3 files changed, 26 insertions(+), 64 deletions(-)

diff --git a/raster_loader/errors.py b/raster_loader/errors.py
index cc00634..2c07c43 100644
--- a/raster_loader/errors.py
+++ b/raster_loader/errors.py
@@ -1,54 +1,17 @@
 def import_error_bigquery():  # pragma: no cover
     msg = (
-        "Google Cloud BigQuery is not installed.\n"
-        "Please install Google Cloud BigQuery to use this function.\n"
-        "See https://googleapis.dev/python/bigquery/latest/index.html\n"
-        "for installation instructions.\n"
-        "OR, run `pip install google-cloud-bigquery` to install from pypi."
+        "Google Cloud BigQuery client is not installed.\n"
+        "Please install Google Cloud BigQuery dependencies to use this function.\n"
+        "run `pip install -U raster-loader[bigquery]` to install from pypi."
     )
     raise ImportError(msg)
 
 
 def import_error_snowflake():  # pragma: no cover
     msg = (
-        "Google Snowflake is not installed.\n"
-        "Please install Snowflake to use this function.\n"
-        "See https://docs.snowflake.com/en/developer-guide/python-connector\n"
-        "for installation instructions.\n"
-        "OR, run `pip install snowflake-connector-python` to install from pypi."
-    )
-    raise ImportError(msg)
-
-
-def import_error_rasterio():  # pragma: no cover
-    msg = (
-        "Rasterio is not installed.\n"
-        "Please install rasterio to use this function.\n"
-        "See https://rasterio.readthedocs.io/en/latest/installation.html\n"
-        "for installation instructions.\n"
-        "Alternatively, run `pip install rasterio` to install from pypi."
-    )
-    raise ImportError(msg)
-
-
-def import_error_rio_cogeo():  # pragma: no cover
-    msg = (
-        "Rasterio is not installed.\n"
-        "Please install rio_cogeo to use this function.\n"
-        "See https://cogeotiff.github.io/rio-cogeo/\n"
-        "for installation instructions.\n"
-        "Alternatively, run `pip install rio-cogeo` to install from pypi."
-    )
-    raise ImportError(msg)
-
-
-def import_error_quadbin():  # pragma: no cover
-    msg = (
-        "Quadbin is not installed.\n"
-        "Please install quadbin to use this function.\n"
-        "See https://github.com/CartoDB/quadbin-py\n"
-        "for installation instructions.\n"
-        "Alternatively, run `pip install quadbin` to install from pypi."
+        "Snowflake client is not installed.\n"
+        "Please install Snowflake dependencies to use this function.\n"
+        "run `pip install -U raster-loader[snowflake]` to install from pypi."
     )
     raise ImportError(msg)
 
diff --git a/raster_loader/io/bigquery.py b/raster_loader/io/bigquery.py
index 4a40de2..287a392 100644
--- a/raster_loader/io/bigquery.py
+++ b/raster_loader/io/bigquery.py
@@ -29,17 +29,28 @@
 
 from raster_loader.io.datawarehouse import DataWarehouseConnection
 
+if _has_bigquery:
 
-class AccessTokenCredentials(Credentials):
-    def __init__(self, access_token):
-        super(AccessTokenCredentials, self).__init__()
-        self._access_token = access_token
+    class AccessTokenCredentials(Credentials):
+        def __init__(self, access_token):
+            super(AccessTokenCredentials, self).__init__()
+            self._access_token = access_token
 
-    def refresh(self, request):
-        pass
+        def refresh(self, request):
+            pass
 
-    def apply(self, headers, token=None):
-        headers["Authorization"] = f"Bearer {self._access_token}"
+        def apply(self, headers, token=None):
+            headers["Authorization"] = f"Bearer {self._access_token}"
+
+else:
+
+    class Credentials:
+        def __init__(self):
+            import_error_bigquery()
+
+    class AccessTokenCredentials:
+        def __init__(self, access_token):
+            import_error_bigquery()
 
 
 class BigQueryConnection(DataWarehouseConnection):
diff --git a/raster_loader/io/common.py b/raster_loader/io/common.py
index 0a5af7f..ba0729d 100644
--- a/raster_loader/io/common.py
+++ b/raster_loader/io/common.py
@@ -9,20 +9,12 @@
 from typing import List
 from typing import Tuple
 from affine import Affine
-
-try:
-    import rio_cogeo
-except ImportError:  # pragma: no cover
-    _has_rio_cogeo = False
-else:
-    _has_rio_cogeo = True
-
+import rio_cogeo
 import rasterio
 import quadbin
 
 from raster_loader.geo import raster_bounds
 from raster_loader.errors import (
-    import_error_rio_cogeo,
     error_not_google_compatible,
 )
 
@@ -99,10 +91,6 @@ def rasterio_metadata(
     bands_info: List[Tuple[int, str]],
     band_rename_function: Callable,
 ):
-    """Requires rio_cogeo."""
-    if not _has_rio_cogeo:  # pragma: no cover
-        import_error_rio_cogeo()
-
     """Open a raster file with rasterio."""
     raster_info = rio_cogeo.cog_info(file_path).dict()