datajunction-server 0.0.1a70.dev0__tar.gz → 0.0.1a70.dev1__tar.gz

{datajunction_server-0.0.1a70.dev0 → datajunction_server-0.0.1a70.dev1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: datajunction-server
-Version: 0.0.1a70.dev0
+Version: 0.0.1a70.dev1
 Summary: DataJunction server library for running to a DataJunction server
 Project-URL: repository, https://github.com/DataJunction/dj
 Author-email: DataJunction Authors <yian.shang@gmail.com>

datajunction_server-0.0.1a70.dev1/alembic/versions/2024_10_05_1707-f953dc7ec5a5_add_indexes_on_history.py

@@ -0,0 +1,58 @@
+"""
+Add indexes on history and node tables
+
+Revision ID: f953dc7ec5a5
+Revises: f3c9b40deb6f
+Create Date: 2024-10-05 17:07:17.320467+00:00
+"""
+# pylint: disable=no-member, invalid-name, missing-function-docstring, unused-import, no-name-in-module
+
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "f953dc7ec5a5"
+down_revision = "f3c9b40deb6f"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    with op.batch_alter_table("history", schema=None) as batch_op:
+        batch_op.create_index("ix_history_entity_name", ["entity_name"], unique=False)
+
+    with op.batch_alter_table("history", schema=None) as batch_op:
+        batch_op.create_index("ix_history_user", ["user"], unique=False)
+
+    with op.batch_alter_table("node", schema=None) as batch_op:
+        batch_op.create_index(
+            "cursor_index",
+            ["created_at", "id"],
+            unique=False,
+            postgresql_using="btree",
+        )
+
+    with op.batch_alter_table("node", schema=None) as batch_op:
+        batch_op.create_index(
+            "namespace_index",
+            ["namespace"],
+            unique=False,
+            postgresql_using="btree",
+            postgresql_ops={"identifier": "varchar_pattern_ops"},
+        )
+
+
+def downgrade():
+    with op.batch_alter_table("node", schema=None) as batch_op:
+        batch_op.drop_index("namespace_index", postgresql_using="text_pattern_ops")
+
+    with op.batch_alter_table("node", schema=None) as batch_op:
+        batch_op.drop_index("cursor_index", postgresql_using="btree")
+
+    with op.batch_alter_table("history", schema=None) as batch_op:
+        batch_op.drop_index("ix_history_user")
+
+    with op.batch_alter_table("history", schema=None) as batch_op:
+        batch_op.drop_index("ix_history_entity_name")

datajunction_server-0.0.1a70.dev1/datajunction_server/__about__.py

@@ -0,0 +1,4 @@
+"""
+Version for Hatch
+"""
+__version__ = "0.0.1a70.dev1"

datajunction_server-0.0.1a70.dev1/datajunction_server/api/graphql/main.py

@@ -0,0 +1,157 @@
+"""DJ graphql"""
+from typing import Annotated, List, Optional
+
+import strawberry
+from fastapi import Depends
+from strawberry.fastapi import GraphQLRouter
+from strawberry.types import Info
+
+from datajunction_server.api.graphql.catalogs import CatalogInfo, list_catalogs
+from datajunction_server.api.graphql.engines import EngineInfo, list_engines
+from datajunction_server.api.graphql.resolvers.nodes import find_nodes_by
+from datajunction_server.api.graphql.scalars import Connection
+from datajunction_server.api.graphql.scalars.node import Node
+from datajunction_server.models.node import NodeCursor, NodeType
+from datajunction_server.utils import get_session, get_settings
+
+
+async def get_context(
+    session=Depends(get_session),
+    settings=Depends(get_settings),
+):
+    """
+    Provides the context for graphql requests
+    """
+    return {"session": session, "settings": settings}
+
+
+@strawberry.type
+class Query:  # pylint: disable=R0903
+    """
+    Parent of all DJ graphql queries
+    """
+
+    list_catalogs: List[CatalogInfo] = strawberry.field(  # noqa: F811
+        resolver=list_catalogs,
+    )
+    list_engines: List[EngineInfo] = strawberry.field(  # noqa: F811
+        resolver=list_engines,
+    )
+
+    @strawberry.field(description="Find nodes based on the search parameters.")
+    async def find_nodes(
+        self,
+        fragment: Annotated[
+            Optional[str],
+            strawberry.argument(
+                description="A fragment of a node name to search for",
+            ),
+        ] = None,
+        names: Annotated[
+            Optional[List[str]],
+            strawberry.argument(
+                description="Filter to nodes with these names",
+            ),
+        ] = None,
+        node_types: Annotated[
+            Optional[List[NodeType]],
+            strawberry.argument(
+                description="Filter nodes to these node types",
+            ),
+        ] = None,
+        tags: Annotated[
+            Optional[List[str]],
+            strawberry.argument(
+                description="Filter to nodes tagged with these tags",
+            ),
+        ] = None,
+        *,
+        info: Info,
+    ) -> List[Node]:
+        """
+        Find nodes based on the search parameters.
+        """
+        return await find_nodes_by(info, names, fragment, node_types, tags)  # type: ignore
+
+    @strawberry.field(
+        description="Find nodes based on the search parameters with pagination",
+    )
+    async def find_nodes_paginated(
+        self,
+        fragment: Annotated[
+            Optional[str],
+            strawberry.argument(
+                description="A fragment of a node name to search for",
+            ),
+        ] = None,
+        names: Annotated[
+            Optional[List[str]],
+            strawberry.argument(
+                description="Filter to nodes with these names",
+            ),
+        ] = None,
+        node_types: Annotated[
+            Optional[List[NodeType]],
+            strawberry.argument(
+                description="Filter nodes to these node types",
+            ),
+        ] = None,
+        tags: Annotated[
+            Optional[List[str]],
+            strawberry.argument(
+                description="Filter to nodes tagged with these tags",
+            ),
+        ] = None,
+        edited_by: Annotated[
+            Optional[str],
+            strawberry.argument(
+                description="Filter to nodes edited by this user",
+            ),
+        ] = None,
+        namespace: Annotated[
+            Optional[str],
+            strawberry.argument(
+                description="Filter to nodes in this namespace",
+            ),
+        ] = None,
+        after: Optional[str] = None,
+        before: Optional[str] = None,
+        limit: Annotated[
+            Optional[int],
+            strawberry.argument(description="Limit nodes"),
+        ] = 100,
+        *,
+        info: Info,
+    ) -> Connection[Node]:
+        """
+        Find nodes based on the search parameters.
+        """
+        if not limit or limit < 0:
+            limit = 100
+        nodes_list = await find_nodes_by(
+            info,
+            names,
+            fragment,
+            node_types,
+            tags,
+            edited_by,
+            namespace,
+            limit + 1,
+            before,
+            after,
+        )
+        return Connection.from_list(
+            items=nodes_list,
+            before=before,
+            after=after,
+            limit=limit,
+            encode_cursor=lambda dj_node: NodeCursor(
+                created_at=dj_node.created_at,
+                id=dj_node.id,
+            ),
+        )
+
+
+schema = strawberry.Schema(query=Query)
+
+graphql_app = GraphQLRouter(schema, context_getter=get_context)

{datajunction_server-0.0.1a70.dev0 → datajunction_server-0.0.1a70.dev1}/datajunction_server/api/graphql/resolvers/nodes.py

@@ -20,6 +20,11 @@ async def find_nodes_by(
     fragment: Optional[str] = None,
     node_types: Optional[List[NodeType]] = None,
     tags: Optional[List[str]] = None,
+    edited_by: Optional[str] = None,
+    namespace: Optional[str] = None,
+    limit: Optional[int] = 100,
+    before: Optional[str] = None,
+    after: Optional[str] = None,
 ) -> List[DBNode]:
     """
     Finds nodes based on the search parameters. This function also tries to optimize
@@ -27,13 +32,24 @@ async def find_nodes_by(
     """
     session = info.context["session"]  # type: ignore
     fields = extract_fields(info)
-    options = load_node_options(fields)
+    options = load_node_options(
+        fields["nodes"]
+        if "nodes" in fields
+        else fields["edges"]["node"]
+        if "edges" in fields
+        else fields,
+    )
     return await DBNode.find_by(
         session,
         names,
         fragment,
         node_types,
         tags,
+        edited_by,
+        namespace,
+        limit,
+        before,
+        after,
         *options,
     )
 

datajunction_server-0.0.1a70.dev1/datajunction_server/api/graphql/scalars/__init__.py

@@ -0,0 +1,143 @@
+"""
+GraphQL scalars
+"""
+import dataclasses
+import datetime
+import json
+from base64 import b64decode, b64encode
+from dataclasses import dataclass
+from typing import Callable, Generic, List, Optional, TypeVar, Union
+
+import strawberry
+from strawberry import field
+
+BigInt = strawberry.scalar(
+    Union[int, str],  # type: ignore
+    serialize=int,
+    parse_value=str,
+    description="BigInt field",
+)
+
+GenericItem = TypeVar("GenericItem")
+GenericItemNode = TypeVar("GenericItemNode")
+
+
+class DateTimeJSONEncoder(json.JSONEncoder):
+    """
+    JSON encoder that handles datetime objects
+    """
+
+    def default(self, obj):  # pylint: disable=arguments-renamed
+        """
+        Check if there are datetime objects and serialize them as ISO
+        format strings.
+        """
+        if isinstance(obj, (datetime.datetime, datetime.date, datetime.time)):
+            return obj.isoformat()
+        return super().default(obj)  # pragma: no cover
+
+
+class DateTimeJSONDecoder(json.JSONDecoder):
+    """
+    JSON decoder that handles ISO format datetime strings
+    """
+
+    def __init__(self, *args, **kwargs):
+        json.JSONDecoder.__init__(self, object_hook=self.object_hook, *args, **kwargs)
+
+    def object_hook(self, source):  # pylint: disable=method-hidden
+        """
+        Check if the string is in ISO 8601 format and convert to a datetime
+        object if it is.
+        """
+        for k, v in source.items():
+            if isinstance(v, str):
+                try:
+                    source[k] = datetime.datetime.fromisoformat(str(v))
+                except ValueError:  # pragma: no cover
+                    pass  # pragma: no cover
+        return source
+
+
+@dataclass
+class Cursor:
+    """
+    Dataclass that serializes into a string and back.
+    """
+
+    def encode(self) -> str:
+        """Serialize this cursor into a string."""
+        json_str = json.dumps(dataclasses.asdict(self), cls=DateTimeJSONEncoder)
+        return b64encode(json_str.encode()).decode()
+
+    @classmethod
+    def decode(cls, serialized: str) -> "Cursor":
+        """Parse the string into an instance of this Cursor class."""
+        json_str = b64decode(serialized.encode()).decode()
+        json_obj = json.loads(json_str, cls=DateTimeJSONDecoder)
+        return cls(**json_obj)
+
+
+@strawberry.type
+class PageInfo:  # pylint: disable=too-few-public-methods
+    """Metadata about a page in a connection."""
+
+    has_next_page: bool = field(
+        description="When paginating forwards, are there more nodes?",
+    )
+    has_prev_page: bool = field(
+        description="When paginating forwards, are there more nodes?",
+    )
+    start_cursor: str | None = field(
+        description="When paginating back, the cursor to continue.",
+    )
+    end_cursor: str | None = field(
+        description="When paginating forwards, the cursor to continue.",
+    )
+
+
+@strawberry.type
+class Edge(Generic[GenericItemNode]):  # pylint: disable=too-few-public-methods
+    """Metadata about an item in a connection."""
+
+    node: GenericItemNode
+
+
+@strawberry.type
+class Connection(Generic[GenericItemNode]):  # pylint: disable=too-few-public-methods
+    """
+    Pagination for a list of items.
+    """
+
+    page_info: PageInfo
+    edges: list[Edge[GenericItemNode]]
+
+    @classmethod
+    def from_list(
+        cls,
+        items: List[GenericItem],
+        before: Optional[str],
+        after: Optional[str],
+        limit: int,
+        encode_cursor: Callable[[GenericItem], Cursor],
+    ) -> "Connection":
+        """
+        Construct a Connection from a list of items.
+        """
+        has_next_page = len(items) > limit or (
+            before is not None and items[0] is not None
+        )
+        has_prev_page = (before is not None and len(items) > limit) or (
+            after is not None and items[0] is not None
+        )
+        start_cursor = encode_cursor(items[0]).encode() if items else None
+        end_cursor = encode_cursor(items[-1]).encode() if items else None
+        return Connection(  # type: ignore
+            page_info=PageInfo(  # type: ignore
+                has_prev_page=has_prev_page,
+                start_cursor=start_cursor,
+                has_next_page=has_next_page,
+                end_cursor=end_cursor,
+            ),
+            edges=[Edge(node=item) for item in items[:limit]],  # type: ignore
+        )

{datajunction_server-0.0.1a70.dev0 → datajunction_server-0.0.1a70.dev1}/datajunction_server/api/graphql/scalars/node.py

@@ -1,5 +1,4 @@
 """Node-related scalars."""
-
 import datetime
 from typing import List, Optional
 
@@ -189,4 +188,4 @@ class Node:  # pylint: disable=too-few-public-methods
         """
         The users who edited this node
         """
-        return list({entry.user for entry in root.history})
+        return root.edited_by

{datajunction_server-0.0.1a70.dev0 → datajunction_server-0.0.1a70.dev1}/datajunction_server/api/graphql/utils.py

@@ -2,6 +2,8 @@
 import re
 from typing import Any, Dict
 
+CURSOR_SEPARATOR = "-"
+
 
 def convert_camel_case(name):
     """
@@ -31,8 +33,8 @@ def extract_fields(query_fields) -> Dict[str, Any]:
     """
     fields = {}
 
-    for field in query_fields.selected_fields:
-        for selection in field.selections:
+    for query_field in query_fields.selected_fields:
+        for selection in query_field.selections:
             field_name = convert_camel_case(selection.name)
             if selection.selections:
                 subfield = extract_subfields(selection)

{datajunction_server-0.0.1a70.dev0 → datajunction_server-0.0.1a70.dev1}/datajunction_server/database/history.py

@@ -3,7 +3,7 @@ from datetime import datetime, timezone
 from functools import partial
 from typing import Any, Dict, Optional
 
-from sqlalchemy import JSON, BigInteger, DateTime, Enum, Integer, String
+from sqlalchemy import JSON, BigInteger, DateTime, Enum, Index, Integer, String
 from sqlalchemy.orm import Mapped, mapped_column
 
 from datajunction_server.database.base import Base
@@ -53,6 +53,10 @@ class History(Base):  # pylint: disable=too-few-public-methods
     """
 
     __tablename__ = "history"
+    __table_args__ = (
+        Index("ix_history_entity_name", "entity_name", postgresql_using="btree"),
+        Index("ix_history_user", "user", postgresql_using="btree"),
+    )
 
     id: Mapped[int] = mapped_column(
         BigInteger().with_variant(Integer, "sqlite"),

{datajunction_server-0.0.1a70.dev0 → datajunction_server-0.0.1a70.dev1}/datajunction_server/database/node.py

@@ -12,6 +12,7 @@ from sqlalchemy import (
     DateTime,
     Enum,
     ForeignKey,
+    Index,
     Integer,
     String,
     UniqueConstraint,
@@ -38,6 +39,7 @@ from datajunction_server.models.base import labelize
 from datajunction_server.models.node import (
     DEFAULT_DRAFT_VERSION,
     BuildCriteria,
+    NodeCursor,
     NodeMode,
     NodeStatus,
 )
@@ -165,6 +167,13 @@ class Node(Base):  # pylint: disable=too-few-public-methods
     __tablename__ = "node"
     __table_args__ = (
         UniqueConstraint("name", "namespace", name="unique_node_namespace_name"),
+        Index("cursor_index", "created_at", "id", postgresql_using="btree"),
+        Index(
+            "namespace_index",
+            "namespace",
+            postgresql_using="btree",
+            postgresql_ops={"identifier": "varchar_pattern_ops"},
+        ),
     )
 
     id: Mapped[int] = mapped_column(
@@ -390,13 +399,18 @@ class Node(Base):  # pylint: disable=too-few-public-methods
         return result.unique().scalars().all()
 
     @classmethod
-    async def find_by(  # pylint: disable=keyword-arg-before-vararg
+    async def find_by(  # pylint: disable=keyword-arg-before-vararg,too-many-locals
         cls,
         session: AsyncSession,
         names: Optional[List[str]] = None,
         fragment: Optional[str] = None,
         node_types: Optional[List[NodeType]] = None,
         tags: Optional[List[str]] = None,
+        edited_by: Optional[str] = None,
+        namespace: Optional[str] = None,
+        limit: Optional[int] = 100,
+        before: Optional[str] = None,
+        after: Optional[str] = None,
         *options: ExecutableOption,  # pylint: disable=keyword-arg-before-vararg
     ) -> List["Node"]:
         """
@@ -416,10 +430,14 @@ class Node(Base):  # pylint: disable=too-few-public-methods
                 return []
 
         statement = select(Node).where(is_(Node.deactivated_at, None))
+        if namespace:
+            statement = statement.where(
+                (Node.namespace.like(f"{namespace}.%")) | (Node.namespace == namespace),
+            )
         if nodes_with_tags:
-            statement = statement.where(  # pragma: no cover
+            statement = statement.where(
                 Node.id.in_(nodes_with_tags),
-            )
+            )  # pragma: no cover
         if names:
             statement = statement.where(
                 Node.name.in_(names),  # type: ignore  # pylint: disable=no-member
@@ -430,8 +448,44 @@ class Node(Base):  # pylint: disable=too-few-public-methods
             )
         if node_types:
             statement = statement.where(Node.type.in_(node_types))
+        if edited_by:
+            edited_node_subquery = (
+                select(History.entity_name)
+                .where((History.user == edited_by))
+                .distinct()
+                .subquery()
+            )
+
+            statement = statement.join(
+                edited_node_subquery,
+                onclause=(edited_node_subquery.c.entity_name == Node.name),
+            ).distinct()
+
+        if after:
+            cursor = NodeCursor.decode(after)
+            statement = statement.where(
+                (Node.created_at, Node.id)
+                <= (cursor.created_at, cursor.id),  # pylint: disable=no-member
+            ).order_by(Node.created_at.desc(), Node.id.desc())
+        elif before:
+            cursor = NodeCursor.decode(before)
+            statement = statement.where(
+                (Node.created_at, Node.id)
+                >= (cursor.created_at, cursor.id),  # pylint: disable=no-member
+            )
+            statement = statement.order_by(Node.created_at.asc(), Node.id.asc())
+        else:
+            statement = statement.order_by(Node.created_at.desc(), Node.id.desc())
+
+        limit = limit if limit and limit > 0 else 100
+        statement = statement.limit(limit)
         result = await session.execute(statement.options(*options))
-        return result.unique().scalars().all()
+        nodes = result.unique().scalars().all()
+
+        # Reverse for backward pagination
+        if before:
+            nodes.reverse()
+        return nodes
 
 
 class NodeRevision(

{datajunction_server-0.0.1a70.dev0 → datajunction_server-0.0.1a70.dev1}/datajunction_server/models/node.py

@@ -5,7 +5,7 @@ Model for nodes.
 import enum
 import sys
 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union, cast
 
 from pydantic import BaseModel, Extra, root_validator, validator
 from pydantic.fields import Field
@@ -15,6 +15,7 @@ from sqlalchemy.sql.schema import Column as SqlaColumn
 from sqlalchemy.types import Enum
 from typing_extensions import TypedDict
 
+from datajunction_server.api.graphql.scalars import Cursor
 from datajunction_server.enum import StrEnum
 from datajunction_server.errors import DJError
 from datajunction_server.models.base import labelize
@@ -959,3 +960,15 @@ class NodeIndegreeOutput(BaseModel):
 
     name: str
     indegree: int
+
+
+@dataclass
+class NodeCursor(Cursor):
+    """Cursor that represents a node in a paginated list."""
+
+    created_at: UTCDatetime
+    id: int
+
+    @classmethod
+    def decode(cls, serialized: str) -> "NodeCursor":
+        return cast(NodeCursor, super().decode(serialized))