[DOP-23867] Change API schema for lineage endpoints

Author: dolfinus

data_rentgen/consumer/openlineage/run.py

@@ -5,13 +5,13 @@
 
 from data_rentgen.consumer.openlineage.base import OpenLineageBase
 from data_rentgen.consumer.openlineage.run_facets import OpenLineageRunFacets
-from data_rentgen.utils import UUID
+from data_rentgen.utils import UUIDv6Plus
 
 
 class OpenLineageRun(OpenLineageBase):
     """Run model.
     See [Run](https://github.com/OpenLineage/OpenLineage/blob/main/spec/OpenLineage.json).
     """
 
-    runId: UUID
+    runId: UUIDv6Plus
     facets: OpenLineageRunFacets = Field(default_factory=OpenLineageRunFacets)

data_rentgen/consumer/openlineage/run_facets/parent_run.py

@@ -3,7 +3,7 @@
 
 from data_rentgen.consumer.openlineage.base import OpenLineageBase
 from data_rentgen.consumer.openlineage.run_facets.base import OpenLineageRunFacet
-from data_rentgen.utils import UUID
+from data_rentgen.utils import UUIDv6Plus
 
 
 class OpenLineageParentJob(OpenLineageBase):
@@ -20,7 +20,7 @@ class OpenLineageParentRun(OpenLineageBase):
     See [ParentRunFacet](https://github.com/OpenLineage/OpenLineage/blob/main/spec/facets/ParentRunFacet.json).
     """
 
-    runId: UUID
+    runId: UUIDv6Plus
 
 
 class OpenLineageParentRunFacet(OpenLineageRunFacet):

data_rentgen/db/repositories/operation.py

@@ -3,6 +3,7 @@
 
 from datetime import datetime, timezone
 from typing import Sequence
+from uuid import UUID
 
 from sqlalchemy import Row, any_, func, select
 from sqlalchemy.dialects.postgresql import insert
@@ -11,7 +12,6 @@
 from data_rentgen.db.repositories.base import Repository
 from data_rentgen.db.utils.uuid import extract_timestamp_from_uuid
 from data_rentgen.dto import OperationDTO, PaginationDTO
-from data_rentgen.utils import UUID
 
 
 class OperationRepository(Repository[Operation]):

data_rentgen/server/providers/auth/dummy_provider.py

@@ -31,7 +31,9 @@ def __init__(
 
     @classmethod
     def setup(cls, app: FastAPI) -> FastAPI:
-        settings = DummyAuthProviderSettings.model_validate(app.state.settings.auth.dict(exclude={"provider"}))
+        settings = DummyAuthProviderSettings.model_validate(
+            app.state.settings.auth.model_dump(exclude={"provider"}),
+        )
         logger.info("Using %s provider with settings:\n%s", cls.__name__, pformat(settings))
         app.dependency_overrides[AuthProvider] = cls
         app.dependency_overrides[DummyAuthProviderSettings] = lambda: settings

data_rentgen/server/providers/auth/keycloak_provider.py

@@ -38,7 +38,7 @@ def __init__(
     @classmethod
     def setup(cls, app: FastAPI) -> FastAPI:
         settings = KeycloakAuthProviderSettings.model_validate(
-            app.state.settings.auth.dict(exclude={"provider"}),
+            app.state.settings.auth.model_dump(exclude={"provider"}),
         )
         logger.info("Using %s provider with settings:\n%s", cls.__name__, settings)
         app.dependency_overrides[AuthProvider] = cls

data_rentgen/server/schemas/v1/__init__.py

@@ -18,8 +18,12 @@
     LineageEntityKindV1,
     LineageEntityV1,
     LineageInputRelationV1,
+    LineageIORelationSchemaFieldV1,
+    LineageIORelationSchemaV1,
+    LineageNodesResponseV1,
     LineageOutputRelationV1,
     LineageParentRelationV1,
+    LineageRelationsResponseV1,
     LineageResponseV1,
     LineageSymlinkRelationV1,
     OperationLineageQueryV1,
@@ -68,12 +72,16 @@
     "DatasetPaginateQueryV1",
     "DatasetResponseV1",
     "LineageDirectionV1",
-    "LineageEntityV1",
     "LineageEntityKindV1",
-    "LineageResponseV1",
+    "LineageEntityV1",
     "LineageInputRelationV1",
+    "LineageIORelationSchemaFieldV1",
+    "LineageIORelationSchemaV1",
+    "LineageNodesResponseV1",
     "LineageOutputRelationV1",
     "LineageParentRelationV1",
+    "LineageRelationsResponseV1",
+    "LineageResponseV1",
     "LineageSymlinkRelationV1",
     "LocationDetailedResponseV1",
     "LocationPaginateQueryV1",

data_rentgen/server/schemas/v1/dataset.py

@@ -2,8 +2,6 @@
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
-from typing import Literal
-
 from fastapi import Query
 from pydantic import BaseModel, ConfigDict, Field
 
@@ -12,8 +10,7 @@
 
 
 class DatasetResponseV1(BaseModel):
-    kind: Literal["DATASET"] = "DATASET"
-    id: int = Field(description="Dataset id")
+    id: str = Field(description="Dataset id", coerce_numbers_to_str=True)
     location: LocationResponseV1 = Field(description="Corresponding Location")
     name: str = Field(description="Dataset name")
     format: str | None = Field(description="Data format", default=None)

data_rentgen/server/schemas/v1/job.py

@@ -1,6 +1,5 @@
 # SPDX-FileCopyrightText: 2024-2025 MTS PJSC
 # SPDX-License-Identifier: Apache-2.0
-from typing import Literal
 
 from fastapi import Query
 from pydantic import BaseModel, ConfigDict, Field
@@ -12,8 +11,7 @@
 class JobResponseV1(BaseModel):
     """Job response"""
 
-    kind: Literal["JOB"] = "JOB"
-    id: int = Field(description="Job id")
+    id: str = Field(description="Job id", coerce_numbers_to_str=True)
     location: LocationResponseV1 = Field(description="Corresponding Location")
     name: str = Field(description="Job name")
     type: str = Field(description="Job type")

data_rentgen/server/schemas/v1/lineage.py

@@ -3,14 +3,15 @@
 from datetime import datetime
 from enum import Enum
 from typing import Literal
+from uuid import UUID
 
 from pydantic import BaseModel, ConfigDict, Field, ValidationInfo, field_validator
 
 from data_rentgen.server.schemas.v1.dataset import DatasetResponseV1
 from data_rentgen.server.schemas.v1.job import JobResponseV1
 from data_rentgen.server.schemas.v1.operation import OperationResponseV1
 from data_rentgen.server.schemas.v1.run import RunResponseV1
-from data_rentgen.utils import UUID
+from data_rentgen.utils import UUIDv6Plus
 
 
 class LineageEntityKindV1(str, Enum):
@@ -34,9 +35,9 @@ def __str__(self) -> str:
 
 class LineageEntityV1(BaseModel):
     kind: LineageEntityKindV1 = Field(description="Type of Lineage entity")
-    id: int | UUID = Field(description="Id of Lineage entity")
+    id: str | UUID = Field(description="Id of Lineage entity")
 
-    model_config = ConfigDict(from_attributes=True, use_enum_values=True)
+    model_config = ConfigDict(from_attributes=True)
 
 
 class BaseLineageQueryV1(BaseModel):
@@ -94,11 +95,11 @@ class JobLineageQueryV1(BaseLineageQueryV1):
 
 
 class OperationLineageQueryV1(BaseLineageQueryV1):
-    start_node_id: UUID = Field(description="Operation id", examples=["00000000-0000-0000-0000-000000000000"])
+    start_node_id: UUIDv6Plus = Field(description="Operation id", examples=["00000000-0000-0000-0000-000000000000"])
 
 
 class RunLineageQueryV1(BaseLineageQueryV1):
-    start_node_id: UUID = Field(description="Run id", examples=["00000000-0000-0000-0000-000000000000"])
+    start_node_id: UUIDv6Plus = Field(description="Run id", examples=["00000000-0000-0000-0000-000000000000"])
     granularity: Literal["OPERATION", "RUN"] = Field(
         default="RUN",
         description="Granularity of the run lineage",
@@ -107,36 +108,34 @@ class RunLineageQueryV1(BaseLineageQueryV1):
 
 
 class LineageParentRelationV1(BaseModel):
-    kind: Literal["PARENT"] = "PARENT"
     from_: LineageEntityV1 = Field(description="Start point of relation", serialization_alias="from")
     to: LineageEntityV1 = Field(description="End point of relation")
 
 
-class LineageOutputRelationSchemaFieldV1(BaseModel):
+class LineageIORelationSchemaFieldV1(BaseModel):
     name: str
     type: str | None = Field(default=None)
     description: str | None = Field(default=None)
-    fields: list["LineageOutputRelationSchemaFieldV1"] = Field(description="Nested fields", default_factory=list)
+    fields: list["LineageIORelationSchemaFieldV1"] = Field(description="Nested fields", default_factory=list)
 
     model_config = ConfigDict(from_attributes=True)
 
 
-class LineageOutputRelationSchemaV1(BaseModel):
-    id: int = Field(description="Schema id")
-    fields: list[LineageOutputRelationSchemaFieldV1] = Field(description="Schema fields")
+class LineageIORelationSchemaV1(BaseModel):
+    id: str = Field(description="Schema id", coerce_numbers_to_str=True)
+    fields: list[LineageIORelationSchemaFieldV1] = Field(description="Schema fields")
 
     model_config = ConfigDict(from_attributes=True)
 
 
 class LineageInputRelationV1(BaseModel):
-    kind: Literal["INPUT"] = "INPUT"
     from_: LineageEntityV1 = Field(description="Start point of relation", serialization_alias="from")
     to: LineageEntityV1 = Field(description="End point of relation")
     last_interaction_at: datetime = Field(description="Last interaction at", examples=["2008-09-15T15:53:00+05:00"])
     num_bytes: int | None = Field(description="Number of bytes", examples=[42], default=None)
     num_rows: int | None = Field(description="Number of rows", examples=[42], default=None)
     num_files: int | None = Field(description="Number of files", examples=[42], default=None)
-    i_schema: LineageOutputRelationSchemaV1 | None = Field(
+    i_schema: LineageIORelationSchemaV1 | None = Field(
         description="Schema",
         default=None,
         # pydantic models have reserved "schema" attribute, using alias
@@ -145,15 +144,14 @@ class LineageInputRelationV1(BaseModel):
 
 
 class LineageOutputRelationV1(BaseModel):
-    kind: Literal["OUTPUT"] = "OUTPUT"
     from_: LineageEntityV1 = Field(description="Start point of relation", serialization_alias="from")
     to: LineageEntityV1 = Field(description="End point of relation")
     type: str | None = Field(description="Type of relation", examples=["CREATE", "APPEND"], default=None)
     last_interaction_at: datetime = Field(description="Last interaction at", examples=["2008-09-15T15:53:00+05:00"])
     num_bytes: int | None = Field(description="Number of bytes", examples=[42], default=None)
     num_rows: int | None = Field(description="Number of rows", examples=[42], default=None)
     num_files: int | None = Field(description="Number of files", examples=[42], default=None)
-    o_schema: LineageOutputRelationSchemaV1 | None = Field(
+    o_schema: LineageIORelationSchemaV1 | None = Field(
         description="Schema",
         default=None,
         # pydantic models have reserved "schema" attribute, using alias
@@ -162,17 +160,28 @@ class LineageOutputRelationV1(BaseModel):
 
 
 class LineageSymlinkRelationV1(BaseModel):
-    kind: Literal["SYMLINK"] = "SYMLINK"
     from_: LineageEntityV1 = Field(description="Start point of relation", serialization_alias="from")
     to: LineageEntityV1 = Field(description="End point of relation")
     type: str = Field(description="Type of relation between datasets", examples=["METASTORE", "WAREHOUSE"])
 
 
+class LineageRelationsResponseV1(BaseModel):
+    parents: list[LineageParentRelationV1] = Field(description="Parent relations", default_factory=list)
+    symlinks: list[LineageSymlinkRelationV1] = Field(description="Symlink relations", default_factory=list)
+    inputs: list[LineageInputRelationV1] = Field(description="Input relations", default_factory=list)
+    outputs: list[LineageOutputRelationV1] = Field(description="Input relations", default_factory=list)
+
+
+class LineageNodesResponseV1(BaseModel):
+    datasets: dict[str, DatasetResponseV1] = Field(description="Dataset nodes", default_factory=dict)
+    jobs: dict[str, JobResponseV1] = Field(description="Job nodes", default_factory=dict)
+    runs: dict[UUID, RunResponseV1] = Field(description="Run nodes", default_factory=dict)
+    operations: dict[UUID, OperationResponseV1] = Field(description="Operation nodes", default_factory=dict)
+
+
 class LineageResponseV1(BaseModel):
-    relations: list[
-        LineageParentRelationV1 | LineageInputRelationV1 | LineageOutputRelationV1 | LineageSymlinkRelationV1
-    ] = Field(description="List of relations", default_factory=list)
-    nodes: list[RunResponseV1 | OperationResponseV1 | JobResponseV1 | DatasetResponseV1] = Field(
-        description="List of nodes",
-        default_factory=list,
+    relations: LineageRelationsResponseV1 = Field(
+        description="Lineage relations",
+        default_factory=LineageRelationsResponseV1,
     )
+    nodes: LineageNodesResponseV1 = Field(description="Lineage nodes", default_factory=LineageNodesResponseV1)

data_rentgen/server/schemas/v1/location.py

@@ -8,7 +8,7 @@
 
 
 class LocationResponseV1(BaseModel):
-    id: int = Field(description="Location id")
+    id: str = Field(description="Location id", coerce_numbers_to_str=True)
     type: str = Field(description="Location type, e.g kafka, hdfs, postgres")
     name: str = Field(description="Location name, e.g. cluster name")
     addresses: list[AddressResponseV1] = Field(description="List of addresses")