def test_SnowflakeCredentialYamlHelper_key_pair_auth(
mock_emit, mock_prompt, empty_data_context_stats_enabled
):
helper = SnowflakeCredentialYamlHelper("my_datasource")
mock_prompt.side_effect = ["3"]
helper.prompt()
assert helper.auth_method == SnowflakeAuthMethod.KEY_PAIR
expected_credentials_snippet = """\
host = "YOUR_HOST" # The account name (include region -- ex 'ABCD.us-east-1')
username = "******"
database = "" # The database name
schema = "" # The schema name
warehouse = "" # The warehouse name
role = "" # The role name
private_key_path = "YOUR_KEY_PATH" # Path to the private key used for authentication
private_key_passphrase = "" # Passphrase for the private key used for authentication (optional -- leave blank for none)"""
assert helper.credentials_snippet() == expected_credentials_snippet
assert (
helper.yaml_snippet()
== '''f"""
name: {datasource_name}
class_name: Datasource
execution_engine:
class_name: SqlAlchemyExecutionEngine
credentials:
host: {host}
username: {username}
database: {database}
query:
schema: {schema}
warehouse: {warehouse}
role: {role}
private_key_path: {private_key_path}
private_key_passphrase: {private_key_passphrase}
drivername: snowflake
data_connectors:
default_runtime_data_connector_name:
class_name: RuntimeDataConnector
batch_identifiers:
- default_identifier_name
default_inferred_data_connector_name:
class_name: InferredAssetSqlDataConnector
include_schema_name: True"""'''
)
helper.send_backend_choice_usage_message(empty_data_context_stats_enabled)
_snowflake_usage_stats_assertions(mock_emit)
renderer = helper.get_notebook_renderer(empty_data_context_stats_enabled)
assert renderer.sql_credentials_code_snippet == expected_credentials_snippet
def test_SnowflakeCredentialYamlHelper_password_auth(
mock_emit, mock_prompt, empty_data_context_stats_enabled
):
helper = SnowflakeCredentialYamlHelper("my_datasource")
mock_prompt.side_effect = ["1"]
helper.prompt()
assert helper.auth_method == SnowflakeAuthMethod.USER_AND_PASSWORD
expected_credentials_snippet = '''\
host = "YOUR_HOST" # The account name (include region -- ex 'ABCD.us-east-1')
username = "******"
database = "" # The database name
schema = "" # The schema name
warehouse = "" # The warehouse name
role = "" # The role name
password = "******"'''
assert helper.credentials_snippet() == expected_credentials_snippet
assert (
helper.yaml_snippet()
== '''f"""
name: {datasource_name}
class_name: Datasource
execution_engine:
class_name: SqlAlchemyExecutionEngine
credentials:
host: {host}
username: {username}
database: {database}
query:
schema: {schema}
warehouse: {warehouse}
role: {role}
password: {password}
drivername: snowflake
data_connectors:
default_runtime_data_connector_name:
class_name: RuntimeDataConnector
batch_identifiers:
- default_identifier_name
default_inferred_data_connector_name:
class_name: InferredAssetSqlDataConnector
include_schema_name: True"""'''
)
helper.send_backend_choice_usage_message(empty_data_context_stats_enabled)
_snowflake_usage_stats_assertions(mock_emit)
renderer = helper.get_notebook_renderer(empty_data_context_stats_enabled)
assert renderer.sql_credentials_code_snippet == expected_credentials_snippet
def test_SnowflakeCredentialYamlHelper_sso_auth(
mock_emit, mock_prompt, empty_data_context_stats_enabled):
helper = SnowflakeCredentialYamlHelper("my_datasource")
mock_prompt.side_effect = ["2"]
helper.prompt()
assert helper.auth_method == SnowflakeAuthMethod.SSO
expected_credentials_snippet = """\
host = "YOUR_HOST" # The account name (include region -- ex 'ABCD.us-east-1')
username = "******"
database = "" # The database name (optional -- leave blank for none)
schema = "" # The schema name (optional -- leave blank for none)
warehouse = "" # The warehouse name (optional -- leave blank for none)
role = "" # The role name (optional -- leave blank for none)
authenticator_url = "externalbrowser" # A valid okta URL or 'externalbrowser' used to connect through SSO"""
assert helper.credentials_snippet() == expected_credentials_snippet
assert (helper.yaml_snippet() == '''f"""
name: {datasource_name}
class_name: Datasource
execution_engine:
class_name: SqlAlchemyExecutionEngine
credentials:
host: {host}
username: {username}
database: {database}
query:
schema: {schema}
warehouse: {warehouse}
role: {role}
connect_args:
authenticator: {authenticator_url}
drivername: snowflake
data_connectors:
default_runtime_data_connector_name:
class_name: RuntimeDataConnector
batch_identifiers:
- default_identifier_name
default_inferred_data_connector_name:
class_name: InferredAssetSqlDataConnector
name: whole_table"""''')
helper.send_backend_choice_usage_message(empty_data_context_stats_enabled)
_snowflake_usage_stats_assertions(mock_emit)
renderer = helper.get_notebook_renderer(empty_data_context_stats_enabled)
assert renderer.sql_credentials_code_snippet == expected_credentials_snippet
def test_SnowflakeCredentialYamlHelper_key_pair_auth(
mock_emit, mock_prompt, empty_data_context_stats_enabled):
helper = SnowflakeCredentialYamlHelper("my_datasource")
mock_prompt.side_effect = ["3"]
helper.prompt()
assert helper.auth_method == SnowflakeAuthMethod.KEY_PAIR
expected_credentials_snippet = """host = "YOUR_HOST" # The account name (include region -- ex 'ABCD.us-east-1')
username = "******"
database = "" # The database name (optional -- leave blank for none)
schema = "" # The schema name (optional -- leave blank for none)
warehouse = "" # The warehouse name (optional -- leave blank for none)
role = "" # The role name (optional -- leave blank for none)
private_key_path = "YOUR_KEY_PATH" # Path to the private key used for authentication
private_key_passphrase = "" # Passphrase for the private key used for authentication (optional -- leave blank for none)"""
assert helper.credentials_snippet() == expected_credentials_snippet
assert (helper.yaml_snippet() == '''f"""
name: {datasource_name}
class_name: SimpleSqlalchemyDatasource
introspection:
whole_table:
data_asset_name_suffix: __whole_table
credentials:
host: {host}
username: {username}
database: {database}
query:
schema: {schema}
warehouse: {warehouse}
role: {role}
private_key_path: {private_key_path}
private_key_passphrase: {private_key_passphrase}
drivername: snowflake"""''')
helper.send_backend_choice_usage_message(empty_data_context_stats_enabled)
_snowflake_usage_stats_assertions(mock_emit)
renderer = helper.get_notebook_renderer(empty_data_context_stats_enabled)
assert renderer.sql_credentials_code_snippet == expected_credentials_snippet