def get_timestamp(path: str) -> datetime:
if path.startswith('s3://'):
return S3Config().get_s3_client().get_key(path).last_modified
elif path.startswith('gs://'):
# for gcs object
# should PR to luigi
bucket, obj = GCSConfig().get_gcs_client()._path_to_bucket_and_key(path)
result = GCSConfig().get_gcs_client().client.objects().get(bucket=bucket, object=obj).execute()
return result['updated']
else:
raise
def get_object_storage_target(path: str, format: Format) -> luigi.Target:
if path.startswith('s3://'):
return luigi.contrib.s3.S3Target(path, client=S3Config().get_s3_client(), format=format)
elif path.startswith('gs://'):
return luigi.contrib.gcs.GCSTarget(path, client=GCSConfig().get_gcs_client(), format=format)
else:
raise
def exists(path: str) -> bool:
if path.startswith('s3://'):
return S3Config().get_s3_client().exists(path)
elif path.startswith('gs://'):
return GCSConfig().get_gcs_client().exists(path)
else:
raise
def test_get_gcs_client_with_json(self):
mock = MagicMock()
json_str = '{"test": 1}'
os.environ['env_name'] = json_str
with patch('luigi.contrib.gcs.GCSClient'):
with patch('google.oauth2.service_account.Credentials.from_service_account_info', mock):
GCSConfig(gcs_credential_name='env_name')._get_gcs_client()
self.assertEqual(dict(test=1), mock.call_args[0][0])
def __init__(self,
workspace_directory: str = '',
task_filters: List[str] = [],
tqdm_disable: bool = False):
"""must set $GCS_CREDENTIAL"""
self.workspace_directory = workspace_directory
self.task_filters = task_filters
self.tqdm_disable = tqdm_disable
self.gcs_client = GCSConfig().get_gcs_client()
def test_get_gcs_client_with_file_path(self):
mock = MagicMock()
file_path = 'test.json'
os.environ['env_name'] = file_path
with patch('luigi.contrib.gcs.GCSClient'):
with patch('google.oauth2.service_account.Credentials.from_service_account_file', mock):
with patch('os.path.isfile', return_value=True):
GCSConfig(gcs_credential_name='env_name')._get_gcs_client()
self.assertEqual(file_path, mock.call_args[0][0])
def test_get_gcs_client_without_gcs_credential_name(self):
mock = MagicMock()
discover_path = 'discover_cache.json'
os.environ['env_name'] = ''
os.environ['discover_path'] = discover_path
with open(f'{discover_path}', 'w') as f:
f.write('{}')
with patch('luigi.contrib.gcs.GCSClient', mock):
with patch('fcntl.flock'):
GCSConfig(gcs_credential_name='env_name', discover_cache_local_path=discover_path).get_gcs_client()
self.assertEqual(dict(oauth_credentials=None, descriptor='{}'), mock.call_args[1])
def __init__(self,
workspace_directory: str = '',
task_filters: List[str] = [],
tqdm_disable: bool = False,
use_cache: bool = True):
"""must set $GCS_CREDENTIAL"""
self.workspace_directory = workspace_directory
self.task_filters = task_filters
self.tqdm_disable = tqdm_disable
self.gcs_client = GCSConfig().get_gcs_client()
self.local_cache = LocalCache(workspace_directory, use_cache)
self.use_cache = use_cache
def test_get_gcs_client_with_json(self):
mock = MagicMock()
json_str = '{"test": 1}'
discover_path = 'discover_cache.json'
os.environ['env_name'] = json_str
os.environ['discover_path'] = discover_path
with open(f'{discover_path}', 'w') as f:
f.write('{}')
with patch('luigi.contrib.gcs.GCSClient'):
with patch('google.oauth2.service_account.Credentials.from_service_account_info', mock):
GCSConfig(gcs_credential_name='env_name', discover_cache_local_path=discover_path).get_gcs_client()
self.assertEqual(dict(test=1), mock.call_args[0][0])
def test_get_gcs_client_with_file_path(self):
mock = MagicMock()
file_path = 'test.json'
discover_path = 'discover_cache.json'
os.environ['env_name'] = file_path
os.environ['discover_path'] = discover_path
with open(f'{discover_path}', 'w') as f:
f.write('{}')
with patch('luigi.contrib.gcs.GCSClient'):
with patch('google.oauth2.service_account.Credentials.from_service_account_file', mock):
with patch('os.path.isfile', return_value=True):
GCSConfig(gcs_credential_name='env_name', discover_cache_local_path=discover_path).get_gcs_client()
self.assertEqual(file_path, mock.call_args[0][0])
def test_get_gcs_client_without_gcs_credential_name(self):
mock = MagicMock()
os.environ['env_name'] = ''
with patch('luigi.contrib.gcs.GCSClient', mock):
GCSConfig(gcs_credential_name='env_name')._get_gcs_client()
self.assertEqual(dict(oauth_credentials=None), mock.call_args[1])
def __init__(self, file_path: str, temporary_directory: str) -> None:
self._file_path = file_path
self._temporary_directory = temporary_directory
self._client = GCSConfig().get_gcs_client()