datajunction-server 0.0.1a45.dev4__tar.gz → 0.0.1a45.dev5__tar.gz

{datajunction_server-0.0.1a45.dev4 → datajunction_server-0.0.1a45.dev5}/.coveragerc

@@ -1,5 +1,6 @@
 # .coveragerc to control coverage.py
 [run]
+concurrency = thread,greenlet
 branch = True
 source = datajunction_server
 omit =

{datajunction_server-0.0.1a45.dev4 → datajunction_server-0.0.1a45.dev5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: datajunction-server
-Version: 0.0.1a45.dev4
+Version: 0.0.1a45.dev5
 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>
@@ -14,9 +14,11 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Requires-Python: <4.0,>=3.8
 Requires-Dist: accept-types<1.0.0,>=0.4.1
+Requires-Dist: aiosqlite>=0.20.0
 Requires-Dist: alembic>=1.10.3
 Requires-Dist: antlr4-python3-runtime==4.12.0
 Requires-Dist: asciidag<1.0.0,>=0.2.0
+Requires-Dist: asyncpg>=0.29.0
 Requires-Dist: bcrypt>=4.0.1
 Requires-Dist: cachelib<1.0.0,>=0.10.2
 Requires-Dist: cachetools>=5.3.1

datajunction_server-0.0.1a45.dev5/datajunction_server/__about__.py

@@ -0,0 +1,4 @@
+"""
+Version for Hatch
+"""
+__version__ = "0.0.1a45.dev5"

{datajunction_server-0.0.1a45.dev4 → datajunction_server-0.0.1a45.dev5}/datajunction_server/api/access/authentication/basic.py

@@ -8,7 +8,7 @@ from fastapi import APIRouter, Depends, Form
 from fastapi.responses import JSONResponse, Response
 from fastapi.security import OAuth2PasswordRequestForm
 from sqlalchemy import select
-from sqlalchemy.orm import Session
+from sqlalchemy.ext.asyncio import AsyncSession
 
 from datajunction_server.constants import AUTH_COOKIE, LOGGED_IN_FLAG_COOKIE
 from datajunction_server.database.user import OAuthProvider, User
@@ -24,20 +24,17 @@ router = APIRouter(tags=["Basic OAuth2"])
 
 
 @router.post("/basic/user/")
-def create_a_user(
+async def create_a_user(
     email: str = Form(),
     username: str = Form(),
     password: str = Form(),
-    session: Session = Depends(get_session),
+    session: AsyncSession = Depends(get_session),
 ) -> JSONResponse:
     """
     Create a new user
     """
-    if (
-        session.execute(select(User).where(User.username == username))
-        .scalars()
-        .one_or_none()
-    ):
+    user_result = await session.execute(select(User).where(User.username == username))
+    if user_result.scalar_one_or_none():
         raise DJException(
             http_status_code=HTTPStatus.CONFLICT,
             errors=[
@@ -54,8 +51,8 @@ def create_a_user(
         oauth_provider=OAuthProvider.BASIC,
     )
     session.add(new_user)
-    session.commit()
-    session.refresh(new_user)
+    await session.commit()
+    await session.refresh(new_user)
     return JSONResponse(
         content={"message": "User successfully created"},
         status_code=HTTPStatus.CREATED,
@@ -63,14 +60,14 @@ def create_a_user(
 
 
 @router.post("/basic/login/")
-def login(
+async def login(
     form_data: OAuth2PasswordRequestForm = Depends(),
-    session: Session = Depends(get_session),
+    session: AsyncSession = Depends(get_session),
 ):
     """
     Get a JWT token and set it as an HTTP only cookie
     """
-    user = validate_user_password(
+    user = await validate_user_password(
         username=form_data.username,
         password=form_data.password,
         session=session,

{datajunction_server-0.0.1a45.dev4 → datajunction_server-0.0.1a45.dev5}/datajunction_server/api/access/authentication/google.py

@@ -14,7 +14,7 @@ import requests
 from fastapi import APIRouter, Depends, Request
 from google.oauth2 import id_token
 from sqlalchemy import select
-from sqlalchemy.orm import Session
+from sqlalchemy.ext.asyncio import AsyncSession
 from starlette.responses import RedirectResponse
 
 from datajunction_server.constants import AUTH_COOKIE, LOGGED_IN_FLAG_COOKIE
@@ -46,11 +46,11 @@ def login(target: Optional[str] = None):
 
 
 @router.get("/google/token/")
-def get_access_token(
+async def get_access_token(
     request: Request,
     state: Optional[str] = None,
     error: Optional[str] = None,
-    session: Session = Depends(get_session),
+    session: AsyncSession = Depends(get_session),
     setting: Settings = Depends(get_settings),
 ):
     """
@@ -75,8 +75,10 @@ def get_access_token(
         audience=setting.google_oauth_client_id,
     )
 
-    existing_user = session.execute(
-        select(User).where(User.email == user_data["email"]),
+    existing_user = (
+        await session.execute(
+            select(User).where(User.email == user_data["email"]),
+        )
     ).scalar()
     if existing_user:
         _logger.info("OAuth user found")
@@ -91,8 +93,8 @@ def get_access_token(
             oauth_provider=OAuthProvider.GOOGLE,
         )
         session.add(new_user)
-        session.commit()
-        session.refresh(new_user)
+        await session.commit()
+        await session.refresh(new_user)
         user = new_user
     response = RedirectResponse(url=urljoin(settings.frontend_host, state))  # type: ignore
     response.set_cookie(

{datajunction_server-0.0.1a45.dev4 → datajunction_server-0.0.1a45.dev5}/datajunction_server/api/access/authentication/whoami.py

@@ -18,7 +18,7 @@ router = SecureAPIRouter(tags=["Who am I?"])
 
 
 @router.get("/whoami/", response_model=UserOutput)
-def get_user(current_user: User = Depends(get_current_user)) -> UserOutput:
+async def get_user(current_user: User = Depends(get_current_user)) -> UserOutput:
     """
     Returns the current authenticated user
     """
@@ -26,7 +26,7 @@ def get_user(current_user: User = Depends(get_current_user)) -> UserOutput:
 
 
 @router.get("/token/")
-def get_short_lived_token(request: Request) -> JSONResponse:
+async def get_short_lived_token(request: Request) -> JSONResponse:
     """
     Returns a token that expires in 24 hours
     """

{datajunction_server-0.0.1a45.dev4 → datajunction_server-0.0.1a45.dev5}/datajunction_server/api/attributes.py

@@ -7,7 +7,7 @@ from typing import List
 
 from fastapi import Depends
 from sqlalchemy import select
-from sqlalchemy.orm import Session
+from sqlalchemy.ext.asyncio import AsyncSession
 
 from datajunction_server.database.attributetype import AttributeType
 from datajunction_server.errors import DJAlreadyExistsException, DJException
@@ -26,13 +26,13 @@ router = SecureAPIRouter(tags=["attributes"])
 
 
 @router.get("/attributes/", response_model=List[AttributeTypeBase])
-def list_attributes(
-    *, session: Session = Depends(get_session)
+async def list_attributes(
+    *, session: AsyncSession = Depends(get_session)
 ) -> List[AttributeTypeBase]:
     """
     List all available attribute types.
     """
-    attributes = session.execute(select(AttributeType)).scalars().all()
+    attributes = await AttributeType.get_all(session)
     return [AttributeTypeBase.from_orm(attr) for attr in attributes]
 
 
@@ -42,8 +42,8 @@ def list_attributes(
     status_code=201,
     name="Add an Attribute Type",
 )
-def add_attribute_type(
-    data: MutableAttributeTypeFields, *, session: Session = Depends(get_session)
+async def add_attribute_type(
+    data: MutableAttributeTypeFields, *, session: AsyncSession = Depends(get_session)
 ) -> AttributeTypeBase:
     """
     Add a new attribute type
@@ -52,26 +52,16 @@ def add_attribute_type(
         raise DJException(
             message="Cannot use `system` as the attribute type namespace as it is reserved.",
         )
-    statement = select(AttributeType).where(AttributeType.name == data.name)
-    attribute_type = session.execute(statement).unique().one_or_none()
+    attribute_type = await AttributeType.get_by_name(session, data.name)
     if attribute_type:
         raise DJAlreadyExistsException(
             message=f"Attribute type `{data.name}` already exists!",
         )
-    attribute_type = AttributeType(
-        namespace=data.namespace,
-        name=data.name,
-        description=data.description,
-        allowed_node_types=data.allowed_node_types,
-        uniqueness_scope=data.uniqueness_scope if data.uniqueness_scope else [],
-    )
-    session.add(attribute_type)
-    session.commit()
-    session.refresh(attribute_type)
-    return attribute_type
+    attribute_type = await AttributeType.create(session, data)
+    return AttributeTypeBase.from_orm(attribute_type)
 
 
-def default_attribute_types(session: Session = Depends(get_session)):
+async def default_attribute_types(session: AsyncSession = Depends(get_session)):
     """
     Loads all the column attribute types that are supported by the system
     by default into the database.
@@ -105,7 +95,7 @@ def default_attribute_types(session: Session = Depends(get_session)):
             set(default_attribute_type_names.keys()),
         ),
     )
-    attribute_types = session.execute(statement).scalars().all()
+    attribute_types = (await session.execute(statement)).scalars().all()
     for type_ in attribute_types:
         # if type_:  # pragma: no cover
         type_.name = default_attribute_type_names[type_.name].name
@@ -123,7 +113,7 @@ def default_attribute_types(session: Session = Depends(get_session)):
     new_types = set(default_attribute_type_names.keys()) - {
         type_.name for type_ in attribute_types if type_
     }
-    session.bulk_save_objects(
+    session.add_all(
         [default_attribute_type_names[name] for name in new_types],
     )
-    session.commit()
+    await session.commit()

{datajunction_server-0.0.1a45.dev4 → datajunction_server-0.0.1a45.dev5}/datajunction_server/api/catalogs.py

@@ -8,7 +8,8 @@ from typing import List
 
 from fastapi import Depends, HTTPException
 from sqlalchemy import select
-from sqlalchemy.orm import Session
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.orm import joinedload
 
 from datajunction_server.api.engines import EngineInfo
 from datajunction_server.api.helpers import get_catalog_by_name
@@ -28,22 +29,27 @@ UNKNOWN_CATALOG_ID = 0
 
 
 @router.get("/catalogs/", response_model=List[CatalogInfo])
-def list_catalogs(*, session: Session = Depends(get_session)) -> List[CatalogInfo]:
+async def list_catalogs(
+    *, session: AsyncSession = Depends(get_session)
+) -> List[CatalogInfo]:
     """
     List all available catalogs
     """
+    statement = select(Catalog).options(joinedload(Catalog.engines))
     return [
         CatalogInfo.from_orm(catalog)
-        for catalog in session.execute(select(Catalog)).scalars()
+        for catalog in (await session.execute(statement)).unique().scalars()
     ]
 
 
 @router.get("/catalogs/{name}/", response_model=CatalogInfo, name="Get a Catalog")
-def get_catalog(name: str, *, session: Session = Depends(get_session)) -> CatalogInfo:
+async def get_catalog(
+    name: str, *, session: AsyncSession = Depends(get_session)
+) -> CatalogInfo:
     """
     Return a catalog by name
     """
-    return get_catalog_by_name(session, name)
+    return await get_catalog_by_name(session, name)
 
 
 @router.post(
@@ -52,16 +58,16 @@ def get_catalog(name: str, *, session: Session = Depends(get_session)) -> Catalo
     status_code=201,
     name="Add A Catalog",
 )
-def add_catalog(
+async def add_catalog(
     data: CatalogInfo,
     *,
-    session: Session = Depends(get_session),
+    session: AsyncSession = Depends(get_session),
 ) -> CatalogInfo:
     """
     Add a Catalog
     """
     try:
-        get_catalog_by_name(session, data.name)
+        await get_catalog_by_name(session, data.name)
     except DJException:
         pass
     else:
@@ -83,15 +89,15 @@ def add_catalog(
         ],
     )
     catalog.engines.extend(
-        list_new_engines(
+        await list_new_engines(
             session=session,
             catalog=catalog,
             create_engines=data.engines,  # type: ignore
         ),
     )
     session.add(catalog)
-    session.commit()
-    session.refresh(catalog)
+    await session.commit()
+    await session.refresh(catalog, ["engines"])
 
     return CatalogInfo.from_orm(catalog)
 
@@ -102,27 +108,27 @@ def add_catalog(
     status_code=201,
     name="Add Engines to a Catalog",
 )
-def add_engines_to_catalog(
+async def add_engines_to_catalog(
     name: str,
     data: List[EngineInfo],
     *,
-    session: Session = Depends(get_session),
+    session: AsyncSession = Depends(get_session),
 ) -> CatalogInfo:
     """
     Attach one or more engines to a catalog
     """
-    catalog = get_catalog_by_name(session, name)
+    catalog = await get_catalog_by_name(session, name)
     catalog.engines.extend(
-        list_new_engines(session=session, catalog=catalog, create_engines=data),
+        await list_new_engines(session=session, catalog=catalog, create_engines=data),
     )
     session.add(catalog)
-    session.commit()
-    session.refresh(catalog)
+    await session.commit()
+    await session.refresh(catalog)
     return CatalogInfo.from_orm(catalog)
 
 
-def list_new_engines(
-    session: Session,
+async def list_new_engines(
+    session: AsyncSession,
     catalog: Catalog,
     create_engines: List[EngineInfo],
 ) -> List[Engine]:
@@ -132,7 +138,7 @@ def list_new_engines(
     new_engines = []
     for engine_ref in create_engines:
         already_set = False
-        engine = get_engine(session, engine_ref.name, engine_ref.version)
+        engine = await get_engine(session, engine_ref.name, engine_ref.version)
         for set_engine in catalog.engines:
             if engine.name == set_engine.name and engine.version == set_engine.version:
                 already_set = True
@@ -141,17 +147,17 @@ def list_new_engines(
     return new_engines
 
 
-def default_catalog(session: Session = Depends(get_session)):
+async def default_catalog(session: AsyncSession = Depends(get_session)):
     """
     Loads a default catalog for nodes that are pure SQL and don't belong in any
     particular catalog. This typically applies to on-the-fly user-defined dimensions.
     """
     statement = select(Catalog).filter(Catalog.id == UNKNOWN_CATALOG_ID)
-    catalogs = session.execute(statement).all()
+    catalogs = (await session.execute(statement)).all()
     if not catalogs:
         unknown = Catalog(
             id=UNKNOWN_CATALOG_ID,
             name="unknown",
         )
         session.add(unknown)
-        session.commit()
+        await session.commit()

{datajunction_server-0.0.1a45.dev4 → datajunction_server-0.0.1a45.dev5}/datajunction_server/api/client.py

@@ -6,12 +6,13 @@ import json
 import logging
 
 from fastapi import Depends
-from sqlalchemy.orm import Session
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.orm import joinedload
 
 from datajunction_server.api.helpers import get_node_by_name
+from datajunction_server.database import Node, NodeRevision
 from datajunction_server.internal.access.authentication.http import SecureAPIRouter
 from datajunction_server.models.materialization import MaterializationJobTypeEnum
-from datajunction_server.models.node import NodeOutput
 from datajunction_server.models.node_type import NodeType
 from datajunction_server.utils import get_session, get_settings
 
@@ -21,62 +22,58 @@ router = SecureAPIRouter(tags=["client"])
 
 
 @router.get("/datajunction-clients/python/new_node/{node_name}", response_model=str)
-def client_code_for_creating_node(
-    node_name: str, *, session: Session = Depends(get_session)
+async def client_code_for_creating_node(
+    node_name: str, *, session: AsyncSession = Depends(get_session)
 ) -> str:
     """
     Generate the Python client code used for creating this node
     """
     node_short_name = node_name.split(".")[-1]
-    node = get_node_by_name(session, node_name)
-
-    # Generic user-configurable node creation params
-    params = NodeOutput.from_orm(node).dict(
-        exclude={
-            "id",
-            "version",
-            "type",
-            "catalog_id",
-            "lineage",
-            "status",
-            "metric_metadata_id",
-            "mode",
-            "node_id",
-            "updated_at",
-            "materializations",
-            "columns",
-            "catalog",
-            "parents",
-            "metric_metadata",
-            "query" if node.type == NodeType.CUBE else "",
-            "dimension_links",
-            "created_at",
-            "current_version",
-            "missing_table",
-            "namespace",
-            "tags",
-        },
-        exclude_none=True,
+    node = await Node.get_by_name(
+        session,
+        node_name,
+        options=[
+            joinedload(Node.current).options(
+                *NodeRevision.default_load_options(),
+                joinedload(NodeRevision.cube_elements),
+            ),
+        ],
+        raise_if_not_exists=True,
     )
+    # Generic user-configurable node creation params
+    params = {
+        "name": node.name,  # type: ignore
+        "display_name": node.current.display_name,  # type: ignore
+        "description": node.current.description,  # type: ignore
+        "mode": node.current.mode,  # type: ignore
+        "query": node.current.query,  # type: ignore
+        "schema_": node.current.schema_,  # type: ignore
+        "table": node.current.table,  # type: ignore
+        "primary_key": [col.name for col in node.current.primary_key()],  # type: ignore
+    }
 
-    params["primary_key"] = [col.name for col in node.current.primary_key()]
-
-    for key in params:
-        if not isinstance(params[key], list) and key != "query" and key != "lineage":
+    for key in params:  # pylint: disable=consider-using-dict-items
+        if (
+            not isinstance(params[key], list)
+            and key != "query"
+            and key != "lineage"
+            and params[key]
+        ):
             params[key] = f'"{params[key]}"'
-        if key == "query":
+        if key == "query" and params[key]:
             params[key] = f'"""{params[key]}"""'
 
     # Cube-specific params
     cube_params = []
-    if node.type == NodeType.CUBE:
+    if node.type == NodeType.CUBE:  # type: ignore
         ordering = {
-            col.name: col.order or idx for idx, col in enumerate(node.current.columns)
+            col.name: col.order or idx
+            for idx, col in enumerate(node.current.columns)  # type: ignore
         }
         metrics_list = sorted(
             [
                 elem.node_revisions[-1].name
-                for elem in node.current.cube_elements
+                for elem in node.current.cube_elements  # type: ignore
                 if elem.node_revisions[-1].type == NodeType.METRIC
             ],
             key=lambda x: ordering[x],
@@ -84,7 +81,7 @@ def client_code_for_creating_node(
         dimensions_list = sorted(
             [
                 elem.node_revisions[-1].name + "." + elem.name
-                for elem in node.current.cube_elements
+                for elem in node.current.cube_elements  # type: ignore
                 if elem.node_revisions[-1].type == NodeType.DIMENSION
             ],
             key=lambda x: ordering[x],
@@ -97,12 +94,14 @@ def client_code_for_creating_node(
         ]
 
     formatted_params = ",\n".join(
-        [f"    {k}={params[k]}" for k in sorted(params.keys())] + cube_params,
+        [f"    {k}={params[k]}" for k in sorted(params.keys()) if params[k]]
+        + cube_params,
     )
 
+    node_type = node.type  # type: ignore
     client_code = f"""dj = DJBuilder(DJ_URL)
 
-{node_short_name} = dj.create_{node.type}(
+{node_short_name} = dj.create_{node_type}(
 {formatted_params}
 )"""
     return client_code  # type: ignore
@@ -112,17 +111,17 @@ def client_code_for_creating_node(
     "/datajunction-clients/python/add_materialization/{node_name}/{materialization_name}",
     response_model=str,
 )
-def client_code_for_adding_materialization(
+async def client_code_for_adding_materialization(
     node_name: str,
     materialization_name: str,
     *,
-    session: Session = Depends(get_session),
+    session: AsyncSession = Depends(get_session),
 ) -> str:
     """
     Generate the Python client code used for adding this materialization
     """
     node_short_name = node_name.split(".")[-1]
-    node = get_node_by_name(session, node_name)
+    node = await get_node_by_name(session, node_name)
     materialization = [
         materialization
         for materialization in node.current.materializations
@@ -160,18 +159,18 @@ materialization = MaterializationConfig(
     "/datajunction-clients/python/link_dimension/{node_name}/{column}/{dimension}/",
     response_model=str,
 )
-def client_code_for_linking_dimension_to_node(
+async def client_code_for_linking_dimension_to_node(
     node_name: str,
     column: str,
     dimension: str,
     *,
-    session: Session = Depends(get_session),
+    session: AsyncSession = Depends(get_session),
 ) -> str:
     """
     Generate the Python client code used for linking this node's column to a dimension
     """
     node_short_name = node_name.split(".")[-1]
-    node = get_node_by_name(session, node_name)
+    node = await get_node_by_name(session, node_name)
     client_code = f"""dj = DJBuilder(DJ_URL)
 {node_short_name} = dj.{node.type}(
     "{node.name}"