# Copyright (C) 2026 The Software Heritage developers
# See the AUTHORS file at the top-level directory of this distribution
# License: GNU General Public License version 3, or any later version
# See top-level LICENSE file for more information
from datetime import datetime
from typing import Any, Dict, List, Optional
import psycopg
from psycopg.rows import dict_row
from psycopg.types.json import Jsonb
import psycopg_pool
from swh.core.db import BaseDb
from swh.core.db.common import db_transaction
from swh.model.swhids import CoreSWHID
from swh.vulns.osv.model import OSVVulnerabilityEvent
from swh.vulns.storage.interface import VulnerabilitiesStorageInterface
[docs]
class VulnerabilitiesStorage(VulnerabilitiesStorageInterface):
current_version = 1
def __init__(
self,
db,
min_pool_conns: int = 1,
max_pool_conns: int = 10,
):
"""
Args:
db: either a libpq connection string, or a psycopg connection
"""
self._db = None
if isinstance(db, psycopg.Connection):
self._pool = None
self._db = BaseDb(db)
else:
self._pool = psycopg_pool.ConnectionPool(
conninfo=db,
min_size=min_pool_conns,
max_size=max_pool_conns,
kwargs={"row_factory": dict_row},
)
# Wait for the first connection to be ready, and raise the
# appropriate exception if connection fails
self._pool.open(wait=True, timeout=1)
[docs]
def get_db(self):
if self._db:
return self._db
assert self._pool
db = BaseDb.from_pool(self._pool)
return db
[docs]
def put_db(self, db):
if db is not self._db:
db.put_conn()
[docs]
@db_transaction()
def osv_vulnerability_report_add(
self,
vulnerability_id: str,
processing_date: datetime,
json_report: Dict[str, Any],
db,
cur,
):
report = Jsonb(json_report)
cur.execute(
"""
INSERT INTO osv_report (vulnerability_id, last_processed, json_report)
VALUES (%s, %s, %s)
ON CONFLICT(vulnerability_id) DO UPDATE SET last_processed = %s, json_report = %s
""",
(vulnerability_id, processing_date, report, processing_date, report),
)
[docs]
@db_transaction()
def osv_vulnerability_report_last_processing_date(
self,
vulnerability_id: str,
db,
cur,
) -> Optional[datetime]:
row = cur.execute(
"""
SELECT last_processed from osv_report WHERE vulnerability_id = %s
""",
(vulnerability_id,),
).fetchone()
return row["last_processed"] if row else None
[docs]
@db_transaction()
def osv_vulnerability_report_get(
self,
vulnerability_id: str,
db,
cur,
) -> Optional[Dict[str, Any]]:
row = cur.execute(
"""
SELECT json_report from osv_report WHERE vulnerability_id = %s
""",
(vulnerability_id,),
).fetchone()
return row["json_report"] if row else None
vulnerability_event_cols = [
"vulnerability_id",
"vulnerability_severity",
"event_type",
"origin_url",
"swhid",
"version",
]
[docs]
@db_transaction()
def osv_vulnerability_event_add(
self, vulnerabilities_events: List[OSVVulnerabilityEvent], db=None, cur=None
):
for vulnerability_event in vulnerabilities_events:
vulnerability_event_d = vulnerability_event.to_dict()
cur.execute(
"INSERT INTO osv_vulnerability_event "
f"({', '.join(self.vulnerability_event_cols)}) "
"VALUES (%s, %s, %s, %s, %s, %s) ON CONFLICT DO NOTHING",
tuple(
vulnerability_event_d[key] for key in self.vulnerability_event_cols
),
)
[docs]
@db_transaction()
def osv_vulnerability_event_get_by_id(
self, vulnerabilities_ids: List[str], db=None, cur=None
) -> List[OSVVulnerabilityEvent]:
query = (
f"SELECT {', '.join(self.vulnerability_event_cols)} "
"FROM osv_vulnerability_event "
"WHERE vulnerability_id = ANY(%s)"
)
cur.execute(query, (vulnerabilities_ids,))
return [OSVVulnerabilityEvent.from_dict(row) for row in cur]
[docs]
@db_transaction()
def osv_vulnerability_event_get_by_swhid(
self, swhids: List[CoreSWHID], db=None, cur=None
) -> List[OSVVulnerabilityEvent]:
query = (
f"SELECT {', '.join(self.vulnerability_event_cols)} "
"FROM osv_vulnerability_event "
"WHERE swhid = ANY(%s)"
)
cur.execute(query, ([str(swhid) for swhid in swhids],))
return [OSVVulnerabilityEvent.from_dict(row) for row in cur]
[docs]
@db_transaction()
def osv_vulnerability_event_get_by_origin_url(
self, origin_urls: List[str], db=None, cur=None
) -> List[OSVVulnerabilityEvent]:
query = (
f"SELECT {', '.join(self.vulnerability_event_cols)} "
"FROM osv_vulnerability_event "
"WHERE origin_url = ANY(%s)"
)
cur.execute(query, (origin_urls,))
return [OSVVulnerabilityEvent.from_dict(row) for row in cur]
