"""Earth Engine workload-tag helpers.
EE workload tags surface in GCP Billing under the
``goog-earth-engine-workload-tag`` label, so tagged calls can be broken
down per user / session / source. This module builds well-formed tags
from arbitrary string parts (sanitizing each so the result is always
accepted by EE).
EE constraints (from ``ee/_state.py`` validation):
- 1 - 63 characters
- begins and ends with a lowercase alphanumeric ``[a-z0-9]``
- middle characters: ``[a-z0-9_-]`` (lowercase alphanumeric, dash, underscore)
No uppercase, no ``.``, no other punctuation. Anything outside that set
(``@``, spaces, slashes, dots, uppercase, etc.) gets sanitized to ``-``.
**Separator: ``__`` (double underscore).** Single ``-`` already appears
inside sanitized parts (e.g. ``ihousman-redcastleresources-com``), so we
reserve double underscore as the between-parts delimiter. That makes
tags trivially parseable with ``tag.split("__")``::
agent__run_code__ihousman-redcastleresources-com__db208a06-1c49
To keep ``__`` an unambiguous separator, runs of ``_`` *within* a part
get collapsed to a single ``_`` during sanitization (so an input like
``run__code`` becomes ``run_code``). Underscores from sources like tool
names — ``run_code``, ``map_control`` — pass through intact because
they're already singletons.
Tag stores
----------
63 characters isn't enough to spell out a real identity tuple (user
email + session id + action + tenant), so :func:`mint_workload_tag`
hashes the parts into a short ``wl_<hex>`` tag instead. That tag is not
reversible on its own — recovering who a billing row belongs to needs
the ``tag -> parts`` mapping written down somewhere. That "somewhere" is
this module's store layer:
- :class:`TagStore` — the ``put`` / ``lookup`` Protocol every store
implements. Implementations must be thread-safe and idempotent on
``put``.
- :class:`InMemoryTagStore` — process-local dict. Fine for one-shot
scripts and tests; the mapping dies with the process.
- :class:`SQLiteTagStore` — single file, typically
``~/.geeViz/workload_tags.db``. Survives kernel restarts and reruns.
Not suitable for multi-instance deployments, where each instance
would keep its own file and miss the others' tags.
- :class:`ChainedTagStore` — several stores treated as one.
- :func:`default_tag_store` — the store used when nobody configured one.
:class:`ChainedTagStore` is the newest piece and the least obvious,
because its two operations deliberately disagree about scope:
- ``put`` writes to the **primary (first) store only**. Attribution
needs exactly one writer; fanning writes out to every backend would
create diverging partial copies and turn a lookup miss into "which
copy is right?".
- ``lookup`` reads from **every backend in order**, first hit wins (a
failing backend is logged and skipped rather than breaking the
chain). Reads have no such problem — tags are content-addressed, so
the same parts and secret always produce the same tag and it does not
matter which backend answered.
That asymmetry exists because more than one store can be live on a
machine at once — an agent writing to Postgres beside a notebook using
the sqlite default. Without chaining, a tag minted under one is simply
invisible to the other's ``lookup``, and the caller gets a bare ``None``
that reads as "never minted" when the truth is "you asked the wrong
store".
"""
from __future__ import annotations
import hashlib
import json
import logging
import os
import re
import sqlite3
import threading
from pathlib import Path
from typing import Any, Optional, Protocol, runtime_checkable
logger = logging.getLogger(__name__)
_TAG_ALLOWED_CHAR = re.compile(r"[^a-z0-9_\-]")
_TAG_MAX_LEN = 63
SEPARATOR = "__"
_TAG_PREFIX = "wl"
[docs]
def sanitize_workload_tag_part(s: str) -> str:
"""Sanitize a single component of a workload tag.
- Lowercases.
- Replaces disallowed characters with ``-``.
- Collapses runs of ``-`` to a single ``-``.
- Collapses runs of ``_`` to a single ``_`` so the ``__`` separator
stays unambiguous when parts are joined.
- Strips leading/trailing ``-`` and ``_`` (EE rejects tags that don't
begin and end with an alphanumeric).
"""
if not s:
return ""
s = s.lower()
s = _TAG_ALLOWED_CHAR.sub("-", s)
s = re.sub(r"-{2,}", "-", s)
s = re.sub(r"_{2,}", "_", s)
s = s.strip("-_")
return s
[docs]
def build_workload_tag(*parts: str) -> str:
"""Join sanitized parts with ``__`` and clamp to EE's 63-char limit.
Empty / falsy parts are dropped. The final tag is guaranteed to satisfy
EE's regex: ``[a-z0-9][a-z0-9_\\-]{0,61}[a-z0-9]``. Returns an empty
string if everything was dropped — callers should treat empty as "no
tag" and skip the workload-tag header / body field entirely.
"""
clean = [sanitize_workload_tag_part(p) for p in parts]
clean = [c for c in clean if c]
if not clean:
return ""
tag = SEPARATOR.join(clean)[:_TAG_MAX_LEN]
# Re-strip in case the truncation left a dangling separator at the end.
tag = tag.rstrip("-_")
return tag
# ---------------------------------------------------------------------------
# Deterministic short-tag minting + reversible lookup
# ---------------------------------------------------------------------------
#
# EE workload tags are capped at 63 chars, which isn't enough to encode a
# realistic identity tuple (user email + session id + action + tenant). The
# recovery pattern is:
#
# tag = mint_workload_tag(parts, secret) # short deterministic hash
# store.put(tag, parts) # remember the mapping
# ...
# parts = store.lookup(tag) # get parts back later
#
# `mint_workload_tag` is deterministic: same (parts, secret) always yields
# the same tag, so re-minting during a session collapses onto the same row
# rather than growing the store.
def _canonical_parts(parts: dict[str, Any]) -> str:
"""Stable serialization of a parts dict — sorted by key, string-cast
values. So {'user':'x','tenant':'a'} and {'tenant':'a','user':'x'}
hash to the SAME tag."""
if not isinstance(parts, dict):
raise TypeError(f"parts must be a dict, got {type(parts).__name__}")
items = sorted((str(k), "" if v is None else str(v)) for k, v in parts.items())
return "|".join(f"{k}={v}" for k, v in items)
[docs]
def mint_workload_tag(
parts: dict[str, Any], *, secret: str, digest_size: int = 8
) -> str:
"""Deterministic short tag from a parts dict + secret.
Returns ``wl_<hex>`` where the hex is a ``digest_size``-byte blake2b of
the canonicalized parts (default 16 hex chars → collision probability
~10⁻⁹ at millions of tags). Same input always yields the same tag.
The tag is NOT reversible on its own — pair with a ``TagStore`` that
records ``tag → parts`` at mint time so lookups can recover identity
later.
"""
if not secret:
raise ValueError("mint_workload_tag: secret is required")
canonical = _canonical_parts(parts)
h = hashlib.blake2b(
(canonical + "|" + secret).encode("utf-8"),
digest_size=digest_size,
).hexdigest()
tag = f"{_TAG_PREFIX}_{h}"
# Belt-and-suspenders: minted tags are guaranteed valid but pass them
# through build_workload_tag anyway so any future format tweak stays
# consistent with the sanitizer.
return build_workload_tag(tag)
# ---------------------------------------------------------------------------
# TagStore Protocol + default impls
# ---------------------------------------------------------------------------
import contextlib
@contextlib.contextmanager
def _sqlite_conn(path: str | Path):
"""Open a short-lived sqlite3 connection and guarantee close on exit.
``with sqlite3.connect(...)`` commits on exit but does NOT close the
connection — on Windows that leaves the file locked and breaks
``TemporaryDirectory`` cleanup and any process that wants to reopen.
"""
conn = sqlite3.connect(str(path), isolation_level=None)
try:
yield conn
finally:
conn.close()
[docs]
def default_tag_store() -> TagStore:
"""Return the process-wide default store. First call constructs
``SQLiteTagStore()`` at ``~/.geeViz/workload_tags.db``. Callers who
want in-memory or a custom store should set it explicitly via
``eeCreds.setTagStore(...)`` before ``eeCreds.start()``."""
global _DEFAULT_STORE
if _DEFAULT_STORE is None:
with _DEFAULT_STORE_LOCK:
if _DEFAULT_STORE is None:
_DEFAULT_STORE = SQLiteTagStore()
return _DEFAULT_STORE
_DEFAULT_STORE: Optional[TagStore] = None
_DEFAULT_STORE_LOCK = threading.Lock()
def _default_secret() -> str:
"""Read the workload-tag secret from ``WORKLOAD_TAG_SECRET`` env var.
Falls back to a per-machine constant so mints are stable across
process restarts on the same box; this fallback is fine for local
dev / notebooks but production callers should set the env var."""
env = os.environ.get("WORKLOAD_TAG_SECRET")
if env:
return env
return "geeviz-workload-tag-local-dev-secret"