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()