From eb107a99d7cb1740af20fdb36346c80ad7052087 Mon Sep 17 00:00:00 2001 From: Sebastian Wagner Date: Tue, 28 Apr 2026 21:10:20 +0200 Subject: [PATCH 1/2] fix: psql initdb: handle existing severity_enum postgres does not have a CREATE TYPE IF NOT EXISTS do as a workaround, we need to encapuslate the CREATE TYPE in a kind of try-catch block --- CHANGELOG.md | 2 ++ intelmq/bin/intelmq_psql_initdb.py | 21 +++++++++++++-------- intelmq/tests/bin/initdb.sql | 20 ++++++++++++-------- 3 files changed, 27 insertions(+), 16 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 77cba765b..c1b8a4e24 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -48,6 +48,8 @@ Please refer to the [NEWS](NEWS.md) for a list of changes which have an affect o ### Tools - `intelmq.lib.bot_debugger`: Optionally read input messages from stdin instead of parameter value (PR#2678 by Sebastian Wager). +- `intelmq.bin.intelmq_psql_initdb`: + - Create the custom type `severity_enum` only when it does not yet exists, allows easy re-execution (PR#2701 by Sebastian Wagner). ### Contrib diff --git a/intelmq/bin/intelmq_psql_initdb.py b/intelmq/bin/intelmq_psql_initdb.py index 82359bdb7..bf8b0420c 100644 --- a/intelmq/bin/intelmq_psql_initdb.py +++ b/intelmq/bin/intelmq_psql_initdb.py @@ -138,15 +138,20 @@ def generate(harmonization_file=HARMONIZATION_CONF_FILE, skip_events=False, FIELDS = {} # ENUM for severity does not only save space, it first and foremost allows for easy sorting by severity (ascending sorting is critical to undefined) + # PostgreSQL has no CREATE TYPE IF NOT EXISTS, so we use a DO block as workaround sql_lines = dedent(""" - CREATE TYPE severity_enum AS ENUM ( - 'critical', - 'high', - 'medium', - 'low', - 'info', - 'undefined' - );""").strip().splitlines() + DO $$ BEGIN + CREATE TYPE severity_enum AS ENUM ( + 'critical', + 'high', + 'medium', + 'low', + 'info', + 'undefined' + ); + EXCEPTION + WHEN duplicate_object THEN NULL; + END $$;""").strip().splitlines() try: print("INFO - Reading %s file" % harmonization_file) diff --git a/intelmq/tests/bin/initdb.sql b/intelmq/tests/bin/initdb.sql index 3d3b4a2fc..09a8e90e7 100644 --- a/intelmq/tests/bin/initdb.sql +++ b/intelmq/tests/bin/initdb.sql @@ -1,11 +1,15 @@ -CREATE TYPE severity_enum AS ENUM ( - 'critical', - 'high', - 'medium', - 'low', - 'info', - 'undefined' -); +DO $$ BEGIN + CREATE TYPE severity_enum AS ENUM ( + 'critical', + 'high', + 'medium', + 'low', + 'info', + 'undefined' + ); +EXCEPTION + WHEN duplicate_object THEN NULL; +END $$; CREATE TABLE events ( "id" BIGSERIAL UNIQUE PRIMARY KEY, "classification.identifier" text, From ecfe816d78212eaedf1712cb5257f4e6f3612864 Mon Sep 17 00:00:00 2001 From: Sebastian Wagner Date: Tue, 28 Apr 2026 21:52:58 +0200 Subject: [PATCH 2/2] fix: psql initdb: use argparse.FileType for output file this way, the tool supports writing to stdout and also allows overwriting the output file, which is necessary to get your own predictable file name, e.g. in orchestration 1. create a temp file 2. set permisions 3. add the content with intelmq_psql_initdb, writing to the empty file 4. delete it --- CHANGELOG.md | 3 ++- intelmq/bin/intelmq_psql_initdb.py | 29 ++++++++--------------------- 2 files changed, 10 insertions(+), 22 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c1b8a4e24..6a1ddd9f0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -49,7 +49,8 @@ Please refer to the [NEWS](NEWS.md) for a list of changes which have an affect o ### Tools - `intelmq.lib.bot_debugger`: Optionally read input messages from stdin instead of parameter value (PR#2678 by Sebastian Wager). - `intelmq.bin.intelmq_psql_initdb`: - - Create the custom type `severity_enum` only when it does not yet exists, allows easy re-execution (PR#2701 by Sebastian Wagner). + - Create the custom type `severity_enum` only when it does not yet exists, allows easy re-execution. + - Use `argparse.FileType` and support stdout and file overwriting for predictable output file naming (PR#2701 by Sebastian Wagner). ### Contrib diff --git a/intelmq/bin/intelmq_psql_initdb.py b/intelmq/bin/intelmq_psql_initdb.py index bf8b0420c..21ef4722d 100644 --- a/intelmq/bin/intelmq_psql_initdb.py +++ b/intelmq/bin/intelmq_psql_initdb.py @@ -12,9 +12,7 @@ """ import argparse import json -import os import sys -import tempfile from textwrap import dedent from intelmq import HARMONIZATION_CONF_FILE @@ -154,12 +152,12 @@ def generate(harmonization_file=HARMONIZATION_CONF_FILE, skip_events=False, END $$;""").strip().splitlines() try: - print("INFO - Reading %s file" % harmonization_file) + print(f"INFO - Reading {harmonization_file} file", file=sys.stderr) with open(harmonization_file) as fp: DATA = json.load(fp)['event'] except OSError: - print("ERROR - Could not find %s" % harmonization_file) - print("ERROR - Make sure that you have intelmq installed.") + print(f"ERROR - Could not find {harmonization_file}", file=sys.stderr) + print("ERROR - Make sure that you have intelmq installed.", file=sys.stderr) sys.exit(2) for field in DATA.keys(): @@ -215,8 +213,9 @@ def main(): formatter_class=argparse.RawDescriptionHelpFormatter ) parser.add_argument('-o', '--outputfile', - help='Defines the Ouputfile', - default='/tmp/initdb.sql' + help='Defines the Outputfile', + default='/tmp/initdb.sql', + type=argparse.FileType('w') ) parser.add_argument("--no-events", action="store_true", default=False, help="Skip generating the events table schema") @@ -235,16 +234,7 @@ def main(): help="Do not use JSONB but JSON type to represent dictionary fields") args = parser.parse_args() - OUTPUTFILE = args.outputfile - fp = None - try: - if os.path.exists(OUTPUTFILE): - print(f'INFO - File {OUTPUTFILE} exists, generating temporary file.') - os_fp, OUTPUTFILE = tempfile.mkstemp(suffix='.initdb.sql', - text=True) - fp = os.fdopen(os_fp, 'wt') - else: - fp = open(OUTPUTFILE, 'w') + with args.outputfile as fp: psql = generate(args.harmonization, skip_events=args.no_events, separate_raws=args.separate_raws, @@ -252,11 +242,8 @@ def main(): skip_or_replace=args.skip_or_replace, no_jsonb=args.no_jsonb, ) - print("INFO - Writing %s file" % OUTPUTFILE) + print(f"INFO - Writing {fp.name} file", file=sys.stderr) fp.write(psql) - finally: - if fp: - fp.close() if __name__ == '__main__': # pragma: no cover