{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "23f38f64",
   "metadata": {},
   "source": [
    "# Sharing On-Demand Logs with RAI\n",
    "\n",
    "## Overview\n",
    "\n",
    "This notebook creates and shares a [secure-share view](https://docs.snowflake.com/en/user-guide/data-sharing-intro) with RAI to provide access to on-demand logs.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "885bf349",
   "metadata": {
    "language": "python"
   },
   "outputs": [],
   "source": [
    "import snowflake.connector\n",
    "from datetime import datetime, timedelta, date"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02635de2",
   "metadata": {
    "language": "python"
   },
   "outputs": [],
   "source": [
    "def get_date_range(start_date, end_date):\n",
    "    start = datetime.strptime(start_date, \"%Y-%m-%d\")\n",
    "    end = datetime.strptime(end_date, \"%Y-%m-%d\")\n",
    "    dates = [(start + timedelta(days=i)).strftime(\"%Y-%m-%d\")\n",
    "             for i in range((end - start).days + 1)]\n",
    "    return \"|\".join(dates)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9dad0e8b",
   "metadata": {},
   "source": [
    "Edit the values in the next cell with your specific values:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "153ef9e5",
   "metadata": {
    "language": "python"
   },
   "outputs": [],
   "source": [
    "# Override this with the account name in the format ORG-account_nameL\n",
    "snowflake_account = \"\"\n",
    "snowflake_user = \"\"\n",
    "snowflake_password = \"\"\n",
    "\n",
    "# Override this to the date you want to get logs starting from in the format YYYY-MM-DD:\n",
    "start_date = date.today().strftime(\"%Y-%m-%d\")\n",
    "\n",
    "# Override this to the date you want to get logs until in the format YYYY-MM-DD\n",
    "end_date = date.today().strftime(\"%Y-%m-%d\")\n",
    "date_range = get_date_range(start_date, end_date)\n",
    "\n",
    "# Override this to True if you want to share spcs_control_plane logs\n",
    "include_erps_logs = False\n",
    "\n",
    "# Override this to the engine name (ex, 'testEngine') or leave it as it is if you want to get logs for all engines\n",
    "engine_name = \".*\"\n",
    "warehouse = \"\"\n",
    "\n",
    "# Override this to a unique id and share it with RAI\n",
    "id = \"\"\n",
    "\n",
    "# The account you want to share the logs with\n",
    "event_sharing_account = \"\"\n",
    "\n",
    "# Your native app name (usually relationalai)\n",
    "native_app_name = \"relationalai\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a26b67a7",
   "metadata": {
    "language": "python"
   },
   "outputs": [],
   "source": [
    "engine_file_pattern = f'{engine_name}/clientlogs-.*({date_range}).*\\\\.json|{engine_name}/clientlogs-engine.json'\n",
    "erp_file_pattern = f'|clientlogs-cp-({date_range}).*\\\\.json|.*clientlogs-cp.json'\n",
    "if include_erps_logs:\n",
    "    file_pattern = f'.*({engine_file_pattern}{erp_file_pattern}).*'\n",
    "else:\n",
    "    file_pattern = f'.*({engine_file_pattern}).*'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56270549",
   "metadata": {
    "language": "python"
   },
   "outputs": [],
   "source": [
    "sql_query = \"\"\"\n",
    "USE ROLE ACCOUNTADMIN;\n",
    "USE WAREHOUSE &warehouse; -- Update to use another warehouse if necessary.\n",
    "\n",
    "CREATE DATABASE IF NOT EXISTS TELEMETRY_SHARING;\n",
    "USE DATABASE TELEMETRY_SHARING;\n",
    "CREATE SCHEMA IF NOT EXISTS LOGS;\n",
    "USE SCHEMA LOGS;\n",
    "\n",
    "--*****--\n",
    "-- Load staged data files to temporary tables\n",
    "--*****--\n",
    "CREATE OR REPLACE TABLE TELEMETRY_SHARING.LOGS.TELEMETRY_LOAD_TABLE_&id (\n",
    "  LOG_RECORD VARCHAR\n",
    ");\n",
    "\n",
    "CREATE OR REPLACE FILE FORMAT json_format TYPE = 'json';\n",
    "\n",
    "CREATE OR REPLACE TABLE TELEMETRY_SHARING.LOGS.TELEMETRY_SHARE_TABLE_&id  (\n",
    "  TIMESTAMP TIMESTAMP,\n",
    "  OBSERVED_TIMESTAMP TIMESTAMP,\n",
    "  SPAN_ID VARCHAR,\n",
    "  TRACE_ID VARCHAR,\n",
    "  MESSAGE VARCHAR,\n",
    "  LOG_RECORD VARCHAR\n",
    ");\n",
    "\n",
    "--*****--\n",
    "-- Create secure view from table with target log records\n",
    "--*****--\n",
    "CREATE OR REPLACE SECURE VIEW TELEMETRY_SHARING.LOGS.TELEMETRY_SHARE_VIEW_&id\n",
    "COMMENT = 'View containing telemetry records to share with the RAI provider account'\n",
    "AS\n",
    "    SELECT *\n",
    "    FROM TELEMETRY_SHARING.LOGS.TELEMETRY_SHARE_TABLE_&id ;\n",
    "\n",
    "--*****--\n",
    "-- Share secure view with the RAI provider account\n",
    "--*****--\n",
    "CREATE OR REPLACE SHARE TELEMETRY_SHARE_&id;\n",
    "\n",
    "CREATE DATABASE ROLE IF NOT EXISTS TELEMETRY_SHARE_ROLE;\n",
    "GRANT USAGE ON DATABASE TELEMETRY_SHARING TO SHARE TELEMETRY_SHARE_&id;\n",
    "GRANT USAGE ON SCHEMA TELEMETRY_SHARING.LOGS TO DATABASE ROLE TELEMETRY_SHARE_ROLE;\n",
    "GRANT SELECT ON VIEW TELEMETRY_SHARE_VIEW_&id TO DATABASE ROLE TELEMETRY_SHARE_ROLE;\n",
    "GRANT DATABASE ROLE TELEMETRY_SHARE_ROLE TO SHARE TELEMETRY_SHARE_&id;\n",
    "\n",
    "ALTER SHARE TELEMETRY_SHARE_&id ADD ACCOUNTS = NDSOEBE.&event_sharing_account;\n",
    "\n",
    "COPY INTO TELEMETRY_SHARING.LOGS.TELEMETRY_LOAD_TABLE_&id\n",
    "FROM (\n",
    "    SELECT\n",
    "        $1 AS log_record\n",
    "    FROM @&native_app_name.app_state.client_log_stage\n",
    ")\n",
    "PATTERN =  '&file_pattern'\n",
    "FILE_FORMAT = (TYPE = JSON)\n",
    "ON_ERROR = CONTINUE; -- This will skip any log records that are invalid JSON.\n",
    "-- The output of the query will indicate if any records were skipped due to errors.\n",
    "\n",
    "\n",
    "-- Copy from TELEMETRY_LOAD_TABLE_&id into TELEMETRY_SHARE_TABLE_&id and remove safe logs\n",
    "INSERT INTO TELEMETRY_SHARING.LOGS.TELEMETRY_SHARE_TABLE_&id\n",
    "    SELECT\n",
    "        to_timestamp(timeUnixNano) as timestamp,\n",
    "        to_timestamp(observedTimeUnixNano) as observed_timestamp,\n",
    "        spanId,\n",
    "        traceId,\n",
    "        a.value:value:stringValue as message,\n",
    "        log_record\n",
    "    FROM (SELECT\n",
    "            value:timeUnixNano as timeUnixNano,\n",
    "            value:observedTimeUnixNano as observedTimeUnixNano,\n",
    "            value:spanId as spanId,\n",
    "            value:traceId as traceId,\n",
    "            value as log_record\n",
    "        FROM TELEMETRY_SHARING.LOGS.TELEMETRY_LOAD_TABLE_&id, LATERAL FLATTEN( INPUT => TRY_PARSE_JSON($1):resourceLogs[0]:scopeLogs[0]:logRecords, OUTER => TRUE )),\n",
    "        LATERAL FLATTEN( INPUT => log_record:body:kvlistValue:values, OUTER => TRUE) a, LATERAL FLATTEN( INPUT => log_record:attributes, OUTER => TRUE) b\n",
    "    WHERE a.VALUE:key = 'message'\n",
    "    and\n",
    "        (\n",
    "            ( -- engine unsafe logs\n",
    "                b.value:key = 'log.file.name'\n",
    "                and b.value:value:stringValue ='engine-unsafe.log'\n",
    "            )\n",
    "            or\n",
    "            (\n",
    "                -- erps unsafe logs\n",
    "                log_record not like '%___safe_to_log%'\n",
    "                and log_record not like '%engine-safe.log%'\n",
    "                and log_record like '%spcs_control_plane%'\n",
    "\n",
    "            )\n",
    "        )\n",
    ";\n",
    "\"\"\"\n",
    "sql_query = (\n",
    "    sql_query\n",
    "        .replace('&warehouse', warehouse)\n",
    "        .replace('&id', id)\n",
    "        .replace('&event_sharing_account', event_sharing_account)\n",
    "        .replace('&date_range', date_range)\n",
    "        .replace('&native_app_name', native_app_name)\n",
    "        .replace('&file_pattern', file_pattern)\n",
    ")\n",
    "\n",
    "sql_statements = sql_query.split(';')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88410031",
   "metadata": {
    "language": "python"
   },
   "outputs": [],
   "source": [
    "conn = snowflake.connector.connect(\n",
    "    account=snowflake_account,\n",
    "    user=snowflake_user,\n",
    "    password= snowflake_password,\n",
    ")\n",
    "cur = conn.cursor()\n",
    "try:\n",
    "    for statement in sql_statements:\n",
    "        statement = statement.strip()\n",
    "        if statement:\n",
    "            cur.execute(statement)\n",
    "            cur.fetchall()\n",
    "finally:\n",
    "    cur.close()\n",
    "    conn.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3b16651",
   "metadata": {},
   "source": [
    "Run the following cell to confirm that the logs are copied to the secure share:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4faabc1",
   "metadata": {
    "language": "python"
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "query = \"\"\"\n",
    "USE ROLE ACCOUNTADMIN;\n",
    "USE WAREHOUSE &warehouse;\n",
    "SELECT * FROM TELEMETRY_SHARING.LOGS.TELEMETRY_SHARE_TABLE_&id limit 5;\n",
    "\"\"\"\n",
    "query = query.replace('&id', id)\n",
    "query = query.replace('&warehouse', warehouse)\n",
    "conn = snowflake.connector.connect(\n",
    "    account=snowflake_account,\n",
    "    user=snowflake_user,\n",
    "    password= snowflake_password,\n",
    ")\n",
    "cur = conn.cursor()\n",
    "try:\n",
    "    sql_statements = query.split(';')\n",
    "    for statement in sql_statements:\n",
    "        statement = statement.strip()\n",
    "        if statement:\n",
    "            cur.execute(statement)\n",
    "            result = cur.fetchall()\n",
    "\n",
    "    columns = [desc[0] for desc in cur.description]\n",
    "    df = pd.DataFrame(result, columns=columns)\n",
    "    print(df)\n",
    "finally:\n",
    "    cur.close()\n",
    "    conn.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Streamlit Notebook",
   "name": "streamlit"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}