Add get_catalog_store_urls + get_github_commit_url (#23)

* Add get_catalog_store_urls + get_github_commit_url

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix pre-commit stuff

* Bugfix in get_catalog_store_urls

* Make inject_attrs optional

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

leap_data_management_utils/data_management_transforms.py

@@ -1,14 +1,52 @@
 # Note: All of this code was written by Julius Busecke and copied from this feedstock:
 # https://github.com/leap-stc/cmip6-leap-feedstock/blob/main/feedstock/recipe.py#L262
 
-import datetime
+import subprocess
 from dataclasses import dataclass
+from datetime import datetime, timezone
 from typing import Optional
 
 import apache_beam as beam
 import zarr
 from google.api_core.exceptions import NotFound
 from google.cloud import bigquery
+from ruamel.yaml import YAML
+
+yaml = YAML(typ='safe')
+
+
+def get_github_commit_url() -> Optional[str]:
+    """Get the GitHub commit URL for the current commit"""
+    # Get GitHub Server URL
+    github_server_url = 'https://github.com'
+
+    # Get the repository's remote origin URL
+    try:
+        repo_origin_url = subprocess.check_output(
+            ['git', 'config', '--get', 'remote.origin.url'], text=True
+        ).strip()
+
+        # Extract the repository path from the remote URL
+        repository_path = repo_origin_url.split('github.com/')[-1].replace('.git', '')
+
+        # Get the current commit SHA
+        commit_sha = subprocess.check_output(['git', 'rev-parse', 'HEAD'], text=True).strip()
+
+        # Construct the GitHub commit URL
+        git_url_hash = f'{github_server_url}/{repository_path}/commit/{commit_sha}'
+
+        # Output the GitHub commit URL
+        return git_url_hash
+
+    except subprocess.CalledProcessError as e:
+        print('Error executing Git command:', e)
+        return None
+
+
+def get_catalog_store_urls(catalog_yaml_path: str) -> dict[str, str]:
+    with open(catalog_yaml_path) as f:
+        catalog_meta = yaml.load(f)
+    return {d['id']: d['url'] for d in catalog_meta['stores']}
 
 
 @dataclass
@@ -54,7 +92,7 @@ def _get_table(self) -> bigquery.table.Table:
         return self.client.get_table(self.table_id)
 
     def insert(self, fields: dict = {}):
-        timestamp = datetime.datetime.now().isoformat()
+        timestamp = datetime.now().isoformat()
 
         rows_to_insert = [
             fields | {'timestamp': timestamp}  # timestamp is always overridden
@@ -120,6 +158,8 @@ def expand(self, pcoll: beam.PCollection) -> beam.PCollection:
 
 @dataclass
 class Copy(beam.PTransform):
+    """Copy a store to a new location. If the target input is False, do nothing."""
+
     target: str
 
     def _copy(self, store: zarr.storage.FSStore) -> zarr.storage.FSStore:
@@ -147,7 +187,22 @@ def expand(self, pcoll: beam.PCollection) -> beam.PCollection:
 
 @dataclass
 class InjectAttrs(beam.PTransform):
-    inject_attrs: dict
+    inject_attrs: dict = None
+    add_provenance: bool = True
+
+    # add a post_init method to add the provenance attributes
+    def __post_init__(self):
+        if self.inject_attrs is None:
+            self.inject_attrs = {}
+
+        if self.add_provenance:
+            git_url_hash = get_github_commit_url()
+            timestamp = datetime.now(timezone.utc).isoformat()
+            provenance_dict = {
+                'pangeo_forge_build_git_hash': git_url_hash,
+                'pangeo_forge_build_timestamp': timestamp,
+            }
+            self.inject_attrs.update(provenance_dict)
 
     def _update_zarr_attrs(self, store: zarr.storage.FSStore) -> zarr.storage.FSStore:
         # TODO: Can we get a warning here if the store does not exist?

leap_data_management_utils/tests/test_data_management_transforms.py

@@ -1,4 +1,36 @@
+from ruamel.yaml import YAML
+
+from leap_data_management_utils.data_management_transforms import (
+    get_catalog_store_urls,
+    get_github_commit_url,
+)
+
+yaml = YAML(typ='safe')
+
+
 def test_smoke_test():
     assert True
     # This is a bit dumb, but it at least checks the the imports are working
     # again super hard to test code involving bigquery here.
+
+
+def test_get_github_commit_url():
+    url = get_github_commit_url()
+    assert url.startswith('https://github.com/leap-stc/leap-data-management-utils')
+
+
+def test_get_catalog_store_urls(tmp_path):
+    # Create a temporary text file
+    temp_file = tmp_path / 'some-name.yaml'
+    data = {
+        'stores': [{'id': 'a', 'url': 'a-url', 'some_other': 'stuff'}, {'id': 'b', 'url': 'b-url'}]
+    }
+    with open(temp_file, 'w') as f:
+        yaml.dump(data, f)
+
+    # Call the function to read the file
+    content = get_catalog_store_urls(temp_file)
+
+    # Assertions
+    assert content['a'] == 'a-url'
+    assert content['b'] == 'b-url'

pyproject.toml

@@ -32,6 +32,7 @@ dependencies = [
     "pydantic-core",
     "pydantic>=2",
     "pyyaml",
+    "ruamel.yaml",
     "universal-pathlib",
     "xarray",
     "zarr",