fix: use the cloud catalog when doing save_as_table in a cloud session

Author: german-typedef

src/fenic/_backends/cloud/catalog.py

@@ -4,6 +4,7 @@
 from dataclasses import dataclass
 from datetime import datetime
 from typing import Any, Coroutine, Dict, List, Optional
+from urllib.parse import urlparse
 from uuid import UUID
 
 import polars as pl
@@ -56,12 +57,12 @@
 
 logger = logging.getLogger(__name__)
 
-
 @dataclass(frozen=True)
 class CatalogKey:
     catalog_name: str
     catalog_id: UUID
 
+CLOUD_SUPPORTED_SCHEMES = ["s3"]
 
 class CloudCatalog(BaseCatalog):
     """A catalog for cloud execution mode. Implements the BaseCatalog -
@@ -112,11 +113,18 @@ def list_catalogs(self) -> List[str]:
             remote_catalogs.append(DEFAULT_CATALOG_NAME)
             return remote_catalogs
 
-    def create_catalog(self, catalog_name: str, ignore_if_exists: bool = True) -> bool:
+    def create_catalog(
+            self,
+            catalog_name: str,
+            location: str,
+            ignore_if_exists: bool = True) -> bool:
         """Create a new catalog."""
         if compare_object_names(catalog_name, DEFAULT_CATALOG_NAME):
             raise CatalogError("Cannot create a catalog with the default name")
 
+        if urlparse(location).scheme not in CLOUD_SUPPORTED_SCHEMES:
+            raise CatalogError(f"Unsupported scheme: {urlparse(location).scheme}")
+
         with self.lock:
             if self._does_catalog_exist(catalog_name):
                 if ignore_if_exists:
@@ -131,7 +139,7 @@ def create_catalog(self, catalog_name: str, ignore_if_exists: bool = True) -> bo
                     created_by_user_id=UUID(self.user_id),
                     parent_organization_id=UUID(self.organization_id),
                     catalog_type=TypedefCatalogTypeReferenceEnum.INTERNAL_TYPEDEF,
-                    catalog_warehouse="",
+                    catalog_warehouse=location,
                 )
             )
             return True
@@ -267,14 +275,12 @@ def create_table(
         self,
         table_name: str,
         schema: Schema,
-        location: str,
         ignore_if_exists: bool = True,
-        file_format: Optional[str] = None,
     ) -> bool:
         """Create a new table in the current database."""
         with self.lock:
             return self._create_table(
-                table_name, schema, location, ignore_if_exists, file_format
+                table_name, schema, ignore_if_exists
             )
 
 
@@ -526,9 +532,7 @@ def _create_table(
         self,
         table_name: str,
         schema: Schema,
-        location: str,
         ignore_if_exists: bool = True,
-        file_format: Optional[str] = None,
     ) -> bool:
         table_identifier = TableIdentifier.from_string(table_name).enrich(
             self.current_catalog_name, self.current_database_name
@@ -557,11 +561,6 @@ def _create_table(
                 raise TableAlreadyExistsError(table_identifier.table, table_identifier.db)
 
         catalog_id = self._get_catalog_id(table_identifier.catalog)
-        fixed_file_format = (
-            FileFormat.PARQUET
-            if file_format is None
-            else FileFormat(file_format.upper())
-        )
         self._execute_catalog_command(
             self.user_client.sc_create_table(
                 dispatch=self._get_catalog_dispatch_input(catalog_id),
@@ -571,8 +570,8 @@ def _create_table(
                     canonical_name=table_identifier.table.casefold(),
                     description=None,
                     external=False,
-                    location=location,
-                    file_format=fixed_file_format,
+                    location=self._get_table_location_from_table_identifier(table_identifier),
+                    file_format=FileFormat.PARQUET,
                     partition_field_names=[],
                     schema_=self._get_schema_input_from_schema(schema),
                 ),
@@ -661,3 +660,8 @@ def _get_schema_type_to_pyarrow(schema_type: str):
             return pa.float64()
         else:
             return schema_type
+
+    @staticmethod
+    def _get_table_location_from_table_identifier(table_identifier: TableIdentifier) -> str:
+        """Gets the key in the s3 bucket for the table based on its database and name."""
+        return f"{table_identifier.db}/{table_identifier.table}"

src/fenic/_backends/cloud/execution.py

@@ -28,19 +28,23 @@
     SaveToFileExecutionRequest,
     ShowExecutionRequest,
     StartExecutionRequest,
-    TableIdentifier,
+)
+from fenic_cloud.protos.engine.v1.engine_pb2 import (
+    TableIdentifier as TableIdentifierProto,
 )
 from fenic_cloud.protos.engine.v1.engine_pb2_grpc import EngineServiceStub
 
 from fenic._backends.cloud.metrics import get_query_execution_metrics
 from fenic._backends.schema_serde import deserialize_schema, serialize_schema
+from fenic._backends.utils.catalog_utils import TableIdentifier
 from fenic.core._interfaces import BaseExecution
 from fenic.core._logical_plan.serde import LogicalPlanSerde
 from fenic.core.error import (
     CloudExecutionError,
     CloudSessionError,
     ExecutionError,
     InternalError,
+    PlanError,
     ValidationError,
 )
 from fenic.core.metrics import LMMetrics, PhysicalPlanRepr, QueryMetrics, RMMetrics
@@ -164,16 +168,57 @@ def save_as_table(
     ) -> QueryMetrics:
         """Execute the logical plan and save the result as a table."""
         logger.debug(f"Saving plan {logical_plan} as table: {table_name}")
+
+        table_identifier = TableIdentifier.from_string(table_name).enrich(
+            self.session_state.catalog.get_current_catalog(),
+            self.session_state.catalog.get_current_database(),
+        )
+
+        # If the table doesn't exist, create it, this has to be done in the user's context.
+        table_identifier_str = str(table_identifier)
+        table_exists = self.session_state.catalog.does_table_exist(table_identifier_str)
+        if table_exists:
+            if mode == "error":
+                raise PlanError(
+                    f"Cannot save to table '{table_name}' - it already exists and mode is 'error'. "
+                    f"Choose a different approach: "
+                    f"1) Use mode='overwrite' to replace the existing table, "
+                    f"2) Use mode='append' to add data to the existing table, "
+                    f"3) Use mode='ignore' to skip saving if table exists, "
+                    f"4) Use a different table name.")
+            if mode == "ignore":
+                logger.warning(f"Table {table_name} already exists, ignoring write.")
+                return QueryMetrics()
+            if mode == "append":
+                saved_schema = self.session_state.catalog.describe_table(table_identifier_str)
+                plan_schema = logical_plan.schema()
+                if saved_schema != plan_schema:
+                    raise PlanError(
+                        f"Cannot append to table '{table_name}' - schema mismatch detected. "
+                        f"The existing table has a different schema than your DataFrame. "
+                        f"Existing schema: {saved_schema} "
+                        f"Your DataFrame schema: {plan_schema} "
+                        f"To fix this: "
+                        f"1) Use mode='overwrite' to replace the table with your DataFrame's schema, "
+                        f"2) Modify your DataFrame to match the existing table's schema, "
+                        f"3) Use a different table name.")
+        else:
+            raise CloudExecutionError(
+                f"Cannot save to table '{table_name}' - it does not exist. "
+                f"Choose a different approach: "
+                f"1) Create the table in question "
+                f"2) Use a different table name.")
+
         # TODO (DY): check that current catalog and schema (if specified in table_name) match session state
-        table_identifier = TableIdentifier(
-            catalog=self.session_state.catalog,
-            schema=self.session_state.schema,
+        table_identifier_proto = TableIdentifierProto(
+            catalog=table_identifier.catalog,
+            schema=table_identifier.db,
             table=table_name,
         )
         request = StartExecutionRequest(
             save_as_table=SaveAsTableExecutionRequest(
                 serialized_plan=LogicalPlanSerde.serialize(logical_plan),
-                table_identifier=table_identifier,
+                table_identifier=table_identifier_proto,
                 mode=mode,
             )
         )

src/fenic/_backends/local/physical_plan/sink.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import logging
-from typing import TYPE_CHECKING, List, Literal, Tuple
+from typing import TYPE_CHECKING, List, Literal, Optional, Tuple
 
 if TYPE_CHECKING:
     from fenic._backends.local.session_state import LocalSessionState

src/fenic/_backends/utils/catalog_utils.py

@@ -103,6 +103,13 @@ def enrich(self, catalog_name: str, db_name: str) -> "TableIdentifier":
             table=self.table,
         )
 
+    def __str__(self) -> str:
+        str_identifier = self.table
+        if self.db:
+            str_identifier = f"{self.db}.{str_identifier}"
+        if self.catalog:
+            str_identifier = f"{self.catalog}.{str_identifier}"
+        return str_identifier
 
 @dataclass(frozen=True)
 class DBIdentifier(BaseIdentifier):
@@ -132,6 +139,12 @@ def enrich(self, catalog_name: str) -> "DBIdentifier":
             return self
         return DBIdentifier(catalog=catalog_name, db=self.db)
 
+    def __str__(self) -> str:
+        str_identifier = self.db
+        if self.catalog:
+            str_identifier = f"{self.catalog}.{str_identifier}"
+        return str_identifier
+
 
 def compare_object_names(object_name_1: str, object_name_2: str) -> bool:
     """Compare two object names, ignoring case."""

src/fenic/api/io/writer.py

@@ -4,7 +4,7 @@
 
 import logging
 from pathlib import Path
-from typing import TYPE_CHECKING, Literal, Union
+from typing import TYPE_CHECKING, Literal, Optional, Union
 
 if TYPE_CHECKING:
     from fenic.api.dataframe import DataFrame

src/fenic/core/_logical_plan/plans/sink.py

@@ -91,6 +91,7 @@ def __init__(
                  - append: Appends data to table if it exists
                  - overwrite: Overwrites existing table
                  - ignore: Silently ignores operation if table exists
+            location: location where the table will be saved (only available for cloud execution)
         """
         self.child = child
         self.table_name = table_name

tests/_backends/cloud/catalog/test_cloud_catalog.py

@@ -531,14 +531,12 @@ def test_create_table(cloud_catalog, schema): # noqa: D103
         cloud_catalog.create_table(
             TEST_TABLE_NAME_1,
             schema=schema,
-            location=TEST_SAMPLE_LOCATION,
             ignore_if_exists=False,
         )
     with pytest.raises(CatalogError):
         cloud_catalog.create_table(
             "some_catalog.some_database.some_table",
             schema=schema,
-            location=TEST_SAMPLE_LOCATION,
         )
 
 
@@ -566,16 +564,16 @@ def test_drop_table(cloud_catalog): # noqa: D103
 
 def test_create_catalog(cloud_catalog): # noqa: D103
     with pytest.raises(CatalogError):
-        cloud_catalog.create_catalog(DEFAULT_CATALOG_NAME)
+        cloud_catalog.create_catalog(DEFAULT_CATALOG_NAME, TEST_SAMPLE_LOCATION)
 
-    assert cloud_catalog.create_catalog(TEST_NEW_CATALOG_NAME)
+    assert cloud_catalog.create_catalog(TEST_NEW_CATALOG_NAME, TEST_SAMPLE_LOCATION)
 
     # The catalog already exists, so we should return False (default for ignore_if_exists is True)
-    assert not cloud_catalog.create_catalog(TEST_CATALOG_NAME)
+    assert not cloud_catalog.create_catalog(TEST_CATALOG_NAME, TEST_SAMPLE_LOCATION)
 
     # The catalog already exists, so we should raise an error if ignore_if_exists is False
     with pytest.raises(CatalogAlreadyExistsError):
-        cloud_catalog.create_catalog(TEST_CATALOG_NAME, ignore_if_exists=False)
+        cloud_catalog.create_catalog(TEST_CATALOG_NAME, TEST_SAMPLE_LOCATION, ignore_if_exists=False)
 
 
 def test_drop_catalog(cloud_catalog): # noqa: D103