def get_variables(self,
to_dict: bool = False) -> List[Union[dict, Variable]]:
variables_raw = dynamodb_helper.scan_dynamodb_table(
ddb_resource=self.resource,
table_name=Settings.variables_table_name,
)
return [
Variable(**var).dict_contents() if to_dict else Variable(**var)
for var in variables_raw
]
def test_create_variables_table(ddb_store):
test_var = Variable(id='bar',
type=VariableType.JSON,
contents='[ {"aa": 1}, true]')
ddb_store.set_variable(test_var)
assert test_var == ddb_store.get_variable(test_var.id)
assert test_var.get_contents() == [{'aa': 1}, True]
ddb_store.delete_variable(test_var.id)
with pytest.raises(TyphoonResourceNotFoundError):
ddb_store.get_variable(test_var.id)
def load_variable(remote: Optional[str], file: str):
"""Read variable from file and add it to the metadata store"""
var = Variable.from_file(file)
set_settings_from_remote(remote)
metadata_store = Settings.metadata_store(Remotes.aws_profile(remote))
metadata_store.set_variable(var)
print(f'Variable {var.id} added')
def add_variable(remote: Optional[str], var_id: str, var_type: str, contents):
"""Add variable to the metadata store"""
set_settings_from_remote(remote)
metadata_store = Settings.metadata_store(Remotes.aws_profile(remote))
var = Variable(var_id, VariableType[var_type.upper()], contents)
metadata_store.set_variable(var)
print(f'Variable {var_id} added')
def get_variables(self, to_dict: bool = False) -> List[Union[dict, Variable]]:
result = []
with set_airflow_db(self.db_path, self.fernet_key) as db:
for af_var in db.get_variables():
if af_var.key.startswith('typhoon:info#'):
info = json.loads(af_var.val)
var_id = info['id']
value_var = typhoon_airflow_variable_name_for_value(var_id)
contents = db.get_variable(value_var)
if contents is None:
raise MetadataObjectNotFound(f'Variable "{var_id}" does not have a value set')
var = Variable(id=var_id, type=VariableType(info['type']), contents=contents)
if to_dict:
var = var.dict_contents()
result.append(var)
return result
def set_variable(self, variable: Variable):
dynamodb_helper.dynamodb_put_item(
ddb_client=self.client,
table_name=Settings.variables_table_name,
item={
'id': variable.id,
**variable.dict_contents(),
})
def get_variable(self, variable_id: str) -> Variable:
info_var = typhoon_airflow_variable_name_for_info(variable_id)
value_var = typhoon_airflow_variable_name_for_value(variable_id)
with set_airflow_db(self.db_path, self.fernet_key) as db:
info = db.get_variable(info_var)
if info is None:
raise MetadataObjectNotFound(f'Variable "{variable_id}" is not set')
info = json.loads(info)
value = db.get_variable(value_var)
return Variable(id=info['id'], type=VariableType(info['type']), contents=value)
def test_sqlite_metadata_store(typhoon_home):
store = Settings.metadata_store()
assert isinstance(store, SQLiteMetadataStore)
conn = Connection(conn_id='foo', conn_type='s3')
store.set_connection(conn)
assert store.get_connection('foo') == conn
var = Variable(id='bar', type=VariableType.STRING, contents='lorem ipsum')
store.set_variable(var)
assert store.get_variable('bar') == var
def get_variable(self, variable_id: str) -> Variable:
try:
item = dynamodb_helper.dynamodb_get_item(
ddb_client=self.client,
table_name=Settings.variables_table_name,
key_name='id',
key_value=variable_id,
)
except TyphoonResourceNotFoundError:
raise MetadataObjectNotFound(
f'Variable "{variable_id}" is not set')
return Variable(**item)
def set_variable(self, variable: Variable):
"""
Sets Airflow variables in JSON like:
{
"id": "variable_name",
"type": "variable_type",
"contents": "contents",
}
Following the same schema as the typhoon variable
"""
info_var = typhoon_airflow_variable_name_for_info(variable.id)
value_var = typhoon_airflow_variable_name_for_value(variable.id)
info = variable.dict_contents()
del info['contents']
with set_airflow_db(self.db_path, self.fernet_key) as db:
db.set_variable(info_var, json.dumps(info))
db.set_variable(value_var, variable.contents)
import pytest
from typhoon.connections import Connection
from typhoon.core.metadata_store_interface import MetadataObjectNotFound
from typhoon.core.settings import Settings
from typhoon.variables import Variable, VariableType
sample_conn = Connection(conn_id='foo', conn_type='some_type')
other_sample_conn = Connection(conn_id='bar', conn_type='some_other_type')
sample_var = Variable(id='foo',
type=VariableType.STRING,
contents='hello world')
other_sample_var = Variable(id='bar',
type=VariableType.STRING,
contents='hello world')
@pytest.fixture
def cfg_path(tmp_path):
Settings.typhoon_home = tmp_path
Settings.project_name = 'unittests'
db_path = tmp_path / 'test.db'
Settings.metadata_db_url = f'sqlite:{db_path}'
return str(cfg_path)
def test_set_and_get_connection(cfg_path):
store = Settings.metadata_store()
store.set_connection(sample_conn)
assert store.get_connection(sample_conn.conn_id) == sample_conn
def get_variable(variable_id: str) -> Variable:
env_var_value = os.environ.get(f'TYPHOON_VARIABLE_{variable_id}')
if env_var_value:
return Variable(variable_id, VariableType.STRING, env_var_value)
else:
return Settings.metadata_store().get_variable(variable_id)