kumoai 2.14.0.dev202601011731__cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl

kumoai/graph/graph.py

@@ -0,0 +1,948 @@
+import copy
+import io
+import logging
+import time
+from dataclasses import dataclass
+from importlib.util import find_spec
+from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Union
+
+if TYPE_CHECKING:
+    import graphviz
+
+import kumoapi.data_snapshot as snapshot_api
+import kumoapi.graph as api
+from kumoapi.common import JobStatus
+from kumoapi.data_snapshot import GraphSnapshotID
+from tqdm.auto import tqdm
+from typing_extensions import Self
+
+from kumoai import global_state
+from kumoai.client.graph import GraphID
+from kumoai.graph.table import Table
+from kumoai.mixin import CastMixin
+
+logger = logging.getLogger(__name__)
+
+_DEFAULT_INTERVAL_S = 20
+
+
+@dataclass(frozen=True)
+class Edge(CastMixin, api.Edge):
+    r"""An edge represents a relationship between two tables in a
+    :class:`~kumoai.graph.Graph`. Note that edges are **always** bidirectional
+    within the Kumo platform.
+
+    Args:
+        src_table: The name of the source table of the edge. This table must
+            have a foreign key with name :obj:`fkey` that links to the primary
+            key in the destination table.
+        fkey: The name of the foreign key in the source table.
+        dst_table: The name of the destination table in the graph. This table
+            must have a primary key that links to the
+            source table's foreign key.
+
+    Example:
+        >>> import kumoai
+        >>> edge = kumoai.Edge("table_with_fkey", "fkey", "table_with_pkey")
+    """
+    def __iter__(self) -> Iterator[str]:
+        # Allows unwrapping an edge via `src_table, fkey, dst_table = edge`
+        return iter((self.src_table, self.fkey, self.dst_table))
+
+    def __hash__(self) -> int:
+        return hash((self.src_table, self.fkey, self.dst_table))
+
+    @property
+    def _fully_qualified_name(self) -> str:
+        return f"{self.src_table}.{self.fkey}.{self.dst_table}"
+
+
+@dataclass
+class GraphHealthStats:
+    r"""Graph health statistics contain important statistics that represent the
+    healthiness of each defined edge in a graph. These statistics are computed
+    as part of a :class:`~kumoai.graph.Graph` snapshot, and can be fetched by
+    indexing with an :class:`~kumoai.graph.graph.Edge` object.
+    """
+    _stats: Dict[str, api.EdgeHealthStatistics]
+
+    def __init__(self, stats: Dict[str, api.EdgeHealthStatistics]):
+        self._stats = stats
+
+    def __getitem__(self, key: Edge) -> api.EdgeHealthStatistics:
+        return self._stats[key._fully_qualified_name]
+
+    def __repr__(self) -> str:
+        representation = "GraphHealthStats\n"
+        for key, stats in self._stats.items():
+            src, fkey, dst = key.split('.')
+            representation += (f" - Edge({src} ({fkey})-> {dst}) \n")
+            representation += (f"    - {stats.total_num_edges} total edges\n")
+            representation += (
+                f"    - {int(stats.absolute_match_stats.src_in_dst)} "
+                f"({round(stats.percent_match_stats.src_in_dst, 2)}%) rows "
+                f"in {src} have valid edges to {dst}\n")
+            representation += (
+                f"    - {int(stats.absolute_match_stats.dst_in_src)} "
+                f"({round(stats.percent_match_stats.dst_in_src, 2)}%) rows "
+                f"in {dst} have valid edges to {src}\n")
+        return representation
+
+
+class Graph:
+    r"""A graph defines the relationships between a set of Kumo tables, akin
+    to relationships between tables in a relational database. Creating a graph
+    is the final step of data definition in Kumo; after a graph is created, you
+    are ready to write a :class:`~kumoai.pquery.PredictiveQuery` and train a
+    predictive model.
+
+
+    .. code-block:: python
+
+        import kumoai
+
+        # Define connector to source data:
+        connector = kumoai.S3Connector('s3://...')
+
+        # Create Kumo Tables. See Table documentation for more information:
+        customer = kumoai.Table(...)
+        article = kumoai.Table(...)
+        transaction = kumoai.Table(...)
+
+        # Create a graph:
+        graph = kumo.Graph(
+            # These are the tables that participate in the graph: the keys of this
+            # dictionary are the names of the tables, and the values are the Table
+            # objects that correspond to these names:
+            tables={
+                'customer': customer,
+                'stock': stock,
+                'transaction': transaction,
+            },
+
+            # These are the edges that define the primary key / foreign key
+            # relationships between the tables defined above. Here, `src_table`
+            # is the table that has the foreign key `fkey`, which maps to the
+            # table `dst_table`'s primary key:`
+            edges=[
+                dict(src_table='transaction', fkey='StockCode', dst_table='stock'),
+                dict(src_table='transaction', fkey='CustomerID', dst_table='customer'),
+            ],
+        )
+
+        # Validate the graph configuration, for use in Kumo downstream models:
+        graph.validate(verbose=True)
+
+        # Visualize the graph:
+        graph.visualize()
+
+        # Fetch the statistics of the tables in this graph (this method will
+        # take a graph snapshot, and as a result may have high latency):
+        graph.get_table_stats(wait_for="minimal")
+
+        # Fetch link health statistics (this method will
+        # take a graph snapshot, and as a result may have high latency):
+        graph.get_edge_stats(non_blocking=Falsej)
+
+    Args:
+        tables: The tables in the graph, represented as a dictionary mapping
+            unique table names (within the context of this graph) to the
+            :class:`~kumoai.graph.Table` definition for the table.
+        edges: The edges (relationships) between the :obj:`tables` in the
+            graph. Edges must specify the source table, foreign key, and
+            destination table that they link.
+
+    .. # noqa: E501
+    """
+    def __init__(
+        self,
+        tables: Optional[Dict[str, Table]] = None,
+        edges: Optional[List[Edge]] = None,
+    ) -> None:
+        self._tables: Dict[str, Table] = {}
+        self._edges: List[Edge] = []
+
+        for name, table in (tables or {}).items():
+            self.add_table(name, table)
+
+        for edge in (edges or []):
+            self.link(Edge._cast(edge))
+
+        # Cached from backend:
+        self._graph_snapshot_id: Optional[GraphSnapshotID] = None
+
+    def print_definition(self) -> None:
+        r"""Prints the full definition for this graph; the definition uses
+        placeholder names in place of `kumoai.graph.Table` variables. Copy and
+        paste this definition, modify the table variable names to re-create
+        the original graph.
+
+        Example:
+            >>> import kumoai
+            >>> graph = kumoai.Graph(...)  # doctest: +SKIP
+            >>> graph.print_definition()  # doctest: +SKIP
+            Graph(
+              tables={
+                'table-1' : <table-1>,
+                'table-2' : <table-2>,
+                ...
+                'table-N' : <table-N>,
+              },
+              edges=[
+                Edge(src_table='table-A', fkey='fkey-AD', dst_table='table-D'),
+                Edge(src_table='table-B', fkey='fkey-BE', dst_table='table-E'),
+                ...
+                Edge(src_table='table-C', fkey='fkey-CF', dst_table='table-F'),
+              ],
+            )
+        """
+        definition = f"{self.__class__.__name__}(\n"
+        definition += "  tables={"
+        for table in self._tables.keys():
+            definition += f"\n    '{table}' : <{table}>,"
+        definition += "\n  },\n"
+        definition += "  edges=["
+        for edge in self._edges:
+            src_table, fkey, dst_table = edge
+            definition += (
+                f"\n    {edge.__class__.__name__}(src_table='{src_table}', "
+                f"fkey='{fkey}', dst_table='{dst_table}'),")
+        definition += "\n  ],\n)"
+        print(definition)
+
+    # Properties ##############################################################
+
+    @property
+    def id(self) -> str:
+        r"""Returns the unique ID for this graph, determined from its
+        schema and the schemas of the tables and columns that it contains. Two
+        graphs with any differences in their constituent tables or columns are
+        guaranteed to have unique identifiers.
+        """
+        return self.save()
+
+    # Save / load #############################################################
+
+    def _to_api_graph_definition(self) -> api.GraphDefinition:
+        col_groups_by_dst_table: Dict[str, List[api.ColumnKey]] = dict()
+        for edge in self.edges:
+            dst_pkey = self[edge.dst_table].primary_key
+            if dst_pkey is None:
+                raise ValueError(
+                    f"The destination table {edge.dst_table} of edge "
+                    f"{edge} does not have a primary key.")
+            if edge.dst_table not in col_groups_by_dst_table:
+                col_groups_by_dst_table[edge.dst_table] = [
+                    api.ColumnKey(edge.dst_table, dst_pkey.name)
+                ]
+            col_groups_by_dst_table[edge.dst_table].append(
+                api.ColumnKey(edge.src_table, edge.fkey))
+
+        return api.GraphDefinition(
+            tables={
+                table_name: table._to_api_table_definition()
+                for table_name, table in self.tables.items()
+            },
+            col_groups=[
+                api.ColumnKeyGroup(columns=tuple(col_keys))
+                for col_keys in col_groups_by_dst_table.values()
+            ],
+        )
+
+    @staticmethod
+    def _edges_from_api_graph_definition(
+            graph_definition: api.GraphDefinition) -> List[Edge]:
+        edges: List[Edge] = []
+        for col_group in graph_definition.col_groups:
+            pkey_col = None
+            for col in col_group.columns:
+                table_def = graph_definition.tables[col.table_name]
+                if col.col_name == table_def.pkey:
+                    pkey_col = col
+                    break
+            assert pkey_col is not None
+            for col in col_group.columns:
+                if col != pkey_col:
+                    edges.append(
+                        Edge(src_table=col.table_name, fkey=col.col_name,
+                             dst_table=pkey_col.table_name))
+
+        return edges
+
+    @staticmethod
+    def _from_api_graph_definition(
+            graph_definition: api.GraphDefinition) -> 'Graph':
+        tables = {
+            k: Table._from_api_table_definition(v)
+            for k, v in graph_definition.tables.items()
+        }
+        edges = Graph._edges_from_api_graph_definition(graph_definition)
+        return Graph(tables, edges)
+
+    def save(
+        self,
+        name: Optional[str] = None,
+        skip_validation: bool = False,
+    ) -> Union[GraphID, str]:
+        r"""Associates this graph with a unique name, that can later be
+        used to fetch the graph either in the Kumo UI or in the Kumo SDK
+        with method :meth:`~kumoai.Graph.load`.
+
+        Args:
+            name: The name to associate with this table definition. If the
+                name is already associated with another table, that table will
+                be overridden.
+            skip_validation: Whether to skip validation of the graph. If
+                :obj:`True`, validation will be skipped, but saving an invalid
+                graph may result in undefined behavior.
+                If :obj:`False`, the graph will be validated before saving.
+
+        Example:
+            >>> import kumoai
+            >>> graph = kumoai.Graph(...)  # doctest: +SKIP
+            >>> graph.save()  # doctest: +SKIP
+            graph-xxx
+            >>> graph.save("template_name")  # doctest: +SKIP
+            >>> loaded = kumoai.Graph.load("template_name")  # doctest: +SKIP
+        """
+        if not skip_validation:
+            self.validate(verbose=False)
+
+        template_resource = (global_state.client.graph_api.get_graph_if_exists(
+            graph_id_or_name=name)) if name else None
+        if template_resource is not None:
+            config = self._from_api_graph_definition(template_resource.graph)
+            logger.warning(
+                ("Graph template %s already exists, with configuration %s. "
+                 "This template will be overridden with configuration %s."),
+                name, str(config), str(self))
+
+        # Save as named template
+        return global_state.client.graph_api.create_graph(
+            graph_def=self._to_api_graph_definition(),
+            force_rename=True if name else False,
+            name_alias=name,
+        )
+
+    @classmethod
+    def load(cls, graph_id_or_template: str) -> 'Graph':
+        r"""Loads a graph from either a graph ID or a named template. Returns a
+        :class:`Graph` object that contains the loaded graph along with its
+        associated tables, columns, etc.
+        """
+        api = global_state.client.graph_api
+        res = api.get_graph_if_exists(graph_id_or_template)
+        if not res:
+            raise ValueError(f"Graph {graph_id_or_template} was not found.")
+        out = cls._from_api_graph_definition(res.graph)
+        return out
+
+    # Snapshot ################################################################
+
+    @property
+    def snapshot_id(self) -> Optional[snapshot_api.GraphSnapshotID]:
+        r"""Returns the snapshot ID of this graph's snapshot, if a snapshot
+        has been taken. Returns `None` otherwise.
+
+        .. warning::
+            This function currently only returns a snapshot ID if a snapshot
+            has been taken *in this session.*
+        """
+        return self._graph_snapshot_id
+
+    def snapshot(
+        self,
+        *,
+        force_refresh: bool = False,
+        non_blocking: bool = False,
+    ) -> snapshot_api.GraphSnapshotID:
+        r"""Takes a *snapshot* of this graph's underlying data, and returns a
+        unique identifier for this snapshot.
+
+        This is equivalent to taking a snapshot for each constituent table in
+        the graph.  For more information, please see the documentation for
+        :meth:`~kumoai.graph.Table.snapshot`.
+
+        .. warning::
+            Please familiarize yourself with the warnings for this method in
+            :class:`~kumoai.graph.Table` before proceeding.
+
+        Args:
+            force_refresh: Indicates whether a snapshot should be taken, if one
+                already exists in Kumo. If :obj:`False`, a previously existing
+                snapshot may be re-used. If :obj:`True`, a new snapshot is
+                always taken.
+            non_blocking: Whether this operation should return immediately
+                after creating the snapshot, or await completion of the
+                snapshot. If :obj:`True`, the snapshot will proceed in the
+                background, and will be used for any downstream job.
+
+        Raises:
+            RuntimeError: if ``non_blocking`` is set to :obj:`False` and the
+                graph snapshot fails.
+        """
+        if self._graph_snapshot_id is None or force_refresh:
+            self.save()
+            if not force_refresh:
+                snapshotted_table_names: List[str] = []
+                for table_name, table in self.tables.items():
+                    if table.snapshot_id is not None:
+                        snapshotted_table_names.append(table_name)
+                if len(snapshotted_table_names) > 0:
+                    logger.warning(
+                        "Tables %s have already been snapshot, and will not "
+                        "be refreshed. If you would like to refresh all "
+                        "tables, please set 'force_refresh=True'.",
+                        snapshotted_table_names)
+
+            self._graph_snapshot_id = (
+                global_state.client.graph_api.create_snapshot(
+                    graph_id=self.id,
+                    refresh_source=True,
+                ))
+            logger.info("Graph snapshot with identifier %s created.",
+                        self._graph_snapshot_id)
+
+            # Perform initial GET to update table snapshot IDs:
+            graph_resource: snapshot_api.GraphSnapshotResource = (
+                global_state.client.graph_api.get_snapshot(
+                    snapshot_id=self._graph_snapshot_id))
+            for table_name, table_id in graph_resource.table_ids.items():
+                self[table_name]._table_snapshot_id = table_id
+
+            # NOTE we do not use a `KumoFuture` here as we do not want to treat
+            # a graph refresh as having its own state; since we only ever
+            # operate on the latest graph version (and do not let users to time
+            # travel), there is no need for a separate Future object:
+            if not non_blocking:
+                stage = snapshot_api.GraphSnapshotStage.INGEST
+                table_status: Dict[str, JobStatus] = {
+                    table_name: JobStatus.NOT_STARTED
+                    for table_name in self.tables
+                }
+
+                # Increment progress bar with table refresh stages:
+                done = [status.is_terminal for status in table_status.values()]
+                graph_done = False
+                if logger.isEnabledFor(logging.INFO):
+                    pbar = tqdm(total=len(done), unit="table",
+                                desc="Ingesting")
+                while not (all(done) and graph_done):
+                    graph_resource = (
+                        global_state.client.graph_api.get_snapshot(
+                            snapshot_id=self._graph_snapshot_id))
+                    for table_name, table_id in graph_resource.table_ids.items(
+                    ):
+                        resource = (global_state.client.table_api.get_snapshot(
+                            snapshot_id=table_id))
+                        table_status[table_name] = resource.stages[
+                            stage].status
+                    done = [
+                        status.is_terminal for status in table_status.values()
+                    ]
+                    graph_done = graph_resource.stages[
+                        stage].status.is_terminal
+                    if logger.isEnabledFor(logging.INFO):
+                        pbar.update(sum(done) - pbar.n)
+                    time.sleep(_DEFAULT_INTERVAL_S)
+                if logger.isEnabledFor(logging.INFO):
+                    pbar.update(len(done) - pbar.n)
+                    pbar.close()
+
+                state = graph_resource.stages[stage]
+                status = state.status
+                warnings = "\n".join([
+                    f"{i}. {message}"
+                    for i, message in enumerate(state.warnings)
+                ])
+                error = state.error
+                if status == JobStatus.FAILED:
+                    raise RuntimeError(
+                        f"Graph snapshot with identifier "
+                        f"{self._graph_snapshot_id} failed, with error "
+                        f"{error} and warnings {warnings}")
+                if len(state.warnings) > 0:
+                    logger.warning(
+                        "Graph snapshot completed with the following "
+                        "warnings: %s", warnings)
+        else:
+            logger.warning(
+                "Graph snapshot with identifier %s already exists, and will "
+                "not be refreshed.", self._graph_snapshot_id)
+
+        # <prefix>@<data_version>:
+        assert self._graph_snapshot_id is not None
+        return self._graph_snapshot_id
+
+    # Statistics ##############################################################
+
+    def get_table_stats(
+        self,
+        wait_for: Optional[str] = None,
+    ) -> Dict[str, Dict[str, Any]]:
+        r"""Returns all currently computed statistics on the latest snapshot of
+        this graph. If a snapshot on this graph has not been taken, this method
+        will take a snapshot.
+
+        .. note::
+            Graph statistics are computed in multiple stages after ingestion is
+            complete. These stages are called *minimal* and *full*; minimal
+            statistics are always computed before full statistics.
+
+        Args:
+            wait_for: Whether this operation should block on the existence of
+                statistics availability. This argument can take one of three
+                values: :obj:`None`, which indicates that the method should
+                return immediately with whatever statistics are present,
+                :obj:`"minimal"`, which indicates that the method should return
+                the when the minimum, maximum, and fraction of NA values
+                statistics are present, or :obj:`"full"`, which indicates that
+                the method should return when all computed statistics are
+                present.
+        """
+        assert wait_for is None or wait_for in {"minimal", "full"}
+
+        # Wait for graph ingestion to be done:
+        if not self._graph_snapshot_id:
+            self.snapshot(force_refresh=False, non_blocking=False)
+        assert self._graph_snapshot_id is not None
+
+        # Wait for all table snapshots to match the `wait_for` stage, if
+        # we support that:
+        if wait_for:
+            if wait_for == "minimal":
+                stage = snapshot_api.TableSnapshotStage.MIN_COL_STATS
+            else:
+                stage = snapshot_api.TableSnapshotStage.FULL_COL_STATS
+
+            table_status: Dict[str, JobStatus] = {
+                table_name: JobStatus.NOT_STARTED
+                for table_name in self.tables
+            }
+            done = [status.is_terminal for status in table_status.values()]
+            if logger.isEnabledFor(logging.INFO):
+                pbar = tqdm(total=len(done), unit="table",
+                            desc="Computing Statistics")
+            while not all(done):
+                for table_name, table in self.tables.items():
+                    resource = (global_state.client.table_api.get_snapshot(
+                        snapshot_id=table._table_snapshot_id))
+                    table_status[table_name] = resource.stages[stage].status
+                done = [status.is_terminal for status in table_status.values()]
+                if logger.isEnabledFor(logging.INFO):
+                    pbar.update(sum(done) - pbar.n)
+                time.sleep(_DEFAULT_INTERVAL_S)
+            if logger.isEnabledFor(logging.INFO):
+                pbar.update(len(done) - pbar.n)
+                pbar.close()
+
+        # Write out statistics:
+        out = {}
+        for table_name, table in self.tables.items():
+            resource = (global_state.client.table_api.get_snapshot(
+                snapshot_id=table._table_snapshot_id))
+            out[table_name] = {
+                stat.column_name: stat.stats
+                for stat in resource.column_stats
+            }
+        return out
+
+    def get_edge_stats(
+        self,
+        *,
+        non_blocking: bool = False,
+    ) -> Optional[GraphHealthStats]:
+        """Retrieves edge health statistics for the edges in a graph, if these
+        statistics have been computed by a graph snapshot.
+
+        Edge health statistics are returned in a
+        :class:`~kumoai.graph.GraphHealthStats` object, and contain information
+        about the match rate between primary key / foreign key relationships
+        between the tables in the graph.
+
+        Args:
+            non_blocking: Whether this operation should return immediately
+                after querying edge statistics (returning `None` if statistics
+                are not available), or await completion of statistics
+                computation.
+        """
+        if self._graph_snapshot_id is None:
+            raise ValueError('In order to calculate edge health statistics, '
+                             'you must first create a snapshot of the graph '
+                             'on which to calculate match statistics for each '
+                             'edge. Please call Graph.snapshot() and then '
+                             'this function.')
+
+        edge_health_response = global_state.client.graph_api.get_edge_stats(
+            graph_snapshot_id=self._graph_snapshot_id)
+
+        if non_blocking:
+            if not edge_health_response.is_ready:
+                return None
+        else:
+            while not edge_health_response.is_ready:
+                edge_health_response = (
+                    global_state.client.graph_api.get_edge_stats(
+                        graph_snapshot_id=self._graph_snapshot_id))
+
+        return GraphHealthStats(edge_health_response.statistics)
+
+    # Tables ##################################################################
+
+    def has_table(self, name: str) -> bool:
+        r"""Returns True if a table by `name` is present in this Graph."""
+        return name in self._tables
+
+    def table(self, name: str) -> Table:
+        r"""Returns a table in this Kumo Graph.
+
+        Raises:
+            KeyError: if no such table is present.
+        """
+        if name not in self._tables:
+            raise KeyError(f"Table '{name}' not found in this graph.")
+        return self._tables[name]
+
+    def add_table(self, name: str, table: Table) -> 'Graph':
+        r"""Adds a table to this Kumo Graph.
+
+        Raises:
+            KeyError: if a table with the same name already exists in this
+                graph.
+        """
+        if name in self._tables:
+            raise KeyError(
+                f"Cannot add table with name '{name}' to this graph; names "
+                f"must be globally unique within a graph.")
+        self._tables[name] = table
+        return self
+
+    def remove_table(self, name: str) -> Self:
+        r"""Removes a table from this graph.
+
+        Raises:
+            KeyError: if no such table is present.
+        """
+        if not self.has_table(name):
+            raise KeyError(f"Table '{name}' not found in this graph.'")
+
+        del self._tables[name]
+        self._edges = [
+            edge for edge in self._edges
+            if edge.src_table != name and edge.dst_table != name
+        ]
+        return self
+
+    @property
+    def tables(self) -> Dict[str, Table]:
+        r"""Returns a list of all :class:`~kumoai.graph.Table` objects that
+        are contained in this graph.
+        """
+        return self._tables
+
+    def infer_metadata(self, inplace: bool = True) -> 'Graph':
+        r"""Infers metadata for the tables in this Graph, by inferring the
+        metadata of each table in the graph. For more information, please
+        see the documentation for
+        :meth:`~kumoai.table.Table.infer_metadata`.
+        """
+        out = self
+        if not inplace:
+            out = copy.deepcopy(self)
+
+        for table in out.tables.values():
+            table.infer_metadata(inplace=True)
+        return out
+
+    # Edges ###################################################################
+
+    def infer_links(self) -> 'Graph':
+        r"""Infers edges for the tables in this Graph. It adds edges to the
+        graph.
+
+        Note that the function only works if the graph edges are empty.
+        """
+        if self._edges is not None and len(self._edges) > 0:
+            raise ValueError(
+                "Cannot infer links if graph edges are not empty.")
+
+        graph_def_with_col_groups = global_state.client.graph_api.infer_links(
+            graph=self._to_api_graph_definition())
+
+        edges = Graph._edges_from_api_graph_definition(
+            graph_def_with_col_groups)
+
+        for edge in (edges or []):
+            logger.info("Inferring edge: %s", edge)
+            self.link(Edge._cast(edge))
+        return self
+
+    def link(self, *args: Optional[Union[str, Edge]],
+             **kwargs: str) -> 'Graph':
+        r"""Links two tables (:obj:`src_table` and :obj:`dst_table`) from the
+        foreign key :obj:`fkey` in the source table to the primary key in the
+        destination table. These edges are treated bidirectionally in Kumo.
+
+        Args:
+            *args: Any arguments to construct a
+                :class:`kumoai.graph.Edge`, or a :class:`kumoai.graph.Edge`
+                itself.
+            **kwargs: Any keyword arguments to construct a
+                :class:`kumoai.graph.Edge`.
+
+        Raises:
+            ValueError: if the edge is already present in the graph, if the
+                source table does not exist in the graph, if the destination
+                table does not exist in the graph, if the source key does not
+                exist in the source table, or if the primary key of the source
+                table is being treated as a foreign key.
+        """
+        edge = Edge._cast(*args, **kwargs)
+        if edge is None:
+            raise ValueError("Cannot add a 'None' edge to a graph.")
+
+        (src_table, fkey, dst_table) = edge
+
+        if edge in self._edges:
+            raise ValueError(f"Cannot add edge {edge} to graph; edge is "
+                             f"already present.")
+
+        if src_table not in self._tables:
+            raise ValueError(
+                f"Source table '{src_table}' does not exist in the graph. "
+                f"Please add it via `Graph.add_table(...)` before proceeding.")
+
+        if dst_table not in self._tables:
+            raise ValueError(
+                f"Destination table '{dst_table}' does not exist in the "
+                f"graph. Please add it via `Graph.add_table(...)` before "
+                f"proceeding.")
+
+        if fkey not in self._tables[src_table]:
+            raise ValueError(
+                f"Source key '{fkey}' does not exist in source table "
+                f"'{src_table}'; please check that you have added it as a "
+                f"column.")
+
+        # Backend limitations: ensure the source is not its primary key:
+        src_pkey = self.table(src_table).primary_key
+        src_is_pkey = src_pkey is not None and src_pkey.name == fkey
+        if src_is_pkey:
+            raise ValueError(f"Cannot treat the primary key of table "
+                             f"'{src_table}' as a foreign key; please "
+                             f"select a different key.")
+
+        self._edges.append(edge)
+        return self
+
+    def unlink(self, *args: Optional[Union[str, Edge]],
+               **kwargs: str) -> 'Graph':
+        r"""Removes an edge added to a Kumo Graph.
+
+        Args:
+            *args: Any arguments to construct a
+                :class:`~kumoai.graph.Edge`, or a :class:`~kumoai.graph.Edge`
+                itself.
+            **kwargs: Any keyword arguments to construct a
+                :class:`~kumoai.graph.Edge`.
+
+        Raises:
+            ValueError: if the edge is not present in the graph.
+        """
+        edge = Edge._cast(*args, **kwargs)
+        if edge not in self._edges:
+            raise ValueError(f"Edge {edge} is not present in {self._edges}")
+        self._edges.remove(edge)
+        return self
+
+    @property
+    def edges(self) -> List[Edge]:
+        r"""Returns a list of all :class:`~kumoai.graph.Edge` objects that
+        represent links in this graph.
+        """
+        return self._edges
+
+    def validate(self, verbose: bool = True) -> Self:
+        r"""Validates a Graph to ensure that all relevant metadata is specified
+        for its Tables and Edges.
+
+        Concretely, validation ensures that all tables are valid (see
+        :meth:`~kumoai.graph.table.validate` for more information), and that
+        edges properly link primary keys and foreign keys between valid tables.
+        It additionally ensures that primary and foreign keys between tables
+        in an edge are of the same data type, so that unexpected mismatches do
+        not occur within the Kumo platform.
+
+        Example:
+            >>> import kumoai
+            >>> graph = kumoai.Graph(...)  # doctest: +SKIP
+            >>> graph.validate()  # doctest: +SKIP
+            ValidationResponse(warnings=[], errors=[])
+
+        Args:
+            verbose: Whether to log non-error output of this validation.
+
+        Raises:
+            ValueError:
+                if validation fails.
+        """
+        # Validate table definitions, so we can properly create a graph
+        # definition:
+        for table_name, table in self.tables.items():
+            try:
+                table.validate(verbose=verbose)
+            except ValueError as e:
+                raise ValueError(
+                    f"Validation of table {table_name} failed. {e}") from e
+
+        resp = global_state.client.graph_api.validate_graph(
+            api.GraphValidationRequest(self._to_api_graph_definition()))
+        if not resp.ok:
+            raise ValueError(resp.error_message())
+        if verbose:
+            if resp.empty():
+                logger.info("Graph is configured correctly.")
+            else:
+                logger.warning(resp.message())
+        return self
+
+    def visualize(
+        self,
+        path: Optional[Union[str, io.BytesIO]] = None,
+        show_cols: bool = True,
+    ) -> 'graphviz.Graph':
+        r"""Visualizes the tables and edges in this graph using the
+        ``graphviz`` library.
+
+        Args:
+            path: An optional local path to write the produced image to. If
+                None, the image will not be written to disk.
+            show_cols: Whether to show all columns of every table in the graph.
+                If False, will only show the primary key, foreign key(s),
+                time column, and end time column of each table.
+
+        Returns:
+            A ``graphviz.Graph`` instance representing the visualized graph.
+        """
+        def has_graphviz_executables() -> bool:
+            import graphviz
+            try:
+                graphviz.Digraph().pipe()
+            except graphviz.backend.ExecutableNotFound:
+                return False
+
+            return True
+
+        # Check basic dependency:
+        if not find_spec('graphviz'):
+            raise ModuleNotFoundError(
+                "The `graphviz` Python package is required for visualization.")
+        elif not has_graphviz_executables():
+            raise RuntimeError(
+                "Could not visualize graph as `graphviz` executables have not "
+                "been installed. These dependencies are required in addition "
+                "to the `graphviz` Python package. Please install them to "
+                "continue. Instructions at https://graphviz.org/download/.")
+        else:
+            import graphviz
+
+        fmt = None
+        if isinstance(path, str):
+            fmt = path.split('.')[-1]
+        elif isinstance(path, io.BytesIO):
+            fmt = 'svg'
+        graph = graphviz.Graph(format=fmt)
+
+        def left_align(list_of_text: List[str]) -> str:
+            return '\\l'.join(list_of_text) + '\\l'
+
+        table_to_fkey: Dict[str, List[str]] = {}
+        for edge in self.edges:
+            src, fkey, dst = edge
+            if src not in table_to_fkey:
+                table_to_fkey[src] = []
+            table_to_fkey[src].append(fkey)
+
+        for table_name, table in self.tables.items():
+            keys = []
+            if table.has_primary_key():
+                assert table.primary_key is not None
+                keys += [f'{table.primary_key.name} (PK)']
+            if table_name in table_to_fkey:
+                keys += [f'{fkey} (FK)' for fkey in table_to_fkey[table_name]]
+            if table.has_time_column():
+                assert table.time_column is not None
+                keys += [f'{table.time_column.name} (Time)']
+            if table.has_end_time_column():
+                assert table.end_time_column is not None
+                keys += [f'{table.end_time_column.name} (End Time)']
+
+            keys_aligned = left_align(keys)
+
+            cols = []
+            cols_aligned = ""
+            if show_cols and len(table.columns) > 0:
+                cols += [
+                    f'{col.name}: {col.stype or "???"} ({col.dtype or "???"})'
+                    for col in table.columns
+                ]
+                cols_aligned = left_align(cols)
+
+            if cols:
+                label = f'{{{table_name}|{keys_aligned}|{cols_aligned}}}'
+            else:
+                label = f'{{{table_name}|{keys_aligned}}}'
+
+            graph.node(table_name, shape='record', label=label)
+
+        for edge in self.edges:
+            src, fkey, dst = edge
+            pkey_obj = self[dst].primary_key
+            assert pkey_obj is not None
+            pkey = pkey_obj.name
+            # Print both key names if different:
+            if fkey != pkey:
+                label = f' {fkey}\n< >\n{pkey} '
+            else:
+                label = f' {fkey} '
+            headlabel, taillabel = '1', '*'
+            graph.edge(src, dst, label=label, headlabel=headlabel,
+                       taillabel=taillabel, minlen='2', fontsize='11pt',
+                       labeldistance='1.5')
+
+        if isinstance(path, str):
+            path = '.'.join(path.split('.')[:-1])
+            graph.render(path, cleanup=True)
+        elif isinstance(path, io.BytesIO):
+            path.write(graph.pipe())
+        else:
+            try:
+                graph.view()
+            except Exception as e:
+                logger.warning(
+                    "Could not visualize graph due to an unexpected error in "
+                    "`graphviz`. If you are in a notebook environment, "
+                    "consider calling `display()` on the returned object "
+                    "from `visualize()`. Error: %s", e)
+        return graph
+
+    # Class properties ########################################################
+
+    def __hash__(self) -> int:
+        return hash((tuple(self.edges), self.tables.values()))
+
+    def __contains__(self, name: str) -> bool:
+        return self.has_table(name)
+
+    def __getitem__(self, name: str) -> Table:
+        return self.table(name)
+
+    def __delitem__(self, name: str) -> None:
+        self.remove_table(name)
+
+    def __repr__(self) -> str:
+        table_names = str(list(self._tables.keys())).replace("'", "")
+        return (f'{self.__class__.__name__}(\n'
+                f'  tables={table_names},\n'
+                f'  edges={self._edges},\n'
+                f')')