Refactor topology to use UUIDs instead of integers

Updated the topology models, migrations, and tests to replace integer-based IDs with UUIDs for better scalability and uniqueness. Introduced an `external_id` field where necessary to preserve compatibility with existing data. Adjusted related logic, schemas, and checks accordingly.

Author: VladimirFilonov

keep/api/models/db/migrations/versions/2025-03-07-17-49_aaec81b991bd.py

@@ -21,13 +21,27 @@ def transfer_data():
     session = Session(bind=op.get_bind())
     dialect = session.bind.dialect.name
 
-    session.execute(sa.text("""
+    uuid_generation_func = "replace(uuid(),'-','')"
+    if dialect == "sqlite":
+        uuid_generation_func = """
+        lower(
+            hex(randomblob(4)) || '-' || hex(randomblob(2)) || '-' || '4' ||
+            substr(hex( randomblob(2)), 2) || '-' ||
+            substr('AB89', 1 + (abs(random()) % 4) , 1)  ||
+            substr(hex(randomblob(2)), 2) || '-' ||
+            hex(randomblob(6))
+          )
+        """
+    elif dialect == "postgresql":
+        uuid_generation_func = "gen_random_uuid()"
+
+    session.execute(sa.text(f"""
         INSERT INTO topologyservice (
-            id, tenant_id, source_provider_id, repository, tags, service, environment, display_name, description, 
+            id, external_id, tenant_id, source_provider_id, repository, tags, service, environment, display_name, description, 
             team, email, slack, ip_address, mac_address, category, manufacturer, namespace, is_manual
         )  
         SELECT 
-            id, tenant_id, source_provider_id, repository, tags, service, environment, display_name, description, 
+            {uuid_generation_func} as id, id as external_id, tenant_id, source_provider_id, repository, tags, service, environment, display_name, description, 
             team, email, slack, ip_address, mac_address, category, manufacturer, namespace, is_manual 
         FROM topologyservice_tmp
     """))
@@ -39,37 +53,17 @@ def transfer_data():
 
     session.execute(sa.text("""
         INSERT INTO topologyserviceapplication (service_id, application_id, tenant_id) 
-        SELECT tsa.service_id, tsa.application_id, ts.tenant_id FROM topologyserviceapplication_tmp as tsa
-        JOIN topologyservice_tmp as ts ON tsa.service_id = ts.id
+        SELECT ts.id, tsa.application_id, ts.tenant_id FROM topologyserviceapplication_tmp as tsa
+        JOIN topologyservice as ts ON tsa.service_id = ts.external_id
     """))
 
-    if dialect == "sqlite":
-        session.execute(sa.text("""
-            INSERT INTO topologyservicedependency (id, service_id, depends_on_service_id, updated_at, protocol, tenant_id) 
-            SELECT lower(
-                hex(randomblob(4)) || '-' || hex(randomblob(2)) || '-' || '4' ||
-                substr(hex( randomblob(2)), 2) || '-' ||
-                substr('AB89', 1 + (abs(random()) % 4) , 1)  ||
-                substr(hex(randomblob(2)), 2) || '-' ||
-                hex(randomblob(6))
-              ) as id, tsd.service_id, tsd.depends_on_service_id, tsd.updated_at, tsd.protocol, ts.tenant_id
-            FROM topologyservicedependency_tmp as tsd
-            JOIN topologyservice_tmp as ts ON tsd.service_id = ts.id
-        """))
-    elif dialect == "postgres":
-        session.execute(sa.text("""
-            INSERT INTO topologyservicedependency (id, service_id, depends_on_service_id, updated_at, protocol, tenant_id) 
-            SELECT gen_random_uuid() as id, tsd.service_id, tsd.depends_on_service_id, tsd.updated_at, tsd.protocol, ts.tenant_id
-            FROM topologyservicedependency_tmp as tsd
-            JOIN topologyservice_tmp as ts ON tsd.service_id = ts.id
-        """))
-    elif dialect == "mysql":
-        session.execute(sa.text("""
-            INSERT INTO topologyservicedependency (id, service_id, depends_on_service_id, updated_at, protocol, tenant_id) 
-            SELECT uuid() as id, tsd.service_id, tsd.depends_on_service_id, tsd.updated_at, tsd.protocol, ts.tenant_id
-            FROM topologyservicedependency_tmp as tsd
-            JOIN topologyservice_tmp as ts ON tsd.service_id = ts.id
-        """))
+    session.execute(sa.text(f"""
+        INSERT INTO topologyservicedependency (id, service_id, depends_on_service_id, updated_at, protocol, tenant_id) 
+        SELECT {uuid_generation_func} as id, ts.id as service_id, ts_dep.id as depends_on_service_id, tsd.updated_at, tsd.protocol, ts.tenant_id
+        FROM topologyservicedependency_tmp as tsd
+        JOIN topologyservice as ts ON tsd.service_id = ts.external_id
+        JOIN topologyservice as ts_dep ON tsd.depends_on_service_id = ts_dep.external_id
+    """))
 
 
 def transfer_data_back():
@@ -80,7 +74,7 @@ def transfer_data_back():
             id, tenant_id, source_provider_id, repository, tags, service, environment, display_name, description, 
             team, email, slack, ip_address, mac_address, category, manufacturer, namespace, is_manual
         )
-        SELECT id, tenant_id, source_provider_id, repository, tags, service, environment, display_name, description, 
+        SELECT external_id as id, tenant_id, source_provider_id, repository, tags, service, environment, display_name, description, 
         team, email, slack, ip_address, mac_address, category, manufacturer, namespace, is_manual 
         FROM topologyservice_tmp
     """))
@@ -92,12 +86,17 @@ def transfer_data_back():
 
     session.execute(sa.text("""
         INSERT INTO topologyserviceapplication (service_id, application_id) 
-        SELECT service_id, application_id FROM topologyserviceapplication_tmp
+        SELECT ts.external_id, tsa.application_id FROM topologyserviceapplication_tmp as tsa
+        JOIN topologyservice_tmp as ts ON tsa.service_id = ts.id
+        
     """))
 
     session.execute(sa.text("""
         INSERT INTO topologyservicedependency (service_id, depends_on_service_id, updated_at, protocol) 
-        SELECT service_id, depends_on_service_id, updated_at, protocol FROM topologyservicedependency_tmp
+        SELECT ts.external_id, ts_dep.external_id, tsd.updated_at, tsd.protocol 
+        FROM topologyservicedependency_tmp as tsd
+        JOIN topologyservice_tmp as ts ON tsd.service_id = ts.id
+        JOIN topologyservice_tmp as ts_dep ON tsd.depends_on_service_id = ts_dep.id
     """))
 
 def upgrade():
@@ -121,7 +120,8 @@ def upgrade():
     )
     op.create_table(
         "topologyservice",
-        sa.Column("id", sa.Integer(), nullable=False),
+        sa.Column("id", sa.Uuid(), nullable=False),
+        sa.Column("external_id", sa.Integer(), nullable=True),
         sa.Column("tenant_id", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
         sa.Column(
             "source_provider_id", sqlmodel.sql.sqltypes.AutoString(), nullable=False
@@ -156,7 +156,7 @@ def upgrade():
     op.create_table(
         "topologyserviceapplication",
         sa.Column("tenant_id", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
-        sa.Column("service_id", sa.Integer(), nullable=False),
+        sa.Column("service_id", sa.Uuid(), nullable=False),
         sa.Column("application_id", sa.Uuid(), nullable=False),
         sa.ForeignKeyConstraint(
             ["application_id", "tenant_id"],
@@ -176,8 +176,8 @@ def upgrade():
         "topologyservicedependency",
         sa.Column("id", sa.Uuid(), nullable=False),
         sa.Column("tenant_id", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
-        sa.Column("service_id", sa.Integer(), nullable=False),
-        sa.Column("depends_on_service_id", sa.Integer(), nullable=False),
+        sa.Column("service_id", sa.Uuid(), nullable=False),
+        sa.Column("depends_on_service_id", sa.Uuid(), nullable=False),
         sa.Column("protocol", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
         sa.Column(
             "updated_at",

keep/api/models/db/topology.py

@@ -9,7 +9,7 @@
 
 class TopologyServiceApplication(SQLModel, table=True):
     tenant_id: str = Field(sa_column=Column(ForeignKey("tenant.id"), primary_key=True))
-    service_id: int = Field(primary_key=True)
+    service_id: UUID = Field(primary_key=True)
     application_id: UUID = Field(primary_key=True)
 
     service: "TopologyService" = Relationship(
@@ -51,7 +51,8 @@ class TopologyApplication(SQLModel, table=True):
 
 
 class TopologyService(SQLModel, table=True):
-    id: int
+    id: UUID = Field(default_factory=uuid4, primary_key=True)
+    external_id: Optional[int]
     tenant_id: str = Field(sa_column=Column(ForeignKey("tenant.id")))
     source_provider_id: str = "unknown"
     repository: Optional[str]
@@ -103,8 +104,8 @@ class Config:
 class TopologyServiceDependency(SQLModel, table=True):
     id: UUID = Field(default_factory=uuid4)
     tenant_id: str = Field(sa_column=Column(ForeignKey("tenant.id")))
-    service_id: int
-    depends_on_service_id: int
+    service_id: UUID
+    depends_on_service_id: UUID
     protocol: Optional[str] = "unknown"
     updated_at: Optional[datetime] = Field(
         sa_column=Column(
@@ -169,7 +170,7 @@ class TopologyServiceInDto(TopologyServiceDtoBase):
 
 class TopologyServiceDependencyDto(BaseModel, extra="ignore"):
     id: str | UUID
-    serviceId: str
+    serviceId: str | UUID
     serviceName: str
     protocol: Optional[str] = "unknown"
 
@@ -194,7 +195,7 @@ class TopologyApplicationDto(BaseModel, extra="ignore"):
 
 
 class TopologyServiceDtoIn(BaseModel, extra="ignore"):
-    id: int
+    id: UUID
 
 
 class TopologyApplicationDtoIn(BaseModel, extra="ignore"):
@@ -239,7 +240,7 @@ def from_orm(
 
 
 class TopologyServiceDtoOut(TopologyServiceDtoBase):
-    id: str
+    id: UUID | str
     dependencies: List[TopologyServiceDependencyDto]
     application_ids: List[UUID]
     updated_at: Optional[datetime]
@@ -298,28 +299,28 @@ class TopologyServiceCreateRequestDTO(BaseModel, extra="ignore"):
 
 
 class TopologyServiceUpdateRequestDTO(TopologyServiceCreateRequestDTO, extra="ignore"):
-    id: int
+    id: UUID | str
 
 
 class TopologyServiceDependencyCreateRequestDto(BaseModel, extra="ignore"):
-    service_id: int
-    depends_on_service_id: int
+    service_id: UUID | str
+    depends_on_service_id: UUID | str
     protocol: Optional[str] = "unknown"
 
 
 class TopologyServiceDependencyUpdateRequestDto(
     TopologyServiceDependencyCreateRequestDto, extra="ignore"
 ):
-    service_id: Optional[int]
-    depends_on_service_id: Optional[int]
-    id: int
+    service_id: Optional[UUID | str]
+    depends_on_service_id: Optional[UUID | str]
+    id: UUID | str
 
 
 class DeleteServicesRequest(BaseModel, extra="ignore"):
     service_ids: List[int]
 
 
 class TopologyServiceYAML(TopologyServiceCreateRequestDTO, extra="ignore"):
-    id: int
-    source_provider_id: Optional[str] = None
+    id: UUID | str
+    source_provider_id: Optional[UUID | str] = None
     is_manual: Optional[bool] = None

tests/test_enrichments.py

@@ -557,7 +557,7 @@ def test_disposable_enrichment(db_session, client, test_app, mock_alert_dto):
 def test_topology_mapping_rule_enrichment(mock_session, mock_alert_dto):
     # Mock a TopologyService with dependencies to simulate the DB structure
     mock_topology_service = TopologyService(
-        id=1, tenant_id="keep", service="test-service", display_name="Test Service"
+        external_id=1, tenant_id="keep", service="test-service", display_name="Test Service"
     )
 
     # Create a mock MappingRule for topology

tests/test_topology.py

@@ -23,12 +23,12 @@
 VALID_API_KEY = "valid_api_key"
 
 
-def create_service(db_session, tenant_id, service_id):
+def create_service(db_session, tenant_id, external_id):
     service = TopologyService(
-        id=service_id,
+        external_id=external_id,
         tenant_id=tenant_id,
-        service="test_service_" + service_id,
-        display_name=service_id,
+        service="test_service_" + external_id,
+        display_name=external_id,
         repository="test_repository",
         tags=["test_tag"],
         description="test_description",
@@ -105,7 +105,7 @@ def test_create_application_by_tenant_id(db_session):
             SINGLE_TENANT_UUID, application_dto, db_session
         )
 
-    application_dto.services.append(TopologyServiceDtoIn(id=123))
+    application_dto.services.append(TopologyServiceDtoIn(id=uuid.uuid4()))
     with pytest.raises(ServiceNotFoundException):
         TopologiesService.create_application_by_tenant_id(
             SINGLE_TENANT_UUID, application_dto, db_session
@@ -131,8 +131,10 @@ def test_create_application_by_tenant_id(db_session):
     assert len(result) == 1
     assert result[0].name == "New Application"
     assert len(result[0].services) == 2
-    assert result[0].services[0].service == "test_service_1"
-    assert result[0].services[1].service == "test_service_2"
+
+    services_names = [s.service for s in result[0].services]
+    assert "test_service_1" in services_names
+    assert "test_service_2" in services_names
 
 
 def test_update_application_by_id(db_session):
@@ -213,7 +215,7 @@ def test_create_application(db_session, client, test_app):
 
     service = create_service(db_session, SINGLE_TENANT_UUID, "1")
 
-    application_data = {"name": "New Application", "services": [{"id": service.id}]}
+    application_data = {"name": "New Application", "services": [{"id": str(service.id)}]}
 
     response = client.post(
         "/topology/applications",
@@ -256,7 +258,7 @@ def test_update_application(db_session, client, test_app):
     assert response.status_code == 200
     assert response.json()["name"] == "Updated Application"
 
-    invalid_update_data = {"name": "Invalid Application", "services": [{"id": "123"}]}
+    invalid_update_data = {"name": "Invalid Application", "services": [{"id": str(random_uuid)}]}
 
     response = client.put(
         f"/topology/applications/{application.id}",
@@ -335,18 +337,24 @@ def test_import_to_db(db_session):
 
     # Do same operation twice - import and re-import
     for i in range(2):
+
+        s1_id = str(uuid.uuid4())
+        s2_id = str(uuid.uuid4())
+
         topology_data = {
             "services": [
                 {
-                    "id": 1,
+                    "id": s1_id,
+                    "external_id": "1",
                     "service": "test_service_1",
                     "display_name": "Service 1",
                     "tags": ["tag1"],
                     "team": "team1",
                     "email": "[email protected]",
                 },
                 {
-                    "id": 2,
+                    "id": s2_id,
+                    "external_id": "2",
                     "service": "test_service_2",
                     "display_name": "Service 2",
                     "tags": ["tag2"],
@@ -358,18 +366,18 @@ def test_import_to_db(db_session):
                 {
                     "name": "Test Application 1",
                     "description": "Application 1 description",
-                    "services": [1],
+                    "services": [s1_id],
                 },
                 {
                     "name": "Test Application 2",
                     "description": "Application 2 description",
-                    "services": [2],
+                    "services": [s2_id],
                 },
             ],
             "dependencies": [
                 {
-                    "service_id": 1,
-                    "depends_on_service_id": 2,
+                    "service_id": s1_id,
+                    "depends_on_service_id": s2_id,
                 }
             ],
         }
@@ -388,5 +396,5 @@ def test_import_to_db(db_session):
 
         dependencies = db_session.exec(select(TopologyServiceDependency)).all()
         assert len(dependencies) == 1
-        assert dependencies[0].service_id == 1
-        assert dependencies[0].depends_on_service_id == 2
+        assert str(dependencies[0].service_id) == s1_id
+        assert str(dependencies[0].depends_on_service_id) == s2_id