def get_connections(self,
to_dict: bool = False
) -> List[Union[dict, Connection]]:
connections_raw = dynamodb_helper.scan_dynamodb_table(
ddb_resource=self.resource,
table_name=Settings.connections_table_name,
)
return [
Connection(**conn).__dict__ if to_dict else Connection(**conn)
for conn in connections_raw
]
def set_connection(self, conn: Connection):
self.delete_connection(conn)
with set_airflow_db(self.db_path, self.fernet_key) as db:
db.set_connection(
conn_id=typhoon_airflow_conn_name(conn.conn_id),
**conn.get_connection_params().__dict__
)
def add_connection(remote: Optional[str], conn_id: str, conn_env: str):
"""Add connection to the metadata store"""
set_settings_from_remote(remote)
metadata_store = Settings.metadata_store(Remotes.aws_profile(remote))
conn_params = connections.get_connection_local(conn_id, conn_env)
metadata_store.set_connection(
Connection(conn_id=conn_id, **asdict(conn_params)))
print(f'Connection {conn_id} added')
def set_connection(self, conn: Connection):
dynamodb_helper.dynamodb_put_item(
ddb_client=self.client,
table_name=Settings.connections_table_name,
item={
'conn_id': conn.conn_id,
**conn.get_connection_params().__dict__,
})
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 test_create_connections_table(ddb_store):
assert ddb_store.exists()
test_conn = Connection(conn_id='foo', conn_type='s3')
ddb_store.set_connection(test_conn)
assert test_conn == ddb_store.get_connection(conn_id=test_conn.conn_id)
ddb_store.delete_connection(conn=test_conn)
with pytest.raises(TyphoonResourceNotFoundError):
ddb_store.get_connection(test_conn.conn_id)
def get_connection(self, conn_id: str) -> Connection:
try:
item = dynamodb_helper.dynamodb_get_item(
ddb_client=self.client,
table_name=Settings.connections_table_name,
key_name='conn_id',
key_value=conn_id,
)
except TyphoonResourceNotFoundError:
raise MetadataObjectNotFound(f'Connection "{conn_id}" is not set')
return Connection(**item)
def get_connection(self, conn_id: str) -> Connection:
af_name = typhoon_airflow_conn_name(conn_id)
with set_airflow_db(self.db_path, self.fernet_key) as db:
af_conn = db.get_connection(af_name)
if af_conn is None:
raise MetadataObjectNotFound(f'Connection "{conn_id}" is not set')
conn = Connection(
conn_id=conn_id,
conn_type=af_conn.conn_type,
host=af_conn.host,
port=af_conn.port,
login=af_conn.login,
password=af_conn.password,
schema=af_conn.schema,
extra=af_conn.extra_dejson,
)
return conn
def get_connections(self, to_dict: bool = False) -> List[Union[dict, Connection]]:
result = []
with set_airflow_db(self.db_path, self.fernet_key) as db:
for af_conn in db.get_connections():
if af_conn.conn_id.startswith('typhoon#'):
conn = Connection(
conn_id=af_conn.conn_id.split('#')[1],
conn_type=af_conn.conn_type,
host=af_conn.host,
port=af_conn.port,
login=af_conn.login,
password=af_conn.password,
schema=af_conn.schema,
extra=af_conn.extra_dejson,
)
if to_dict:
conn = conn.__dict__
result.append(conn)
return result
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