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