datajunction-server 0.0.1a99.dev2__tar.gz → 0.0.1a100__tar.gz

{datajunction_server-0.0.1a99.dev2 → datajunction_server-0.0.1a100}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: datajunction-server
-Version: 0.0.1a99.dev2
+Version: 0.0.1a100
 Summary: DataJunction server library for running to a DataJunction server
 Project-URL: Homepage, https://datajunction.io
 Project-URL: Repository, https://github.com/DataJunction/dj

datajunction_server-0.0.1a100/datajunction_server/__about__.py

@@ -0,0 +1,5 @@
+"""
+Version for Hatch
+"""
+
+__version__ = "0.0.1a100"

datajunction_server-0.0.1a100/datajunction_server/alembic/versions/2025_05_08_0601-ee0eb65f743b_convert_dialect.py

@@ -0,0 +1,38 @@
+"""
+Convert dialect
+
+Revision ID: ee0eb65f743b
+Revises: 51547dcccb10
+Create Date: 2025-05-08 06:01:38.039079+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
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = "ee0eb65f743b"
+down_revision = "51547dcccb10"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    with op.batch_alter_table("engine", schema=None) as batch_op:
+        batch_op.alter_column(
+            "dialect",
+            existing_type=postgresql.ENUM("SPARK", "TRINO", "DRUID", name="dialect"),
+            type_=sa.String(),
+            existing_nullable=True,
+        )
+
+
+def downgrade():
+    with op.batch_alter_table("engine", schema=None) as batch_op:
+        batch_op.alter_column(
+            "dialect",
+            existing_type=sa.String(),
+            type_=postgresql.ENUM("SPARK", "TRINO", "DRUID", name="dialect"),
+            existing_nullable=True,
+        )

{datajunction_server-0.0.1a99.dev2 → datajunction_server-0.0.1a100}/datajunction_server/api/graphql/queries/sql.py

@@ -84,7 +84,7 @@ async def measures_sql(
     """
     Get measures SQL for a set of metrics with dimensions and filters
     """
-    session, settings = info.context["session"], info.context["settings"]
+    session = info.context["session"]
     queries = await get_measures_query(
         session=session,
         metrics=list(OrderedDict.fromkeys(cube.metrics)),  # type: ignore
@@ -94,7 +94,6 @@ async def measures_sql(
         engine_name=engine.name if engine else None,
         engine_version=engine.version if engine else None,
         include_all_columns=include_all_columns,
-        sql_transpilation_library=settings.sql_transpilation_library,
         use_materialized=use_materialized,
         preagg_requested=preaggregate,
     )

{datajunction_server-0.0.1a99.dev2 → datajunction_server-0.0.1a100}/datajunction_server/api/graphql/scalars/sql.py

@@ -72,7 +72,6 @@ class GeneratedSQL:
 
     node: Node
     sql: str
-    sql_transpilation_library: str | None = None
     columns: list[ColumnMetadata]
     dialect: Dialect  # type: ignore
     upstream_tables: list[str]
@@ -91,7 +90,6 @@ class GeneratedSQL:
                 name=obj.node.name,
             ),
             sql=obj.sql,
-            sql_transpilation_library=obj.sql_transpilation_library,
             columns=[
                 ColumnMetadata(  # type: ignore
                     name=col.name,

{datajunction_server-0.0.1a99.dev2 → datajunction_server-0.0.1a100}/datajunction_server/api/helpers.py

@@ -56,6 +56,11 @@ from datajunction_server.models.history import status_change_history
 from datajunction_server.models.metric import TranslatedSQL
 from datajunction_server.models.node import NodeStatus
 from datajunction_server.models.node_type import NodeType
+from datajunction_server.models.materialization import (
+    MaterializationConfigInfoUnified,
+    MaterializationStrategy,
+    MaterializationConfigOutput,
+)
 from datajunction_server.models.query import ColumnMetadata, QueryWithResults
 from datajunction_server.naming import LOOKUP_CHARS
 from datajunction_server.service_clients import QueryServiceClient
@@ -717,7 +722,7 @@ async def build_sql_for_multiple_metrics(
         ]
         engine = materialized_cube_catalog.engines[0]
         return (
-            TranslatedSQL(
+            TranslatedSQL.create(
                 sql=str(query_ast),
                 columns=query_metric_columns + query_dimension_columns,
                 dialect=materialized_cube_catalog.engines[0].dialect,
@@ -745,7 +750,7 @@ async def build_sql_for_multiple_metrics(
     for tbl in upstream_tables:
         await refresh_if_needed(session, tbl.dj_node, ["availability"])
     return (
-        TranslatedSQL(
+        TranslatedSQL.create(
             sql=str(query_ast),
             columns=columns,
             dialect=engine.dialect if engine else None,
@@ -875,7 +880,7 @@ async def build_sql_for_dj_query(  # pragma: no cover
     ]
 
     return (  # pragma: no cover
-        TranslatedSQL(
+        TranslatedSQL.create(
             sql=str(query_ast),
             columns=columns,
             dialect=engine.dialect if engine else None,
@@ -940,3 +945,70 @@ async def get_save_history(notify: Callable = Depends(get_notifier)) -> Callable
         await save_history(event, session, notify)
 
     return save_history_with_notify
+
+
+def get_materialization_info(
+    query_service_client: QueryServiceClient,
+    node: Node,
+    include_all_revisions: bool,
+    show_inactive: bool,
+    request_headers: Optional[Dict[str, str]] = None,
+):
+    """Get materializations for a node
+
+    If include_all_revisions=true this pulls materializations for each node revision and
+    combines them into a single list. Otherwise it just pulls materializations for the
+    current node revision.
+
+    If show_inactive=true this also returns materializations with a deactivated_at timestamp
+    """
+    return (
+        [
+            materialization
+            for node_revision in node.revisions  # type: ignore
+            for materialization in get_node_revision_materialization(
+                query_service_client=query_service_client,
+                node_revision=node_revision,
+                show_inactive=show_inactive,
+                request_headers=request_headers,
+            )
+        ]
+        if include_all_revisions
+        else get_node_revision_materialization(
+            query_service_client=query_service_client,
+            node_revision=node.current,  # type: ignore
+            show_inactive=show_inactive,
+            request_headers=request_headers,
+        )
+    )
+
+
+def get_node_revision_materialization(
+    query_service_client: QueryServiceClient,
+    node_revision: NodeRevision,
+    show_inactive: bool,
+    request_headers: Optional[Dict[str, str]] = None,
+) -> list[MaterializationConfigInfoUnified]:
+    """Merge in materialization info from the query service for a node revision"""
+    materializations = []
+    for materialization in node_revision.materializations:
+        if not materialization.deactivated_at or show_inactive:
+            info = query_service_client.get_materialization_info(
+                node_revision.name,
+                node_revision.version,
+                node_revision.type,
+                materialization.name,
+                request_headers=request_headers,
+            )
+            if materialization.strategy != MaterializationStrategy.INCREMENTAL_TIME:
+                info.urls = [info.urls[0]]
+            materialization_config_output = MaterializationConfigOutput.from_orm(
+                materialization,
+            )
+            materializations.append(
+                MaterializationConfigInfoUnified(
+                    **materialization_config_output.dict(),
+                    **info.dict(),
+                ),
+            )
+    return materializations

{datajunction_server-0.0.1a99.dev2 → datajunction_server-0.0.1a100}/datajunction_server/api/main.py

@@ -2,14 +2,10 @@
 Main DJ server app.
 """
 
-# All the models need to be imported here so that SQLModel can define their
-# relationships at runtime without causing circular imports.
-# See https://sqlmodel.tiangolo.com/tutorial/code-structure/#make-circular-imports-work.
-
 import logging
+from datajunction_server.api import setup_logging  # noqa
+
 from http import HTTPStatus
-from logging import config
-from os import path
 from typing import TYPE_CHECKING
 
 from fastapi import Depends, FastAPI, Request
@@ -41,6 +37,7 @@ from datajunction_server.api import (
     tags,
     users,
 )
+
 from datajunction_server.api.access.authentication import basic, whoami
 from datajunction_server.api.attributes import default_attribute_types
 from datajunction_server.api.catalogs import default_catalog
@@ -55,11 +52,6 @@ if TYPE_CHECKING:  # pragma: no cover
 _logger = logging.getLogger(__name__)
 settings = get_settings()
 
-config.fileConfig(
-    path.join(path.dirname(path.abspath(__file__)), "logging.conf"),
-    disable_existing_loggers=False,
-)
-
 dependencies = [Depends(default_attribute_types), Depends(default_catalog)]
 
 app = FastAPI(

{datajunction_server-0.0.1a99.dev2 → datajunction_server-0.0.1a100}/datajunction_server/api/materializations.py

@@ -12,7 +12,7 @@ from fastapi.responses import JSONResponse
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import joinedload, selectinload
 
-from datajunction_server.api.helpers import get_save_history
+from datajunction_server.api.helpers import get_save_history, get_materialization_info
 from datajunction_server.database import Node, NodeRevision
 from datajunction_server.database.backfill import Backfill
 from datajunction_server.database.column import Column, ColumnAttribute
@@ -32,7 +32,6 @@ from datajunction_server.models.base import labelize
 from datajunction_server.models.cube_materialization import UpsertCubeMaterialization
 from datajunction_server.models.materialization import (
     MaterializationConfigInfoUnified,
-    MaterializationConfigOutput,
     MaterializationInfo,
     MaterializationJobTypeEnum,
     MaterializationStrategy,
@@ -277,7 +276,8 @@ async def upsert_materialization(
 )
 async def list_node_materializations(
     node_name: str,
-    show_deleted: bool = False,
+    show_inactive: bool = False,
+    include_all_revisions: bool = False,
     *,
     session: AsyncSession = Depends(get_session),
     request: Request,
@@ -286,30 +286,32 @@ async def list_node_materializations(
     """
     Show all materializations configured for the node, with any associated metadata
     like urls from the materialization service, if available.
+
+    show_inactive: bool - Show materializations that have a deactivated_at timestamp set
+    include_all_revisions: bool - Show  materializations for all revisions of the node
     """
     request_headers = dict(request.headers)
-    node = await Node.get_by_name(session, node_name)
-    materializations = []
-    for materialization in node.current.materializations:  # type: ignore
-        if not materialization.deactivated_at or show_deleted:  # pragma: no cover
-            info = query_service_client.get_materialization_info(
-                node_name,
-                node.current.version,  # type: ignore
-                node.type,  # type: ignore
-                materialization.name,  # type: ignore
-                request_headers=request_headers,
-            )
-            if materialization.strategy != MaterializationStrategy.INCREMENTAL_TIME:
-                info.urls = [info.urls[0]]
-            materialization_config_output = MaterializationConfigOutput.from_orm(
-                materialization,
-            )
-            materialization = MaterializationConfigInfoUnified(
-                **materialization_config_output.dict(),
-                **info.dict(),
-            )
-            materializations.append(materialization)
-    return materializations
+
+    # If materializations from all revisions are requested,
+    # this includes the joined load to pull all node revisions
+    node = await Node.get_by_name(
+        session,
+        node_name,
+        options=[
+            joinedload(Node.revisions).options(*NodeRevision.default_load_options()),
+        ]
+        if include_all_revisions
+        else [],
+        raise_if_not_exists=True,
+    )
+
+    return get_materialization_info(
+        query_service_client=query_service_client,
+        node=node,  # type: ignore
+        include_all_revisions=include_all_revisions,
+        show_inactive=show_inactive,
+        request_headers=request_headers,
+    )
 
 
 @router.delete(

datajunction_server-0.0.1a100/datajunction_server/api/setup_logging.py

@@ -0,0 +1,7 @@
+from logging import config
+from os import path
+
+config.fileConfig(
+    path.join(path.dirname(path.abspath(__file__)), "logging.conf"),
+    disable_existing_loggers=False,
+)

{datajunction_server-0.0.1a99.dev2 → datajunction_server-0.0.1a100}/datajunction_server/api/sql.py

@@ -94,7 +94,7 @@ async def get_measures_sql_for_cube_v2(
     )
 
     metrics = list(OrderedDict.fromkeys(metrics))
-    measures_query = await get_measures_query(
+    return await get_measures_query(
         session=session,
         metrics=metrics,
         dimensions=dimensions,
@@ -105,12 +105,10 @@ async def get_measures_sql_for_cube_v2(
         current_user=current_user,
         validate_access=validate_access,
         include_all_columns=include_all_columns,
-        sql_transpilation_library=settings.sql_transpilation_library,
         use_materialized=use_materialized,
         preagg_requested=preaggregate,
         query_parameters=json.loads(query_params),
     )
-    return measures_query
 
 
 async def build_and_save_node_sql(
@@ -297,7 +295,7 @@ async def get_node_sql(
             query_parameters=query_parameters,
         )
         return (
-            TranslatedSQL(
+            TranslatedSQL.create(
                 sql=query_request.query,
                 columns=query_request.columns,
                 dialect=engine.dialect if engine else None,
@@ -319,7 +317,7 @@ async def get_node_sql(
         query_parameters=query_parameters,
     )
     return (
-        TranslatedSQL(
+        TranslatedSQL.create(
             sql=query_request.query,
             columns=query_request.columns,
             dialect=engine.dialect if engine else None,
@@ -452,7 +450,7 @@ async def get_sql_for_metrics(
             if engine_name
             else None
         )
-        return TranslatedSQL(  # pragma: no cover
+        return TranslatedSQL.create(  # pragma: no cover
             sql=query_request.query,
             columns=query_request.columns,
             dialect=engine.dialect if engine else None,

{datajunction_server-0.0.1a99.dev2 → datajunction_server-0.0.1a100}/datajunction_server/config.py

@@ -66,8 +66,8 @@ class Settings(BaseSettings):  # pragma: no cover
     # DJ UI host, used for OAuth redirection
     frontend_host: Optional[str] = "http://localhost:3000"
 
-    # Library to use when transpiling SQL to other dialects
-    sql_transpilation_library: Optional[str] = None
+    # Enabled transpilation plugin names
+    transpilation_plugins: List[str] = ["default", "sqlglot"]
 
     # 128 bit DJ secret, used to encrypt passwords and JSON web tokens
     secret: str = "a-fake-secretkey"

{datajunction_server-0.0.1a99.dev2 → datajunction_server-0.0.1a100}/datajunction_server/construction/build_v2.py

@@ -101,7 +101,6 @@ async def get_measures_query(
     current_user: Optional[User] = None,
     validate_access: access.ValidateAccessFn = None,
     include_all_columns: bool = False,
-    sql_transpilation_library: Optional[str] = None,
     use_materialized: bool = True,
     preagg_requested: bool = False,
     query_parameters: dict[str, Any] = None,
@@ -126,9 +125,7 @@ async def get_measures_query(
     )
 
     engine = (
-        await get_engine(session, engine_name, engine_version)
-        if engine_name and engine_version
-        else None
+        await get_engine(session, engine_name, engine_version) if engine_name else None  # type: ignore
     )
     build_criteria = BuildCriteria(
         dialect=engine.dialect if engine and engine.dialect else Dialect.SPARK,
@@ -245,9 +242,8 @@ async def get_measures_query(
             for col in final_query.select.projection
         ]
         measures_queries.append(
-            GeneratedSQL(
+            GeneratedSQL.create(
                 node=parent_node.current,
-                sql_transpilation_library=sql_transpilation_library,
                 sql=str(final_query),
                 columns=columns_metadata,
                 dialect=build_criteria.dialect,

{datajunction_server-0.0.1a99.dev2 → datajunction_server-0.0.1a100}/datajunction_server/construction/dimensions.py

@@ -114,7 +114,7 @@ async def build_dimensions_from_cube_query(
         amenable_name(elem.node_revision().name + SEPARATOR + elem.name): elem.type  # type: ignore
         for elem in cube.cube_elements
     }
-    return TranslatedSQL(
+    return TranslatedSQL.create(
         sql=str(query_ast),
         columns=[
             ColumnMetadata(

{datajunction_server-0.0.1a99.dev2 → datajunction_server-0.0.1a100}/datajunction_server/database/engine.py

@@ -3,11 +3,22 @@
 from typing import Optional
 
 import sqlalchemy as sa
-from sqlalchemy import Enum
 from sqlalchemy.orm import Mapped, mapped_column
 
 from datajunction_server.database.base import Base
-from datajunction_server.models.engine import Dialect
+from datajunction_server.models.dialect import Dialect
+
+
+class DialectType(sa.TypeDecorator):
+    impl = sa.String
+
+    def process_bind_param(self, value, dialect):
+        if isinstance(value, Dialect):
+            return value.name
+        return value
+
+    def process_result_value(self, value, dialect):
+        return Dialect(value) if value else None
 
 
 class Engine(Base):
@@ -24,6 +35,4 @@ class Engine(Base):
     name: Mapped[str]
     version: Mapped[str]
     uri: Mapped[Optional[str]]
-    dialect: Mapped[Optional[Dialect]] = mapped_column(
-        Enum(Dialect),
-    )
+    dialect: Mapped[Optional[Dialect]] = mapped_column(DialectType())

datajunction_server-0.0.1a100/datajunction_server/models/dialect.py

@@ -0,0 +1,114 @@
+import logging
+from datajunction_server.enum import StrEnum
+from datajunction_server.utils import get_settings
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from datajunction_server.transpilation import SQLTranspilationPlugin
+
+
+logger = logging.getLogger(__name__)
+settings = get_settings()
+
+
+class Dialect(StrEnum):
+    """
+    SQL dialect
+    """
+
+    SPARK = "spark"
+    TRINO = "trino"
+    DRUID = "druid"
+
+    @classmethod
+    def _missing_(cls, value: object) -> "Dialect":
+        """
+        Create a Dialect enum member from a string, but only if it corresponds to a
+        registered transpilation plugin in the plugin registry. If it has a corresponding
+        plugin, a dynamic enum member is created for the dialect. If not, a ValueError is
+        raised to prevent the use of unsupported dialects.
+        """
+        if isinstance(value, str):
+            key = value.lower()
+            if key in DialectRegistry._registry:
+                new_member = str.__new__(cls, key)
+                new_member._name_ = key.upper()
+                new_member._value_ = key
+                return new_member
+            raise ValueError(
+                f"Dialect '{value}' does not have a registered SQL transpilation plugin that "
+                "supports it. Please configure a plugin on the server to use this dialect.",
+            )
+        raise TypeError(f"{value!r} is not a valid string for {cls.__name__}")
+
+
+class DialectRegistry:
+    """
+    Registry for SQL dialect plugins.
+    """
+
+    _registry: dict[str, type["SQLTranspilationPlugin"]] = {}
+
+    @classmethod
+    def register(cls, dialect_name: str, plugin: type["SQLTranspilationPlugin"]):
+        """
+        Registers a new dialect plugin.
+        """
+        logger.info(
+            "Registering plugin %s for dialect %s",
+            plugin.__name__,
+            dialect_name,
+        )
+        cls._registry[dialect_name.lower()] = plugin
+
+    @classmethod
+    def get_plugin(cls, dialect_name: str) -> type["SQLTranspilationPlugin"] | None:
+        """
+        Retrieves the plugin class for a given dialect name.
+        """
+        return cls._registry.get(dialect_name.lower())
+
+    @classmethod
+    def list(cls) -> list[str]:
+        """
+        A list of supported dialects.
+        """
+        return list(cls._registry.keys())
+
+
+def register_dialect_plugin(name: str, plugin_cls: type["SQLTranspilationPlugin"]):
+    """
+    Registers a plugin for the given dialect. Skips registration if the plugin is not
+    configured in the settings.
+    """
+    if plugin_cls.package_name not in settings.transpilation_plugins:
+        logger.warning(
+            "Skipping plugin registration for '%s' (%s) (not in configured transpilation plugins: %s)",
+            name,
+            plugin_cls.package_name,
+            settings.transpilation_plugins,
+        )
+        return
+
+    DialectRegistry.register(name, plugin_cls)
+
+
+def dialect_plugin(name: str):
+    """
+    Decorator to register a dialect plugin. This decorator should be used on a class that
+    inherits from SQLTranspilationPlugin.
+
+    Example usage:
+
+    @dialect_plugin("my_dialect")
+    class MyDialectPlugin(SQLTranspilationPlugin):
+        def transpile_sql(self, query: str, input_dialect: str = None, output_dialect: str = None):
+            # Custom transpilation logic here
+            return query
+    """
+
+    def wrapper(cls):
+        register_dialect_plugin(name, cls)
+        return cls
+
+    return wrapper

{datajunction_server-0.0.1a99.dev2 → datajunction_server-0.0.1a100}/datajunction_server/models/engine.py

@@ -5,18 +5,7 @@ Models for columns.
 from typing import Optional
 
 from pydantic.main import BaseModel
-
-from datajunction_server.enum import StrEnum
-
-
-class Dialect(StrEnum):
-    """
-    SQL dialect
-    """
-
-    SPARK = "spark"
-    TRINO = "trino"
-    DRUID = "druid"
+from datajunction_server.models.dialect import Dialect
 
 
 class EngineInfo(BaseModel):

{datajunction_server-0.0.1a99.dev2 → datajunction_server-0.0.1a100}/datajunction_server/models/metric.py

@@ -4,7 +4,6 @@ Models for metrics.
 
 from typing import List, Optional
 
-from pydantic.class_validators import root_validator
 from pydantic.main import BaseModel
 
 from datajunction_server.database.node import Node
@@ -15,11 +14,11 @@ from datajunction_server.models.node import (
     MetricMetadataOutput,
 )
 from datajunction_server.models.query import ColumnMetadata
+from datajunction_server.models.sql import TranspiledSQL
 from datajunction_server.sql.decompose import MeasureExtractor
 from datajunction_server.sql.parsing.backends.antlr4 import ast, parse
-from datajunction_server.transpilation import get_transpilation_plugin
+from datajunction_server.transpilation import transpile_sql
 from datajunction_server.typing import UTCDatetime
-from datajunction_server.utils import get_settings
 
 
 class Metric(BaseModel):
@@ -85,7 +84,7 @@ class Metric(BaseModel):
         )
 
 
-class TranslatedSQL(BaseModel):
+class TranslatedSQL(TranspiledSQL):
     """
     Class for SQL generated from a given metric.
     """
@@ -97,21 +96,11 @@ class TranslatedSQL(BaseModel):
     dialect: Optional[Dialect] = None
     upstream_tables: Optional[List[str]] = None
 
-    @root_validator(pre=False)
-    def transpile_sql(
-        cls,
-        values,
-    ) -> "TranslatedSQL":
-        """
-        Transpiles SQL to the specified dialect with the configured transpilation plugin.
-        If no plugin is configured, it will just return the original generated query.
-        """
-        settings = get_settings()
-        if settings.sql_transpilation_library:  # pragma: no cover
-            plugin = get_transpilation_plugin(settings.sql_transpilation_library)
-            values["sql"] = plugin.transpile_sql(
-                values["sql"],
-                input_dialect=Dialect.SPARK,
-                output_dialect=values["dialect"],
-            )
-        return values
+    @classmethod
+    def create(cls, *, dialect: Dialect | None = None, **kwargs):
+        sql = transpile_sql(kwargs["sql"], dialect)
+        return cls(
+            sql=sql,
+            dialect=dialect,
+            **{k: v for k, v in kwargs.items() if k not in {"sql", "dialect"}},
+        )