Parse zipped PostgreSQL's logs and save them in a parquet file.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Parse zipped PostgreSQL's logs and save them in a parquet file\n",
    "\n",
    "I'm administrating a large number of PostgreSQL's servers and I get their logs zipped. To analyze them I've done a Spark task for:\n",
    "\n",
    "1. Unzip the files\n",
    "2. Parse then logs of PostgreSQL\n",
    "3. Save (append) the data into a parquet file\n",
    "\n",
    "In a following notebook I will show how to query them to get usefull information.\n",
    "\n",
    "\n",
    "### PostgreSQL logs\n",
    "\n",
    "The log format specified in the PostgreSQL's config file is the following:\n",
    "\n",
    "`log_line_prefix = '%t %a %u %d %c '`\n",
    "\n",
    "Special values:\n",
    "\n",
    "* %a = application name\n",
    "* %u = user name\n",
    "* %d = database name\n",
    "* %t = timestamp without milliseconds\n",
    "* %c = session ID"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Importing needed packages:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import datetime\n",
    "import io\n",
    "import re\n",
    "import zipfile\n",
    "\n",
    "from pyspark.sql.types import Row\n",
    "from pyspark.sql import SparkSession"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create Spark Session to work with:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spark = (SparkSession\n",
    "         .builder\n",
    "         .appName('PostgreSQL_logs_into_Parquet')\n",
    "         .config(\"spark.sql.execution.arrow.enabled\", \"true\")\n",
    "         .getOrCreate())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Functions \n",
    "\n",
    "### Read ZIP files\n",
    "\n",
    "To read zip files, there's no native funcition in Spark, so I had to create a function for this. I got the idea from [StackOverflow](\n",
    "https://stackoverflow.com/questions/28569788/how-to-open-stream-zip-files-through-spark)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_postgresql_zip(zip_filename: str, zip_bytes: bytes):\n",
    "    \"\"\"Unzip a zip bytes (already in memory, not a file) and add additional fields based on the filename of the zip file\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    zip_filename : str\n",
    "        Full path of the zip filename\n",
    "    zip_bytes : bytes\n",
    "        Bytes of the zip file\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    result : list\n",
    "        List of tuples (info_file, raw_log_file ) \n",
    "        where info_file is a dict with the zip_filename and zip_filename break into properties\n",
    "        and raw_log_file is a str\n",
    "    Example:\n",
    "        [({'type': 'POSTGRESQL',\n",
    "        'servername': 'myserver',\n",
    "        'date': '2018-10-02_120852',\n",
    "        'filename': 'postgresql-2018-01-01_000000.log',\n",
    "        'zip_filename': 'file:/zip_files/POSTGRESQL__myserver__2018-10-02_120852.zip'},\n",
    "        'raw_postgresql_log_data_HERE'), ({},)]\n",
    "    \"\"\"\n",
    "    raw_log_files = zip_extract(zip_bytes)\n",
    "    zip_filename_splitted = zip_filename.split('/')[-1].split('.')[0].split('__')\n",
    "    zip_filename_parts = {\n",
    "        'type': zip_filename_splitted[0],\n",
    "        'servername': zip_filename_splitted[1]\n",
    "    }\n",
    "    \n",
    "    return [({**zip_filename_parts, 'log_filename': file, 'zip_filename': zip_filename }, content)\n",
    "                       for file, content in raw_log_files]\n",
    "\n",
    "\n",
    "def zip_extract(zip_bytes: bytes):\n",
    "    \"\"\"Unzip a zip bytes (already in memory, not a file)\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    zipbytes : bytes\n",
    "        Bytes of the zip file\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    result : list\n",
    "        List of tuples (postgresql_log_filename, raw_log_content)\n",
    "    \"\"\"\n",
    "    in_memory_data = io.BytesIO(zip_bytes)\n",
    "    file_obj = zipfile.ZipFile(in_memory_data, \"r\")\n",
    "    zipped_files = [i for i in file_obj.namelist()]\n",
    "    files_content = []\n",
    "    for file in zipped_files:\n",
    "        try:\n",
    "            files_content.append((file, file_obj.open(file).read().decode('UTF-8')))\n",
    "        except:\n",
    "            print('Failed to process file {file}'.format(file=file))\n",
    "    return files_content"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### To parse PostgreSQL logs\n",
    "\n",
    "`parse_postgresql_log_file()` detects every log entry. It's important to note that a log entry can have multiple lines, for example when PostgreSQL logs DDL it writes in the log file all SQL executed to create tables, functions...\n",
    "\n",
    "Then each log entry is parsed with `parse_postgresql_log_entry()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_postgresql_log_file(info_file: dict, raw_file: str):\n",
    "    \"\"\"Parse a PostgreSQL log file with the folowing log line prefix:\n",
    "                log_line_prefix = '%t %a %u %d %c '\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    info_file : dict\n",
    "        Information related to the raw_file (date, etc..)\n",
    "    raw_file : str\n",
    "        The PostgreSQL log to parse\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    result : list\n",
    "        List of all entries parsed. So, it's a list of dicts\n",
    "    \"\"\"\n",
    "    reg_expression = '^(?P<date>\\d{4}-\\d{2}-\\d{2}) (?P<time>\\d{2}:\\d{2}:\\d{2}) CEST '\n",
    "    pattern = re.compile(reg_expression, re.MULTILINE)\n",
    "    raw_entries = pattern.finditer(raw_file)\n",
    "    prior_entry = None\n",
    "    log_entries = []\n",
    "    for entry in raw_entries:\n",
    "        if prior_entry is not None:\n",
    "            log_entries.append(parse_postgresql_log_entry(info_file, raw_file[prior_entry.start():entry.start()]))\n",
    "        prior_entry = entry\n",
    "    return log_entries\n",
    "\n",
    "def parse_postgresql_log_entry(info_file: dict, log_entry: str):\n",
    "    \"\"\"Parse an entry of a PostgreSQL log file. Usuarlly one entry is one line, \n",
    "            but can be multiple lines in case of DDL.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    info_file : dict\n",
    "        Information related to the raw_file (date, etc..)\n",
    "    log_entry : str\n",
    "        Log entry to parse. Must have the folowing log line prefix:\n",
    "                log_line_prefix = '%t %a %u %d %c '\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    result : dict\n",
    "        A dictionary with the parsed information\n",
    "    \"\"\"    \n",
    "    try:\n",
    "        reg_expression = '(?P<timestamp>\\d{4}-\\d{2}-\\d{2} \\d{2}:\\d{2}:\\d{2}) CEST (?P<app>.*) (?P<user>\\S*) (?P<db>\\S*) (?P<trans_id>\\w{8}\\.\\w*) (?P<msg_type>INFO|STATEMENT|ERROR|LOG|DETAIL|FATAL):  (?P<msg>(.|\\s)*)'\n",
    "        pattern = re.compile(reg_expression)\n",
    "        match = pattern.match(log_entry)\n",
    "        dict_mappings = match.groupdict()\n",
    "        dict_mappings['log_type'] = info_file['type']\n",
    "        dict_mappings['servername'] = info_file['servername']\n",
    "        dict_mappings['timestamp'] = datetime.strptime(dict_mappings['timestamp'], '%Y-%m-%d %H:%M:%S') #'2018-03-27 09:25:22'\n",
    "        return dict_mappings\n",
    "    except:\n",
    "        reg_expression = '(?P<timestamp>\\d{4}-\\d{2}-\\d{2} \\d{2}:\\d{2}:\\d{2}) CEST (?P<msg>.*)'\n",
    "        pattern = re.compile(reg_expression)\n",
    "        match = pattern.match(log_entry)\n",
    "        if match:\n",
    "            dict_mappings = match.groupdict()\n",
    "            dict_mappings['log_type'] = info_file['type']\n",
    "            dict_mappings['servername'] = info_file['servername']\n",
    "            dict_mappings['timestamp'] = datetime.strptime(dict_mappings['timestamp'], '%Y-%m-%d %H:%M:%S')\n",
    "            dict_mappings['msg_type'] = 'NOT_PARSED'\n",
    "            dict_mappings['app'] = ''\n",
    "            dict_mappings['user'] = ''\n",
    "            dict_mappings['db'] = ''\n",
    "            dict_mappings['trans_id'] = ''\n",
    "        else:\n",
    "            dict_mappings = { 'log_type': info_file['type'], 'servername': info_file['servername'],\n",
    "                             'msg_type': 'NOT_PARSED', 'app': '', 'user': '', 'db': '',\n",
    "                             'timestamp': None, 'trans_id': '', 'msg': log_entry }\n",
    "        return dict_mappings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "##  Process zip files\n",
    "\n",
    "`binaryFiles()` returns a RDD where each element is a tuple with the content: `(zip_filename, zip_content)`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "zips = spark.sparkContext.binaryFiles(\"/zips/POSTGRESQL__*.zip\")\n",
    "zips.count()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I use `flatMap()` the zip extraction because we want that each row is a log file, instead of being grouped by the original zip files.\n",
    "\n",
    "Then `raw_pg_logs` is a RDD of tuples. An example of the tuple:\n",
    "\n",
    "`({'type': 'POSTGRESQL',\n",
    "        'servername': 'myserver',\n",
    "        'date': '2018-10-02_120852',\n",
    "        'filename': 'postgresql-2018-01-01_000000.log',\n",
    "        'zip_filename': 'file:/zip_files/POSTGRESQL__myserver__2018-10-02_120852.zip'},\n",
    "        'raw_postgresql_log_data_HERE')`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_pg_logs = zips.flatMap(lambda x: extract_postgresql_zip(x[0], x[1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`flatMap()` the parsing of the log files to have that each raw is a log entry:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pg_logs = raw_pg_logs.flatMap(lambda log: parse_postgresql_log_file(log[0], log[1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Save logs to a Parquet fine\n",
    "\n",
    "When saving a Parquet file, Spark enables us to partition it by a field. Databricks Support [recomendation](https://forums.databricks.com/questions/101/what-is-an-optimal-size-for-file-partitions-using.html) about Parquet partitions size:\n",
    "> This all depends on the dataset size and specific use cases, but, in general, we've seen that Parquet partitions of about 1GB are optimal.\n",
    "\n",
    "To save the RDD to a parquet file, I convert it to a Spark DataFrame:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "df = pg_logs.map(lambda r: Row(**r)).toDF()\n",
    "df.write.partitionBy(['servername']).mode(\"append\").parquet('pg_logs.parquet')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}