#!/usr/bin/env python3

# -*- coding: UTF-8 -*-

import argparse

import collections

import logging

import re

import requests

logger = logging.getLogger(__name__)

logger.setLevel(logging.INFO)

ERR_CODES_PAGE = 'https://www.postgresql.org/docs/12/static/errcodes-appendix.html'

HTML_REGEX = r'(?:class=\"bold\"[\S]*>([^<]+))' '|' r'(?:class=\"literal\">([^<]+).*?class=.*?\"symbol\">([^<]+))'

DEF_OFFSET = ' '

LIT_OFFSET = ' '

CLASSES_ENUM = """

// PostgreSQL error classes

/// Docs: https://www.postgresql.org/docs/12/static/errcodes-appendix.html

/// Enumeration was generated by userver/scripts/postgres/pg_sql_codes.py

enum class SqlStateClass : std::int64_t {"""

ERRORS_ENUM = """

/// PostgreSQL error codes

/// Docs: https://www.postgresql.org/docs/12/static/errcodes-appendix.html

/// Enumeration was generated by userver/scripts/postgres/pg_sql_codes.py

enum class SqlState : std::int64_t {

kUnknownState, ///!< Unknown state, not in PostgreSQL docs"""

ERROR_LITERALS = """

// This goes to the cpp file in an anonymous namespace

// Data was generated by userver/scripts/postgres/pg_sql_codes.py

const std::unordered_map<std::string_view, SqlState> kCodeStrToState{"""

DISAMBIGUATIONS = {

'kWarning': 'Warn',

'kExternalRoutineException': 'Ex',

'kExternalRoutineInvocationException': 'Ex',

}

def static_vars(**kwargs):

def decorate(func):

for k in kwargs:

setattr(func, k, kwargs[k])

return func

return decorate

@static_vars(value=0)

def class_value(prev_count):

if class_value.value == 0:

class_value.value = 1

else:

val = class_value.value * 2

while class_value.value + prev_count > val:

val *= 2

class_value.value = val

return class_value.value

def gen_enum_value(snake_str):

components = snake_str.split('_')

return 'k' + ''.join(x.title() for x in components)

def print_with_offset(stream, offset, *args):

lines = '\n'.join(args).split('\n')

print(offset + ('\n' + offset).join(lines), file=stream)

class ErrorClass:

error_cnt_by_symbol: dict = collections.Counter()

class ErrorCode:

def __init__(self, klass, symbol, literal):

self.klass = klass

self.symbol = symbol

self.literal = literal

def disambiguated_symbol(self):

disambiguation = ''

if ErrorClass.error_cnt_by_symbol[self.symbol] > 1:

disambiguation = DISAMBIGUATIONS.get(self.klass, '')

return self.symbol + disambiguation

def print_decl(self, stream):

if self.symbol == self.klass:

print_with_offset(

stream,

DEF_OFFSET,

'{symbol} = static_cast<std::int64_t>(SqlStateClass::{symbol}), //!< {literal}'.format(

symbol=self.disambiguated_symbol(),

literal=self.literal,

),

)

else:

print_with_offset(

stream,

DEF_OFFSET,

'{symbol}, //!< {literal}'.format(

symbol=self.disambiguated_symbol(),

literal=self.literal,

),

)

def print_literal(self, stream):

print_with_offset(

stream,

LIT_OFFSET,

'{{"{literal}", SqlState::{symbol}}},'.format(

literal=self.literal,

symbol=self.disambiguated_symbol(),

),

)

def print_class_test(self, stream):

print_with_offset(

stream,

LIT_OFFSET,

'EXPECT_EQ(pg::SqlStateClass::{klass}, pg::GetSqlStateClass(pg::SqlState::{symbol}));'.format(

klass=self.klass,

symbol=self.disambiguated_symbol(),

),

)

def print_parse_test(self, stream):

print_with_offset(

stream,

LIT_OFFSET,

'EXPECT_EQ(pg::SqlState::{symbol}, pg::SqlStateFromString("{literal}"));'.format(

symbol=self.disambiguated_symbol(),

literal=self.literal,

),

)

def __init__(self, desc, prev_count):

self.desc = desc

self.value = class_value(prev_count)

self.errors = []

self.name = None

def add_error(self, symbol, literal):

if self.name is None:

self.name = symbol

self.errors.append(ErrorClass.ErrorCode(self.name, symbol, literal))

ErrorClass.error_cnt_by_symbol[symbol] += 1

def print_class(self, stream):

print_with_offset(

stream,

DEF_OFFSET,

'{name} = 0x{value:02x},'.format(name=self.name, value=self.value),

)

def print_symbols(self, stream):

print_with_offset(

stream,

DEF_OFFSET,

'//@{',

'/** @name {desc} */'.format(desc=self.desc),

)

for err_class in self.errors:

err_class.print_decl(stream)

print_with_offset(stream, DEF_OFFSET, '//@}')

def print_literals(self, stream):

print_with_offset(

stream,

LIT_OFFSET,

'//@{',

'/** @name {desc} */'.format(desc=self.desc),

)

for err_class in self.errors:

err_class.print_literal(stream)

print_with_offset(stream, LIT_OFFSET, '//@}')

def print_test(self, stream):

print_with_offset(

stream,

LIT_OFFSET,

'// {desc}'.format(desc=self.desc),

)

for err_class in self.errors:

err_class.print_parse_test(stream)

err_class.print_class_test(stream)

def main():

parser = argparse.ArgumentParser(description='SqlCode enum generator')

parser.add_argument(

'-o',

'--header',

help='Output header',

metavar='out.hpp',

)

parser.add_argument(

'-s',

'--source',

help='Output source',

metavar='out.cpp',

)

parser.add_argument(

'-t',

'--test',

help='Generate test',

metavar='test.cpp',

)

args = parser.parse_args()

logger.info('Retrieving documentation')

page = requests.get(ERR_CODES_PAGE)

logger.info('Cleaning html')

page_text = re.sub(

r'.*?<table[^>]+summary=\"PostgreSQL Error Codes\"[^>]*>',

'',

page.text,

flags=re.MULTILINE | re.DOTALL,

)

page_text = re.sub(

r'</table>.*',

'',

page_text,

flags=re.MULTILINE | re.DOTALL,

)

logger.info('Parsing html')

matches = re.finditer(HTML_REGEX, page_text, re.MULTILINE | re.DOTALL)

err_count = 0

current_class = None

error_classes = []

logger.info('Collecting definitions')

for match in matches:

if match.group(1):

class_desc = match.group(1).replace('\n', ' ')

current_class = ErrorClass(class_desc, err_count)

error_classes.append(current_class)

err_count = 0

elif match.group(2):

literal = match.group(2)

symbol = gen_enum_value(match.group(3))

current_class.add_error(symbol, literal)

err_count += 1

logger.info('Printing sources')

if args.header:

with open(args.header, 'w') as header:

print(CLASSES_ENUM, file=header)

for err_class in error_classes:

err_class.print_class(header)

print('};\n\n', file=header)

print(ERRORS_ENUM, file=header)

for err_class in error_classes:

err_class.print_symbols(header)

print('};\n', file=header)

if args.source:

with open(args.source, 'w') as source:

print(ERROR_LITERALS, file=source)

for err_class in error_classes:

err_class.print_literals(source)

print('};\n', file=source)

if args.test:

with open(args.test, 'w') as test:

print(

'// Test was generated by userver/scripts/postgres/pg_sql_codes.py',

'#include <userver/utest/utest.hpp>',

'',

'#include <userver/storages/postgres/exceptions.hpp>',

'',

'namespace pg = storages::postgres;',

'',

sep='\n',

file=test,

)

for err_class in error_classes:

print(

'TEST(PostgreError, SqlState' + err_class.name[1:] + ') {',

file=test,

)

err_class.print_test(test)

print('}\n', file=test)

logger.info('Done')

if __name__ == '__main__':

main()