def test_run_query_executes_desired_query(self, mocker):
mocker.patch("sqlalchemy.create_engine")
conn = SnowflakeConnector()
query = "MY FUN TESTING QUERY"
conn.run_query(query)
conn.engine.assert_has_calls([mocker.call.connect().__enter__().execute(query)])
def test_run_query_returns_results(self, mocker):
mocker.patch("sqlalchemy.create_engine")
conn = SnowflakeConnector()
expectedResult = "MY DATABASE RESULT"
mocker.patch.object(conn.engine.connect().__enter__(),
"execute",
return_value=expectedResult)
result = conn.run_query("query")
assert result is expectedResult
def test_snowflaky(self):
db1 = "analytics.schema.table"
db2 = "1234raw.schema.table"
db3 = '"123-with-quotes".schema.table'
db4 = "1_db-9-RANDOM.schema.table"
assert SnowflakeConnector.snowflaky(db1) == "analytics.schema.table"
assert SnowflakeConnector.snowflaky(db2) == "1234raw.schema.table"
assert SnowflakeConnector.snowflaky(db3) == '"123-with-quotes".schema.table'
assert SnowflakeConnector.snowflaky(db4) == '"1_db-9-RANDOM".schema.table'
def test_get_current_user(self, mocker):
mocker.patch("sqlalchemy.create_engine")
conn = SnowflakeConnector()
conn.run_query = mocker.MagicMock()
mocker.patch.object(
conn.run_query(), "fetchone", return_value={"user": "******"}
)
user = conn.get_current_user()
conn.run_query.assert_has_calls([mocker.call("SELECT CURRENT_USER() AS USER")])
assert user == "test_user"
def test_get_current_role(self, mocker):
mocker.patch("sqlalchemy.create_engine")
conn = SnowflakeConnector()
conn.run_query = mocker.MagicMock()
mocker.patch.object(
conn.run_query(), "fetchone", return_value={"role": "TEST_ROLE"}
)
role = conn.get_current_role()
conn.run_query.assert_has_calls([mocker.call("SELECT CURRENT_ROLE() AS ROLE")])
assert role == "test_role"
def test_uses_username_password_by_default(
selfself, mocker, snowflake_connector_env
):
mocker.patch("sqlalchemy.create_engine")
SnowflakeConnector()
sqlalchemy.create_engine.assert_called_with(
"snowflake://*****:*****@TEST/TEST?role=TEST&warehouse=TEST"
)
def test_uses_oauth_if_available(selfself, mocker, snowflake_connector_env):
mocker.patch("sqlalchemy.create_engine")
os.environ["PERMISSION_BOT_OAUTH_TOKEN"] = "TEST"
SnowflakeConnector()
del os.environ["PERMISSION_BOT_OAUTH_TOKEN"]
sqlalchemy.create_engine.assert_called_with(
"snowflake://TEST:@TEST/?authenticator=oauth&token=TEST&warehouse=TEST"
)
def test_show_roles(self, mocker):
mocker.patch("sqlalchemy.create_engine")
conn = SnowflakeConnector()
conn.run_query = mocker.MagicMock()
mocker.patch.object(
conn.run_query(),
"fetchall",
return_value=[
{"name": "TEST_ROLE", "owner": "SUPERADMIN"},
{"name": "SUPERADMIN", "owner": "SUPERADMIN"},
],
)
roles = conn.show_roles()
conn.run_query.assert_has_calls([mocker.call("SHOW ROLES")])
assert roles["test_role"] == "superadmin"
assert roles["superadmin"] == "superadmin"
def grant(spec, dry, diff):
"""Grant the permissions provided in the provided specification file."""
try:
spec_loader = SnowflakeSpecLoader(spec)
sql_grant_queries = spec_loader.generate_permission_queries()
click.secho()
if diff:
click.secho(
"SQL Commands generated for given spec file (Full diff with both new and already granted commands):"
)
else:
click.secho("SQL Commands generated for given spec file:")
click.secho()
conn = SnowflakeConnector()
for query in sql_grant_queries:
if not dry:
status = None
if not query.get("already_granted"):
try:
result = conn.run_query(query.get("sql"))
outcome = result.fetchall()
status = True
except:
status = False
ran_query = query
ran_query["run_status"] = status
print_command(ran_query, diff)
# If already granted, print command
else:
print_command(query, diff)
# If dry, print commands
else:
print_command(query, diff)
except SpecLoadingError as exc:
for line in str(exc).splitlines():
click.secho(line, fg="red")
sys.exit(1)
def check_permissions_on_snowflake_server(self,
conn: SnowflakeConnector = None
) -> None:
if conn is None:
conn = SnowflakeConnector()
error_messages = []
click.secho(f" Current user is: {conn.get_current_user()}.",
fg="green")
current_role = conn.get_current_role()
if "securityadmin" != current_role:
error_messages.append(
f"Current role is not securityadmin! "
"Permifrost expects to run as securityadmin, please update your connection settings."
)
click.secho(f" Current role is: {current_role}.", fg="green")
if error_messages:
raise SpecLoadingError("\n".join(error_messages))
def test_uses_key_pair_if_available(selfself, mocker, snowflake_connector_env):
mocker.patch("sqlalchemy.create_engine")
test_private_key = "TEST_PK"
mocker.patch.object(
SnowflakeConnector, "generate_private_key", return_value=test_private_key
)
os.environ["PERMISSION_BOT_KEY_PATH"] = "TEST"
os.environ["PERMISSION_BOT_KEY_PASSPHRASE"] = "TEST"
SnowflakeConnector()
del os.environ["PERMISSION_BOT_KEY_PATH"]
del os.environ["PERMISSION_BOT_KEY_PASSPHRASE"]
sqlalchemy.create_engine.assert_called_with(
"snowflake://TEST:@TEST/TEST?role=TEST&warehouse=TEST",
connect_args={"private_key": test_private_key},
)
def get_privileges_from_snowflake_server(self,
conn: SnowflakeConnector = None
) -> None:
"""
Get the privileges granted to users and roles in the Snowflake account
Gets the future privileges granted in all database and schema objects
Consolidates role and future privileges into a single object for self.grants_to_role
"""
if conn is None:
conn = SnowflakeConnector()
future_grants = {}
for database in self.entities["database_refs"]:
grant_results = conn.show_future_grants(database=database)
for role in grant_results:
for privilege in grant_results[role]:
for grant_on in grant_results[role][privilege]:
(future_grants.setdefault(role, {}).setdefault(
privilege, {}).setdefault(grant_on, []).extend(
grant_results[role][privilege][grant_on]))
# Get all schemas in all ref'd databases. Not all schemas will be
# ref'd in the spec.
for schema in conn.show_schemas(database=database):
grant_results = conn.show_future_grants(schema=schema)
for role in grant_results:
for privilege in grant_results[role]:
for grant_on in grant_results[role][privilege]:
(future_grants.setdefault(role, {}).setdefault(
privilege, {}).setdefault(grant_on, []).extend(
grant_results[role][privilege][grant_on]))
for role in self.entities["roles"]:
grant_results = conn.show_grants_to_role(role)
for privilege in grant_results:
for grant_on in grant_results[privilege]:
(future_grants.setdefault(role, {}).setdefault(
privilege, {}).setdefault(grant_on, []).extend(
grant_results[privilege][grant_on]))
self.grants_to_role = future_grants
for user in self.entities["users"]:
self.roles_granted_to_user[user] = conn.show_roles_granted_to_user(
user)
def check_entities_on_snowflake_server(self,
conn: SnowflakeConnector = None
) -> None:
"""
Make sure that all [warehouses, dbs, schemas, tables, users, roles]
referenced in the spec are defined in Snowflake.
Raises a SpecLoadingError with all the errors found while checking
Snowflake for missinf entities.
"""
error_messages = []
if conn is None:
conn = SnowflakeConnector()
if len(self.entities["warehouses"]) > 0:
warehouses = conn.show_warehouses()
for warehouse in self.entities["warehouses"]:
if warehouse not in warehouses:
error_messages.append(
f"Missing Entity Error: Warehouse {warehouse} was not found on"
" Snowflake Server. Please create it before continuing."
)
else:
logging.debug(
"`warehouses` not found in spec, skipping SHOW WAREHOUSES call."
)
if len(self.entities["databases"]) > 0:
databases = conn.show_databases()
for db in self.entities["databases"]:
if db not in databases:
error_messages.append(
f"Missing Entity Error: Database {db} was not found on"
" Snowflake Server. Please create it before continuing."
)
else:
logging.debug(
"`databases` not found in spec, skipping SHOW DATABASES call.")
if len(self.entities["schema_refs"]) > 0:
schemas = conn.show_schemas()
for schema in self.entities["schema_refs"]:
if "*" not in schema and schema not in schemas:
error_messages.append(
f"Missing Entity Error: Schema {schema} was not found on"
" Snowflake Server. Please create it before continuing."
)
else:
logging.debug(
"`schemas` not found in spec, skipping SHOW SCHEMAS call.")
if len(self.entities["table_refs"]) > 0:
tables = conn.show_tables()
views = conn.show_views()
for table in self.entities["table_refs"]:
if "*" not in table and table not in tables and table not in views:
error_messages.append(
f"Missing Entity Error: Table/View {table} was not found on"
" Snowflake Server. Please create it before continuing."
)
else:
logging.debug(
"`tables` not found in spec, skipping SHOW TABLES/VIEWS call.")
if len(self.entities["roles"]) > 0:
roles = conn.show_roles()
for role in self.entities["roles"]:
if role not in roles:
error_messages.append(
f"Missing Entity Error: Role {role} was not found on"
" Snowflake Server. Please create it before continuing."
)
else:
logging.debug(
"`roles` not found in spec, skipping SHOW ROLES call.")
if len(self.entities["users"]) > 0:
users = conn.show_users()
for user in self.entities["users"]:
if user not in users:
error_messages.append(
f"Missing Entity Error: User {user} was not found on"
" Snowflake Server. Please create it before continuing."
)
else:
logging.debug(
"`users` not found in spec, skipping SHOW USERS call.")
if error_messages:
raise SpecLoadingError("\n".join(error_messages))