def snowflake_get_first(
hook: SnowflakeHook, sql: str, parameters: Union[Mapping, Iterable, None] = None
) -> Tuple[Any, int, str]:
"""
Executes the sql and returns the first resulting row alongside with
session_id and query_id generated by Snowflake.
:param hook: Configured SnowflakeHook with Snoflake connection details.
:param sql: the sql statement to be executed (str) or a list of
sql statements to execute
:param parameters: The parameters to render the SQL query with.
:return: Query result alongside with session_id and query_id
assigned by Snowflake to each submitted query
"""
if sys.version_info[0] < 3:
sql = sql.encode("utf-8")
with closing(hook.get_conn()) as conn:
session_id = conn.session_id
with closing(conn.cursor()) as cur:
if parameters is not None:
cur.execute(sql, parameters)
else:
cur.execute(sql)
query_id = cur.sfqid
return cur.fetchone(), session_id, query_id
def fetch_data_from_snowflake():
hook = SnowflakeHook("snowflake_conn")
conn = hook.get_conn()
roles = []
users = []
with conn.cursor() as cursor:
cursor.execute("USE DATABASE MY_AUDIT_DB;")
cursor.execute("USE ROLE ACCOUNTADMIN;")
# Queries for Roles and Role Grants
cursor.execute("SHOW roles")
rec_set = cursor.fetchall()
for rec in rec_set:
roles.append(Role(rec[1], rec[9]))
for role in roles:
cursor.execute("SHOW GRANTS TO ROLE " + role.name)
grant_set = cursor.fetchall()
for cur_grant in grant_set:
role.add_grant(
RoleGrant(cur_grant[1], cur_grant[2], cur_grant[3]), roles
)
# Queries for User and User Roles
cursor.execute("SHOW users")
user_set = cursor.fetchmany(1000)
for user in user_set:
users.append(User(user))
while len(user_set) > 0:
user_set = cursor.fetchmany(1000)
for user in user_set:
users.append(User(user))
for user in users:
cursor.execute("SHOW GRANTS TO USER " + user.user_name)
user.get_roles(cursor.fetchall())
with open(ROLES_PATH, "w") as roles_file, open(
ROLE_GRANTS_PATH, "w"
) as role_grants_file:
for role in roles:
role.write_roles(roles_file)
role.write_grants(role.name, "ROOT", role_grants_file)
with open(USERS_PATH, "w") as users_file, open(
USER_ROLES_PATH, "w"
) as user_roles_file:
for user in users:
user.write_user_record(users_file)
user.write_roles(user_roles_file)
def snowflake_run(
hook: SnowflakeHook,
sql: Union[str, List[str]],
autocommit: bool = False,
parameters: Union[Mapping, Iterable, None] = None,
) -> Tuple[int, List[str]]:
"""
Runs a command or a list of commands and returns session_id and query_id(s).
Pass a list of sql statements to the sql parameter to get them to execute sequentially.
:param hook: Configured SnowflakeHook with Snoflake connection details.
:param sql: the sql statement to be executed (str) or a list of
sql statements to execute
:param autocommit: What to set the connection's autocommit setting to
before executing the query.
:param parameters: The parameters to render the SQL query with.
:return: session_id and query_id(s) assigned by Snowflake to each submitted query
"""
if isinstance(sql, six.string_types):
sql = [sql]
with closing(hook.get_conn()) as conn:
session_id, query_ids = conn.session_id, []
if hook.supports_autocommit:
hook.set_autocommit(conn, autocommit)
with closing(conn.cursor()) as cur:
for s in sql:
if sys.version_info[0] < 3:
s = s.encode("utf-8")
if parameters is not None:
hook.log.info("{} with parameters {}".format(s, parameters))
res = cur.execute(s, parameters)
else:
hook.log.info(s)
res = cur.execute(s)
query_ids.append(res.sfqid)
# If autocommit was set to False for db that supports autocommit,
# or if db does not supports autocommit, we do a manual commit.
if not hook.get_autocommit(conn):
conn.commit()
logger.info(
"Executed queries '{}', got session_id {}, query_id {}".format(
sql, session_id, query_ids
)
)
return (session_id, query_ids)
def pivot_data(**kwargs):
#Make connection to Snowflake
hook = SnowflakeHook(snowflake_conn_id='snowflake')
conn = hook.get_conn()
#Define SQL query
query = 'SELECT DATE, STATE, POSITIVE FROM STATE_DATA;'
#Read data into pandas dataframe
df = pd.read_sql(query, conn)
#Pivot dataframe into new format
pivot_df = df.pivot(index='DATE', columns='STATE', values='POSITIVE').reset_index()
#Save dataframe to S3
s3_hook = S3Hook(aws_conn_id=S3_CONN_ID)
s3_hook.load_string(pivot_df.to_csv(index=False),
'{0}.csv'.format(filename),
bucket_name=BUCKET,
replace=True)
def get_conn():
hook = SnowflakeHook(snowflake_conn_id=conn_id)
return hook.get_conn()
