datajunction 0.0.149__tar.gz → 0.0.151__tar.gz

{datajunction-0.0.149 → datajunction-0.0.151}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datajunction
-Version: 0.0.149
+Version: 0.0.151
 Summary: DataJunction client library for connecting to a DataJunction server
 Project-URL: repository, https://github.com/DataJunction/dj
 Author-email: DataJunction Authors <yian.shang@gmail.com>

{datajunction-0.0.149 → datajunction-0.0.151}/datajunction/__about__.py

@@ -2,4 +2,4 @@
 Version for Hatch
 """
 
-__version__ = "0.0.149"
+__version__ = "0.0.151"

{datajunction-0.0.149 → datajunction-0.0.151}/datajunction/_internal.py

@@ -683,12 +683,28 @@ class DJClient:
         response = self._session.get(f"/namespaces/{namespace}/export/")
         return response.json()
 
-    def _export_namespace_spec(self, namespace):
+    def _export_namespace_yaml_zip(
+        self,
+        namespace: str,
+        existing_zip_bytes: Optional[bytes] = None,
+    ) -> bytes:
         """
-        Export a deployment spec for a namespace
+        Export a namespace as a ZIP of YAML files.
+
+        If `existing_zip_bytes` is provided, the server merges new content into
+        those files, preserving key ordering and comments. Otherwise it falls back
+        to the configured git branch or produces a fresh export.
         """
-        response = self._session.get(f"/namespaces/{namespace}/export/spec/")
-        return response.json()
+        files = (
+            {"existing_zip": ("existing.zip", existing_zip_bytes, "application/zip")}
+            if existing_zip_bytes is not None
+            else None
+        )
+        response = self._session.post(
+            f"/namespaces/{namespace}/export/yaml",
+            files=files,
+        )
+        return response.content
 
     #
     # Methods for Tags

{datajunction-0.0.149 → datajunction-0.0.151}/datajunction/deployment.py

@@ -56,123 +56,49 @@ class DeploymentService:
         self.client = client
         self.console = console or Console()
 
-    @staticmethod
-    def clean_dict(d: dict) -> dict:
-        """
-        Recursively remove None, empty list, and empty dict values.
-        """
-        result = {}
-        for k, v in d.items():
-            if v is None:
-                continue
-            if isinstance(v, (list, dict)) and not v:
-                continue
-            if isinstance(v, dict):
-                nested = DeploymentService.clean_dict(v)
-                if nested:  # only include if not empty after cleaning
-                    result[k] = nested
-            else:
-                result[k] = v  # type: ignore
-        return result
-
-    @staticmethod
-    def filter_node_for_export(node: dict) -> dict:
-        """
-        Filter a node dict for export to YAML.
-
-        For columns:
-        - Cubes: columns are always excluded (they're inferred from metrics/dimensions)
-        - Other nodes: only includes columns with meaningful customizations
-          (display_name different from name, attributes, description, or partition).
-          Column types are excluded - let DJ infer them from the query/source.
-        """
-        result = DeploymentService.clean_dict(node)
-
-        # Cubes should never have columns in export - they're inferred from metrics/dimensions
-        if result.get("node_type") == "cube":
-            result.pop("columns", None)
-        # For other nodes, filter columns to only include meaningful customizations
-        elif "columns" in result and result["columns"]:
-            filtered_columns = []
-            for col in result["columns"]:
-                # Check for meaningful customizations
-                has_custom_display = col.get("display_name") and col.get(
-                    "display_name",
-                ) != col.get("name")
-                has_attributes = bool(col.get("attributes"))
-                has_description = bool(col.get("description"))
-                has_partition = bool(col.get("partition"))
-
-                if (
-                    has_custom_display
-                    or has_attributes
-                    or has_description
-                    or has_partition
-                ):
-                    # Include column but exclude type (let DJ infer)
-                    filtered_col = {
-                        k: v
-                        for k, v in col.items()
-                        if k != "type" and v  # Exclude type and empty values
-                    }
-                    filtered_columns.append(filtered_col)
-
-            if filtered_columns:
-                result["columns"] = filtered_columns
-            else:
-                # Remove columns entirely if none have customizations
-                del result["columns"]
-
-        return result
-
     def pull(
         self,
         namespace: str,
         target_path: Union[str, Path],
-        ignore_existing_files: bool = False,
     ):
         """
         Export a namespace to a local project.
+
+        Pulls the YAML files from the server's `/export/yaml` endpoint, which
+        runs every node through the same serializer (`node_spec_to_yaml`) used
+        by the UI sync-to-git flow. This way `dj pull` and a UI export produce
+        identical YAML for the same node state.
+
+        When the target directory already contains YAML files, they are uploaded
+        to the server so it can merge new content into them — preserving key
+        ordering, inline comments, and scalar styles. This means a `dj pull`
+        against an already-populated directory produces minimal diffs.
         """
-        path = Path(target_path)
-        if any(path.iterdir()) and not ignore_existing_files:
-            raise DJClientException("The target path must be empty")
-        deployment_spec = self.client._export_namespace_spec(namespace)
+        import io
+        import zipfile
 
-        namespace = deployment_spec["namespace"]
-        nodes: list[dict[str, Any]] = deployment_spec.get("nodes", [])
         base_path = Path(target_path)
         base_path.mkdir(parents=True, exist_ok=True)
 
-        # Create a YAML for each node in the appropriate namespace folder
-        for node in nodes:
-            node_name = node["name"]
-            # Namespace folder is everything except the last part of the node
-            node_parts = node_name.replace("${prefix}", "").split(".")
-            node_namespace_path = base_path.joinpath(*node_parts[:-1])
-            node_namespace_path.mkdir(parents=True, exist_ok=True)
-
-            # File name is the last part of the node
-            file_name = node_parts[-1] + ".yaml"
-            file_path = node_namespace_path / file_name
-
-            # Write YAML for this node (filter columns for cleaner output)
-            with open(file_path, "w") as yaml_file:
-                yaml.dump(
-                    DeploymentService.filter_node_for_export(node),
-                    yaml_file,
-                    sort_keys=False,
-                )
+        existing_zip_bytes = None
+        existing_yaml_files = list(base_path.rglob("*.yaml"))
+        if existing_yaml_files:
+            buf = io.BytesIO()
+            with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
+                for yaml_file in existing_yaml_files:
+                    zf.writestr(
+                        str(yaml_file.relative_to(base_path)),
+                        yaml_file.read_bytes(),
+                    )
+            existing_zip_bytes = buf.getvalue()
 
-        # Write top-level dj.yaml with full deployment info
-        dj_yaml_path = base_path / "dj.yaml"
-        with open(dj_yaml_path, "w") as yaml_file:
-            project_spec = {
-                "name": f"Project {namespace} (Autogenerated)",
-                "description": f"This is an autogenerated project for namespace {namespace}",
-                "namespace": namespace,
-            }
-            yaml.safe_dump(project_spec, yaml_file, sort_keys=False)
+        zip_bytes = self.client._export_namespace_yaml_zip(
+            namespace,
+            existing_zip_bytes,
+        )
+
+        with zipfile.ZipFile(io.BytesIO(zip_bytes)) as zf:
+            zf.extractall(base_path)
 
     def push(
         self,

{datajunction-0.0.149 → datajunction-0.0.151}/tests/conftest.py

@@ -192,14 +192,14 @@ def module__query_service_client(
     qs_client = QueryServiceClient(uri="query_service:8001")
     qs_client.query_state = QueryState.RUNNING  # type: ignore
 
-    def mock_get_columns_for_table(
+    async def mock_get_columns_for_table(
         catalog: str,
         schema: str,
         table: str,
-        engine: Optional[Engine] = None,  # pylint: disable=unused-argument
         request_headers: Optional[  # pylint: disable=unused-argument
             Dict[str, str]
         ] = None,
+        engine: Optional[Engine] = None,  # pylint: disable=unused-argument
     ) -> List[Column]:
         return COLUMN_MAPPINGS[f"{catalog}.{schema}.{table}"]
 
@@ -209,23 +209,27 @@ def module__query_service_client(
         mock_get_columns_for_table,
     )
 
-    def mock_submit_query(
+    async def mock_submit_query(
         query_create: QueryCreate,
         request_headers: Optional[  # pylint: disable=unused-argument
             Dict[str, str]
         ] = None,
     ) -> QueryWithResults:
-        normalized_query = (
-            query_create.submitted_query.strip()
-            .replace('"', "")
-            .replace("\n", "")
-            .replace(" ", "")
-        )
-        # Strip LIMIT clause for matching (allows queries with/without LIMIT to match)
-        normalized_query = re.sub(r"LIMIT\d+$", "", normalized_query)
+        def _normalize(s: str) -> str:
+            out = (
+                s.strip()
+                .replace('"', "")
+                .replace("\n", "")
+                .replace("\t", "")
+                .replace(" ", "")
+            )
+            return re.sub(r"LIMIT\d+$", "", out)
+
+        normalized_query = _normalize(query_create.submitted_query)
 
-        if normalized_query in QUERY_DATA_MAPPINGS:
-            results = QUERY_DATA_MAPPINGS[normalized_query]
+        normalized_mappings = {_normalize(k): v for k, v in QUERY_DATA_MAPPINGS.items()}
+        if normalized_query in normalized_mappings:
+            results = normalized_mappings[normalized_query]
         else:
             # SQL-agnostic fallback: return a generic successful result.
             # Client tests should not depend on the exact SQL generated — they
@@ -262,7 +266,7 @@ def module__query_service_client(
         mock_submit_query,
     )
 
-    def mock_create_view(
+    async def mock_create_view(
         view_name: str,
         query_create: QueryCreate,  # pylint: disable=unused-argument
         request_headers: Optional[  # pylint: disable=unused-argument

{datajunction-0.0.149 → datajunction-0.0.151}/tests/test_deploy.py

@@ -2,6 +2,7 @@ import importlib.metadata
 import io
 from pathlib import Path
 import time
+import zipfile
 from unittest import mock
 import pytest
 from unittest.mock import MagicMock, patch
@@ -18,126 +19,46 @@ import yaml
 from rich.console import Console
 
 
-def test_clean_dict_removes_nones_and_empty():
-    dirty = {
-        "a": None,
-        "b": [],
-        "c": {},
-        "d": {"x": None, "y": {"z": []}, "k": "keep"},
-        "e": [1, 2],
-    }
-    cleaned = DeploymentService.clean_dict(dirty)
-    assert cleaned == {"d": {"k": "keep"}, "e": [1, 2]}
-
-
-def test_filter_node_for_export_removes_columns_without_customizations():
-    """Columns without meaningful customizations should be removed."""
-    node = {
-        "name": "test.node",
-        "query": "SELECT * FROM foo",
-        "columns": [
-            # Should be kept: has custom display_name
-            {"name": "user_id", "type": "INT", "display_name": "User ID"},
-            # Should be removed: display_name same as name
-            {"name": "created_at", "type": "TIMESTAMP", "display_name": "created_at"},
-            # Should be kept: has attributes
-            {"name": "id", "type": "BIGINT", "attributes": ["primary_key"]},
-            # Should be removed: no customizations
-            {"name": "plain_col", "type": "VARCHAR"},
-            # Should be kept: has description
-            {"name": "desc_col", "type": "TEXT", "description": "A useful column"},
-        ],
-    }
-
-    filtered = DeploymentService.filter_node_for_export(node)
-
-    # Only columns with customizations should remain
-    assert len(filtered["columns"]) == 3
-
-    # Type should be excluded from all columns
-    for col in filtered["columns"]:
-        assert "type" not in col
-
-    # Check correct columns were kept
-    col_names = [c["name"] for c in filtered["columns"]]
-    assert "user_id" in col_names
-    assert "id" in col_names
-    assert "desc_col" in col_names
-    assert "created_at" not in col_names
-    assert "plain_col" not in col_names
-
-
-def test_filter_node_for_export_removes_columns_key_when_empty():
-    """If no columns have customizations, the columns key should be removed."""
-    node = {
-        "name": "test.node",
-        "query": "SELECT * FROM foo",
-        "columns": [
-            {"name": "a", "type": "INT"},
-            {"name": "b", "type": "VARCHAR", "display_name": "b"},  # same as name
-        ],
-    }
-
-    filtered = DeploymentService.filter_node_for_export(node)
-
-    assert "columns" not in filtered
-
+def _make_zip(files: dict[str, str]) -> bytes:
+    buf = io.BytesIO()
+    with zipfile.ZipFile(buf, "w") as zf:
+        for name, content in files.items():
+            zf.writestr(name, content)
+    return buf.getvalue()
 
-def test_filter_node_for_export_always_removes_columns_for_cubes():
-    """Cube columns should always be removed - they're inferred from metrics/dimensions."""
-    node = {
-        "name": "test.cube",
-        "node_type": "cube",
-        "metrics": ["test.metric1", "test.metric2"],
-        "dimensions": ["test.dim1"],
-        "columns": [
-            {"name": "metric1", "type": "BIGINT"},
-            {"name": "dim1", "type": "VARCHAR", "display_name": "Dimension 1"},
-        ],
-    }
-
-    filtered = DeploymentService.filter_node_for_export(node)
-
-    # Columns should be removed regardless of customizations
-    assert "columns" not in filtered
-    # Other fields should remain
-    assert filtered["metrics"] == ["test.metric1", "test.metric2"]
-    assert filtered["dimensions"] == ["test.dim1"]
 
-
-def test_pull_writes_yaml_files(tmp_path):
-    # fake client returning a minimal deployment spec
+def test_pull_extracts_zip_from_server(tmp_path):
+    zip_bytes = _make_zip(
+        {
+            "dj.yaml": "namespace: foo.bar\n",
+            "foo/bar/baz.yaml": "name: foo.bar.baz\nquery: SELECT 1\n",
+        },
+    )
     client = MagicMock()
-    client._export_namespace_spec.return_value = {
-        "namespace": "foo.bar",
-        "nodes": [
-            {"name": "foo.bar.baz", "query": "SELECT 1"},
-            {"name": "foo.bar.qux", "query": "SELECT 2"},
-        ],
-    }
+    client._export_namespace_yaml_zip.return_value = zip_bytes
     svc = DeploymentService(client)
 
     svc.pull("foo.bar", tmp_path)
 
-    # project-level yaml
-    project_yaml = yaml.safe_load((tmp_path / "dj.yaml").read_text())
-    assert project_yaml["namespace"] == "foo.bar"
-
-    # node files
-    baz_file = tmp_path / "foo" / "bar" / "baz.yaml"
-    assert baz_file.exists()
-    assert yaml.safe_load(baz_file.read_text())["query"] == "SELECT 1"
+    client._export_namespace_yaml_zip.assert_called_once_with("foo.bar", None)
+    assert yaml.safe_load((tmp_path / "dj.yaml").read_text())["namespace"] == "foo.bar"
+    assert (tmp_path / "foo" / "bar" / "baz.yaml").exists()
 
-    qux_file = tmp_path / "foo" / "bar" / "qux.yaml"
-    assert qux_file.exists()
 
-
-def test_pull_raises_if_target_not_empty(tmp_path):
-    (tmp_path / "something.txt").write_text("not empty")
+def test_pull_uploads_existing_yaml_files(tmp_path):
+    (tmp_path / "old.yaml").write_text("name: ns.old\n")
+    updated_zip = _make_zip({"old.yaml": "name: ns.old\nquery: SELECT 2\n"})
     client = MagicMock()
+    client._export_namespace_yaml_zip.return_value = updated_zip
     svc = DeploymentService(client)
-    with pytest.raises(DJClientException):
-        svc.pull("ns", tmp_path)
+
+    svc.pull("ns", tmp_path)
+
+    _, call_kwargs = client._export_namespace_yaml_zip.call_args
+    existing_bytes = client._export_namespace_yaml_zip.call_args[0][1]
+    assert existing_bytes is not None
+    with zipfile.ZipFile(io.BytesIO(existing_bytes)) as zf:
+        assert "old.yaml" in zf.namelist()
 
 
 def test_print_results_success():