"""Client-side helpers for routing the Earth Engine Python SDK through a
token-injecting proxy.
The pattern:
1. Run a proxy server (see ``geeViz.eeAuth.server``) that holds the SA
credentials and substitutes the right bearer token per request based
on a tenant header / query param.
2. Tell the EE SDK to send all REST calls through that proxy instead of
directly to Google. Pass anonymous credentials — the proxy supplies
the real ones.
3. Switch tenants on the client side by setting a ``ContextVar``; the
custom HTTP transport reads it and stamps the routing header on
every outbound request.
That gives you full multi-tenant concurrency in a single Python process,
which the bare EE SDK can't do because ``ee.Initialize()`` stores
credentials in module-global state.
Quick start
-----------
::
from geeViz.eeAuth import initialize_via_proxy, tenant_context
import ee
initialize_via_proxy("http://localhost:8888/ee-api")
# Now ee.X calls go through the proxy with the default tenant
with tenant_context("training"):
ee.Image(1).getInfo() # uses the training SA
Attribution context
-------------------
Routing says *which credential* pays; attribution says *who* to charge
it to. Four ContextVars carry that second answer to the proxy:
:data:`CURRENT_USER_EMAIL`, :data:`CURRENT_SESSION_ID`,
:data:`CURRENT_ACTION` and :data:`CURRENT_BILLING_TENANT`. Whatever is
set when an EE call goes out is stamped onto the outbound request as
``X-Agent-*`` headers, and the proxy's workload-tag builder folds those
parts into the tag it mints — so the resulting EE spend can be traced
back to a real person, session and operation.
Setting them is optional and additive. A standalone geeViz install
leaves all four empty and gets exactly the historic behavior: no
attribution headers, and a tag built from credential/pid/source alone.
They exist for hosts that already know the caller's identity (an agent's
MCP tool wrapper, a web backend) and would otherwise see every
programmatic EE call — thumbnails, GIFs, charts, zonal stats — collapse
into ``anonymous``, because those requests carry no browser Referer,
session cookie or IAP header to identify them by.
``GEEVIZ_AGENT_ATTRIB_SECRET``
------------------------------
The attribution headers are sent **only** when the
``GEEVIZ_AGENT_ATTRIB_SECRET`` environment variable is set, and the
value is sent alongside them as ``X-Agent-Attrib-Secret``. The gate is a
security boundary, not a convenience toggle: those headers name who gets
billed for EE compute, and the proxy's ``/ee-api`` path skips auth for
callers that appear to be on loopback. Under
``forwarded_allow_ips="*"`` "appears to be loopback" is a
client-controlled judgement — a remote caller can present as loopback
via ``X-Forwarded-For`` — so an ungated header would let anyone bill
their usage to a victim's account. The server that spawns geeViz puts
the secret in the child's environment; when it is absent we send no
attribution headers at all and the call attributes to nobody, which is
the safe failure.
"""
from __future__ import annotations
import contextvars
import logging
import os
import sys
from contextlib import contextmanager
from typing import Optional
logger = logging.getLogger(__name__)
#: Email of the user the current EE call is being made on behalf of.
#: Propagated from the calling agent to the outbound EE proxy: set by the
#: MCP tool wrapper (or the caller of ``initialize_via_proxy`` in a
#: plain-Python setup) from the same workload-tag builder the agent uses.
#: ``TenantAwareHttp`` reads it on every outbound request and stamps it as
#: an HTTP header so the proxy's own workload-tag builder can attribute the
#: call to the right user even when the request has no browser Referer,
#: session cookie, or IAP header (i.e., every MCP-generated EE call — thumb,
#: gif, chart, zonal stats, etc.). Without this thread-through, those calls
#: all collapsed into ``anonymous`` in Cloud Monitoring and were invisible
#: in the per-user spend table.
CURRENT_USER_EMAIL: contextvars.ContextVar[str] = contextvars.ContextVar(
"current_user_email", default=""
)
#: Session the current EE call belongs to — the companion to
#: :data:`CURRENT_USER_EMAIL`, stamped by the same code path. Lets spend be
#: sliced per conversation/run, not just per user. Empty means no session
#: known; the tag is still minted, just without a session part.
CURRENT_SESSION_ID: contextvars.ContextVar[str] = contextvars.ContextVar(
"current_session_id", default=""
)
#: ContextVar holding the current tenant id. ``TenantAwareHttp`` reads it
#: on every outbound EE REST call and stamps the routing header. Empty
#: string means "use the proxy's default tenant" (no header sent).
CURRENT_TENANT: contextvars.ContextVar[str] = contextvars.ContextVar(
"geeViz_ee_tenant", default="",
)
#: Billing tenant slug — distinct from :data:`CURRENT_TENANT` (which is the
#: eeCreds credential-name, e.g. "ee-persistent" for ADC or "acme-prod"
#: for a named SA). Agents set this to their fixed deploy-time tenant
#: (e.g. "askterra") so the workload-tag mint carries the right
#: attribution regardless of which cred is currently active. Empty means
#: no separate billing tenant known — the default builder falls back to
#: :data:`CURRENT_TENANT` (matches standalone-geeViz behavior).
CURRENT_BILLING_TENANT: contextvars.ContextVar[str] = contextvars.ContextVar(
"geeViz_billing_tenant", default="",
)
#: Action label for the current EE call (e.g. "map", "compute", "thumb",
#: "export"). Set by the caller so ``_default_workload_tag_builder`` can
#: emit that as the ``action`` part instead of a generic "proxy-default"
#: marker. Empty means no action known — builder still mints, just
#: without an action label.
CURRENT_ACTION: contextvars.ContextVar[str] = contextvars.ContextVar(
"geeViz_ee_action", default="",
)
# Header name the proxy looks at to pick a credential. Matches the
# default in ``geeViz.eeAuth.server`` — both sides must agree on the
# string for routing to work. Override at init time if you've deployed
# the proxy with a different header (e.g. the AskTerra agent uses
# ``X-AskTerra-Tenant``).
DEFAULT_TENANT_HEADER = "X-geeViz-Creds"
[docs]
def set_tenant(tenant: str):
"""Set the current tenant for subsequent EE calls in this context.
Returns a token that can be passed to ``reset_tenant`` to restore
the previous value. Prefer ``tenant_context()`` for scoped use.
"""
return CURRENT_TENANT.set(tenant or "")
[docs]
def reset_tenant(token) -> None:
"""Restore the tenant to what it was before the matching ``set_tenant`` call.
Args:
token: The token object returned by the paired ``set_tenant()`` call.
Passing a token from a different scope raises ``ValueError``.
"""
CURRENT_TENANT.reset(token)
[docs]
@contextmanager
def tenant_context(tenant: str):
"""Scoped tenant switch.
.. code-block:: python
with tenant_context("training"):
ee.Image(1).getInfo()
# back to previous tenant here
"""
token = set_tenant(tenant)
try:
yield
finally:
reset_tenant(token)
[docs]
class TenantAwareHttp:
"""``httplib2.Http`` subclass that stamps the tenant header on every
outbound request and strips whatever Authorization the SDK injected.
Subclassed at first instantiation so ``httplib2`` is only imported
when actually used — keeps unit tests that mock the EE init path
from needing the dependency. The header name is set per-instance so
you can run multiple proxies with different conventions in the same
process if you really need to.
Thread safety
-------------
``httplib2.Http`` is NOT thread-safe — its per-host connection cache
(``self.connections``) is a plain ``dict`` mutated from inside
``request()`` and ``socket.HTTPConnection`` objects hold per-instance
socket state. When the EE SDK shares one transport across a
``ThreadPoolExecutor`` (as ``Map.testLayers()`` does with 8 workers),
concurrent threads tear down each other's sockets mid-request,
surfacing as ``'NoneType' object has no attribute 'close'`` and
Windows ``WinError 10038/10057`` socket errors.
Workaround: route each thread's ``request()`` call to its OWN
``httplib2.Http`` instance stored in ``threading.local()``. EE's SDK
only consults the transport for ``request()`` — it doesn't reach into
``self.connections`` directly — so per-thread instances are a safe
drop-in.
Attribution headers
-------------------
Besides the tenant header, ``request()`` stamps the ``X-Agent-*``
attribution headers built from this module's four ContextVars — but
only when ``GEEVIZ_AGENT_ATTRIB_SECRET`` is set in the environment,
in which case its value goes out as ``X-Agent-Attrib-Secret`` to
prove the headers came from a process the server itself spawned.
Those headers decide who gets billed for EE compute and who has CDU
caps consumed, and the proxy's ``/ee-api`` path treats
apparently-loopback callers as trusted — a judgement the client can
influence via ``X-Forwarded-For`` when uvicorn runs with
``forwarded_allow_ips="*"``. Ungated, that would let a remote caller
bill their usage to anyone they can name. With no secret present we
send no attribution headers and the call attributes to nobody, which
is the safe failure. See the module docstring for the full rationale.
"""
_impl_cls = None # lazily-defined subclass keyed by header name
def __new__(cls, tenant_header: str = DEFAULT_TENANT_HEADER):
if cls._impl_cls is None:
import httplib2
import threading as _threading
class _Impl(httplib2.Http):
_tenant_header = tenant_header
_thread_local = _threading.local()
def _thread_http(self):
"""Lazy per-thread ``httplib2.Http`` so concurrent
SDK calls don't share socket/connection state."""
h = getattr(self.__class__._thread_local, "http", None)
if h is None:
h = httplib2.Http()
self.__class__._thread_local.http = h
return h
def request(self, uri, method="GET", body=None, headers=None, **kw):
"""Forward an EE SDK HTTP request through the proxy.
Runs the actual HTTP via a per-thread ``httplib2.Http``
to avoid the SDK's thread-unsafe connection cache.
Rewrites the request to:
* Strip any ``Authorization`` header the SDK added
(proxy stamps its own).
* Add the tenant header from the current ContextVar.
* Add user + session workload-tag headers when set.
Args:
uri: Upstream URL the SDK wants to hit.
method: HTTP verb. Default ``"GET"``.
body: Request body bytes.
headers: Dict of headers to forward (after rewrite).
**kw: Passed through to ``httplib2.Http.request``.
Returns:
Same ``(response, content)`` tuple ``httplib2``
normally returns.
"""
headers = dict(headers or {})
# Strip SDK-injected auth — proxy substitutes its own
headers.pop("Authorization", None)
headers.pop("authorization", None)
# Stamp current tenant
tenant = CURRENT_TENANT.get()
if tenant:
headers[self._tenant_header] = tenant
# Stamp user + session so the proxy's workload-tag
# builder has an attribution source when there's no
# Referer (MCP-initiated calls: thumb, gif, chart,
# zonal stats). Values come from ContextVars the
# MCP tool wrapper set from the agent's ``_workload_tag``
# argument on each tool call. Both header names are
# ``X-Agent-*`` so they're clearly distinct from
# any browser/IAP auth headers the proxy already
# inspects.
#
# ``X-Agent-Attrib-Secret`` proves these headers came
# from a process the server itself spawned. Without
# it the proxy MUST NOT believe them: it bills EE
# spend and consumes CDU caps against whatever email
# they name, and the proxy path skips auth for
# apparently-loopback peers — which is a
# client-controlled judgement when uvicorn runs with
# forwarded_allow_ips="*". The server puts the value
# in our environment when it spawns us; when it's
# absent we send no attribution headers at all and
# the call attributes to nobody, which is the safe
# failure.
secret = os.environ.get("GEEVIZ_AGENT_ATTRIB_SECRET", "")
if secret:
user = CURRENT_USER_EMAIL.get()
if user:
headers["X-Agent-User-Email"] = user
session = CURRENT_SESSION_ID.get()
if session:
headers["X-Agent-Session-ID"] = session
headers["X-Agent-Attrib-Secret"] = secret
return self._thread_http().request(
uri, method, body, headers, **kw
)
cls._impl_cls = _Impl
return cls._impl_cls()
[docs]
def initialize_via_proxy(
proxy_url: str,
tenant_header: str = DEFAULT_TENANT_HEADER,
project: Optional[str] = None,
origin: str = "",
) -> bool:
"""Initialize the Earth Engine SDK to route all REST calls through
``proxy_url``.
Uses ``AnonymousCredentials`` since the proxy holds the real SA
credentials. The SDK's bearer-token header is stripped by
``TenantAwareHttp`` before reaching the proxy anyway.
Args:
proxy_url: Base URL of the EE proxy, e.g.
``"http://localhost:8888/ee-api"``. No trailing slash.
tenant_header: Header name the proxy expects for tenant routing.
Default ``X-geeViz-Creds`` matches ``geeViz.eeAuth.server``.
project: Placeholder project id passed to ``ee.Initialize`` (EE
requires one but the proxy overrides per-tenant via
``x-goog-user-project``). Default
``"ee-proxy-placeholder"``.
origin: Human-readable provenance label appended to the
"EE initialized via proxy" line printed on stderr — today
``"spawned new detached proxy"`` or ``"attached to existing
detached proxy"``. Purely diagnostic; it changes no
behavior. Default ``""`` omits the label entirely. Worth
passing: without it the line reads identically whether a
proxy was started or merely reused, which is how a
kill-and-respawn on a different port looked like "two
proxies running" in a notebook.
Returns:
True on success, False if init failed (caller should fall back
to direct ``ee.Initialize`` or surface the error). Prints any
underlying exception to stderr; doesn't re-raise.
"""
try:
import ee
from google.auth.credentials import AnonymousCredentials
ee.Initialize(
credentials=AnonymousCredentials(),
url=proxy_url.rstrip("/"),
http_transport=TenantAwareHttp(tenant_header=tenant_header),
project=project or "ee-proxy-placeholder",
)
# ``origin`` says whether this proxy was newly SPAWNED, an
# existing one we ATTACHED to, or a re-point of an already-live
# session. Without it the line reads identically in all three
# cases, which is how a kill-and-respawn on a different port
# looked like "two proxies running" in a notebook.
print(
f"[geeViz.eeAuth] EE initialized via proxy: {proxy_url} "
f"({origin}, tenant_header={tenant_header})"
if origin else
f"[geeViz.eeAuth] EE initialized via proxy: {proxy_url} "
f"(tenant_header={tenant_header})",
file=sys.stderr,
)
return True
except Exception as e:
msg = str(e)
# Pull the offending project out of the standard EE
# USER_PROJECT_DENIED message so we can name it in the hint.
import re as _re
m_proj = _re.search(
r"permission to use project ([A-Za-z0-9._-]+)", msg
)
# GCP project IDs can contain ``-`` but not trailing ``.``;
# the regex above happens to gobble a sentence-ending period
# from the EE error text, so trim it.
denied_project = m_proj.group(1).rstrip(".") if m_proj else None
# The classic "OAuth credentials have no project" failure mode:
# EE routes the call to earthengine-legacy as the consumer
# project, and the personal Google account can't use that.
if (denied_project == "earthengine-legacy"
and "USER_PROJECT_DENIED" in msg):
print(
"[geeViz.eeAuth] proxy init failed: OAuth credentials have "
"no project set, so EE routed to 'earthengine-legacy' "
"which your account can't use.\n"
" Fix: tell the library which GCP project to bill against "
"via ONE of:\n"
" 1. eeCreds.addCreds(..., project='YOUR_PROJECT')\n"
" 2. export GOOGLE_CLOUD_PROJECT=YOUR_PROJECT\n"
" 3. gcloud config set project YOUR_PROJECT\n"
" 4. ee.Initialize(project='YOUR_PROJECT') # before importing geeViz\n"
" Your project must have the Earth Engine API enabled and "
"be registered for EE.",
file=sys.stderr,
)
# SA registered with a project_id its identity doesn't have
# ``serviceusage.services.use`` on. Discovery 403s but direct
# EE calls (value:compute, etc.) may still work — this is a real
# IAM gap, not a library bug. Surface it actionably.
elif (denied_project
and "USER_PROJECT_DENIED" in msg
and "serviceusage" in msg.lower()):
print(
f"[geeViz.eeAuth] proxy init: the credential can't use "
f"project {denied_project!r} for service-usage discovery.\n"
f" This usually means the service account is missing "
f"`roles/serviceusage.serviceUsageConsumer` on "
f"{denied_project!r}. Fix by EITHER:\n"
f" 1. Granting that role to the SA in GCP Console → "
f"IAM for project {denied_project!r}, OR\n"
f" 2. Overriding the project at registration:\n"
f" eeCreds.addCreds(..., project='OTHER_PROJECT')\n"
f" (use a project the SA HAS got serviceusage.services.use on).\n"
f" Some EE calls may still succeed (they take a different "
f"code path), but discovery-driven operations will fail.",
file=sys.stderr,
)
else:
print(
f"[geeViz.eeAuth] proxy init failed: {type(e).__name__}: {e}",
file=sys.stderr,
)
return False