fix: linting error

datahub-project · Nov 8, 2024 · dc5c84b · dc5c84b
1 parent 267c814
commit dc5c84b
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 5 deletions.
diff --git a/metadata-ingestion/src/datahub/ingestion/source/cassandra/cassandra.py b/metadata-ingestion/src/datahub/ingestion/source/cassandra/cassandra.py
@@ -43,7 +43,6 @@
 from datahub.ingestion.source.state.stateful_ingestion_base import (
     StatefulIngestionSourceBase,
 )
-from datahub.metadata._schema_classes import DatasetPropertiesClass, TimeStampClass
 from datahub.metadata.com.linkedin.pegasus2avro.common import StatusClass
 from datahub.metadata.com.linkedin.pegasus2avro.schema import (
     SchemaField,
@@ -52,8 +51,10 @@
 from datahub.metadata.schema_classes import (
     DataPlatformInstanceClass,
     DatasetLineageTypeClass,
+    DatasetPropertiesClass,
     OtherSchemaClass,
     SubTypesClass,
+    TimeStampClass,
     UpstreamClass,
     UpstreamLineageClass,
 )

diff --git a/metadata-ingestion/src/datahub/ingestion/source/cassandra/cassandra_utils.py b/metadata-ingestion/src/datahub/ingestion/source/cassandra/cassandra_utils.py
@@ -1,6 +1,6 @@
 import json
 import logging
-from typing import Any, Dict, Generator, List, Optional, Type
+from typing import Dict, Generator, List, Optional, Type
 
 from datahub.metadata.com.linkedin.pegasus2avro.schema import (
     SchemaField,
@@ -122,7 +122,7 @@ def _get_cur_field_path(self) -> str:
         return ".".join(self._prefix_name_stack)
 
     def _get_schema_fields(
-        self, cassandra_column_infos: List[dict[str, Any]]
+        self, cassandra_column_infos: List
     ) -> Generator[SchemaField, None, None]:
         # append each schema field (sort so output is consistent)
         for column_info in cassandra_column_infos:
@@ -166,7 +166,7 @@ def _get_schema_fields(
 
     @classmethod
     def get_schema_fields(
-        cls, cassandra_column_infos: List[dict[str, Any]]
+        cls, cassandra_column_infos: List
     ) -> Generator[SchemaField, None, None]:
         converter = cls()
         yield from converter._get_schema_fields(cassandra_column_infos)
diff --git a/metadata-ingestion/tests/unit/test_cassandra_source.py b/metadata-ingestion/tests/unit/test_cassandra_source.py
@@ -56,7 +56,7 @@ def test_cassandra_schema_conversion(
     schema: str, expected_field_paths: List[str]
 ) -> None:
     schema_dict: Dict[str, List[Any]] = json.loads(schema)
-    column_infos: List[dict[str, Any]] = schema_dict["column_infos"]
+    column_infos: List = schema_dict["column_infos"]
     actual_fields = list(
         CassandraToSchemaFieldConverter.get_schema_fields(column_infos)
     )