export area of interest, track state with sqlite (#3)

* add command to export area of interst, track state in sqlite db

* set task id in db

* add simple task tracking with sqlite

* remove job_name

Author: geospatial-jeff

.gitignore

@@ -208,4 +208,7 @@ __marimo__/
 
 # uv
 .python-version
-uv.lock
+uv.lock
+
+# sqlite
+*.db

aef_export/cli.py

@@ -1,8 +1,10 @@
 import click
+import json
 
-from aef_export.embeddings import export_image
+from aef_export.embeddings import export_image, export_aoi
 from aef_export.coverage import export_image_collection
 from aef_export.settings import get_settings
+from aef_export.task_tracking import update_db_state, get_task_summary, BillingTier
 from aef_export.utils import initialize_ee
 
 
@@ -54,3 +56,56 @@ def image(
     initialize_ee(settings.google_cloud_project)
     task_id = export_image(image_id, gcs_bucket_name, gcs_key_prefix, quantize)
     click.echo(f"Task id: {task_id}")
+
+
+@app.command()
+@click.argument(
+    "geojson_filepath", type=click.Path(exists=True, file_okay=True, dir_okay=False)
+)
+@click.argument("bq_dataset_name")
+@click.argument("bq_table_name")
+@click.argument("gcs_bucket_name")
+@click.option("--limit", type=int, required=False, default=None)
+def aoi(
+    geojson_filepath: str,
+    bq_dataset_name: str,
+    bq_table_name: str,
+    gcs_bucket_name: str,
+    limit: int | None = None,
+):
+    settings = get_settings()
+    initialize_ee(settings.google_cloud_project)
+
+    with open(geojson_filepath) as f:
+        data = json.load(f)
+
+    if data["type"] == "Feature":
+        polygon = data["geometry"]
+    elif data["type"] == "Polygon":
+        polygon = data
+    else:
+        raise ValueError(
+            "Input file must be a geojson polygon, either geometry or feature"
+        )
+
+    export_aoi(polygon, bq_dataset_name, bq_table_name, gcs_bucket_name, limit)
+
+
+@app.group()
+def db():
+    pass
+
+
+@db.command()
+def update_task_status():
+    settings = get_settings()
+    initialize_ee(settings.google_cloud_project)
+    update_db_state()
+
+
+@db.command()
+@click.option(
+    "--billing-tier", type=click.Choice(BillingTier), default=BillingTier.tier1
+)
+def summarize(billing_tier: BillingTier = BillingTier.tier1):
+    click.echo(json.dumps(get_task_summary(billing_tier)))

aef_export/coverage.py

@@ -1,5 +1,9 @@
-import ee
 import uuid
+from typing import Generator
+
+import ee
+from google.cloud import bigquery
+from shapely.geometry import shape
 
 from aef_export.utils import set_workload_tag
 
@@ -77,3 +81,27 @@ def export_image_collection(
         task.start()
 
     return task.id
+
+
+def query_coverage(
+    geojson_geometry: dict,
+    bq_dataset_name: str,
+    bq_table_name: str,
+    limit: int | None = None,
+) -> Generator[dict, None, None]:
+    geom = shape(geojson_geometry)
+
+    bq_client = bigquery.Client()
+    query = f"""
+        SELECT
+            `system:id` as system_id,
+            SPLIT(start_date, '-')[OFFSET(0)] as year,
+            UTM_ZONE as utm_zone
+        FROM {bq_dataset_name}.{bq_table_name}
+        WHERE ST_Intersects(ST_GEOGFROMTEXT('{geom.wkt}'), geo)
+    """
+    if limit:
+        query += f" LIMIT {limit}"
+    query_job = bq_client.query(query)
+    for row in query_job.result():
+        yield dict(row)

aef_export/embeddings.py

@@ -1,6 +1,8 @@
 import ee
 
 from aef_export.utils import set_workload_tag
+from aef_export.sqlite import init_database, Row, insert_row
+from aef_export.coverage import query_coverage
 
 
 def _quantize_embeddings(image: ee.Image) -> ee.Image:
@@ -71,3 +73,40 @@ def export_image(
         task.start()
 
     return task.id
+
+
+def export_aoi(
+    geojson_geometry: dict,
+    bq_dataset_name: str,
+    bq_table_name: str,
+    gcs_bucket_name: str,
+    limit: int | None = None,
+):
+    init_database()
+
+    rows_to_export = query_coverage(
+        geojson_geometry, bq_dataset_name, bq_table_name, limit
+    )
+    for row in rows_to_export:
+        # Build the key prefix.
+        system_id = row["system_id"]
+        year = row["year"]
+        utm_zone = row["utm_zone"]
+        key_prefix = "/".join([year, utm_zone, system_id.split("/")[-1]]) + "/"
+
+        # Start the export.
+        # TODO: Protect against 3000+ tasks in the queue.
+        task_id = export_image(system_id, gcs_bucket_name, key_prefix, quantize=True)
+        print("Submitting task: ", task_id)
+
+        # Insert record of this row into sqlite
+        row = Row(
+            task_id=task_id,
+            eecu_seconds=None,
+            runtime_seconds=None,
+            status="queued",
+            image_id=system_id,
+            year=year,
+            s3_path=f"s3://{gcs_bucket_name}/{key_prefix}",
+        )
+        insert_row(row)

aef_export/sqlite.py

@@ -0,0 +1,77 @@
+from dataclasses import dataclass, asdict
+import sqlite3
+
+
+DATABASE_NAME = "sqlite_aef_export.db"
+
+
+@dataclass
+class Row:
+    task_id: str
+    eecu_seconds: float | None
+    runtime_seconds: float | None
+    status: str
+    image_id: str
+    year: int
+    s3_path: str
+
+
+def get_connection():
+    return sqlite3.connect(DATABASE_NAME)
+
+
+def init_database():
+    cur = get_connection()
+    cur.execute("""
+        CREATE TABLE IF NOT EXISTS exports(
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            task_id VARCHAR,
+            eecu_seconds FLOAT,
+            runtime_seconds FLOAT,
+            status VARCHAR,
+            image_id VARCHAR,
+            year INTEGER,
+            s3_path VARCHAR
+        )
+    """)
+
+
+def insert_row(row: Row):
+    d = asdict(row)
+    columns = ", ".join(d.keys())
+    placeholders = ", ".join(["?" for _ in d.values()])
+    query = f"INSERT INTO exports ({columns}) VALUES ({placeholders})"
+
+    cur = get_connection()
+    cur.execute(query, tuple(d.values()))
+    cur.commit()
+
+
+def update_row(
+    task_id: str,
+    status: str,
+    eecu_seconds: float | None = None,
+    runtime_seconds: float | None = None,
+):
+    query = "UPDATE exports SET status = ?, eecu_seconds = ?, runtime_seconds = ? WHERE task_id = ?"
+    cur = get_connection()
+    cur.execute(query, (status, eecu_seconds, runtime_seconds, task_id))
+    cur.commit()
+
+
+def get_summary():
+    query = """
+        SELECT
+            status,
+            COUNT(*) as count,
+            SUM(eecu_seconds) as eecu_seconds,
+            AVG(runtime_seconds) as avg_runtime_seconds
+        FROM exports
+        GROUP BY status;
+    """
+    conn = get_connection()
+    conn.row_factory = sqlite3.Row
+    cur = conn.cursor()
+    cur.execute(query)
+    resp = cur.fetchall()
+    return [dict(row) for row in resp]

aef_export/task_tracking.py

@@ -0,0 +1,57 @@
+from datetime import datetime
+from enum import Enum
+from typing import Generator
+
+import ee
+
+from aef_export.sqlite import update_row, get_summary
+
+
+class BillingTier(Enum):
+    tier1 = "tier1"
+    tier2 = "tier2"
+    tier3 = "tier3"
+
+
+def list_tasks() -> Generator[dict, None, None]:
+    tasks = ee.data.listOperations()
+    for task in tasks:
+        if task["metadata"]["type"] != "EXPORT_IMAGE":
+            continue
+        yield task
+
+
+def update_db_state():
+    for task in list_tasks():
+        task_state = task["metadata"]["state"]
+        task_id = task["name"].split("/")[-1]
+        eecu_seconds = None
+        duration_seconds = None
+        if task_state == "SUCCEEDED":
+            eecu_seconds = task["metadata"]["batchEecuUsageSeconds"]
+            start_time = datetime.fromisoformat(task["metadata"]["startTime"])
+            end_time = datetime.fromisoformat(task["metadata"]["endTime"])
+            duration_seconds = (end_time - start_time).total_seconds()
+
+        update_row(task_id, task_state, eecu_seconds, duration_seconds)
+
+
+def get_task_summary(tier: BillingTier) -> list[dict]:
+    billing_tiers = {
+        BillingTier.tier1: 0.40,
+        BillingTier.tier2: 0.28,
+        BillingTier.tier3: 0.16,
+    }
+
+    summary = get_summary()
+
+    out_rows = []
+    for row in summary:
+        if eecu_seconds := row.get("eecu_seconds"):
+            eecu_hours = eecu_seconds / 3600
+            row["compute_cost"] = billing_tiers[tier] * eecu_hours
+
+        if avg_runtime_seconds := row.pop("avg_runtime_seconds"):
+            row["avg_runtime_minutes"] = avg_runtime_seconds / 60
+        out_rows.append(row)
+    return out_rows

pyproject.toml

@@ -16,9 +16,11 @@ build-backend = "setuptools.build_meta"
 
 [dependency-groups]
 dev = [
+    "google-cloud-bigquery>=3.37.0",
     "pre-commit>=4.3.0",
     "pytest>=8.4.2",
     "ruff>=0.12.12",
+    "shapely>=2.1.1",
 ]
 
 [project.scripts]