datajunction-server 0.0.1a55.dev0__tar.gz → 0.0.1a56.dev0__tar.gz

{datajunction_server-0.0.1a55.dev0 → datajunction_server-0.0.1a56.dev0}/.coveragerc

@@ -16,6 +16,7 @@ source =
     */site-packages/
 
 [report]
+show_missing = True
 # Regexes for lines to exclude from consideration
 exclude_lines =
     # Have to re-enable the standard pragma

{datajunction_server-0.0.1a55.dev0 → datajunction_server-0.0.1a56.dev0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: datajunction-server
-Version: 0.0.1a55.dev0
+Version: 0.0.1a56.dev0
 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.1a55.dev0 → datajunction_server-0.0.1a56.dev0}/alembic/env.py

@@ -11,7 +11,7 @@ from alembic import context
 from datajunction_server.database import (
     AttributeType,
     Catalog,
-    Column,
+    Collection,
     ColumnAttribute,
     Database,
     DimensionLink,
@@ -27,6 +27,7 @@ from datajunction_server.database import (
     User,
 )
 from datajunction_server.database.base import Base
+from datajunction_server.database.column import Column
 from datajunction_server.utils import get_settings
 
 settings = get_settings()

datajunction_server-0.0.1a56.dev0/alembic/versions/2024_06_21_1301-640a814db2d8_add_collection_tables.py

@@ -0,0 +1,56 @@
+"""Add collection tables
+
+Revision ID: 640a814db2d8
+Revises: 57fc93ef6947
+Create Date: 2024-06-21 13:01:48.141719+00:00
+
+"""
+# pylint: disable=no-member, invalid-name, missing-function-docstring, unused-import, no-name-in-module
+
+import sqlalchemy as sa
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "640a814db2d8"
+down_revision = "57fc93ef6947"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    op.create_table(
+        "collection",
+        sa.Column("id", sa.Integer(), nullable=False),
+        sa.Column("name", sa.String(), nullable=True),
+        sa.Column("description", sa.String(), nullable=True),
+        sa.Column("created_at", sa.DateTime(timezone=True), nullable=False),
+        sa.Column("deactivated_at", sa.DateTime(timezone=True), nullable=True),
+        sa.PrimaryKeyConstraint("id"),
+        sa.UniqueConstraint("name"),
+    )
+    op.create_table(
+        "collectionnodes",
+        sa.Column("collection_id", sa.Integer(), nullable=False),
+        sa.Column(
+            "node_id",
+            sa.BigInteger().with_variant(sa.Integer(), "sqlite"),
+            nullable=False,
+        ),
+        sa.ForeignKeyConstraint(
+            ["collection_id"],
+            ["collection.id"],
+            name="fk_collectionnodes_collection_id_collection",
+        ),
+        sa.ForeignKeyConstraint(
+            ["node_id"],
+            ["node.id"],
+            name="fk_collectionnodes_node_id_node",
+        ),
+        sa.PrimaryKeyConstraint("collection_id", "node_id"),
+    )
+
+
+def downgrade():
+    op.drop_table("collectionnodes")
+    op.drop_table("collection")

datajunction_server-0.0.1a56.dev0/datajunction_server/__about__.py

@@ -0,0 +1,4 @@
+"""
+Version for Hatch
+"""
+__version__ = "0.0.1a56.dev0"

datajunction_server-0.0.1a56.dev0/datajunction_server/api/collection.py

@@ -0,0 +1,154 @@
+"""
+Collection related APIs.
+"""
+import logging
+from http import HTTPStatus
+from typing import List
+
+from fastapi import Depends, HTTPException, Response
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.sql.operators import is_
+
+from datajunction_server.database.collection import Collection
+from datajunction_server.database.node import Node
+from datajunction_server.errors import DJException
+from datajunction_server.internal.access.authentication.http import SecureAPIRouter
+from datajunction_server.models.collection import CollectionInfo
+from datajunction_server.utils import get_session, get_settings
+
+_logger = logging.getLogger(__name__)
+settings = get_settings()
+router = SecureAPIRouter(tags=["collections"])
+
+
+@router.post(
+    "/collections/",
+    response_model=CollectionInfo,
+    status_code=HTTPStatus.CREATED,
+)
+async def create_a_collection(
+    data: CollectionInfo,
+    *,
+    session: AsyncSession = Depends(get_session),
+) -> CollectionInfo:
+    """
+    Create a Collection
+    """
+    try:
+        await Collection.get_by_name(session, data.name, raise_if_not_exists=True)
+    except DJException:
+        pass
+    else:
+        raise HTTPException(
+            status_code=HTTPStatus.CONFLICT,
+            detail=f"Collection already exists: `{data.name}`",
+        )
+
+    collection = Collection(
+        name=data.name,
+        description=data.description,
+    )
+    session.add(collection)
+    await session.commit()
+    await session.refresh(collection)
+
+    return CollectionInfo.from_orm(collection)
+
+
+@router.delete("/collections/{name}", status_code=HTTPStatus.NO_CONTENT)
+async def delete_a_collection(
+    name: str,
+    *,
+    session: AsyncSession = Depends(get_session),
+):
+    """
+    Delete a collection
+    """
+    collection = await Collection.get_by_name(session, name)
+    await session.delete(collection)
+    await session.commit()
+    return Response(status_code=HTTPStatus.NO_CONTENT)
+
+
+@router.get("/collections/")
+async def list_collections(
+    *,
+    session: AsyncSession = Depends(get_session),
+) -> List[CollectionInfo]:
+    """
+    List all collections
+    """
+    collections = await session.execute(
+        select(Collection).where(is_(Collection.deactivated_at, None)),
+    )
+    return collections.scalars().all()
+
+
+@router.get("/collections/{name}")
+async def get_collection(
+    name: str,
+    *,
+    session: AsyncSession = Depends(get_session),
+):
+    """
+    Get a collection and its nodes
+    """
+    collection = await Collection.get_by_name(
+        session,
+        name=name,
+        raise_if_not_exists=True,
+    )
+    return collection
+
+
+@router.post(
+    "/collections/{name}/nodes/",
+    status_code=HTTPStatus.NO_CONTENT,
+    name="Add Nodes to a Collection",
+)
+async def add_nodes_to_collection(
+    name: str,
+    data: List[str],
+    *,
+    session: AsyncSession = Depends(get_session),
+):
+    """
+    Add one or more nodes to a collection
+    """
+    collection = await Collection.get_by_name(session, name, raise_if_not_exists=True)
+    nodes = await Node.get_by_names(session=session, names=data)
+    if not nodes:
+        raise HTTPException(
+            status_code=HTTPStatus.NOT_FOUND,
+            detail=f"Cannot add nodes to collection, no nodes found: `{data}`",
+        )
+    collection.nodes.extend(nodes)  # type: ignore
+    session.add(collection)
+    await session.commit()
+    await session.refresh(collection)
+    return Response(status_code=HTTPStatus.NO_CONTENT)
+
+
+@router.post(
+    "/collections/{name}/remove/",
+    status_code=HTTPStatus.NO_CONTENT,
+    name="Delete Nodes from a Collection",
+)
+async def delete_nodes_from_collection(
+    name: str,
+    data: List[str],
+    *,
+    session: AsyncSession = Depends(get_session),
+):
+    """
+    Delete one or more nodes from a collection
+    """
+    collection = await Collection.get_by_name(session, name)
+    nodes = await Node.get_by_names(session=session, names=data)
+    for node in nodes:
+        if node in collection.nodes:  # type: ignore
+            collection.nodes.remove(node)  # type: ignore
+    await session.commit()
+    await session.refresh(collection)
+    return Response(status_code=HTTPStatus.NO_CONTENT)

{datajunction_server-0.0.1a55.dev0 → datajunction_server-0.0.1a56.dev0}/datajunction_server/api/main.py

@@ -27,6 +27,7 @@ from datajunction_server.api import (
     attributes,
     catalogs,
     client,
+    collection,
     cubes,
     data,
     dimensions,
@@ -89,6 +90,7 @@ app.add_middleware(
 )
 
 app.include_router(catalogs.router)
+app.include_router(collection.router)
 app.include_router(engines.router)
 app.include_router(metrics.router)
 app.include_router(djsql.router)

{datajunction_server-0.0.1a55.dev0 → datajunction_server-0.0.1a56.dev0}/datajunction_server/api/materializations.py

@@ -12,8 +12,9 @@ from fastapi.responses import JSONResponse
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import joinedload, selectinload
 
-from datajunction_server.database import Column, ColumnAttribute, Node, NodeRevision
+from datajunction_server.database import Node, NodeRevision
 from datajunction_server.database.backfill import Backfill
+from datajunction_server.database.column import Column, ColumnAttribute
 from datajunction_server.database.history import ActivityType, EntityType, History
 from datajunction_server.database.user import User
 from datajunction_server.errors import DJDoesNotExistException, DJException

{datajunction_server-0.0.1a55.dev0 → datajunction_server-0.0.1a56.dev0}/datajunction_server/database/__init__.py

@@ -3,7 +3,7 @@ __all__ = [
     "AttributeType",
     "ColumnAttribute",
     "Catalog",
-    "Column",
+    "Collection",
     "Database",
     "DimensionLink",
     "Engine",
@@ -21,7 +21,7 @@ __all__ = [
 
 from datajunction_server.database.attributetype import AttributeType, ColumnAttribute
 from datajunction_server.database.catalog import Catalog
-from datajunction_server.database.column import Column
+from datajunction_server.database.collection import Collection
 from datajunction_server.database.database import Database, Table
 from datajunction_server.database.dimensionlink import DimensionLink
 from datajunction_server.database.engine import Engine

datajunction_server-0.0.1a56.dev0/datajunction_server/database/collection.py

@@ -0,0 +1,79 @@
+"""Collection database schema."""
+from datetime import datetime, timezone
+from functools import partial
+from typing import List, Optional
+
+from sqlalchemy import DateTime, ForeignKey, String, select
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.orm import Mapped, mapped_column, relationship
+
+from datajunction_server.database.base import Base
+from datajunction_server.database.node import Node
+from datajunction_server.errors import DJDoesNotExistException
+from datajunction_server.typing import UTCDatetime
+
+
+class Collection(Base):  # pylint: disable=too-few-public-methods
+    """
+    A collection of nodes
+    """
+
+    __tablename__ = "collection"
+
+    id: Mapped[int] = mapped_column(primary_key=True)
+    name: Mapped[Optional[str]] = mapped_column(String, default=None, unique=True)
+    description: Mapped[Optional[str]] = mapped_column(String, default=None)
+    nodes: Mapped[List[Node]] = relationship(
+        secondary="collectionnodes",
+        primaryjoin="Collection.id==CollectionNodes.collection_id",
+        secondaryjoin="Node.id==CollectionNodes.node_id",
+        lazy="selectin",
+    )
+    created_at: Mapped[UTCDatetime] = mapped_column(
+        DateTime(timezone=True),
+        insert_default=partial(datetime.now, timezone.utc),
+    )
+    deactivated_at: Mapped[UTCDatetime] = mapped_column(
+        DateTime(timezone=True),
+        nullable=True,
+        default=None,
+    )
+
+    def __hash__(self) -> int:
+        return hash(self.id)  # pragma: no cover
+
+    @classmethod
+    async def get_by_name(
+        cls,
+        session: AsyncSession,
+        name: str,
+        raise_if_not_exists: bool = False,
+    ) -> Optional["Collection"]:
+        """
+        Get a collection by name
+        """
+        statement = select(Collection).where(Collection.name == name)
+        collection = (await session.execute(statement)).scalar()
+        if not collection and raise_if_not_exists:
+            raise DJDoesNotExistException(
+                message=f"Collection with name `{name}` does not exist.",
+                http_status_code=404,
+            )
+        return collection
+
+
+class CollectionNodes(Base):  # type: ignore  # pylint: disable=too-few-public-methods
+    """
+    Join table for collections and nodes.
+    """
+
+    __tablename__ = "collectionnodes"
+
+    collection_id: Mapped[int] = mapped_column(
+        ForeignKey("collection.id", name="fk_collectionnodes_collection_id_collection"),
+        primary_key=True,
+    )
+    node_id: Mapped[int] = mapped_column(
+        ForeignKey("node.id", name="fk_collectionnodes_node_id_node"),
+        primary_key=True,
+    )

{datajunction_server-0.0.1a55.dev0 → datajunction_server-0.0.1a56.dev0}/datajunction_server/database/column.py

@@ -116,6 +116,15 @@ class Column(Base):  # type: ignore
             for attr in self.attributes  # pylint: disable=not-an-iterable
         )
 
+    def has_attributes_besides(self, attribute_name: str) -> bool:
+        """
+        Whether the column has any attribute besides the one specified.
+        """
+        return any(
+            attr.attribute_type.name != attribute_name
+            for attr in self.attributes  # pylint: disable=not-an-iterable
+        )
+
     def __hash__(self) -> int:
         return hash(self.id)
 

{datajunction_server-0.0.1a55.dev0 → datajunction_server-0.0.1a56.dev0}/datajunction_server/database/database.py

@@ -1,7 +1,7 @@
 """Database schema"""
 from datetime import datetime, timezone
 from functools import partial
-from typing import Dict, List, Optional, Tuple
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
 from uuid import UUID, uuid4
 
 from sqlalchemy import JSON, BigInteger, DateTime, ForeignKey, Integer, String
@@ -9,9 +9,11 @@ from sqlalchemy.orm import Mapped, mapped_column, relationship
 from sqlalchemy_utils import UUIDType
 
 from datajunction_server.database.base import Base
-from datajunction_server.database.column import Column
 from datajunction_server.typing import UTCDatetime
 
+if TYPE_CHECKING:
+    from datajunction_server.database.column import Column
+
 
 class Database(Base):  # pylint: disable=too-few-public-methods
     """

{datajunction_server-0.0.1a55.dev0 → datajunction_server-0.0.1a56.dev0}/datajunction_server/database/materialization.py

@@ -15,9 +15,9 @@ from sqlalchemy import (
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import Mapped, joinedload, mapped_column, relationship
 
-from datajunction_server.database import Column
 from datajunction_server.database.backfill import Backfill
 from datajunction_server.database.base import Base
+from datajunction_server.database.column import Column
 from datajunction_server.models.materialization import (
     DruidMeasuresCubeConfig,
     GenericMaterializationConfig,

{datajunction_server-0.0.1a55.dev0 → datajunction_server-0.0.1a56.dev0}/datajunction_server/internal/client.py

@@ -9,7 +9,8 @@ from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import joinedload, selectinload
 
 from datajunction_server.construction.utils import to_namespaced_name
-from datajunction_server.database import Column, DimensionLink, Node, NodeRevision
+from datajunction_server.database import DimensionLink, Node, NodeRevision
+from datajunction_server.database.column import Column
 from datajunction_server.models.node_type import NodeType
 from datajunction_server.sql.dag import topological_sort
 from datajunction_server.sql.parsing import ast
@@ -81,6 +82,43 @@ def python_client_code_for_linking_complex_dimension(
     )
 
 
+async def python_client_code_for_setting_column_attributes(
+    session: AsyncSession,
+    node_name: str,
+):
+    """
+    Returns the python client code to set column attributes.
+    """
+    node_short_name = node_name.split(SEPARATOR)[-1]
+    node = await Node.get_by_name(
+        session,
+        node_name,
+        options=[
+            joinedload(Node.current).options(
+                selectinload(NodeRevision.columns).options(
+                    selectinload(Column.attributes),
+                ),
+            ),
+        ],
+    )
+
+    template = jinja_env.get_template("set_column_attributes.j2")
+    snippets = [
+        template.render(
+            node_short_name=node_short_name,
+            column_name=col.name,
+            attributes=[
+                attr.attribute_type
+                for attr in col.attributes
+                if attr.attribute_type.name != "primary_key"
+            ],
+        )
+        for col in node.current.columns  # type: ignore
+        if col.has_attributes_besides("primary_key")
+    ]
+    return "\n\n".join(snippets)
+
+
 async def python_client_create_node(
     session: AsyncSession,
     node_name: str,
@@ -227,7 +265,7 @@ async def export_nodes_notebook_cells(session: AsyncSession, nodes: List[Node]):
     A node export means the following:
     - Client code to create the node and set the right tags
     - Client code to link all dimensions set on the node
-    - TODO: Client code to set all column attributes on the node
+    - Client code to set all column attributes on the node
     """
     cells = []
     cells.append(
@@ -248,7 +286,6 @@ async def export_nodes_notebook_cells(session: AsyncSession, nodes: List[Node]):
         namespace = SEPARATOR.join(node.name.split(SEPARATOR)[:-1])
         cells.append(
             new_code_cell(
-                # for cubes: NAMESPACE_MAPPING['{SEPARATOR.join(metric.split(SEPARATOR)[:-1])}']
                 await python_client_create_node(
                     session,
                     node.name,
@@ -275,4 +312,17 @@ async def export_nodes_notebook_cells(session: AsyncSession, nodes: List[Node]):
                 ],
             )
             cells.append(new_code_cell(link_dimensions))
+
+        # Add cell for setting column attributes if needed
+        if any(
+            col.has_attributes_besides("primary_key") for col in node.current.columns
+        ):
+            cells.append(
+                new_code_cell(
+                    await python_client_code_for_setting_column_attributes(
+                        session,
+                        node.name,
+                    ),
+                ),
+            )
     return cells

datajunction_server-0.0.1a56.dev0/datajunction_server/internal/templates/client_setup.j2

@@ -0,0 +1,9 @@
+from datajunction import (
+    DJBuilder, Source, Dimension, Transform, Metric,
+    Namespace, MetricUnit, MetricDirection, ColumnAttribute,
+)
+
+DJ_URL = "{{ request_url }}"
+
+dj = DJBuilder(DJ_URL)
+dj.basic_login("dj", "dj")

datajunction_server-0.0.1a56.dev0/datajunction_server/internal/templates/set_column_attributes.j2

@@ -0,0 +1,7 @@
+{{ node_short_name }}.set_column_attributes(
+    "{{ column_name }}",
+    [{%- for attribute in attributes %}
+        ColumnAttribute(namespace="{{ attribute.namespace }}", name="{{ attribute.name }}"),
+        {%- endfor %}{% if attributes %}
+    {% else %}{% endif %}],
+)

{datajunction_server-0.0.1a55.dev0 → datajunction_server-0.0.1a56.dev0}/datajunction_server/internal/validation.py

@@ -6,7 +6,8 @@ from sqlalchemy.exc import MissingGreenlet
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from datajunction_server.api.helpers import find_bound_dimensions
-from datajunction_server.database import Column, Node, NodeRevision
+from datajunction_server.database import Node, NodeRevision
+from datajunction_server.database.column import Column
 from datajunction_server.errors import DJError, DJException, ErrorCode
 from datajunction_server.models.base import labelize
 from datajunction_server.models.node import NodeRevisionBase, NodeStatus

datajunction_server-0.0.1a56.dev0/datajunction_server/models/collection.py

@@ -0,0 +1,17 @@
+"""
+Models for collections
+"""
+
+from pydantic.main import BaseModel
+
+
+class CollectionInfo(BaseModel):
+    """
+    Class for a collection information
+    """
+
+    name: str
+    description: str
+
+    class Config:  # pylint: disable=missing-class-docstring, too-few-public-methods
+        orm_mode = True

{datajunction_server-0.0.1a55.dev0 → datajunction_server-0.0.1a56.dev0}/datajunction_server/sql/functions.py

@@ -4438,7 +4438,7 @@ class Unnest(TableFunction):
     It will generate a new row for each element in the specified column.
     """
 
-    dialects = [Dialect.SPARK, Dialect.DRUID]
+    dialects = [Dialect.TRINO, Dialect.DRUID]
 
 
 @Unnest.register

{datajunction_server-0.0.1a55.dev0 → datajunction_server-0.0.1a56.dev0}/datajunction_server/sql/parsing/ast.py

@@ -2279,7 +2279,9 @@ class FunctionTable(FunctionTableExpression):
             if self.column_list
             else ""
         )
-        column_list_str = f"({cols})" if len(self.column_list) > 1 else str(cols)
+        column_list_str = (
+            f"{cols}" if self.name.name.upper() != "UNNEST" else f"({cols})"
+        )
         args_str = f"({', '.join(str(col) for col in self.args)})" if self.args else ""
         return f"{self.name}{args_str}{alias}{as_}{column_list_str}"
 

{datajunction_server-0.0.1a55.dev0 → datajunction_server-0.0.1a56.dev0}/tests/api/client_test.py

@@ -229,6 +229,15 @@ async def test_export_namespace_as_notebook(
     """
     Verify exporting all nodes in a namespace as a notebook.
     """
+    response = await module__client_with_roads.post(
+        "/nodes/default.repair_order_details/columns/repair_type_id/attributes/",
+        json=[
+            {
+                "namespace": "system",
+                "name": "dimension",
+            },
+        ],
+    )
     response = await module__client_with_roads.get(
         "/datajunction-clients/python/notebook?namespace=default",
     )
@@ -236,7 +245,7 @@ async def test_export_namespace_as_notebook(
         response.headers["content-disposition"] == 'attachment; filename="export.ipynb"'
     )
     notebook = response.json()
-    assert len(notebook["cells"]) == 51
+    assert len(notebook["cells"]) >= 50
     # Intro cell
     assert notebook["cells"][0]["cell_type"] == "markdown"
     assert (
@@ -282,6 +291,13 @@ NAMESPACE_MAPPING = {
             "notebook.link_dimension.txt",
         ).strip()
     )
+    # Check column attributes
+    assert (
+        notebook["cells"][7]["source"]
+        == load_expected_file(
+            "notebook.set_attribute.txt",
+        ).strip()
+    )
 
 
 @pytest.mark.asyncio
@@ -344,7 +360,7 @@ async def test_export_cube_as_notebook(
         "&include_sources=true&include_dimensions=true",
     )
     notebook = response.json()
-    assert len(notebook["cells"]) == 20
+    assert len(notebook["cells"]) == 21
     assert (
         notebook["cells"][2]["source"]
         == """### Upserting Nodes:
@@ -358,7 +374,7 @@ async def test_export_cube_as_notebook(
 * default.roads_cube"""
     )
     assert (
-        trim_trailing_whitespace(notebook["cells"][19]["source"])
+        trim_trailing_whitespace(notebook["cells"][20]["source"])
         == load_expected_file(
             "notebook.create_cube.txt",
         ).strip()