neo4j-viz 0.4.1__py3-none-any.whl → 0.5.0__py3-none-any.whl

neo4j_viz/snowflake.py

@@ -0,0 +1,344 @@
+from __future__ import annotations
+
+from enum import Enum
+from typing import Annotated, Any, Optional
+
+from pandas import DataFrame
+from pydantic import (
+    AfterValidator,
+    BaseModel,
+    BeforeValidator,
+)
+from pydantic_core.core_schema import ValidationInfo
+from snowflake.snowpark import Session
+from snowflake.snowpark.exceptions import SnowparkSQLException
+from snowflake.snowpark.types import (
+    ArrayType,
+    BooleanType,
+    ByteType,
+    DataType,
+    DateType,
+    DecimalType,
+    DoubleType,
+    FloatType,
+    GeographyType,
+    GeometryType,
+    IntegerType,
+    LongType,
+    MapType,
+    ShortType,
+    StringType,
+    StructField,
+    StructType,
+    TimestampType,
+    TimeType,
+    VariantType,
+    VectorType,
+)
+
+from neo4j_viz import VisualizationGraph
+from neo4j_viz.colors import ColorSpace
+from neo4j_viz.pandas import from_dfs
+
+
+def _data_type_name(type: DataType) -> str:
+    if isinstance(type, StringType):
+        return "VARCHAR"
+    elif isinstance(type, LongType):
+        return "BIGINT"
+    elif isinstance(type, IntegerType):
+        return "INT"
+    elif isinstance(type, DoubleType):
+        return "DOUBLE"
+    elif isinstance(type, DecimalType):
+        return "NUMBER"
+    elif isinstance(type, BooleanType):
+        return "BOOLEAN"
+    elif isinstance(type, ByteType):
+        return "TINYINT"
+    elif isinstance(type, DateType):
+        return "DATE"
+    elif isinstance(type, ShortType):
+        return "SMALLINT"
+    elif isinstance(type, FloatType):
+        return "FLOAT"
+    elif isinstance(type, ArrayType):
+        return "ARRAY"
+    elif isinstance(type, VectorType):
+        return "VECTOR"
+    elif isinstance(type, MapType):
+        return "OBJECT"
+    elif isinstance(type, TimeType):
+        return "TIME"
+    elif isinstance(type, TimestampType):
+        return "TIMESTAMP"
+    elif isinstance(type, VariantType):
+        return "VARIANT"
+    elif isinstance(type, GeographyType):
+        return "GEOGRAPHY"
+    elif isinstance(type, GeometryType):
+        return "GEOMETRY"
+    else:
+        # This actually does the job much of the time anyway
+        return type.simple_string().upper()
+
+
+SUPPORTED_ID_TYPES = [_data_type_name(data_type) for data_type in [StringType(), LongType(), IntegerType()]]
+
+
+def _validate_id_column(schema: StructType, column_name: str, index: int, supported_types: list[str]) -> None:
+    if column_name.lower() not in [name.lower() for name in schema.names]:
+        raise ValueError(f"Schema must contain a `{column_name}` column")
+
+    field: StructField = schema.fields[index]
+
+    if field.name.lower() != column_name.lower():
+        raise ValueError(f"Column `{column_name}` must have column index {index}")
+
+    if _data_type_name(field.datatype) not in supported_types:
+        raise ValueError(
+            f"Column `{column_name}` has invalid type `{_data_type_name(field.datatype)}`. Expected one of [{', '.join(supported_types)}]"
+        )
+
+
+def _validate_viz_node_table(table: str, info: ValidationInfo) -> str:
+    context = info.context
+    if context and context["session"] is not None:
+        session = context["session"]
+        try:
+            schema = session.table(table).schema
+            _validate_id_column(schema, "nodeId", 0, SUPPORTED_ID_TYPES)
+        except SnowparkSQLException as e:
+            raise ValueError(f"Table '{table}' does not exist or is not accessible.") from e
+    return table
+
+
+def _validate_viz_relationship_table(
+    table: str,
+    info: ValidationInfo,
+) -> str:
+    context = info.context
+    if context and context["session"] is not None:
+        session = context["session"]
+        try:
+            schema = session.table(table).schema
+            _validate_id_column(schema, "sourceNodeId", 0, SUPPORTED_ID_TYPES)
+            _validate_id_column(schema, "targetNodeId", 1, SUPPORTED_ID_TYPES)
+        except SnowparkSQLException as e:
+            raise ValueError(f"Table '{table}' does not exist or is not accessible.") from e
+    return table
+
+
+def _parse_identifier_groups(identifier: str) -> list[str]:
+    """
+    Parses a table identifier into a list of individual identifier groups.
+
+    This function handles identifiers that may include double-quoted segments
+    and ensures proper validation of the identifier's structure. It raises
+    errors for invalid formats, such as unbalanced quotes, invalid characters,
+    or improper use of dots.
+
+    Args:
+        identifier (str): The input string identifier to parse.
+
+    Returns:
+        list[str]: A list of parsed identifier groups.
+
+    Raises:
+        ValueError: If the identifier contains:
+            - Empty double quotes.
+            - Consecutive dots outside of double quotes.
+            - Unbalanced double quotes.
+            - Invalid characters in unquoted segments.
+            - Improper placement of dots around double-quoted segments.
+    """
+    inside = False  # Tracks whether the current character is inside double quotes
+    quoted_starts = []  # Stores the start indices of double-quoted segments
+    quoted_ends = []  # Stores the end indices of double-quoted segments
+    remaining = ""  # Stores the unquoted part of the identifier
+    previous_is_dot = False  # Tracks if the previous character was a dot
+
+    for i, c in enumerate(identifier):
+        if c == '"':
+            if not inside:
+                quoted_starts.append(i + 1)  # Mark the start of a quoted segment
+                previous_is_dot = False
+            else:
+                quoted_ends.append(i)  # Mark the end of a quoted segment
+                if quoted_ends[-1] - quoted_starts[-1] == 0:
+                    raise ValueError("Empty double quotes")
+            inside = not inside  # Toggle the inside state
+        else:
+            if not inside:
+                remaining += c  # Append unquoted characters to `remaining`
+                if c == ".":
+                    if previous_is_dot:
+                        raise ValueError("Not ok to have consecutive dots outside of double quote")
+                    previous_is_dot = True
+                else:
+                    previous_is_dot = False
+
+    if len(quoted_starts) != len(quoted_ends):
+        raise ValueError("Unbalanced double quotes")
+
+    for quoted_start in quoted_starts:
+        if quoted_start > 1:
+            if identifier[quoted_start - 2] != ".":
+                raise ValueError("Only dot character may precede before double quoted identifier")
+
+    for quoted_end in quoted_ends:
+        if quoted_end < len(identifier) - 1:
+            if identifier[quoted_end + 1] != ".":
+                raise ValueError("Only dot character may follow double quoted identifier")
+
+    words = remaining.split(".")  # Split the unquoted part by dots
+    for word in words:
+        if len(word) == 0:
+            continue
+        if word.lower()[0] not in "abcdefghijklmnopqrstuvwxyz_":
+            raise ValueError(f"Invalid first character in identifier {word}. Only a-z, A-Z, and _ are allowed.")
+        if not set(word.lower()).issubset(set("abcdefghijklmnopqrstuvwxyz$_0123456789")):
+            raise ValueError(f"Invalid characters in identifier {word}. Only a-z, A-Z, 0-9, _, and $ are allowed.")
+
+    empty_words_idx = [i for i, w in enumerate(words) if w == ""]
+    for i in range(len(quoted_starts)):
+        # Replace empty words with their corresponding quoted segments
+        words[empty_words_idx[i]] = f'"{identifier[quoted_starts[i] : quoted_ends[i]]}"'
+
+    return words
+
+
+def _validate_table_name(table: str) -> str:
+    if not isinstance(table, str):
+        raise TypeError(f"Table name must be a string, got {type(table).__name__}")
+
+    try:
+        words = _parse_identifier_groups(table)
+    except ValueError as e:
+        raise ValueError(f"Invalid table name '{table}'. {str(e)}") from e
+
+    if len(words) not in {1, 3}:
+        raise ValueError(
+            f"Invalid table name '{table}'. Table names must be in the format '<database>.<schema>.<table>' or '<table>'"
+        )
+
+    return table
+
+
+Table = Annotated[str, BeforeValidator(_validate_table_name)]
+
+VizNodeTable = Annotated[Table, AfterValidator(_validate_viz_node_table)]
+VizRelationshipTable = Annotated[Table, AfterValidator(_validate_viz_relationship_table)]
+
+
+class Orientation(Enum):
+    NATURAL = "natural"
+    UNDIRECTED = "undirected"
+    REVERSE = "reverse"
+
+
+def _to_lower(value: str) -> str:
+    return value.lower() if value and isinstance(value, str) else value
+
+
+LowercaseOrientation = Annotated[Orientation, BeforeValidator(_to_lower)]
+
+
+class VizRelationshipTableConfig(BaseModel, extra="forbid"):
+    sourceTable: VizNodeTable
+    targetTable: VizNodeTable
+    orientation: Optional[LowercaseOrientation] = Orientation.NATURAL
+
+
+class VizProjectConfig(BaseModel, extra="forbid"):
+    defaultTablePrefix: Optional[str] = None
+    nodeTables: list[VizNodeTable]
+    relationshipTables: dict[VizRelationshipTable, VizRelationshipTableConfig]
+
+
+def _map_tables(
+    session: Session, project_model: VizProjectConfig
+) -> tuple[list[DataFrame], list[DataFrame], list[str]]:
+    offset = 0
+    to_internal = {}
+    node_dfs = []
+    for table in project_model.nodeTables:
+        df = session.table(table).to_pandas()
+        internal_ids = range(offset, offset + df.shape[0])
+        to_internal[table] = df[["NODEID"]].copy()
+        to_internal[table]["INTERNALID"] = internal_ids
+        offset += df.shape[0]
+
+        df["SNOWFLAKEID"] = df["NODEID"]
+        df["NODEID"] = internal_ids
+
+        node_dfs.append(df)
+
+    rel_dfs = []
+    rel_table_names = []
+    for table, rel_table_config in project_model.relationshipTables.items():
+        df = session.table(table).to_pandas()
+
+        source_table = rel_table_config.sourceTable
+        target_table = rel_table_config.targetTable
+
+        df = df.merge(to_internal[source_table], left_on="SOURCENODEID", right_on="NODEID")
+        df.drop(["SOURCENODEID", "NODEID"], axis=1, inplace=True)
+        df.rename({"INTERNALID": "SOURCENODEID"}, axis=1, inplace=True)
+        df = df.merge(to_internal[target_table], left_on="TARGETNODEID", right_on="NODEID")
+        df.drop(["TARGETNODEID", "NODEID"], axis=1, inplace=True)
+        df.rename({"INTERNALID": "TARGETNODEID"}, axis=1, inplace=True)
+
+        if (
+            rel_table_config.orientation == Orientation.NATURAL
+            or rel_table_config.orientation == Orientation.UNDIRECTED
+        ):
+            rel_dfs.append(df)
+            rel_table_names.append(table)
+
+        if rel_table_config.orientation == Orientation.REVERSE:
+            df_rev = df.rename(columns={"SOURCENODEID": "TARGETNODEID", "TARGETNODEID": "SOURCENODEID"}, copy=False)
+            rel_dfs.append(df_rev)
+            rel_table_names.append(table)
+
+        if rel_table_config.orientation == Orientation.UNDIRECTED:
+            df_rev = df.rename(columns={"SOURCENODEID": "TARGETNODEID", "TARGETNODEID": "SOURCENODEID"}, copy=True)
+            rel_dfs.append(df_rev)
+            rel_table_names.append(table)
+
+    return node_dfs, rel_dfs, rel_table_names
+
+
+def from_snowflake(
+    session: Session,
+    project_config: dict[str, Any],
+    node_radius_min_max: Optional[tuple[float, float]] = (3, 60),
+) -> VisualizationGraph:
+    project_model = VizProjectConfig.model_validate(project_config, strict=False, context={"session": session})
+    node_dfs, rel_dfs, rel_table_names = _map_tables(session, project_model)
+
+    node_caption_present = False
+    for node_df in node_dfs:
+        if "CAPTION" in node_df.columns:
+            node_caption_present = True
+            break
+
+    if not node_caption_present:
+        for i, node_df in enumerate(node_dfs):
+            node_df["caption"] = project_model.nodeTables[i].split(".")[-1]
+
+    rel_caption_present = False
+    for rel_df in rel_dfs:
+        if "CAPTION" in rel_df.columns:
+            rel_caption_present = True
+            break
+
+    if not rel_caption_present:
+        for i, rel_df in enumerate(rel_dfs):
+            rel_df["caption"] = rel_table_names[i].split(".")[-1]
+
+    VG = from_dfs(node_dfs, rel_dfs, node_radius_min_max)
+
+    VG.color_nodes(field="caption", color_space=ColorSpace.DISCRETE)
+
+    return VG

neo4j_viz/visualization_graph.py

@@ -5,6 +5,7 @@ from collections.abc import Iterable
 from typing import Any, Callable, Hashable, Optional, Union
 
 from IPython.display import HTML
+from pydantic.alias_generators import to_snake
 from pydantic_extra_types.color import Color, ColorType
 
 from .colors import NEO4J_COLORS_CONTINUOUS, NEO4J_COLORS_DISCRETE, ColorSpace, ColorsType
@@ -277,7 +278,7 @@ class VisualizationGraph:
                 return node.properties.get(attribute)
         else:
             assert field is not None
-            attribute = field
+            attribute = to_snake(field)
 
             def node_to_attr(node: Node) -> Any:
                 return getattr(node, attribute)

{neo4j_viz-0.4.1.dist-info → neo4j_viz-0.5.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: neo4j-viz
-Version: 0.4.1
+Version: 0.5.0
 Summary: A simple graph visualization tool
 Author-email: Neo4j <team-gds@neo4j.org>
 Project-URL: Homepage, https://neo4j.com/
@@ -34,7 +34,7 @@ Requires-Dist: pydantic-extra-types<3,>=2
 Requires-Dist: enum-tools==0.12.0
 Provides-Extra: dev
 Requires-Dist: ruff==0.11.8; extra == "dev"
-Requires-Dist: mypy==1.15.0; extra == "dev"
+Requires-Dist: mypy==1.17.1; extra == "dev"
 Requires-Dist: pytest==8.3.4; extra == "dev"
 Requires-Dist: selenium==4.32.0; extra == "dev"
 Requires-Dist: ipykernel==6.29.5; extra == "dev"
@@ -55,6 +55,8 @@ Provides-Extra: gds
 Requires-Dist: graphdatascience<2,>=1; extra == "gds"
 Provides-Extra: neo4j
 Requires-Dist: neo4j; extra == "neo4j"
+Provides-Extra: snowflake
+Requires-Dist: snowflake-snowpark-python<2,>=1; extra == "snowflake"
 Provides-Extra: notebook
 Requires-Dist: ipykernel>=6.29.5; extra == "notebook"
 Requires-Dist: pykernel>=0.1.6; extra == "notebook"
@@ -62,7 +64,7 @@ Requires-Dist: neo4j>=5.26.0; extra == "notebook"
 Requires-Dist: ipywidgets>=8.0.0; extra == "notebook"
 Requires-Dist: palettable>=3.3.3; extra == "notebook"
 Requires-Dist: matplotlib>=3.9.4; extra == "notebook"
-Requires-Dist: snowflake-snowpark-python==1.26.0; extra == "notebook"
+Requires-Dist: snowflake-snowpark-python==1.37.0; extra == "notebook"
 
 # Graph Visualization for Python by Neo4j
 

{neo4j_viz-0.4.1.dist-info → neo4j_viz-0.5.0.dist-info}/RECORD

@@ -1,22 +1,23 @@
 neo4j_viz/__init__.py,sha256=Q-VZlJe3_kAow_-F_-9RsHCQfbOfv5on26YD9ihw27o,504
 neo4j_viz/colors.py,sha256=IvOCTmCu7WTMna_wNLZ3GrThTwFyIoKtNkmZYDLdYac,6694
-neo4j_viz/gds.py,sha256=ux41zwbfBoeH-A4lqzTAYbdM4d-4JwmXv4ooVzlFflI,7595
+neo4j_viz/gds.py,sha256=I6G69KmX7gmEBmfN0vfo9K1_5DAM9orJ3HL2-RPnSsg,8170
 neo4j_viz/gql_create.py,sha256=K33cT6dOj8eJPGNNJXiXlCfLIzNxTwcW4n_2AG3_zaY,14751
 neo4j_viz/neo4j.py,sha256=8oNhsEd33wayyNlDi5KirG_vKgvVc3nJgAyvAZuKcNw,7296
 neo4j_viz/node.py,sha256=MiLoghsn2NLs_iV65NuW7u3iaxP8MTKoNy6La9TdreY,3886
 neo4j_viz/node_size.py,sha256=c_sMtQSD8eJ_6Y0Kr6ku0LOs9VoEDxfYCUUzUWZ-1Xo,1197
 neo4j_viz/nvl.py,sha256=ZN3tyWar9ugR88r5N6txW3ThfNEWOt5A1KzrrRnLKwk,5262
 neo4j_viz/options.py,sha256=eOpiLcIfFvUiPoozyT44F9MHGRkqCfBZFmh0u_6DfwY,6400
-neo4j_viz/pandas.py,sha256=x_hD1IAFKKAo-cwgVMXLnavP0zkMIhB28Csw83vztNo,5755
+neo4j_viz/pandas.py,sha256=7ac8kY2GQfLzh64Hn_V7OWdu4UEDG_P5Lb7FWdN24Hk,5776
 neo4j_viz/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 neo4j_viz/relationship.py,sha256=sRgGjNzlqt6wPmiB3WXBTxR_W_5Z40VofpHUZvkRS1I,4143
-neo4j_viz/visualization_graph.py,sha256=Fizmzm8p4p3Fbd3gsiTymkXgI3crpdmr5ipv4gFv3SY,13760
+neo4j_viz/snowflake.py,sha256=9FO7acOHOzLmjMeJxaZW6rbyoJyI08dQEvnbD91LOu8,12176
+neo4j_viz/visualization_graph.py,sha256=aQhItfpPqKlebzldFRAtzamEyw-fxh4C11jO3X_51U4,13817
 neo4j_viz/resources/icons/screenshot.svg,sha256=Ns9Yi2Iq4lIaiFvzc0pXBmjxt4fcmBO-I4cI8Xiu1HE,311
 neo4j_viz/resources/icons/zoom-in.svg,sha256=PsO5yFkA1JnGM2QV_qxHKG13qmoR-RrlWARpaXNp5qU,415
 neo4j_viz/resources/icons/zoom-out.svg,sha256=OQRADAoe2bxbCeFufg6W22nR41q5NlI8QspT9l5pXUw,400
-neo4j_viz/resources/nvl_entrypoint/base.js,sha256=jYzVHzrBe32hPUgznTDRBNri3urSJ083GDnFmXOkAvc,1811792
+neo4j_viz/resources/nvl_entrypoint/base.js,sha256=SQm93kmdN6ZIDlXWgtWPUXQBVzYp1Td4UHkRi9-_fjw,1815362
 neo4j_viz/resources/nvl_entrypoint/styles.css,sha256=JjeTSB9OJT2KMfb8yFUUMLMG7Rzrf3o60hSCD547zTk,1123
-neo4j_viz-0.4.1.dist-info/METADATA,sha256=zlnSq2VYXVOmpXmal4gS8LVM64wbb3Zoh4FFARJXKTg,7074
-neo4j_viz-0.4.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-neo4j_viz-0.4.1.dist-info/top_level.txt,sha256=jPUM3z8MOtxqDanc2VzqkxG4HJn8aaq4S7rnCFNk_Vs,10
-neo4j_viz-0.4.1.dist-info/RECORD,,
+neo4j_viz-0.5.0.dist-info/METADATA,sha256=vBmmB4PcBfuHcj-yypqOWc5IS_Fe8oBeoVBGq1JiLXI,7169
+neo4j_viz-0.5.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+neo4j_viz-0.5.0.dist-info/top_level.txt,sha256=jPUM3z8MOtxqDanc2VzqkxG4HJn8aaq4S7rnCFNk_Vs,10
+neo4j_viz-0.5.0.dist-info/RECORD,,