def get_definitions(service: str) -> dict:
service_account_obj = service_account.Credentials.from_service_account_file(
os.environ['MYSQL_BIG_QUERY_GOOGLE_AUTH'])
storage = StorageAdapter(service_account_obj)
storage.get_client()
key = storage.get_file_as_string(bucket="mysql_sync_keys", file=service)
if key is False:
raise RuntimeError(storage.errors)
crypto = Fernet(key)
bqa = BigQueryAdapter(service_account_obj)
bqa.get_client()
query = f"SELECT * FROM mysql_sync.data_sources WHERE service = '{service}'"
result = bqa.query(query)
if result is False:
raise RuntimeError(bqa.errors)
definitions = {}
for item in result:
definitions['service'] = item['service']
definitions['data_set'] = item['data_set']
definitions['database'] = item['database']
definitions['host'] = crypto.decrypt(
bytes(item['host'], encoding='utf-8'))
definitions['user'] = crypto.decrypt(
bytes(item['user'], encoding='utf-8'))
definitions['password'] = crypto.decrypt(
bytes(item['password'], encoding='utf-8'))
return definitions
def write_df_to_storage(self, df: DataFrame) -> bool:
storage_client = StorageAdapter(self.service_account)
storage_client.get_client()
date_time_obj = datetime.datetime.utcnow()
location = f'error/csv/{self.table}/{date_time_obj.strftime("%m-%d-%Y_%H:%M:%S")}_UTC'
result = storage_client.write_string(
bucket=self.definitions['service'] + '-etl',
destination=location,
string=df.to_csv(),
encoding='text/csv')
if not result:
self.sd_logger.error(
storage_client.errors, {
'class': 'Data',
'method': 'write_df_to_storage',
'table': self.table
})
return False
self.sd_logger.info({'message': f"Failed CSV added to {location}"}, {
'class': 'Data',
'method': 'write_df_to_storage',
'table': self.table
})
return result
def save_key(self):
storage_client = StorageAdapter(self.service_account)
storage_client.get_client()
result = storage_client.write_string(bucket='mysql_sync_keys',
destination=f"{self.service}",
string=self.key)
if result is False:
raise RuntimeError(storage_client.errors)
def create_service_bucket(self, service: str):
client = StorageAdapter(self.service_account)
client.get_client()
result = client.make_bucket(f"{service}-sync")
if result:
self.sd_logger.info({'message': f'Created {service} Bucket'}, {
'class': 'PrefectAddServiceStorage',
'method': 'create_service_bucket'
})
else:
self.sd_logger.warning(
client.errors, {
'class': 'PrefectAddServiceStorage',
'method': 'create_service_bucket'
})
return result
def create_config_bucket(self):
client = StorageAdapter(self.service_account)
client.get_client()
result = client.make_bucket('mysql_sync_keys')
if result:
self.sd_logger.info({'message': f'Created mysql_sync_keys Bucket'},
{
'class': 'PrefectInstallStorage',
'method': 'create_config_bucket'
})
else:
self.sd_logger.warning(client.errors, {
'class': 'PrefectInstallStorage',
'method': 'create_config_bucket'
})
return result
def store_mysql_schema(self, schema: list, table: str) -> bool:
encoded_schema = json.dumps(schema)
storage_client = StorageAdapter(self.service_account)
storage_client.get_client()
date_time_obj = datetime.utcnow()
result = storage_client.write_string(
bucket=self.definitions['service'] + '-etl',
destination=
f'schema/{table}/{date_time_obj.strftime("%m-%d-%Y_%H:%M:%S")}_UTC',
string=encoded_schema)
if not result:
self.sd_logger.error(storage_client.errors, {
'class': 'Schema',
'method': 'store_mysql_schema',
'table': table
})
return False
return result
def test_create_bucket_client_is_none(mock_service_account):
client = StorageAdapter(mock_service_account)
with pytest.raises(RuntimeError):
client.make_bucket('string')
def test_create_client_failure_bad_credentials(mock_service_account):
client = StorageAdapter(mock_service_account)
with pytest.raises(ValueError):
client.get_client()
def mock_adapter(mock_service_account):
adapter = StorageAdapter(mock_service_account)
mock = Mock(spec=google.cloud.storage.Client)
adapter.client = mock
return adapter
def test_write_string_client_is_none(mock_service_account):
client = StorageAdapter(mock_service_account)
with pytest.raises(RuntimeError):
client.write_string('bucket', 'destination', 'string to write')
def test_write_object_client_is_none(mock_service_account):
client = StorageAdapter(mock_service_account)
with pytest.raises(RuntimeError):
client.write_file('bucket', 'no-file', 'no-file', True)