mirror of
https://github.com/google/earthengine-api.git
synced 2025-12-08 19:26:12 +00:00
289 lines
9.6 KiB
Python
289 lines
9.6 KiB
Python
#!/usr/bin/env python
|
|
"""Support utilities used by the Earth Engine command line interface.
|
|
|
|
This module defines the Command class which is the base class of all
|
|
the commands supported by the EE command line tool. It also defines
|
|
the classes for configuration and runtime context management.
|
|
"""
|
|
from __future__ import print_function
|
|
import collections
|
|
from datetime import datetime
|
|
import json
|
|
import os
|
|
import re
|
|
import threading
|
|
import time
|
|
|
|
import urllib
|
|
import httplib2
|
|
|
|
from google.oauth2.credentials import Credentials
|
|
import ee
|
|
|
|
HOMEDIR = os.path.expanduser('~')
|
|
EE_CONFIG_FILE = 'EE_CONFIG_FILE'
|
|
DEFAULT_EE_CONFIG_FILE_RELATIVE = os.path.join(
|
|
'.config', 'earthengine', 'credentials')
|
|
DEFAULT_EE_CONFIG_FILE = os.path.join(
|
|
HOMEDIR, DEFAULT_EE_CONFIG_FILE_RELATIVE)
|
|
|
|
CONFIG_PARAMS = {
|
|
'url': 'https://earthengine.googleapis.com',
|
|
'account': None,
|
|
'private_key': None,
|
|
'refresh_token': None,
|
|
'use_cloud_api': False,
|
|
'cloud_api_key': None,
|
|
'project': None,
|
|
}
|
|
|
|
TASK_FINISHED_STATES = (ee.batch.Task.State.COMPLETED,
|
|
ee.batch.Task.State.FAILED,
|
|
ee.batch.Task.State.CANCELLED)
|
|
|
|
|
|
class CommandLineConfig(object):
|
|
"""Holds the configuration parameters used by the EE command line interface.
|
|
|
|
This class attempts to load the configuration parameters from a file
|
|
specified as a constructor argument. If not provided, it attempts to load
|
|
the configuration from a file specified via the EE_CONFIG_FILE environment
|
|
variable. If the variable is not set, it looks for a JSON file at the
|
|
path ~/.config/earthengine/credentials. If all fails, it falls back to using
|
|
some predefined defaults for each configuration parameter.
|
|
|
|
If --service_account_file is specified, it is used instead.
|
|
"""
|
|
|
|
def __init__(
|
|
self, config_file=None, service_account_file=None, use_cloud_api=False):
|
|
if not config_file:
|
|
config_file = os.environ.get(EE_CONFIG_FILE, DEFAULT_EE_CONFIG_FILE)
|
|
self.config_file = config_file
|
|
config = {}
|
|
if os.path.exists(config_file):
|
|
with open(config_file) as config_file_json:
|
|
config = json.load(config_file_json)
|
|
CONFIG_PARAMS['use_cloud_api'] = use_cloud_api
|
|
for key, default_value in CONFIG_PARAMS.items():
|
|
setattr(self, key, config.get(key, default_value))
|
|
self.service_account_file = service_account_file
|
|
if service_account_file:
|
|
# Load the file to verify that it exists.
|
|
with open(service_account_file) as service_file_json:
|
|
service = json.load(service_file_json)
|
|
for key, value in service.items():
|
|
setattr(self, key, value)
|
|
|
|
def ee_init(self):
|
|
"""Loads the EE credentials and initializes the EE client."""
|
|
if self.service_account_file:
|
|
credentials = ee.ServiceAccountCredentials(
|
|
self.client_email, self.service_account_file)
|
|
elif self.account and self.private_key:
|
|
credentials = ee.ServiceAccountCredentials(self.account, self.private_key)
|
|
elif self.refresh_token:
|
|
credentials = Credentials(
|
|
None,
|
|
refresh_token=self.refresh_token,
|
|
token_uri=ee.oauth.TOKEN_URI,
|
|
client_id=ee.oauth.CLIENT_ID,
|
|
client_secret=ee.oauth.CLIENT_SECRET,
|
|
scopes=ee.oauth.SCOPES)
|
|
else:
|
|
credentials = 'persistent'
|
|
|
|
ee.Initialize(
|
|
credentials=credentials,
|
|
opt_url=self.url,
|
|
use_cloud_api=self.use_cloud_api,
|
|
cloud_api_key=self.cloud_api_key,
|
|
project=self.project)
|
|
|
|
def save(self):
|
|
config = {}
|
|
for key in CONFIG_PARAMS:
|
|
value = getattr(self, key)
|
|
if value is not None:
|
|
config[key] = value
|
|
with open(self.config_file, 'w') as output_file:
|
|
json.dump(config, output_file)
|
|
|
|
|
|
def query_yes_no(msg):
|
|
print('%s (y/n)' % msg)
|
|
while True:
|
|
confirm = raw_input().lower()
|
|
if confirm == 'y':
|
|
return True
|
|
elif confirm == 'n':
|
|
return False
|
|
else:
|
|
print('Please respond with \'y\' or \'n\'.')
|
|
|
|
|
|
def truncate(string, length):
|
|
return (string[:length] + '..') if len(string) > length else string
|
|
|
|
|
|
def wait_for_task(task_id, timeout, log_progress=True):
|
|
"""Waits for the specified task to finish, or a timeout to occur."""
|
|
start = time.time()
|
|
elapsed = 0
|
|
last_check = 0
|
|
while True:
|
|
elapsed = time.time() - start
|
|
status = ee.data.getTaskStatus(task_id)[0]
|
|
state = status['state']
|
|
if state in TASK_FINISHED_STATES:
|
|
error_message = status.get('error_message', None)
|
|
print('Task %s ended at state: %s after %.2f seconds'
|
|
% (task_id, state, elapsed))
|
|
if error_message:
|
|
raise ee.ee_exception.EEException('Error: %s' % error_message)
|
|
return
|
|
if log_progress and elapsed - last_check >= 30:
|
|
print('[{:%H:%M:%S}] Current state for task {}: {}'
|
|
.format(datetime.now(), task_id, state))
|
|
last_check = elapsed
|
|
remaining = timeout - elapsed
|
|
if remaining > 0:
|
|
time.sleep(min(10, remaining))
|
|
else:
|
|
break
|
|
print('Wait for task %s timed out after %.2f seconds' % (task_id, elapsed))
|
|
|
|
|
|
def wait_for_tasks(task_id_list, timeout, log_progress=False):
|
|
"""For each task specified in task_id_list, wait for that task or timeout."""
|
|
|
|
if len(task_id_list) == 1:
|
|
wait_for_task(task_id_list[0], timeout, log_progress)
|
|
return
|
|
|
|
threads = []
|
|
for task_id in task_id_list:
|
|
t = threading.Thread(target=wait_for_task,
|
|
args=(task_id, timeout, log_progress))
|
|
threads.append(t)
|
|
t.start()
|
|
|
|
for thread in threads:
|
|
thread.join()
|
|
|
|
status_list = ee.data.getTaskStatus(task_id_list)
|
|
status_counts = collections.defaultdict(int)
|
|
for status in status_list:
|
|
status_counts[status['state']] += 1
|
|
num_incomplete = (len(status_list) - status_counts['COMPLETED']
|
|
- status_counts['FAILED'] - status_counts['CANCELLED'])
|
|
print('Finished waiting for tasks.\n Status summary:')
|
|
print(' %d tasks completed successfully.' % status_counts['COMPLETED'])
|
|
print(' %d tasks failed.' % status_counts['FAILED'])
|
|
print(' %d tasks cancelled.' % status_counts['CANCELLED'])
|
|
print(' %d tasks are still incomplete (timed-out)' % num_incomplete)
|
|
|
|
|
|
def expand_gcs_wildcards(source_files):
|
|
"""Implements glob-like '*' wildcard completion for cloud storage objects.
|
|
|
|
Args:
|
|
source_files: A list of one or more cloud storage paths of the format
|
|
gs://[bucket]/[path-maybe-with-wildcards]
|
|
|
|
Yields:
|
|
cloud storage paths of the above format with '*' wildcards expanded.
|
|
Raises:
|
|
EEException: If badly formatted source_files
|
|
(e.g., missing gs://) are specified
|
|
"""
|
|
for source in source_files:
|
|
if '*' not in source:
|
|
yield source
|
|
continue
|
|
|
|
# We extract the bucket and prefix from the input path to match
|
|
# the parameters for calling GCS list objects and reduce the number
|
|
# of items returned by that API call
|
|
|
|
# Capture the part of the path after gs:// and before the first /
|
|
bucket_regex = 'gs://([a-z0-9_.-]+)(/.*)'
|
|
bucket_match = re.match(bucket_regex, source)
|
|
if bucket_match:
|
|
bucket, rest = bucket_match.group(1, 2)
|
|
else:
|
|
raise ee.ee_exception.EEException(
|
|
'Badly formatted source file or bucket: %s' % source)
|
|
prefix = rest[:rest.find('*')] # Everything before the first wildcard
|
|
|
|
bucket_files = _gcs_ls(bucket, prefix)
|
|
|
|
# Regex to match the source path with wildcards expanded
|
|
regex = re.escape(source).replace(r'\*', '[^/]*') + '$'
|
|
for gcs_path in bucket_files:
|
|
if re.match(regex, gcs_path):
|
|
yield gcs_path
|
|
|
|
|
|
def _gcs_ls(bucket, prefix=''):
|
|
"""Retrieve a list of cloud storage filepaths from the given bucket.
|
|
|
|
Args:
|
|
bucket: The cloud storage bucket to be queried
|
|
prefix: Optional, a prefix used to select the objects to return
|
|
Yields:
|
|
Cloud storage filepaths matching the given bucket and prefix
|
|
Raises:
|
|
EEException:
|
|
If there is an error in accessing the specified bucket
|
|
"""
|
|
|
|
base_url = 'https://www.googleapis.com/storage/v1/b/%s/o'%bucket
|
|
method = 'GET'
|
|
http = ee.data.authorizeHttp(httplib2.Http(0))
|
|
next_page_token = None
|
|
|
|
# Loop to handle paginated responses from GCS;
|
|
# Exits once no 'next page token' is returned
|
|
while True:
|
|
params = {'fields': 'items/name,nextPageToken'}
|
|
if next_page_token:
|
|
params['pageToken'] = next_page_token
|
|
if prefix:
|
|
params['prefix'] = prefix
|
|
payload = urllib.urlencode(params)
|
|
|
|
url = base_url + '?' + payload
|
|
try:
|
|
response, content = http.request(url, method=method)
|
|
except httplib2.HttpLib2Error as e:
|
|
raise ee.ee_exception.EEException(
|
|
'Unexpected HTTP error: %s' % e.message)
|
|
|
|
if response.status < 100 or response.status >= 300:
|
|
raise ee.ee_exception.EEException(('Error retrieving bucket %s;'
|
|
' Server returned HTTP code: %d' %
|
|
(bucket, response.status)))
|
|
|
|
json_content = json.loads(content)
|
|
if 'error' in json_content:
|
|
json_error = json_content['error']['message']
|
|
raise ee.ee_exception.EEException('Error retrieving bucket %s: %s' %
|
|
(bucket, json_error))
|
|
|
|
if 'items' not in json_content:
|
|
raise ee.ee_exception.EEException(
|
|
'Cannot find items list in the response from GCS: %s' % json_content)
|
|
objects = json_content['items']
|
|
object_names = [str(gc_object['name']) for gc_object in objects]
|
|
|
|
for name in object_names:
|
|
yield 'gs://%s/%s' % (bucket, name)
|
|
|
|
# GCS indicates no more results
|
|
if 'nextPageToken' not in json_content:
|
|
return
|
|
|
|
# Load next page, continue at beginning of while True:
|
|
next_page_token = json_content['nextPageToken']
|