atlas-init 0.4.5__py3-none-any.whl → 0.6.0__py3-none-any.whl

atlas_init/crud/mongo_dao.py

@@ -0,0 +1,296 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+from dataclasses import dataclass, field
+from datetime import datetime
+from functools import cached_property
+from pathlib import Path
+from typing import ClassVar, Self
+
+from model_lib import Entity, dump, field_names, parse_model
+from motor.motor_asyncio import AsyncIOMotorCollection
+from pydantic import model_validator
+from zero_3rdparty.file_utils import ensure_parents_write_text
+from zero_3rdparty.iter_utils import ignore_falsy
+
+from atlas_init.cli_tf.go_test_run import GoTestRun
+from atlas_init.cli_tf.go_test_tf_error import (
+    ErrorClassAuthor,
+    ErrorDetailsT,
+    GoTestAPIError,
+    GoTestError,
+    GoTestErrorClass,
+    GoTestErrorClassification,
+    GoTestResourceCheckError,
+)
+from atlas_init.crud.mongo_client import get_collection, init_mongo
+from atlas_init.crud.mongo_utils import MongoQueryOperation, create_or_replace, dump_with_id
+from atlas_init.repos.path import TFResoure, terraform_resources
+from atlas_init.settings.env_vars import AtlasInitSettings
+
+logger = logging.getLogger(__name__)
+
+
+def crud_dir(settings: AtlasInitSettings) -> Path:
+    return settings.static_root / "crud"
+
+
+@dataclass
+class TFResources:
+    resources: list[TFResoure] = field(default_factory=list)
+
+    def find_test_resources(self, test: GoTestRun) -> list[str]:
+        found_resources = []
+        for resource in self.resources:
+            url = test.package_url
+            if url and url.endswith(resource.package_rel_path):
+                found_resources.append(resource.name)
+        return found_resources
+
+
+def read_tf_resources(settings: AtlasInitSettings, repo_path: Path, branch: str) -> TFResources:
+    return TFResources(resources=terraform_resources(repo_path))
+
+
+class TFErrors(Entity):
+    errors: list[GoTestError] = field(default_factory=list)
+
+    @model_validator(mode="after")
+    def sort_errors(self) -> TFErrors:
+        self.errors.sort()
+        return self
+
+    def look_for_existing_classifications(self, error: GoTestError) -> tuple[GoTestErrorClass, GoTestErrorClass] | None:
+        for candidate in self.errors:
+            if error.match(candidate) and (classifications := candidate.classifications):
+                logger.info(f"found existing classification for {error.run.name}: {classifications}")
+                return classifications
+
+    def classified_errors(self) -> list[GoTestError]:
+        return [error for error in self.errors if error.classifications is not None]
+
+
+def read_tf_errors(settings: AtlasInitSettings) -> TFErrors:
+    path = crud_dir(settings) / "tf_errors.yaml"
+    return parse_model(path, TFErrors) if path.exists() else TFErrors()
+
+
+def read_tf_errors_for_day(settings: AtlasInitSettings, branch: str, date: datetime) -> list[GoTestError]:
+    raise NotImplementedError
+
+
+def store_or_update_tf_errors(settings: AtlasInitSettings, errors: list[GoTestError]) -> None:
+    existing = read_tf_errors(settings)
+    new_error_ids = {error.run.id for error in errors}
+    existing_without_new = [error for error in existing.errors if error.run.id not in new_error_ids]
+    all_errors = existing_without_new + errors
+    yaml_dump = dump(TFErrors(errors=all_errors), "yaml")
+    ensure_parents_write_text(crud_dir(settings) / "tf_errors.yaml", yaml_dump)
+
+
+def read_tf_error_by_run(settings: AtlasInitSettings, run: GoTestRun) -> GoTestError | None:
+    errors = read_tf_errors(settings)
+    return next((error for error in errors.errors if error.run.id == run.id), None)
+
+
+class TFTestRuns(Entity):
+    test_runs: list[GoTestRun] = field(default_factory=list)
+
+    @model_validator(mode="after")
+    def sort_test_runs(self) -> TFTestRuns:
+        self.test_runs.sort()
+        return self
+
+
+def read_tf_test_runs(settings: AtlasInitSettings) -> list[GoTestRun]:
+    path = crud_dir(settings) / "tf_test_runs.yaml"
+    return parse_model(path, TFTestRuns).test_runs if path.exists() else []
+
+
+def read_tf_tests_for_day(settings: AtlasInitSettings, branch: str, date: datetime) -> list[GoTestRun]:
+    start_date = date.replace(hour=0, minute=0, second=0, microsecond=0)
+    end_date = start_date.replace(hour=23, minute=59, second=59, microsecond=999999)
+    return read_tf_tests(settings, branch, start_date, end_date)
+
+
+def read_tf_tests(
+    settings: AtlasInitSettings, branch: str, start_date: datetime, end_date: datetime | None = None
+) -> list[GoTestRun]:
+    raise NotImplementedError
+
+
+async def init_mongo_dao(settings: AtlasInitSettings) -> MongoDao:
+    dao = MongoDao(settings=settings)
+    return await dao.connect()
+
+
+class GoTestRunNotFound(Exception):
+    def __init__(self, run_id: str) -> None:
+        self.run_id = run_id
+        super().__init__(run_id)
+
+
+@dataclass
+class MongoDao:
+    settings: AtlasInitSettings
+    property_keys_run: ClassVar[list[str]] = ["group_name"]
+
+    @cached_property
+    def runs(self) -> AsyncIOMotorCollection:
+        return get_collection(GoTestRun)
+
+    @cached_property
+    def classifications(self) -> AsyncIOMotorCollection:
+        return get_collection(GoTestErrorClassification)
+
+    @cached_property
+    def _field_names_runs(self) -> set[str]:
+        return set(field_names(GoTestRun)) | set(self.property_keys_run)
+
+    async def connect(self) -> Self:
+        await init_mongo(
+            mongo_url=self.settings.mongo_url,
+            db_name=self.settings.mongo_database,
+        )
+        return self
+
+    async def store_tf_test_runs(self, test_runs: list[GoTestRun]) -> list[GoTestRun]:
+        if not test_runs:
+            return []
+        col = self.runs
+        tasks = []
+        loop = asyncio.get_event_loop()
+        for run in test_runs:
+            dumped = dump_with_id(run, id=run.id, dt_keys=["ts", "finish_ts"], property_keys=self.property_keys_run)
+            tasks.append(loop.create_task(create_or_replace(col, dumped)))
+        await asyncio.gather(*tasks)
+        return test_runs
+
+    async def read_tf_tests_for_day(self, branch: str, date: datetime) -> list[GoTestRun]:
+        start_date = date.replace(hour=0, minute=0, second=0, microsecond=0)
+        end_date = start_date.replace(hour=23, minute=59, second=59, microsecond=999999)
+        query = {
+            "branch": branch,
+            "ts": {MongoQueryOperation.gte: start_date, MongoQueryOperation.lte: end_date},
+        }
+        return await self._find_runs(query)
+
+    async def _find_runs(self, query: dict) -> list[GoTestRun]:
+        runs = []
+        async for raw_run in self.runs.find(query):
+            runs.append(self._parse_run(raw_run))
+        return runs
+
+    async def read_error_classifications(
+        self, run_ids: list[str] | None = None
+    ) -> dict[str, GoTestErrorClassification]:
+        run_ids = run_ids or []
+        if not run_ids:
+            return {}
+        query = {"_id": {MongoQueryOperation.in_: run_ids}}
+        return await self._find_classifications(query)
+
+    async def _find_classifications(self, query: dict) -> dict[str, GoTestErrorClassification]:
+        classifications: dict[str, GoTestErrorClassification] = {}
+        async for raw_error in self.classifications.find(query):
+            run_id = raw_error.pop("_id", None)
+            classification = parse_model(raw_error, t=GoTestErrorClassification)
+            classifications[run_id] = classification
+        return classifications
+
+    async def read_similar_error_classifications(
+        self, details: ErrorDetailsT, *, author_filter: ErrorClassAuthor | None = None
+    ) -> dict[str, GoTestErrorClassification]:
+        query = {}
+        if author_filter:
+            query["author"] = {MongoQueryOperation.eq: author_filter}
+        match details:
+            case GoTestAPIError(
+                api_error_code_str=api_error_code_str,
+                api_method=api_method,
+                api_response_code=api_response_code,
+                api_path_normalized=api_path_normalized,
+            ) if api_path_normalized:
+                query |= {
+                    "details.api_error_code_str": {MongoQueryOperation.eq: api_error_code_str},
+                    "details.api_method": {MongoQueryOperation.eq: api_method},
+                    "details.api_response_code": {MongoQueryOperation.eq: api_response_code},
+                    "details.api_path_normalized": {MongoQueryOperation.eq: api_path_normalized},
+                }
+            case GoTestResourceCheckError(
+                tf_resource_name=tf_resource_name,
+                tf_resource_type=tf_resource_type,
+                step_nr=step_nr,
+                check_errors=check_errors,
+                test_name=test_name,
+            ):
+                query |= {
+                    "details.tf_resource_name": {MongoQueryOperation.eq: tf_resource_name},
+                    "details.tf_resource_type": {MongoQueryOperation.eq: tf_resource_type},
+                    "details.step_nr": {MongoQueryOperation.eq: step_nr},
+                    "test_name": {MongoQueryOperation.eq: test_name},
+                }
+                classifications = await self._find_classifications(query)
+                return {
+                    run_id: classification
+                    for run_id, classification in classifications.items()
+                    if isinstance(classification.details, GoTestResourceCheckError)
+                    and classification.details.check_errors_match(check_errors)
+                }
+            case _:
+                return {}  # todo: vector search to match on error output
+        return await self._find_classifications(query)
+
+    async def add_classification(self, classification: GoTestErrorClassification) -> bool:
+        """Returns is_new"""
+        raw = dump_with_id(classification, id=classification.run_id, dt_keys=["ts"])
+        return await create_or_replace(self.classifications, raw)
+
+    async def read_tf_test_run(self, run_id: str) -> GoTestRun:
+        raw = await self.runs.find_one({"_id": run_id})
+        if raw is None:
+            raise GoTestRunNotFound(run_id)
+        return self._parse_run(raw)
+
+    def _parse_run(self, raw: dict) -> GoTestRun:
+        raw.pop("_id")
+        for key in self.property_keys_run:
+            raw.pop(key, None)  # Remove properties that are not part of the model
+        return parse_model(raw, t=GoTestRun)
+
+    async def read_run_history(
+        self,
+        test_name: str,
+        branches: list[str] | None = None,
+        package_url: str | None = None,
+        group_name: str | None = None,
+        start_date: datetime | None = None,
+        end_date: datetime | None = None,
+        envs: list[str] | None = None,
+    ) -> list[GoTestRun]:
+        eq = MongoQueryOperation.eq
+        query = {
+            "name": {eq: test_name},
+        }
+        eq_parts = {
+            "package_url": {eq: package_url} if package_url else None,
+            "group_name": {eq: group_name} if group_name else None,
+        }
+        in_op = MongoQueryOperation.in_
+        in_parts = {
+            "branch": {in_op: branches} if branches else None,
+            "env": {in_op: envs} if envs else None,
+        }
+        date_parts = {
+            "ts": ignore_falsy(
+                **{
+                    MongoQueryOperation.lte: end_date or None,
+                    MongoQueryOperation.gte: start_date or None,
+                }
+            )
+        }
+        query |= ignore_falsy(**eq_parts, **in_parts, **date_parts)
+        if invalid_fields := set(query) - self._field_names_runs:
+            raise ValueError(f"Invalid fields in query: {invalid_fields}")
+        return await self._find_runs(query)

atlas_init/crud/mongo_utils.py

@@ -0,0 +1,239 @@
+import logging
+import re
+from dataclasses import dataclass
+from datetime import datetime
+from functools import wraps
+from typing import Any, AsyncIterable, Iterable, List, Optional, Type, TypeVar
+
+from model_lib import dump_as_dict
+from motor.core import AgnosticCollection
+from pydantic import BaseModel
+from pymongo import ASCENDING, DESCENDING, IndexModel, ReturnDocument
+from pymongo.errors import DuplicateKeyError, PyMongoError
+from pymongo.results import DeleteResult
+from zero_3rdparty.enum_utils import StrEnum
+
+logger = logging.getLogger(__name__)
+ModelT = TypeVar("ModelT", bound=BaseModel)
+
+
+class MongoUpdateOperation(StrEnum):
+    """
+    References:
+        https://docs.mongodb.com/manual/reference/operator/update-array/
+    """
+
+    slice = "$slice"
+    pop = "$pop"
+    pull = "$pull"
+    unset = "$unset"
+    push = "$push"
+    each = "$each"
+    set = "$set"
+    set_on_insert = "$setOnInsert"
+    inc = "$inc"
+
+
+UPDATE_OPERATIONS = set(MongoUpdateOperation)
+
+
+def ensure_mongo_operation(updates: dict):
+    """
+    >>> ensure_mongo_operation({"field1": 2})
+    {'$set': {'field1': 2}}
+    >>> ensure_mongo_operation({MongoUpdateOperation.set: {"field1": 2}})
+    {'$set': {'field1': 2}}
+    >>> ensure_mongo_operation({MongoUpdateOperation.push: {"field1": 2}})
+    {'$push': {'field1': 2}}
+    """
+    if updates.keys() - UPDATE_OPERATIONS == set():
+        return updates
+    return {MongoUpdateOperation.set: updates}
+
+
+class MongoQueryOperation(StrEnum):
+    # must be used when checking if a boolean field is false
+    eq = "$eq"
+    # https://stackoverflow.com/questions/18837486/query-for-boolean-field-as-not-true-e-g-either-false-or-non-existent
+    ne = "$ne"
+    in_ = "$in"
+    # https://www.mongodb.com/docs/manual/reference/operator/query/nin/#mongodb-query-op.-nin
+    nin = "$nin"
+    gt = "$gt"
+    gte = "$gte"
+    lt = "$lt"
+    lte = "$lte"
+    slice = "$slice"
+
+    @classmethod
+    def boolean_or_none(cls, bool_value: bool | None) -> dict | None:
+        if bool_value is None:
+            return None
+        return {cls.eq: True} if bool_value else {cls.ne: True}
+
+    @classmethod
+    def in_or_none(cls, options: Iterable[Any] | None) -> dict | None:
+        return None if options is None else {cls.in_: list(options)}
+
+    @classmethod
+    def nin_or_none(cls, options: Iterable[Any] | None) -> dict | None:
+        return None if options is None else {cls.nin: list(options)}
+
+
+duplicate_key_regex = re.compile(
+    r".*error collection:"
+    r"\s(?P<collection_path>[-\w\d\\.]+)"
+    r"\sindex:\s"
+    r"(?P<index_name>[\w_\\.\d]+)"
+    r"\sdup key.*?"
+    r'(?P<dup_key_value>("?[\\.\w_\d]+"?)|(null))'
+)
+
+
+@dataclass
+class MongoConstraintDetails:
+    collection_path: str
+    index_name: str
+    dup_key_value: Optional[str]
+
+    def __post_init__(self):
+        if self.dup_key_value:
+            self.dup_key_value = self.dup_key_value.strip('"')
+        if self.dup_key_value == "null":
+            self.dup_key_value = None
+
+
+def parse_error(error: PyMongoError) -> Optional[MongoConstraintDetails]:
+    """
+    >>> raw = 'E11000 duplicate key error collection: dev_situation.Robot index: _id_ dup key: { : "mw_wheel_id" }'
+    >>> parse_error(raw)
+    MongoConstraintDetails(collection_path='dev_situation.Robot', index_name='_id_', dup_key_value='mw_wheel_id')
+    ''
+    """
+    error_str = str(error)
+    for m in duplicate_key_regex.finditer(error_str):
+        constraints = MongoConstraintDetails(**m.groupdict())
+        if isinstance(error, DuplicateKeyError):
+            _, constraints.dup_key_value = error.details["keyValue"].popitem()  # type: ignore
+        return constraints
+    logger.warning(f"unknown pymongo error:{error}")
+
+
+class MongoConstraintError(Exception):
+    def __init__(self, details: MongoConstraintDetails):
+        self.details: MongoConstraintDetails = details
+
+
+T = TypeVar("T")
+
+ConstraintSubT = TypeVar("ConstraintSubT", bound=MongoConstraintError)
+
+
+def raise_mongo_constraint_error(f: T = None, *, cls: Type[ConstraintSubT] = MongoConstraintError) -> T:
+    def decorator(f: T):
+        @wraps(f)  # type: ignore
+        async def inner(*args, **kwargs):
+            try:
+                return await f(*args, **kwargs)  # type: ignore
+            except PyMongoError as e:
+                if details := parse_error(e):
+                    raise cls(details) from e
+                raise e
+
+        return inner
+
+    return decorator(f) if f else decorator  # type: ignore
+
+
+def dump_with_id(
+    model: BaseModel,
+    id: str = "",
+    dt_keys: Optional[List[str]] = None,
+    property_keys: Optional[List[str]] = None,
+    exclude: Optional[set[str]] = None,
+) -> dict:
+    """
+    Warning:
+        If you want to index on datetime, you have to set them afterwards
+        As they will be dumped as strings
+    """
+    raw = dump_as_dict(model) if exclude is None else dump_as_dict(model.model_dump(exclude=exclude))
+    if id:
+        raw["_id"] = id
+    if dt_keys:
+        for key in dt_keys:
+            raw[key] = getattr(model, key)
+    if property_keys:
+        for key in property_keys:
+            raw[key] = getattr(model, key)
+    return raw
+
+
+async def create_or_replace(collection: AgnosticCollection, raw: dict) -> bool:
+    """
+    Returns:
+        is_new: bool
+    """
+    result = await collection.replace_one({"_id": raw["_id"]}, raw, upsert=True)
+    return bool(result.upserted_id)
+
+
+async def find_one_and_update(
+    collection: AgnosticCollection,
+    id: str,
+    updates: dict,
+    return_raw_after: bool = True,
+    upsert: bool = False,
+    **query,
+) -> Optional[dict]:
+    """
+    Warning:
+        pops the "_id" from serialize_lib
+    """
+    return_doc = ReturnDocument.AFTER if return_raw_after else ReturnDocument.BEFORE
+    updates = ensure_mongo_operation(updates)
+    raw = await collection.find_one_and_update({"_id": id, **query}, updates, return_document=return_doc, upsert=upsert)
+    if raw:
+        raw.pop("_id", None)
+        return raw
+
+
+def microsecond_compare(mongo_dt: datetime, dt: datetime) -> bool:
+    """Mongo only stores milliseconds since epoch
+    https://stackoverflow.com/questions/39963143/why-is-there-a-difference-
+    between-the-stored-and-queried-time-in-mongo-database."""
+    with_microseconds = mongo_dt.replace(microsecond=dt.microsecond)
+    return with_microseconds == dt and (mongo_dt - dt).total_seconds() < 0.001
+
+
+def safe_key(key: str) -> str:
+    return key.replace(".", "_DOT_")
+
+
+def replace_dot_keys(values: dict) -> dict:
+    """avoid InvalidDocument("key 'dev.amironenko' must not contain '.'")"""
+    return {safe_key(key): value for key, value in values.items()}
+
+
+def decode_delete_count(result: DeleteResult) -> int:
+    return result.deleted_count
+
+
+def push_and_limit_length_update(field_name: str, new_value: Any, max_size: int) -> dict:
+    return {
+        MongoUpdateOperation.push: {
+            field_name: {
+                MongoUpdateOperation.each: [new_value],
+                MongoUpdateOperation.slice: -max_size,
+            }
+        }
+    }
+
+
+def index_dec(column: str) -> IndexModel:
+    return IndexModel([(column, DESCENDING)])
+
+
+def query_and_sort(collection: AgnosticCollection, query: dict, sort_col: str, desc: bool) -> AsyncIterable[dict]:
+    sort_order = DESCENDING if desc else ASCENDING
+    return collection.find(query).sort(sort_col, sort_order)

atlas_init/repos/go_sdk.py

@@ -2,11 +2,11 @@ from collections import defaultdict
 from pathlib import Path
 
 import requests
-from model_lib import parse_model
+from model_lib import Entity, parse_model
 
-from atlas_init.cli_tf.debug_logs_test_data import ApiSpecPath
+from atlas_init.cli_tf.debug_logs_test_data import ApiSpecPath, find_normalized_path
 from atlas_init.cli_tf.schema import logger
-from atlas_init.cli_tf.schema_v2_api_parsing import OpenapiSchema
+from atlas_init.cli_tf.openapi import OpenapiSchema
 
 
 def go_sdk_breaking_changes(repo_path: Path, go_sdk_rel_path: str = "../atlas-sdk-go") -> Path:
@@ -21,6 +21,15 @@ def api_spec_path_transformed(sdk_repo_path: Path) -> Path:
     return sdk_repo_path / "openapi/atlas-api-transformed.yaml"
 
 
+class ApiSpecPaths(Entity):
+    method_paths: dict[str, list[ApiSpecPath]]
+
+    def normalize_path(self, method: str, path: str) -> str:
+        if path.startswith("/api/atlas/v1.0"):
+            return ""
+        return find_normalized_path(path, self.method_paths[method]).path
+
+
 def parse_api_spec_paths(api_spec_path: Path) -> dict[str, list[ApiSpecPath]]:
     model = parse_model(api_spec_path, t=OpenapiSchema)
     paths: dict[str, list[ApiSpecPath]] = defaultdict(list)