add example for data usage

geoengine/__init__.py

@@ -6,6 +6,7 @@
 
 from geoengine_openapi_client.exceptions import BadRequestException, OpenApiException, ApiTypeError, ApiValueError, \
     ApiKeyError, ApiAttributeError, ApiException, NotFoundException
+from geoengine_openapi_client import UsageSummaryGranularity
 from .auth import Session, get_session, initialize, reset
 from .colorizer import Colorizer, ColorBreakpoint, LinearGradientColorizer, PaletteColorizer, \
     LogarithmicGradientColorizer
@@ -30,7 +31,8 @@
     MultiBandRasterColorizer
 
 from .util import clamp_datetime_ms_ns
-from .workflow import WorkflowId, Workflow, workflow_by_id, register_workflow, get_quota, update_quota
+from .workflow import WorkflowId, Workflow, workflow_by_id, register_workflow, get_quota, update_quota, data_usage, \
+    data_usage_summary, plot_data_usage_summary
 from .raster import RasterTile2D
 from .raster_workflow_rio_writer import RasterWorkflowRioWriter
 

geoengine/workflow.py

@@ -967,3 +967,57 @@ def update_quota(user_id: UUID, new_available_quota: int, timeout: int = 60) ->
             ),
             _request_timeout=timeout
         )
+
+
+def data_usage(offset: int = 0, limit: int = 10) -> List[geoengine_openapi_client.DataUsage]:
+    '''
+    Get data usage
+    '''
+
+    session = get_session()
+
+    with geoengine_openapi_client.ApiClient(session.configuration) as api_client:
+        user_api = geoengine_openapi_client.UserApi(api_client)
+        response = user_api.data_usage_handler(offset=offset, limit=limit)
+
+        # create dataframe from response
+        usage_dicts = [data_usage.dict(by_alias=True) for data_usage in response]
+        df = pd.DataFrame(usage_dicts)
+
+    return df
+
+
+def data_usage_summary(granularity: geoengine_openapi_client.UsageSummaryGranularity, dataset: Optional[str] = None, offset: int = 0, limit: int = 10) -> List[geoengine_openapi_client.DataUsage]:
+    '''
+    Get data usage summary
+    '''
+
+    session = get_session()
+
+    with geoengine_openapi_client.ApiClient(session.configuration) as api_client:
+        user_api = geoengine_openapi_client.UserApi(api_client)
+        response = user_api.data_usage_summary_handler(dataset=dataset, granularity=granularity, offset=offset, limit=limit)
+
+        # create dataframe from response
+        usage_dicts = [data_usage.dict(by_alias=True) for data_usage in response]
+        df = pd.DataFrame(usage_dicts)
+
+    return df
+
+
+def plot_data_usage_summary(granularity: geoengine_openapi_client.UsageSummaryGranularity, dataset: Optional[str] = None, offset: int = 0, limit: int = 10):
+    import matplotlib.pyplot as plt
+    import pandas as pd
+
+    df = data_usage_summary(granularity, dataset, offset, limit)
+    df['timestamp'] = pd.to_datetime(df['timestamp']).dt.tz_localize(None)
+
+    pivot_df = df.pivot(index='timestamp', columns='dataset', values='count').fillna(0)
+    pivot_df.plot(kind='bar', figsize=(10, 6))
+
+    plt.title('Data Usage by Dataset over time')
+    plt.xlabel('Timestamp')
+    plt.ylabel('Count')
+    plt.xticks(rotation=45)
+    plt.legend(title='Dataset')
+    plt.show()

setup.cfg

@@ -18,7 +18,7 @@ package_dir =
 packages = find:
 python_requires = >=3.9
 install_requires =
-    geoengine-openapi-client == 0.0.17
+    geoengine-openapi-client @ git+https://github.com/geo-engine/openapi-client@esg-quota#subdirectory=python # TODO update when merged
     geopandas >=0.9,<0.15
     matplotlib >=3.5,<3.8
     numpy >=1.21,<2