Fix authentication error (#5)

main.py

@@ -7,15 +7,119 @@ from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from geemap.ee_tile_layers import _get_tile_url_format, _validate_palette
 from starlette.middleware.cors import CORSMiddleware
+from typing import Any, Dict, List, Optional, Tuple, Union
 
-# Earth Engine auth
-if "EARTHENGINE_TOKEN" not in os.environ:
-    raise RuntimeError("EARTHENGINE_TOKEN environment variable not found")
+# # Earth Engine auth
+# if "EARTHENGINE_TOKEN" not in os.environ:
+#     raise RuntimeError("EARTHENGINE_TOKEN environment variable not found")
 
-try:
-    geemap.ee_initialize()
-except Exception as e:
-    raise RuntimeError(f"Earth Engine authentication failed: {e}")
+# try:
+#     geemap.ee_initialize()
+# except Exception as e:
+#     raise RuntimeError(f"Earth Engine authentication failed: {e}")
+
+
+def get_env_var(key: str) -> Optional[str]:
+    """Retrieves an environment variable or Colab secret for the given key.
+
+    Colab secrets have precedence over environment variables.
+
+    Args:
+        key (str): The key that's used to fetch the environment variable.
+
+    Returns:
+        Optional[str]: The retrieved key, or None if no environment variable was found.
+    """
+    if not key:
+        return None
+
+    return os.environ.get(key)
+
+
+def ee_initialize(
+    token_name: str = "EARTHENGINE_TOKEN",
+    auth_mode: Optional[str] = None,
+    auth_args: Optional[Dict[str, Any]] = None,
+    project: Optional[str] = None,
+    **kwargs: Any,
+) -> None:
+    """Authenticates Earth Engine and initialize an Earth Engine session
+
+    Args:
+        token_name (str, optional): The name of the Earth Engine token.
+            Defaults to "EARTHENGINE_TOKEN". In Colab, you can also set a secret
+            named "EE_PROJECT_ID" to initialize Earth Engine.
+        auth_mode (str, optional): The authentication mode, can be one of colab,
+            notebook, localhost, or gcloud.
+            See https://developers.google.com/earth-engine/guides/auth for more
+            details. Defaults to None.
+        auth_args (dict, optional): Additional authentication parameters for
+            aa.Authenticate(). Defaults to {}.
+        user_agent_prefix (str, optional): If set, the prefix (version-less)
+            value used for setting the user-agent string. Defaults to "geemap".
+        project (str, optional): The Google cloud project ID for Earth Engine.
+            Defaults to None.
+        kwargs (dict, optional): Additional parameters for ee.Initialize().
+            For example, opt_url='https://earthengine-highvolume.googleapis.com'
+            to use the Earth Engine High-Volume platform. Defaults to {}.
+    """
+    import google.oauth2.credentials
+
+    # pylint: disable-next=protected-access
+    if ee.data._get_state().credentials is not None:
+        return
+
+    if get_env_var("EE_SERVICE_ACCOUNT") is not None:
+
+        key_data = get_env_var("EE_SERVICE_ACCOUNT")
+
+        try:
+            email = json.loads(key_data)["client_email"]
+        except json.JSONDecodeError as e:
+            raise ValueError(f"Invalid JSON for key_data: {e}")
+        except KeyError:
+            raise ValueError("key_data JSON does not contain 'client_email'")
+        credentials = ee.ServiceAccountCredentials(email=email, key_data=key_data)
+        ee.Initialize(credentials)
+        return
+
+    ee_token = get_env_var(token_name)
+    if ee_token is not None:
+
+        stored = json.loads(ee_token)
+        credentials = google.oauth2.credentials.Credentials(
+            None,
+            token_uri="https://oauth2.googleapis.com/token",
+            client_id=stored["client_id"],
+            client_secret=stored["client_secret"],
+            refresh_token=stored["refresh_token"],
+            quota_project_id=stored["project"],
+        )
+
+        ee.Initialize(credentials=credentials, **kwargs)
+        return
+
+    if auth_args is None:
+        auth_args = {}
+
+    if project is None:
+        kwargs["project"] = get_env_var("EE_PROJECT_ID")
+    else:
+        kwargs["project"] = project
+
+    if auth_mode is None:
+        # pylint: disable-next=protected-access
+        if ee.data._get_state().credentials is None:
+            ee.Authenticate()
+            ee.Initialize(**kwargs)
+            return
+        else:
+            auth_mode = "notebook"
+
+    auth_args["auth_mode"] = auth_mode
+
+    ee.Authenticate(**auth_args)
+    ee.Initialize(**kwargs)
 
 
 # ---- Shared Tile Logic ----
@@ -90,6 +194,8 @@ def get_tile(asset_id, vis_params=None, start_date=None, end_date=None, bbox=Non
         return f"Error: {str(e)}"
 
 
+ee_initialize()
+
 # ---- FastAPI ----
 app = FastAPI()
 app.add_middleware(
@@ -116,18 +222,50 @@ def get_tile_api(req: TileRequest):
 
 
 # ---- Gradio UI ----
+def get_tile_gradio(asset_id, vis_params, start_date, end_date, bbox_str):
+    """Wrapper for Gradio that converts string inputs to proper types."""
+    # Convert empty strings to None
+    start_date = start_date.strip() if start_date and start_date.strip() else None
+    end_date = end_date.strip() if end_date and end_date.strip() else None
+
+    # Convert bbox string to list
+    bbox = None
+    if bbox_str and bbox_str.strip():
+        try:
+            bbox = [float(x.strip()) for x in bbox_str.split(",")]
+        except ValueError:
+            return "Error: bbox must be comma-separated numbers (west,south,east,north)"
+
+    return get_tile(asset_id, vis_params, start_date, end_date, bbox)
+
+
 gradio_ui = gr.Interface(
-    fn=get_tile,
+    fn=get_tile_gradio,
     inputs=[
         gr.Textbox(label="Earth Engine Asset ID", placeholder="e.g., USGS/SRTMGL1_003"),
         gr.Textbox(
             label="Visualization Parameters (JSON)",
             placeholder='{"min":0,"max":5000,"palette":"terrain"}',
         ),
+        gr.Textbox(
+            label="Start Date (Optional)",
+            placeholder="e.g., 2023-01-01",
+            value="",
+        ),
+        gr.Textbox(
+            label="End Date (Optional)",
+            placeholder="e.g., 2023-12-31",
+            value="",
+        ),
+        gr.Textbox(
+            label="Bounding Box (Optional)",
+            placeholder="e.g., -122.5,37.5,-122.0,38.0 (west,south,east,north)",
+            value="",
+        ),
     ],
     outputs="text",
     title="Earth Engine Tile URL Generator",
-    description="Supports ee.Image, ee.ImageCollection, ee.FeatureCollection. Tile URL is suitable for basemap usage.",
+    description="Supports ee.Image, ee.ImageCollection, ee.FeatureCollection with optional date range and bbox filtering. Tile URL is suitable for basemap usage.",
 )
 
 app = gr.mount_gradio_app(app, gradio_ui, path="/")