Feature/sc 488704/support labeldescription in raster loader from user input

raster_loader/cli/bigquery.py

@@ -4,8 +4,9 @@
 import click
 from functools import wraps, partial
 
-from raster_loader.utils import get_default_table_name
+from raster_loader.lib.utils import get_default_table_name
 from raster_loader.io.bigquery import BigQueryConnection, AccessTokenCredentials
+from raster_loader.lib.valuelabels import validate_band_valuelabels
 
 
 def catch_exception(func=None, *, handle=Exception):
@@ -101,6 +102,18 @@ def bigquery(args=None):
     type=int,
     default=6,
 )
+@click.option(
+    "--band-valuelabels",
+    help="Custom data for valuelabels in JSON format, or 'None'. "
+    "i.e: '{<value_1>: <label_1>, <value_2>: <label_2>, ...}'. "
+    "Could repeat --band-valuelabels to specify multiple bands data. "
+    "They will be considered in the order they appear in the file. "
+    "Note that you can set any value to 'None' to omit valuelabels for that band. ",
+    type=str,
+    default=[],
+    multiple=True,
+    callback=validate_band_valuelabels,
+)
 @catch_exception()
 def upload(
     file_path,
@@ -119,6 +132,7 @@ def upload(
     exact_stats=False,
     basic_stats=False,
     compression_level=6,
+    band_valuelabels=[],
 ):
     from raster_loader.io.common import (
         get_number_of_blocks,
@@ -194,6 +208,7 @@ def upload(
         basic_stats=basic_stats,
         compress=compress,
         compression_level=compression_level,
+        band_valuelabels=band_valuelabels,
     )
 
     click.echo("Raster file uploaded to Google BigQuery")

raster_loader/cli/databricks.py

@@ -4,8 +4,9 @@
 import click
 from functools import wraps, partial
 
-from raster_loader.utils import get_default_table_name
+from raster_loader.lib.utils import get_default_table_name
 from raster_loader.io.databricks import DatabricksConnection
+from raster_loader.lib.valuelabels import validate_band_valuelabels
 
 
 def catch_exception(func=None, *, handle=Exception):
@@ -115,6 +116,18 @@ def databricks(args=None):
     type=int,
     default=6,
 )
+@click.option(
+    "--band-valuelabels",
+    help="Custom data for valuelabels in JSON format, or 'None'. "
+    "i.e: '{<value_1>: <label_1>, <value_2>: <label_2>, ...}'. "
+    "Could repeat --band-valuelabels to specify multiple bands data. "
+    "They will be considered in the order they appear in the file. "
+    "Note that you can set any value to 'None' to omit valuelabels for that band. ",
+    type=str,
+    default=[],
+    multiple=True,
+    callback=validate_band_valuelabels,
+)
 @catch_exception()
 def upload(
     server_hostname,
@@ -136,6 +149,7 @@ def upload(
     exact_stats=False,
     basic_stats=False,
     compression_level=6,
+    band_valuelabels=[],
 ):
     from raster_loader.io.common import (
         get_number_of_blocks,
@@ -213,6 +227,7 @@ def upload(
         basic_stats=basic_stats,
         compress=compress,
         compression_level=compression_level,
+        band_valuelabels=band_valuelabels,
     )
 
     click.echo("Raster file uploaded to Databricks")

raster_loader/cli/snowflake.py

@@ -4,8 +4,9 @@
 import click
 from functools import wraps, partial
 
-from raster_loader.utils import get_default_table_name, check_private_key
+from raster_loader.lib.utils import get_default_table_name, check_private_key
 from raster_loader.io.snowflake import SnowflakeConnection
+from raster_loader.lib.valuelabels import validate_band_valuelabels
 
 
 def catch_exception(func=None, *, handle=Exception):
@@ -123,6 +124,18 @@ def snowflake(args=None):
     type=int,
     default=6,
 )
+@click.option(
+    "--band-valuelabels",
+    help="Custom data for valuelabels in JSON format, or 'None'. "
+    "i.e: '{<value_1>: <label_1>, <value_2>: <label_2>, ...}'. "
+    "Could repeat --band-valuelabels to specify multiple bands data. "
+    "They will be considered in the order they appear in the file. "
+    "Note that you can set any value to 'None' to omit valuelabels for that band. ",
+    type=str,
+    default=[],
+    multiple=True,
+    callback=validate_band_valuelabels,
+)
 @catch_exception()
 def upload(
     account,
@@ -148,6 +161,7 @@ def upload(
     exact_stats=False,
     basic_stats=False,
     compression_level=6,
+    band_valuelabels=[],
 ):
     from raster_loader.io.common import (
         get_number_of_blocks,
@@ -255,6 +269,7 @@ def upload(
         basic_stats=basic_stats,
         compress=compress,
         compression_level=compression_level,
+        band_valuelabels=band_valuelabels,
     )
 
     click.echo("Raster file uploaded to Snowflake")

raster_loader/io/bigquery.py

@@ -6,8 +6,8 @@
 
 from itertools import chain
 from raster_loader import __version__
-from raster_loader.errors import import_error_bigquery, IncompatibleRasterException
-from raster_loader.utils import ask_yes_no_question, batched
+from raster_loader.lib.errors import import_error_bigquery, IncompatibleRasterException
+from raster_loader.lib.utils import ask_yes_no_question, batched
 from raster_loader.io.common import (
     check_metadata_is_compatible,
     get_number_of_blocks,
@@ -18,7 +18,7 @@
     update_metadata,
 )
 
-from typing import Iterable, List, Tuple
+from typing import Dict, Iterable, List, Tuple
 from functools import partial
 
 try:
@@ -111,6 +111,7 @@ def upload_raster(
         basic_stats: bool = False,
         compress: bool = False,
         compression_level: int = 6,
+        band_valuelabels: List[Dict[int, str]] = [],
     ):
         """Write a raster file to a BigQuery table."""
         print("Loading raster file to BigQuery...")
@@ -139,6 +140,7 @@ def upload_raster(
                 exact_stats,
                 basic_stats,
                 compress=compress,
+                band_valuelabels=band_valuelabels,
             )
 
             overviews_records_gen = rasterio_overview_to_records(

raster_loader/io/common.py

@@ -10,16 +10,18 @@
 from typing import Dict, Callable, Iterable, List, Tuple, Union
 from affine import Affine
 from shapely import wkt  # Can not use directly from shapely.wkt
+from raster_loader.lib.valuelabels import get_band_valuelabels
 
 import rio_cogeo
 import rasterio
 import quadbin
 
-from raster_loader.geo import raster_bounds
-from raster_loader.errors import (
+from raster_loader.lib.geo import raster_bounds
+from raster_loader.lib.errors import (
     error_not_google_compatible,
 )
-from raster_loader.utils import warnings
+from raster_loader.lib.utils import warnings
+
 
 DEFAULT_COG_BLOCK_SIZE = 256
 
@@ -210,6 +212,7 @@ def rasterio_metadata(
     exact_stats: bool = False,
     basic_stats: bool = False,
     compress: bool = False,
+    band_valuelabels: List[Dict[int, str]] = [],
 ):
     """Open a raster file with rasterio."""
     raster_info = rio_cogeo.cog_info(file_path).dict()
@@ -278,12 +281,14 @@ def rasterio_metadata(
                 band_nodata = band_value_as_string(
                     raster_dataset, band, band_nodata_value(raster_dataset, band)
                 )
+
             meta = {
                 "band": band,
                 "type": raster_band_type(raster_dataset, band),
                 "name": band_field_name(band_name, band, band_rename_function),
                 "colorinterp": band_colorinterp,
                 "colortable": get_color_table(raster_dataset, band),
+                "valuelabels": get_band_valuelabels(file_path, band, band_valuelabels),
                 "stats": stats,
                 "nodata": band_nodata,
             }
@@ -311,6 +316,7 @@ def rasterio_metadata(
                 "colorinterp": e["colorinterp"],
                 "nodata": e["nodata"],
                 "colortable": e["colortable"],
+                "valuelabels": e["valuelabels"],
             }
             for e in bands_metadata
         ]
@@ -506,7 +512,7 @@ def compute_quantiles(data: List[Union[int, float]], cast_function: Callable) ->
 
 def get_stats(
     raster_dataset: rasterio.io.DatasetReader, band: int
-) -> rasterio.Statistics:
+) -> rasterio._io.Statistics:
     """Get statistics for a raster band."""
     try:
         # stats method is supported since rasterio 1.4.0 and statistics
@@ -937,7 +943,8 @@ def is_valid_raster_dataset(raster_dataset: rasterio.io.DatasetReader) -> bool:
 def band_without_stats(band):
     return {
         k: band[k]
-        for k in set(list(band.keys())) - set(["stats", "colorinterp", "colortable"])
+        for k in set(list(band.keys()))
+        - set(["stats", "colorinterp", "colortable", "valuelabels"])
     }
 
 

raster_loader/io/databricks.py

@@ -3,8 +3,11 @@
 import rasterio
 
 from itertools import chain
-from raster_loader.errors import import_error_databricks, IncompatibleRasterException
-from raster_loader.utils import ask_yes_no_question, batched
+from raster_loader.lib.errors import (
+    import_error_databricks,
+    IncompatibleRasterException,
+)
+from raster_loader.lib.utils import ask_yes_no_question, batched
 from raster_loader.io.common import (
     check_metadata_is_compatible,
     get_number_of_blocks,
@@ -15,7 +18,7 @@
     update_metadata,
 )
 
-from typing import Iterable, List, Tuple
+from typing import Dict, Iterable, List, Tuple
 
 try:
     from databricks.connect import DatabricksSession
@@ -141,6 +144,7 @@ def upload_raster(
         basic_stats: bool = False,
         compress: bool = False,
         compression_level: int = 6,
+        band_valuelabels: List[Dict[int, str]] = [],
     ):
         """Write a raster file to a Databricks table."""
         # Wait for cluster to be ready before starting the upload
@@ -174,6 +178,7 @@ def upload_raster(
                 exact_stats,
                 basic_stats,
                 compress=compress,
+                band_valuelabels=band_valuelabels,
             )
 
             overviews_records_gen = rasterio_overview_to_records(

raster_loader/io/snowflake.py

@@ -3,14 +3,14 @@
 import pandas as pd
 
 from itertools import chain
-from typing import Iterable, List, Tuple
+from typing import Dict, Iterable, List, Tuple
 
-from raster_loader.errors import (
+from raster_loader.lib.errors import (
     IncompatibleRasterException,
     import_error_snowflake,
 )
 
-from raster_loader.utils import ask_yes_no_question, batched
+from raster_loader.lib.utils import ask_yes_no_question, batched
 
 from raster_loader.io.common import (
     rasterio_metadata,
@@ -208,6 +208,7 @@ def upload_raster(
         basic_stats: bool = False,
         compress: bool = False,
         compression_level: int = 6,
+        band_valuelabels: List[Dict[int, str]] = [],
     ) -> bool:
         """Write a raster file to a Snowflake table."""
 
@@ -243,6 +244,7 @@ def band_rename_function(x):
                 exact_stats,
                 basic_stats,
                 compress,
+                band_valuelabels,
             )
 
             overviews_records_gen = rasterio_overview_to_records(

raster_loader/lib/valuelabels.py

@@ -0,0 +1,25 @@
+import json
+import click
+from typing import Dict, List, Optional
+
+
+def get_band_valuelabels(
+    file_path: str, band: int, band_valuelabels: List[Dict[int, str]]
+) -> Optional[Dict[int, str]]:
+    valuelabels = None
+    if len(band_valuelabels) >= band and band_valuelabels[band - 1] is not None:
+        print(f"Using the provided valuelabels for band {band}")
+        valuelabels = band_valuelabels[band - 1]
+    return valuelabels
+
+
+def validate_band_valuelabels(_, __, value):
+    """
+    Validate the band valuelabels parameter for click library callback.
+    """
+    try:
+        return [json.loads(item) if item != "None" else None for item in value]
+    except json.JSONDecodeError:
+        raise click.BadParameter(
+            "Invalid JSON format. Please provide a valid JSON object."
+        )

raster_loader/tests/bigquery/test_cli.py

@@ -308,6 +308,33 @@ def test_bigquery_describe(*args, **kwargs):
     assert result.exit_code == 0
 
 
+@patch(
+    "raster_loader.cli.bigquery.BigQueryConnection.upload_records", return_value=None
+)
+@patch("raster_loader.cli.bigquery.BigQueryConnection.__init__", return_value=None)
+def test_bigquery_upload_custom_valuelabels(*args, **kwargs):
+    runner = CliRunner()
+    result = runner.invoke(
+        main,
+        [
+            "bigquery",
+            "upload",
+            "--file_path",
+            f"{tiff}",
+            "--project",
+            "project",
+            "--dataset",
+            "dataset",
+            "--table",
+            "table",
+            "--band-valuelabels",
+            '{"90": "Woody Wetlands", "52": "Shrub/Scrub"}',
+        ],
+    )
+    assert result.exit_code == 1
+    assert "Using the provided valuelabels for band 1" in result.output
+
+
 def test_info(*args, **kwargs):
     runner = CliRunner()
     result = runner.invoke(main, ["info"])