diff --git a/graphistry/Plottable.py b/graphistry/Plottable.py index b0f88ec83f..94c5c7ae9a 100644 --- a/graphistry/Plottable.py +++ b/graphistry/Plottable.py @@ -7,7 +7,9 @@ from graphistry.models.compute.dbscan import DBSCANEngine from graphistry.models.compute.umap import UMAPEngineConcrete from graphistry.plugins_types.cugraph_types import CuGraphKind -from graphistry.Engine import Engine, EngineAbstract +from graphistry.plugins_types.kusto_types import KustoConfig +from graphistry.plugins_types.spanner_types import SpannerConfig +from graphistry.Engine import EngineAbstract from graphistry.utils.json import JSONVal @@ -65,8 +67,9 @@ class Plottable(object): _complex_encodings : dict _bolt_driver : Any _tigergraph : Any - _spannergraph: Any - _kustograph: Any + + _spanner_config: Optional[SpannerConfig] + _kusto_config: Optional[KustoConfig] _dataset_id: Optional[str] _url: Optional[str] diff --git a/graphistry/PlotterBase.py b/graphistry/PlotterBase.py index f71a6ecadf..6eddaaf200 100644 --- a/graphistry/PlotterBase.py +++ b/graphistry/PlotterBase.py @@ -10,6 +10,8 @@ from .constants import SRC, DST, NODE from .plugins_types import CuGraphKind +from .plugins_types.kusto_types import KustoConfig +from .plugins_types.spanner_types import SpannerConfig from .plugins.igraph import ( to_igraph as to_igraph_base, from_igraph as from_igraph_base, compute_igraph as compute_igraph_base, @@ -177,9 +179,9 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: # Integrations self._bolt_driver : Any = None self._tigergraph : Any = None - self._spannergraph: Any - self._kustograph: Any - + self._kusto_config : Optional[KustoConfig] = None + self._spanner_config : Optional[SpannerConfig] = None + # feature engineering self._node_embedding = None self._node_encoder = None @@ -2283,45 +2285,50 @@ def bolt(self, driver): res._bolt_driver = to_bolt_driver(driver) return res - def spanner_init(self: Plottable, spanner_config: Dict[str, str]) -> Plottable: + + def spanner(self: Plottable, spanner_config: SpannerConfig) -> Plottable: """ - Initializes a SpannerGraph object with the provided configuration and connects to the instance db + Set spanner configuration for this Plottable. - spanner_config dict must contain the include the following keys, credentials_file is optional: - - "project_id": The GCP project ID. + SpannerConfig - "instance_id": The Spanner instance ID. - "database_id": The Spanner database ID. + - "project_id": The GCP project ID. - "credentials_file": json file API key for service accounts - - :param spanner_config A dictionary containing the Spanner configuration. - :type (Dict[str, str]) + + If credentials_file is provided, it will be used to authenticate with the Spanner instance. + Otherwise, project_id and the spanner login process will be used to authenticate. + + :param spanner_config: A dictionary containing the Spanner configuration. + :type (SpannerConfig) :return: Plottable with a Spanner connection :rtype: Plottable - :raises ValueError: If any of the required keys in `spanner_config` are missing or have invalid values. - """ - from .plugins.spannergraph import SpannerGraph - - res = copy.copy(self) - - res._spannergraph = SpannerGraph(res, spanner_config) - logger.debug("Created SpannerGraph object: {res._spannergraph}") - return res + self._spanner_config = spanner_config + return self + - def kusto_init(self: Plottable, kusto_config: Dict[str, str]) -> Plottable: - from .plugins.kustograph import KustoGraph - res = copy.copy(self) - res._kustograph = KustoGraph(res, kusto_config) - return res + def kusto(self: Plottable, kusto_config: KustoConfig) -> Plottable: + """ + Set kusto configuration for this Plottable. + + KustoConfig + - "cluster": The Kusto cluster name. + - "database": The Kusto database name. + For AAD authentication: + - "client_id": The Kusto client ID. + - "client_secret": The Kusto client secret. + - "tenant_id": The Kusto tenant ID. + Otherwise: process will use web browser to authenticate. + + :param kusto_config: A dictionary containing the Kusto configuration. + :type (KustoConfig) + :returns: Plottable with a Kusto connection + :rtype: Plottable + """ + self._kusto_config = kusto_config + return self - def kusto_query_to_df(self: Plottable, query: str) -> pd.DataFrame: - from .pygraphistry import PyGraphistry - if not hasattr(self, '_kustograph') or self._kustograph is None: - kusto_config = PyGraphistry._config["kusto"] - if not kusto_config: - raise ValueError("Missing kusto_config. Use kusto_init() or register() with kusto_config.") - self = self.kusto_init(kusto_config) - return self._kustograph.query_to_df(query) def infer_labels(self): """ @@ -2550,22 +2557,11 @@ def spanner_gql_to_g(self: Plottable, query: str) -> Plottable: g.plot() """ - from .pygraphistry import PyGraphistry - from .plugins.spannergraph import SpannerGraph - + from .plugins.spannergraph import SpannerGraphContext res = copy.copy(self) + with SpannerGraphContext(res._spanner_config) as sg: + return sg.gql_to_graph(query, g=res) - if not hasattr(res, '_spannergraph'): - spanner_config = PyGraphistry._config.get("spanner", None) - - if spanner_config is not None: - logger.debug(f"Spanner Config: {spanner_config}") - else: - raise ValueError('spanner_config not defined. Pass spanner_config via register() and retry query.') - - res = res.spanner_init(spanner_config) # type: ignore[attr-defined] - - return res._spannergraph.gql_to_graph(res, query) def spanner_query_to_df(self: Plottable, query: str) -> pd.DataFrame: """ @@ -2602,22 +2598,65 @@ def spanner_query_to_df(self: Plottable, query: str) -> pd.DataFrame: g.plot() """ + from .plugins.spannergraph import SpannerGraphContext + with SpannerGraphContext(self._spanner_config) as sg: + return sg.query_to_df(query) + - from .pygraphistry import PyGraphistry + def kusto_query(self: Plottable, query: str, unwrap_nested: bool | None = None) -> List[pd.DataFrame]: + """ + Submit a Kusto/Azure Data Explorer *query* and return result tables. - res = copy.copy(self) - - if not hasattr(res, '_spannergraph'): - spanner_config = PyGraphistry._config["spanner"] - if spanner_config is not None: - logger.debug(f"Spanner Config: {spanner_config}") - else: - logger.warning('PyGraphistry._config["spanner"] is None') + Because a Kusto request may emit multiple tables, a **list of + DataFrames** is always returned; most queries yield a single entry. + + unwrap_nested + ------------- + Controls auto-flattening of *dynamic* (JSON) columns: + • True - always try to flatten, raise if it fails + • None - default heuristic: flatten only if table looks nested + • False - leave results untouched + + :param query: Kusto query string + :type query: str + :param unwrap_nested: flatten strategy above + :type unwrap_nested: bool | None + :returns: list of Pandas DataFrames + :rtype: List[pd.DataFrame] + + **Example** + :: + + frames = graphistry.kusto_query("StormEvents | take 100") + df = frames[0] + """ + from .plugins.kustograph import KustoGraphContext + with KustoGraphContext(self._kusto_config) as kg: + return kg.query(query, unwrap_nested=unwrap_nested) - res = res.spanner_init(PyGraphistry._config["spanner"]) # type: ignore[attr-defined] + def kusto_query_graph(self: Plottable, graph_name: str, snap_name: str | None = None) -> Plottable: + """ + Fetch a Kusto *graph* (and optional *snapshot*) as a Graphistry object. + + Under the hood: `graph(..)` + `graph-to-table` to pull **nodes** and + **edges**, then binds them to *self*. - return res._spannergraph.query_to_df(query) + :param graph_name: name of Kusto graph entity + :type graph_name: str + :param snap_name: optional snapshot/version + :type snap_name: str | None + :returns: Plottable ready for `.plot()` or further transforms + :rtype: Plottable + **Example** + :: + + g = graphistry.kusto_query_graph("HoneypotNetwork").plot() + """ + from .plugins.kustograph import KustoGraphContext + with KustoGraphContext(self._kusto_config) as kg: + return kg.query_graph(graph_name, snap_name, g=self) + def nodexl(self, xls_or_url, source='default', engine=None, verbose=False): diff --git a/graphistry/__init__.py b/graphistry/__init__.py index c197de6fe8..3bdc5689cf 100644 --- a/graphistry/__init__.py +++ b/graphistry/__init__.py @@ -31,11 +31,12 @@ bolt, cypher, tigergraph, + spanner, spanner_gql_to_g, spanner_query_to_df, - spanner_init, - kusto_query_to_df, - kusto_init, + kusto, + kusto_query, + kusto_query_graph, gsql, gsql_endpoint, cosmos, diff --git a/graphistry/plugins/kustograph.py b/graphistry/plugins/kustograph.py index 187f2254c0..091b4dd32d 100644 --- a/graphistry/plugins/kustograph.py +++ b/graphistry/plugins/kustograph.py @@ -1,82 +1,236 @@ import pandas as pd import time -import json -from typing import Any, Dict, List, Optional +from typing import Any, List, TYPE_CHECKING +from collections import OrderedDict +from typing import Mapping, List, Dict, Any, Iterable, Tuple, Optional -from azure.kusto.data import KustoClient, KustoConnectionStringBuilder -from azure.kusto.data.exceptions import KustoServiceError +if TYPE_CHECKING: + from azure.kusto.data import KustoClient +else: + KustoClient = Any from graphistry.Plottable import Plottable from graphistry.util import setup_logger +from graphistry.pygraphistry import PyGraphistry +from graphistry.plugins_types.kusto_types import KustoConfig, KustoConnectionError, KustoQueryResult logger = setup_logger(__name__) -class KustoConnectionError(Exception): - pass -class KustoQueryResult: - def __init__(self, data: List[List[Any]], column_names: List[str]): - self.data = data - self.column_names = column_names - class KustoGraph: + def __init__(self, client: KustoClient, database: str): + self.client = client + self.database = database - def __init__(self, g: Plottable, kusto_config: Dict[str, str]): - required_keys = ["cluster", "database"] - for key in required_keys: - if not kusto_config.get(key): - raise ValueError(f"Missing required kusto_config key: '{key}'") - - self.cluster = kusto_config["cluster"] - self.database = kusto_config["database"] - - if all(kusto_config.get(k) for k in ["client_id", "client_secret", "tenant_id"]): - self.credential_mode = "aad_app" - self.client_id = kusto_config["client_id"] - self.client_secret = kusto_config["client_secret"] - self.tenant_id = kusto_config["tenant_id"] - else: - self.credential_mode = "device_code" + @classmethod + def from_config(cls, cfg: KustoConfig) -> "KustoGraph": + from azure.kusto.data import KustoClient, KustoConnectionStringBuilder - self.client = self._connect() + try: + cluster = cfg["cluster"] + database = cfg["database"] + except KeyError as e: + raise ValueError(f"Missing required kusto_config key: '{e}'") from e - def _connect(self) -> KustoClient: try: - if self.credential_mode == "aad_app": + client_id, client_secret, tenant_id = cfg.get("client_id"), cfg.get("client_secret"), cfg.get("tenant_id") + if all((client_id, client_secret, tenant_id)): kcsb = KustoConnectionStringBuilder.with_aad_application_key_authentication( - self.cluster, - self.client_id, - self.client_secret, - self.tenant_id, + cluster, + client_id, + client_secret, + tenant_id, ) else: - kcsb = KustoConnectionStringBuilder.with_aad_device_authentication(self.cluster) + kcsb = KustoConnectionStringBuilder.with_aad_device_authentication(cluster) + + logger.info("Connecting to Kusto cluster %s", cluster) + client = KustoClient(kcsb) + except Exception as exc: + raise KustoConnectionError(f"Failed to connect to Kusto cluster: {exc}") from exc + + return cls(client, database=database) + + + def close(self) -> None: + self.client.close() - logger.info(f"Connecting to Kusto cluster: {self.cluster}") - return KustoClient(kcsb) - except Exception as e: - raise KustoConnectionError(f"Failed to connect to Kusto cluster: {e}") - def execute_query(self, query: str) -> KustoQueryResult: - logger.debug(f"KustoGraph execute_query(): {query}") + def _query(self, query: str) -> List[KustoQueryResult]: + from azure.kusto.data.exceptions import KustoServiceError + logger.debug(f"KustoGraph._query(): {query}") + try: start = time.time() response = self.client.execute(self.database, query) - rows = list(response.primary_results[0]) - col_names = [col.column_name for col in response.primary_results[0].columns] + results = [] + row_lengths = [] + for result in response.primary_results: + rows = result.rows + col_names = [col.column_name for col in result.columns] + col_types = [col.column_type for col in result.columns] + results.append(KustoQueryResult(rows, col_names, col_types)) + row_lengths.append((len(rows), len(col_names))) - logger.info(f"Query returned {len(rows)} rows in {time.time() - start:.3f} sec") - return KustoQueryResult(rows, col_names) + logger.info(f"Query returned {len(results)} results shapes: {row_lengths} in {time.time() - start:.3f} sec") + print(f"Query returned {len(results)} results shapes: {row_lengths} in {time.time() - start:.3f} sec") + return results except KustoServiceError as e: logger.error(f"Kusto query failed: {e}") raise RuntimeError(f"Kusto query failed: {e}") - def query_to_df(self, query: str) -> pd.DataFrame: - result = self.execute_query(query) - return pd.DataFrame(result.data, columns=result.column_names) + + def query( + self, + query: str, + *, + unwrap_nested: bool | None = True + ) -> List[pd.DataFrame]: + """ + Run *query* and return a list of DataFrames. + + unwrap_nested semantics + ----------------------- + • True - Always attempt to unwrap; raise on failure. + • None - Use heuristic: unwrap iff the first result *looks* nested. + • False - Never attempt to unwrap. + """ + results = self._query(query) + if not results: + return [] + + if unwrap_nested is False: + return [pd.DataFrame(r.data, columns=r.column_names) for r in results] + + first = results[0] + do_unwrap = ( + unwrap_nested is True or + (unwrap_nested is None and _should_unwrap(first)) + ) + + if not do_unwrap: + return [pd.DataFrame(r.data, columns=r.column_names) for r in results] + + try: + frames_od = _unwrap_nested(first) + return list(frames_od.values()) + except Exception as exc: + if unwrap_nested is True: + raise RuntimeError(f"_unwrap_nested failed: {exc}") from exc + # Heuristic miss – fall back silently + return [pd.DataFrame(r.data, columns=r.column_names) for r in results] + + + def query_graph(self, graph_name: str, snap_name: str | None = None, g: Plottable = PyGraphistry.bind()) -> Plottable: + if snap_name: + graph_query = f'graph("{graph_name}", "{snap_name}" | graph-to-table nodes as N with_node_id=NodeId, edges as E with_source_id=src with_target_id=dst; N;E' + else: + graph_query = f'graph("{graph_name}") | graph-to-table nodes as N with_node_id=NodeId, edges as E with_source_id=src with_target_id=dst; N;E' + results = self._query(graph_query) + if len(results) != 2: + raise ValueError(f"Expected 2 results, got {len(results)}") + nodes = pd.DataFrame(results[0].data, columns=results[0].column_names) + edges = pd.DataFrame(results[1].data, columns=results[1].column_names) + return g.nodes(nodes, node='NodeId').edges(edges, source='src', destination='dst') + + + +# ================================================================ +# kusto_graph.py – new imports +# ================================================================ +from collections import OrderedDict +from typing import Any, Iterable, Mapping, List, Dict, OrderedDict as OD +import pandas as pd + + +# ================================================================ +# Low‑level utilities +# ================================================================ +def _is_dynamic(val: Any) -> bool: + """ADT check for Kusto 'dynamic' JSON values.""" + return isinstance(val, (dict, list)) + +def _normalize(records: Iterable[Mapping]) -> pd.DataFrame: + """json_normalize + dedup + stable column ordering.""" + if not records: + return pd.DataFrame() + df = pd.json_normalize(list(records)) + return df.loc[:, sorted(df.columns)].drop_duplicates().reset_index(drop=True) + + +# ================================================================ +# Core transformer +# ================================================================ +def _unwrap_nested(result: "KustoQueryResult") -> "OrderedDict[str, pd.DataFrame]": + """ + Transform one Kusto result whose columns contain *dynamic* objects + (typical from `graph-match … project a, edge, v`) into + OrderedDict[col_name -> DataFrame]. + """ + frames: "OD[str, pd.DataFrame]" = OrderedDict() + + for col_idx, col_name in enumerate(result.column_names): + col_vals = [row[col_idx] for row in result.data if row[col_idx] is not None] + + # Non‑dynamic column → trivial DF + if not col_vals or not _is_dynamic(col_vals[0]): + frames[col_name] = pd.DataFrame({col_name: col_vals}) + continue + + # Dynamic column → flatten each JSON object + recs: List[Mapping] = [] + for cell in col_vals: + if isinstance(cell, list): + recs.extend(cell) # explode lists + else: + recs.append(cell) + frames[col_name] = _normalize(recs) + + return frames + + +# ================================================================ +# Heuristic for ``unwrap_nested is None`` +# ================================================================ +def _should_unwrap(result: "KustoQueryResult", sample_rows: int = 5) -> bool: + """ + Decide whether result *looks* like it contains nested/dynamic columns. + Strategy: + 1. Prefer explicit type info (column_type == 'dynamic') if present. + 2. Otherwise inspect up to `sample_rows` rows for dict / list values. + """ + try: + if any(c.lower() == "dynamic" for c in result.column_types): + return True + except AttributeError: + pass # `.column_type` not available in older SDK versions. + + for col_idx in range(len(result.column_names)): + sample = (row[col_idx] for row in result.data[:sample_rows]) + if any(_is_dynamic(v) for v in sample): + return True + return False + + + + + +class KustoGraphContext: + def __init__(self, config: KustoConfig | None = None): + config = config or PyGraphistry._config.get("kusto") + if not config: + raise ValueError("Missing kusto_config. Register globally with kusto_config or use with_kusto().") + self.kusto_graph = KustoGraph.from_config(config) + + def __enter__(self): + return self.kusto_graph + + def __exit__(self, exc_type, exc_value, traceback): + self.kusto_graph.close() + diff --git a/graphistry/plugins/spannergraph.py b/graphistry/plugins/spannergraph.py index 3aeb6a34a1..17ac87c0f8 100644 --- a/graphistry/plugins/spannergraph.py +++ b/graphistry/plugins/spannergraph.py @@ -1,106 +1,80 @@ -import os import pandas as pd import json import time -from typing import Any, List, Dict, Optional +from typing import Any, List, Dict, TYPE_CHECKING -from graphistry.Plottable import Plottable +if TYPE_CHECKING: + from google.cloud.spanner_dbapi.connection import Connection +else: + Connection = Any +from graphistry.pygraphistry import PyGraphistry +from graphistry.Plottable import Plottable +from graphistry.plugins_types.spanner_types import SpannerConfig, SpannerConnectionError, SpannerQueryResult from graphistry.util import setup_logger logger = setup_logger(__name__) -class SpannerConnectionError(Exception): - """Custom exception for errors related to Spanner connection.""" - pass - -class SpannerQueryResult: - """ - Encapsulates the results of a query, including metadata. - - :ivar list data: The raw query results. - """ - - def __init__(self, data: List[Any], column_names: List[str]): - """ - Initializes a SpannerQueryResult instance. - - :param data: The raw query results. - :type List[Any] - :param column_names: a list of the column names from the cursor, defaults to None - :type: List[str] - """ - self.data = data - self.column_names = column_names class SpannerGraph: """ A comprehensive interface for interacting with Google Spanner Graph databases. - :ivar str project_id: The Google Cloud project ID. - :ivar str instance_id: The Spanner instance ID. - :ivar str database_id: The Spanner database ID. :ivar Any connection: The active connection to the Spanner database. - :ivar Any graphistry: The Graphistry parent object. """ - def __init__(self, g: Plottable, spanner_config: Dict[str, str]): + def __init__(self, connection: Connection): + self.connection = connection + + + @classmethod + def from_config(cls, spanner_config: SpannerConfig) -> "SpannerGraph": """ - Initializes the SpannerGraph instance. + Initializes the SpannerGraph instance and establishes a connection to the Spanner database. :param graphistry: The Graphistry parent object. - :param project_id: The Google Cloud project ID. :param instance_id: The Spanner instance ID. :param database_id: The Spanner database ID. """ - - # check if valid - required_keys = ["project_id", "instance_id", "database_id"] - for key in required_keys: - value = spanner_config.get(key) - if not value: # check for None or empty values - raise ValueError(f"Missing or invalid value for required Spanner configuration: '{key}'") - - self.project_id = spanner_config["project_id"] - self.instance_id = spanner_config["instance_id"] - self.database_id = spanner_config["database_id"] - - if spanner_config.get("credentials_file"): - self.credentials_file = spanner_config["credentials_file"] - - self.connection = self.__connect() - - def __connect(self) -> Any: - """ - Establishes a connection to the Spanner database. - - :return: A connection object to the Spanner database. - :rtype: google.cloud.spanner_dbapi.connection - :raises SpannerConnectionError: If the connection to Spanner fails. - """ from google.cloud.spanner_dbapi.connection import connect + + try: + instance_id = spanner_config["instance_id"] + database_id = spanner_config["database_id"] + except KeyError as e: + raise ValueError(f"Missing required Spanner configuration: '{e}'") + + credentials_file = spanner_config.get("credentials_file", None) try: - if hasattr(self, 'credentials_file') and self.credentials_file is not None: - - connection = connect(self.instance_id, self.database_id, credentials=self.credentials_file) + if credentials_file: + connection = connect(instance_id, database_id, credentials=credentials_file) else: - connection = connect(self.instance_id, self.database_id) + project_id = spanner_config.get("project_id") + if not project_id: + raise ValueError("Missing required Spanner configuration: 'project_id' or 'credentials_file'") + connection = connect(project_id, instance_id, database_id) - connection.autocommit = True logger.info("Connected to Spanner database.") - return connection + + connection.read_only = True + # NOTE: or connection.autocommit=True; Required to query INFORMATION_SCHEMA + # Otherwise: "Unsupported concurrency mode in query using INFORMATION_SCHEMA" + + return cls(connection) + except Exception as e: raise SpannerConnectionError(f"Failed to connect to Spanner: {e}") - def close_connection(self) -> None: + + def close(self) -> None: """ Closes the connection to the Spanner database. """ - if self.connection: - self.connection.close() - logger.info("Connection to Spanner database closed.") + self.connection.close() + logger.info("Closed Spanner connection.") + def execute_query(self, query: str) -> SpannerQueryResult: """ @@ -119,13 +93,13 @@ def execute_query(self, query: str) -> SpannerQueryResult: cursor = self.connection.cursor() cursor.execute(query) results = cursor.fetchall() - column_names = [desc[0] for desc in cursor.description] # extract column names + column_names = [desc[0] for desc in cursor.description or []] # extract column names logger.debug(f'column names returned from query: {column_names}') execution_time_s = time.time() - start_time logger.info(f"Query completed in {execution_time_s:.3f} seconds.") return SpannerQueryResult(results, column_names) except Exception as e: - raise RuntimeError(f"Query execution failed: {e}") + raise RuntimeError(f"Query execution failed: {e}") from e @staticmethod def convert_spanner_json(data: List[Any]) -> List[Dict[str, Any]]: @@ -241,11 +215,13 @@ def get_edges_df(json_data: list) -> pd.DataFrame: return SpannerGraph.add_type_from_label_to_df(edges_df) - def gql_to_graph(self, res: Plottable, query: str) -> Plottable: + def gql_to_graph(self, query: str, g: Plottable = PyGraphistry.bind()) -> Plottable: """ Executes a query and constructs a Graphistry graph from the results. :param query: The GQL query to execute. + :param g: The Graphistry graph object to use. + :type g: Plottable :return: A Graphistry graph object constructed from the query results. :rtype: Plottable """ @@ -260,7 +236,8 @@ def gql_to_graph(self, res: Plottable, query: str) -> Plottable: edges_df = self.get_edges_df(json_data) # TODO(tcook): add more error handling here if nodes or edges are empty - return res.nodes(nodes_df, 'identifier').edges(edges_df, 'source', 'destination') + return g.nodes(nodes_df, 'identifier').edges(edges_df, 'source', 'destination') + def query_to_df(self, query: str) -> pd.DataFrame: """ @@ -274,3 +251,17 @@ def query_to_df(self, query: str) -> pd.DataFrame: # create DataFrame from json results, adding column names return pd.DataFrame(query_result.data, columns=query_result.column_names) + + +class SpannerGraphContext: + def __init__(self, config: SpannerConfig | None = None): + config = config or PyGraphistry._config.get("spanner") + if not config: + raise ValueError("Missing spanner_config. Register globally with spanner_config or use with_spanner().") + self.spanner_graph = SpannerGraph.from_config(config) + + def __enter__(self): + return self.spanner_graph + + def __exit__(self, exc_type, exc_value, traceback): + self.spanner_graph.close() \ No newline at end of file diff --git a/graphistry/plugins_types/kusto_types.py b/graphistry/plugins_types/kusto_types.py new file mode 100644 index 0000000000..0977dcf8d6 --- /dev/null +++ b/graphistry/plugins_types/kusto_types.py @@ -0,0 +1,20 @@ +from typing import List, Any +from typing_extensions import TypedDict, NotRequired + +class KustoConnectionError(Exception): + pass + + +class KustoQueryResult: + def __init__(self, data: List[List[Any]], column_names: List[str], column_types: List[str]): + self.data = data + self.column_names = column_names + self.column_types = column_types + + +class KustoConfig(TypedDict): + cluster: str + database: str + client_id: NotRequired[str] + client_secret: NotRequired[str] + tenant_id: NotRequired[str] diff --git a/graphistry/plugins_types/spanner_types.py b/graphistry/plugins_types/spanner_types.py new file mode 100644 index 0000000000..30e6a26341 --- /dev/null +++ b/graphistry/plugins_types/spanner_types.py @@ -0,0 +1,32 @@ +from typing import List, Any +from typing_extensions import TypedDict, NotRequired + +class SpannerConnectionError(Exception): + """Custom exception for errors related to Spanner connection.""" + pass + +class SpannerQueryResult: + """ + Encapsulates the results of a query, including metadata. + + :ivar list data: The raw query results. + """ + + def __init__(self, data: List[Any], column_names: List[str]): + """ + Initializes a SpannerQueryResult instance. + + :param data: The raw query results. + :type List[Any] + :param column_names: a list of the column names from the cursor, defaults to None + :type: List[str] + """ + self.data = data + self.column_names = column_names + + +class SpannerConfig(TypedDict): + project_id: NotRequired[str] + instance_id: str + database_id: str + credentials_file: NotRequired[str] diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py index de46b41ace..67a0ff7ebe 100644 --- a/graphistry/pygraphistry.py +++ b/graphistry/pygraphistry.py @@ -3,6 +3,8 @@ from graphistry.Plottable import Plottable from graphistry.privacy import Mode, Privacy from graphistry.utils.requests import log_requests_error +from graphistry.plugins_types.spanner_types import SpannerConfig +from graphistry.plugins_types.kusto_types import KustoConfig """Top-level import of class PyGraphistry as "Graphistry". Used to connect to the Graphistry server and then create a base plotter.""" import calendar, gzip, io, json, os, numpy as np, pandas as pd, requests, sys, time, warnings @@ -1083,46 +1085,46 @@ def bolt(driver=None): return Plotter().bolt(driver) @staticmethod - def spanner_init(spanner_config: Dict[str, str]) -> Plottable: + def spanner(spanner_config: SpannerConfig) -> Plottable: """ - Initializes a SpannerGraph object with the provided configuration and connects to the instance db + Set spanner configuration for this Plottable. - spanner_config dict must contain the include the following keys, credentials_file is optional: - - "project_id": The GCP project ID. + SpannerConfig - "instance_id": The Spanner instance ID. - "database_id": The Spanner database ID. + - "project_id": The GCP project ID. - "credentials_file": json file API key for service accounts - - :param spanner_config A dictionary containing the Spanner configuration. - :type (Dict[str, str]) + + If credentials_file is provided, it will be used to authenticate with the Spanner instance. + Otherwise, project_id and the spanner login process will be used to authenticate. + + :param spanner_config: A dictionary containing the Spanner configuration. + :type (SpannerConfig) :return: Plottable with a Spanner connection :rtype: Plottable - :raises ValueError: If any of the required keys in `spanner_config` are missing or have invalid values. - - Call this to create a Plotter with a Spanner Graph Connection - - **Example** - - :: - - import graphistry - spanner_CONF = { project_id: "my_project", instance_id: "my_instance", database_id: "my_database"} - g = graphistry.spanner_init(spanner_CONF) - """ - if spanner_config is None: - logger.warn('spanner_init called with spanner_config with None type. Not connected.') - return None - else: - return Plotter().spanner_init(spanner_config) + return Plotter().spanner(spanner_config) @staticmethod - def kusto_init(kusto_config: Dict[str, str]) -> Plottable: - if kusto_config is None: - logger.warn('kusto_init called with kusto_config with None type. Not connected.') - return None - else: - return Plotter().kusto_init(kusto_config) + def kusto(kusto_config: KustoConfig) -> Plottable: + """ + Set kusto configuration for this Plottable. + + KustoConfig + - "cluster": The Kusto cluster name. + - "database": The Kusto database name. + For AAD authentication: + - "client_id": The Kusto client ID. + - "client_secret": The Kusto client secret. + - "tenant_id": The Kusto tenant ID. + Otherwise: process will use web browser to authenticate. + + :param kusto_config: A dictionary containing the Kusto configuration. + :type (KustoConfig) + :returns: Plottable with a Kusto connection + :rtype: Plottable + """ + return Plotter().kusto(kusto_config) @staticmethod def cypher(query, params={}): @@ -1998,8 +2000,56 @@ def spanner_query_to_df(query: str) -> pd.DataFrame: return Plotter().spanner_query_to_df(query) @staticmethod - def kusto_query_to_df(query: str) -> pd.DataFrame: - return Plotter().kusto_query_to_df(query) + def kusto_query(query: str, unwrap_nested: bool | None = None) -> List[pd.DataFrame]: + """ + Submit a Kusto/Azure Data Explorer *query* and return result tables. + + Because a Kusto request may emit multiple tables, a **list of + DataFrames** is always returned; most queries yield a single entry. + + unwrap_nested + ------------- + Controls auto‑flattening of *dynamic* (JSON) columns: + • True  – always try to flatten, raise if it fails + • None  – default heuristic: flatten only if table looks nested + • False – leave results untouched + + :param query: Kusto query string + :type query: str + :param unwrap_nested: flatten strategy above + :type unwrap_nested: bool | None + :returns: list of Pandas DataFrames + :rtype: List[pd.DataFrame] + + **Example** + :: + + frames = graphistry.kusto_query("StormEvents | take 100") + df = frames[0] + """ + return Plotter().kusto_query(query, unwrap_nested) + + @staticmethod + def kusto_query_graph(graph_name: str, snap_name: str | None = None) -> Plottable: + """ + Fetch a Kusto *graph* (and optional *snapshot*) as a Graphistry object. + + Under the hood: `graph(..)` + `graph-to-table` to pull **nodes** and + **edges**, then binds them to *self*. + + :param graph_name: name of Kusto graph entity + :type graph_name: str + :param snap_name: optional snapshot/version + :type snap_name: str | None + :returns: Plottable ready for `.plot()` or further transforms + :rtype: Plottable + + **Example** + :: + + g = graphistry.kusto_query_graph("HoneypotNetwork").plot() + """ + return Plotter().kusto_query_graph(graph_name, snap_name) @staticmethod def gsql_endpoint( @@ -2662,13 +2712,12 @@ def _handle_api_response(response): cypher = PyGraphistry.cypher nodexl = PyGraphistry.nodexl tigergraph = PyGraphistry.tigergraph +spanner = PyGraphistry.spanner spanner_gql_to_g = PyGraphistry.spanner_gql_to_g spanner_query_to_df = PyGraphistry.spanner_query_to_df -spanner_init = PyGraphistry.spanner_init - -kusto_query_to_df = PyGraphistry.kusto_query_to_df -kusto_init = PyGraphistry.kusto_init - +kusto = PyGraphistry.kusto +kusto_query = PyGraphistry.kusto_query +kusto_query_graph = PyGraphistry.kusto_query_graph cosmos = PyGraphistry.cosmos neptune = PyGraphistry.neptune gremlin = PyGraphistry.gremlin diff --git a/pyrightconfig.json b/pyrightconfig.json new file mode 100644 index 0000000000..4548920cc0 --- /dev/null +++ b/pyrightconfig.json @@ -0,0 +1,8 @@ +{ + "typeCheckingMode": "basic", + "ignore": [], + "reportArgumentType": "none", + "reportMissingTypeArgument": "none", + "reportUnknownMemberType": "none", + "exclude": [] +}