Adding tests and updating clients, vendors

Author: james-geiger

gqlalchemy/models.py

@@ -16,7 +16,6 @@
 from collections import defaultdict
 from dataclasses import dataclass
 from datetime import datetime, date, time, timedelta
-from enum import Enum
 import json
 from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union
 from enum import Enum, EnumMeta
@@ -59,34 +58,37 @@ def _format_timedelta(duration: timedelta) -> str:
 
     return f"P{days}DT{hours}H{minutes}M{remainder_sec}S"
 
+
 class GraphEnum(ABC):
     def __init__(self, enum):
 
         if not isinstance(enum, (Enum, EnumMeta)):
             raise TypeError()
-        
+
         self.enum = enum if isinstance(enum, Enum) else None
         self.cls = enum.__class__ if isinstance(enum, Enum) else enum
-    
+
     @property
     def name(self):
         return self.cls.__name__
-    
+
     @property
     def members(self):
         return self.cls.__members__
-    
+
     @abstractmethod
     def _to_cypher(self):
         pass
 
+
 class MemgraphEnum(GraphEnum):
     def _to_cypher(self):
         return f"{{ {', '.join(self.cls._member_names_)} }}"
-        
+
     def __repr__(self):
-        return f"<enum '{self.name}'>" if self.enum is None else f'{self.name}::{self.enum.name}'
-    
+        return f"<enum '{self.name}'>" if self.enum is None else f"{self.name}::{self.enum.name}"
+
+
 class TriggerEventType:
     """An enum representing types of trigger events."""
 
@@ -343,7 +345,7 @@ def __init__(self, **data):
             if issubclass(cls, Enum) and not attrs.get("enum", False):
                 value = data.get(field)
                 if isinstance(value, dict):
-                    member  = value.get("__value").split('::')[1]
+                    member = value.get("__value").split("::")[1]
                     data[field] = cls[member].value
         super().__init__(**data)
 
@@ -583,7 +585,7 @@ def get_base_labels() -> Set[str]:
                 if cls.enums is None:
                     cls.enums = db.get_enums()
                 enum_names = [x.name for x in cls.enums]
-                if(field_cls.__name__ in enum_names):
+                if field_cls.__name__ in enum_names:
                     existing = cls.enums[enum_names.index(field_cls.__name__)]
                     db.sync_enum(existing, MemgraphEnum(field_cls))
                 else:
@@ -742,7 +744,7 @@ def __new__(mcs, name, bases, namespace, **kwargs):  # noqa C901
                 if cls.enums is None:
                     cls.enums = db.get_enums()
                 enum_names = [x.name for x in cls.enums]
-                if(field_type in enum_names):
+                if field_type in enum_names:
                     existing = cls.enums[enum_names.index(field_type)]
                     db.sync_enum(existing, MemgraphEnum(field_cls))
                 else:

gqlalchemy/vendors/database_client.py

@@ -17,13 +17,8 @@
 
 from gqlalchemy.connection import Connection
 from gqlalchemy.exceptions import GQLAlchemyError
-from gqlalchemy.models import (
-    Constraint,
-    Index,
-    GraphEnum,
-    Node,
-    Relationship
-)
+from gqlalchemy.models import Constraint, Index, GraphEnum, Node, Relationship
+
 
 class DatabaseClient(ABC):
     def __init__(
@@ -127,7 +122,7 @@ def ensure_constraints(
             self.drop_constraint(obsolete_constraints)
         for missing_constraint in new_constraints.difference(old_constraints):
             self.create_constraint(missing_constraint)
-    
+
     @abstractmethod
     def create_enum(self, enum: GraphEnum) -> None:
         pass
@@ -136,7 +131,7 @@ def create_enum(self, enum: GraphEnum) -> None:
     def get_enums(self) -> List[GraphEnum]:
         """Returns a list of all enums defined in the database."""
         pass
-    
+
     @abstractmethod
     def sync_enum(self, existing: GraphEnum, new: GraphEnum) -> None:
         """Ensures that database enum matches input enum."""
@@ -146,7 +141,7 @@ def sync_enum(self, existing: GraphEnum, new: GraphEnum) -> None:
     def drop_enum(self, enum: GraphEnum) -> None:
         """Drops a single enum in the database."""
         pass
-    
+
     @abstractmethod
     def drop_enums(self) -> None:
         """Drops all enums in the database"""

gqlalchemy/vendors/memgraph.py

@@ -16,7 +16,6 @@
 import os
 import sqlite3
 from typing import List, Optional, Union
-import warnings
 
 from gqlalchemy.connection import Connection, MemgraphConnection
 from gqlalchemy.disk_storage import OnDiskPropertyDatabase
@@ -25,7 +24,6 @@
     GQLAlchemyFileNotFoundError,
     GQLAlchemyOnDiskPropertyDatabaseNotDefinedError,
     GQLAlchemyUniquenessConstraintError,
-    GQLAlchemyWarning,
 )
 from gqlalchemy.models import (
     MemgraphConstraintExists,
@@ -169,7 +167,7 @@ def get_constraints(
                     )
                 )
         return constraints
-        
+
     def create_enum(self, graph_enum: MemgraphEnum) -> None:
         query = f"CREATE ENUM {graph_enum.name} VALUES {graph_enum._to_cypher()};"
         self.execute(query)
@@ -178,9 +176,9 @@ def get_enums(self) -> List[MemgraphEnum]:
         """Returns a list of all enums defined in the database."""
         enums: List[MemgraphEnum] = []
         for result in self.execute_and_fetch("SHOW ENUMS;"):
-            enums.append(MemgraphEnum(Enum(result['Enum Name'], result['Enum Values'])))
+            enums.append(MemgraphEnum(Enum(result["Enum Name"], result["Enum Values"])))
         return enums
-    
+
     def sync_enum(self, existing: MemgraphEnum, new: MemgraphEnum) -> None:
         """Ensures that database enum matches input enum."""
         for value in new.members:
@@ -190,9 +188,11 @@ def sync_enum(self, existing: MemgraphEnum, new: MemgraphEnum) -> None:
 
     def drop_enum(self, graph_enum: MemgraphEnum):
         raise GQLAlchemyError(f"DROP ENUM not yet implemented. Enum {graph_enum.name} is persisted in the database.")
-    
+
     def drop_enums(self, graph_enums: List[MemgraphEnum]):
-        raise GQLAlchemyError(f"DROP ENUM not yet implemented. Enums {', '.join(graph_enums)} are persisted in the database.")
+        raise GQLAlchemyError(
+            f"DROP ENUM not yet implemented. Enums {', '.join(graph_enums)} are persisted in the database."
+        )
 
     def get_exists_constraints(
         self,

gqlalchemy/vendors/neo4j.py

@@ -14,7 +14,6 @@
 
 import os
 from typing import List, Optional, Union
-from enum import Enum
 
 from gqlalchemy.connection import Connection, Neo4jConnection
 from gqlalchemy.exceptions import (
@@ -107,14 +106,14 @@ def create_enum(self, graph_enum: GraphEnum) -> None:
     def get_enums(self) -> List[GraphEnum]:
         """Returns a list of all enums defined in the database."""
         raise GQLAlchemyError(f"SHOW ENUMS not yet implemented in Neo4j.")
-    
+
     def sync_enum(self, existing: GraphEnum, new: GraphEnum) -> None:
         """Ensures that database enum matches input enum."""
         raise GQLAlchemyError(f"ALTER ENUM not yet implemented in Neo4j.")
-    
+
     def drop_enum(self, graph_enum: GraphEnum):
         raise GQLAlchemyError(f"DROP ENUM not yet implemented in Neo4j.")
-    
+
     def drop_enums(self, graph_enums: List[GraphEnum]):
         raise GQLAlchemyError(f"DROP ENUM not yet implemented in Neo4j.")
 

tests/ogm/test_custom_fields.py

@@ -13,10 +13,13 @@
 
 from pydantic.v1 import Field
 
+from enum import Enum
+
 from gqlalchemy import (
     MemgraphConstraintExists,
     MemgraphConstraintUnique,
     MemgraphIndex,
+    MemgraphEnum,
     Neo4jConstraintUnique,
     Neo4jIndex,
     Node,
@@ -56,6 +59,19 @@ def test_create_index(memgraph):
     assert actual_constraints == [memgraph_index]
 
 
+def test_create_graph_enum(memgraph):
+    enum1 = Enum("MgEnum", (("MEMBER1", "value1"), ("MEMBER2", "value2"), ("MEMBER3", "value3")))
+
+    class Node3(Node):
+        type: enum1
+
+    memgraph_enum = MemgraphEnum(enum1)
+
+    actual_enums = memgraph.get_enums()
+
+    assert actual_enums == [memgraph_enum]
+
+
 def test_create_constraint_unique_neo4j(neo4j):
     class Node2(Node):
         id: int = Field(db=neo4j)