diff --git a/migrations/versions/ac04850e5f68_add_is_ioc_field.py b/migrations/versions/ac04850e5f68_add_is_ioc_field.py index e9bc161ac26b822bd885654f1341a796bedc902a..2c96468d9d60a809d86dcf0ecb5175aec4b2af87 100644 --- a/migrations/versions/ac04850e5f68_add_is_ioc_field.py +++ b/migrations/versions/ac04850e5f68_add_is_ioc_field.py @@ -32,17 +32,18 @@ def upgrade(): conn = op.get_bind() res = conn.execute("SELECT id FROM tag WHERE name = 'IOC';") row = res.fetchone() - ioc_tag_id = row[0] - res.close() - res = conn.execute( - f"""SELECT interface.host_id FROM interface - INNER JOIN interfacetags ON interface.id = interfacetags.interface_id - WHERE interfacetags.tag_id = {ioc_tag_id}; - """ - ) - results = res.fetchall() - for result in results: - op.execute(host.update().where(host.c.id == result[0]).values(is_ioc=True)) + if row is not None: + ioc_tag_id = row[0] + res.close() + res = conn.execute( + f"""SELECT interface.host_id FROM interface + INNER JOIN interfacetags ON interface.id = interfacetags.interface_id + WHERE interfacetags.tag_id = {ioc_tag_id}; + """ + ) + results = res.fetchall() + for result in results: + op.execute(host.update().where(host.c.id == result[0]).values(is_ioc=True)) op.drop_table("interfacetags") op.drop_table("tag")