def _get_connection(
cm: ConnectionManager,
identifier: str,
sleep: Optional[int] = 0,
enabled_alive: Optional[bool] = True,
):
def __connect():
if 0 < sleep:
time.sleep(sleep)
return ConnectionObject()
def __alive(conn):
if enabled_alive:
return True
else:
return False
def __close(conn):
return True
connection = cm.get(
identifier,
connect_method=__connect,
alive_method=__alive,
close_method=__close,
save=True,
)
return connection
def connection_manager():
return ConnectionManager(
name='connection_manager_name',
timeout=3,
wait=0.2,
time_between_clean=1,
time_keep_alive=5,
connection_timeout=60,
)
def connection_manager_with_error():
return ConnectionManager(
name='connection_manager_name',
timeout=3,
wait=0.2,
time_between_clean=1,
time_keep_alive=5,
connection_timeout=60,
tortank='test',
)
def test_init_exception_connectionmanager():
with pytest.raises(KeyError):
ConnectionManager(
name='connection_manager_name',
timeout=3,
wait=0.2,
time_between_clean=1,
time_keep_alive=5,
connection_timeout=60,
tortank='test',
)
def _get_connection_without_alive_close_method(
cm: ConnectionManager,
identifier: str,
):
def __connect():
return ConnectionObject()
connection = cm.get(
identifier,
connect_method=__connect,
alive_method='alive',
close_method='close',
)
return connection
def _get_connection_without_connect_method(
cm: ConnectionManager,
identifier: str,
):
def __alive(conn):
return True
def __close(conn):
return True
connection = cm.get(
identifier,
connect_method='connect',
alive_method=__alive,
close_method=__close,
)
return connection
def _get_connection_exception(cm: ConnectionManager, identifier: str):
def __connect():
return ConnectionObject()
def __alive(conn):
return True
def __close(conn):
raise TimeoutError
connection = cm.get(
identifier,
connect_method=__connect,
alive_method=__alive,
close_method=__close,
)
return connection
def _get_connection_long_closing(cm: ConnectionManager, identifier: str,
time_sleep: int):
def __connect():
return ConnectionObject()
def __alive(conn):
return True
def __close(conn):
time.sleep(time_sleep)
return True
connection = cm.get(
identifier,
connect_method=__connect,
alive_method=__alive,
close_method=__close,
)
return connection
)
from toucan_connectors.toucan_connector import (
Category,
DataSlice,
DiscoverableConnector,
ToucanConnector,
strlist_to_enum,
)
logger = logging.getLogger(__name__)
snowflake_connection_manager = None
if not snowflake_connection_manager:
snowflake_connection_manager = ConnectionManager(name='snowflake',
timeout=10,
wait=0.2,
time_between_clean=10,
time_keep_alive=600)
class Path(str):
@classmethod
def __get_validators__(cls):
yield cls.validate # pragma: no cover
@classmethod
def validate(cls, v):
if not path.exists(v): # pragma: no cover
raise ValueError(f'path does not exists: {v}') # pragma: no cover
return v # pragma: no cover
'user',
'password',
'access_key_id',
'secret_access_key',
'session_token',
'profile',
'region',
]
logger = logging.getLogger(__name__)
redshift_connection_manager = None
if not redshift_connection_manager:
redshift_connection_manager = ConnectionManager(name='redshift',
timeout=25,
wait=0.2,
time_between_clean=10,
time_keep_alive=60)
class AuthenticationMethod(str, Enum):
DB_CREDENTIALS: str = 'db_credentials'
AWS_CREDENTIALS: str = 'aws_credentials'
AWS_PROFILE: str = 'aws_profile'
class AuthenticationMethodError(str, Enum):
DB_CREDENTIALS: str = f'User & Password are required for {AuthenticationMethod.DB_CREDENTIALS}'
AWS_CREDENTIALS: str = f'AccessKeyId, SecretAccessKey & db_user are required for {AuthenticationMethod.AWS_CREDENTIALS}'
AWS_PROFILE: str = f'Profile & db_user are required for {AuthenticationMethod.AWS_PROFILE}'
UNKNOWN: str = 'Unknown AuthenticationMethod'
)
from toucan_connectors.toucan_connector import (
Category,
DataSlice,
DiscoverableConnector,
ToucanConnector,
strlist_to_enum,
)
logger = logging.getLogger(__name__)
connection_manager = None
if not connection_manager:
connection_manager = ConnectionManager(name='snowflake_oauth2',
timeout=5,
wait=0.2,
time_between_clean=3,
time_keep_alive=600)
class SnowflakeoAuth2DataSource(SnowflakeDataSource):
@classmethod
def _get_databases(cls, connector: 'SnowflakeoAuth2Connector'):
return connector._get_databases()
@classmethod
def get_form(cls, connector: 'SnowflakeoAuth2Connector', current_config):
constraints = {}
with suppress(Exception):
databases = connector._get_databases()