fix: make catalog use sc method to implement does_table_exist

Author: german-typedef

src/fenic/_backends/cloud/catalog.py

@@ -17,7 +17,6 @@
     TypedefCatalogTypeReferenceEnum,
 )
 from fenic_cloud.hasura_client.generated_graphql_client.input_types import (
-    CatalogNamespaceInsertInput,
     CreateTableInput,
     NestedFieldInput,
     SchemaInput,
@@ -30,7 +29,6 @@
 )
 
 from fenic._backends.cloud.manager import CloudSessionManager
-from fenic._backends.cloud.session_state import CloudSessionState
 from fenic._backends.local.catalog import (
     DEFAULT_CATALOG_NAME,
     DEFAULT_DATABASE_NAME,
@@ -70,16 +68,21 @@ class CloudCatalog(BaseCatalog):
     all table reads and writes should go through this class for unified table name canonicalization.
     """
 
-    def __init__(self, session_state: CloudSessionState, cloud_session_manager: CloudSessionManager):
+    def __init__(self,
+        ephemeral_catalog_id: str,
+        asyncio_loop: asyncio.AbstractEventLoop,
+        cloud_session_manager: CloudSessionManager):
         """Initialize the remote catalog."""
-        self.session_state = session_state
+        self.cloud_session_manager = cloud_session_manager
         self.lock = threading.Lock()
-        self.current_catalog_id: UUID = UUID(session_state.ephemeral_catalog_id)
+        self.asyncio_loop = asyncio_loop
+        self.ephemeral_catalog_id: UUID = UUID(ephemeral_catalog_id)
+        self.current_catalog_id: UUID = self.ephemeral_catalog_id
         self.current_catalog_name: str = DEFAULT_CATALOG_NAME
         self.current_database_name: str = DEFAULT_DATABASE_NAME
-        self.user_id = cloud_session_manager.user_id
-        self.organization_id = cloud_session_manager.organization_id
-        self.user_client = cloud_session_manager.hasura_user_client
+        self.user_id = self.cloud_session_manager._client_id
+        self.organization_id = self.cloud_session_manager._organization_id
+        self.user_client = self.cloud_session_manager.hasura_user_client
 
     def does_catalog_exist(self, catalog_name: str) -> bool:
         """Checks if a catalog with the specified name exists."""
@@ -129,7 +132,6 @@ def create_catalog(self, catalog_name: str, ignore_if_exists: bool = True) -> bo
                     parent_organization_id=UUID(self.organization_id),
                     catalog_type=TypedefCatalogTypeReferenceEnum.INTERNAL_TYPEDEF,
                     catalog_warehouse="",
-                    catalog_description=None,
                 )
             )
             return True
@@ -342,14 +344,12 @@ def _create_database(
                 raise DatabaseAlreadyExistsError(database_name)
 
         self._execute_catalog_command(
-            self.user_client.create_namespace(
-                namespace=CatalogNamespaceInsertInput(
-                    name=db_identifier.db,
-                    canonical_name=db_identifier.db.casefold(),
-                    parent_organization_id=self.organization_id,
-                    catalog_id=self.current_catalog_id,
-                    created_by_user_id=self.user_id,
-                )
+            self.user_client.sc_create_namespace(
+                dispatch=self._get_catalog_dispatch_input(self.current_catalog_id),
+                name=db_identifier.db,
+                canonical_name=db_identifier.db.casefold(),
+                description=None,
+                properties=[],
             )
         )
         return True
@@ -365,8 +365,8 @@ def _does_table_exist(
         if not self._does_database_exist(catalog_name, db_name):
             return False
 
-        tables = self._get_tables_for_database(catalog_name, db_name)
-        return any(compare_object_names(table, table_name) for table in tables)
+        table = self._get_table(catalog_name, db_name, table_name)
+        return table is not None
 
     def _set_current_catalog(self, catalog_name: str) -> None:
         if not catalog_name:
@@ -376,7 +376,7 @@ def _set_current_catalog(self, catalog_name: str) -> None:
             return
 
         if compare_object_names(catalog_name, DEFAULT_CATALOG_NAME):
-            self.current_catalog_id = UUID(self.session_state.ephemeral_catalog_id)
+            self.current_catalog_id = self.ephemeral_catalog_id
             self.current_catalog_name = DEFAULT_CATALOG_NAME
             return
 
@@ -421,7 +421,7 @@ def _does_database_exist(self, catalog_name: str, database_name: str) -> bool:
 
     def _execute_catalog_command(self, command: Coroutine[Any, Any, Any]) -> Any:
         return asyncio.run_coroutine_threadsafe(
-            command, self.session_state.asyncio_loop
+            command, self.asyncio_loop
         ).result()
 
     def _get_catalog_by_name(self, catalog_name: str) -> Optional[CatalogKey]:
@@ -481,18 +481,35 @@ def _get_tables_for_database(
         )
         return [dataset.name for dataset in result.catalog_dataset]
 
+    def _get_table(
+            self,
+            catalog_name: str,
+            db_name: str,
+            table_name: str,
+            ignore_if_not_exists: bool = True,
+    ) -> LoadTableSimpleCatalogLoadTable:
+        catalog_id = self._get_catalog_id(catalog_name)
+        try:
+            result = self._execute_catalog_command(
+                self.user_client.load_table(
+                    dispatch=self._get_catalog_dispatch_input(catalog_id),
+                    namespace=db_name,
+                    name=table_name,
+                )
+            )
+            return result.simple_catalog.load_table
+        except Exception as e:
+            if ignore_if_not_exists:
+                return None
+            logger.debug(f"Error getting table {table_name} from catalog {catalog_name} and database {db_name}: {e}")
+            raise e
+
+
     def _get_table_details(
         self, catalog_name: str, db_name: str, table_name: str
     ) -> Schema:
-        catalog_id = self._get_catalog_id(catalog_name)
-        result = self._execute_catalog_command(
-            self.user_client.load_table(
-                dispatch=self._get_catalog_dispatch_input(catalog_id),
-                namespace=db_name,
-                name=table_name,
-            )
-        )
-        return self._get_table_schema(result.simple_catalog.load_table)
+        load_table = self._get_table(catalog_name, db_name, table_name, ignore_if_not_exists=False)
+        return self._get_table_schema(load_table)
 
     def _get_catalog_dispatch_input(
         self,

src/fenic/_backends/cloud/session_state.py

@@ -25,6 +25,7 @@
 from fenic._backends.cloud.execution import CloudExecution
 from fenic._backends.cloud.settings import CloudSettings
 from fenic.core._interfaces import BaseSessionState
+from fenic.core._interfaces.catalog import BaseCatalog
 from fenic.core._resolved_session_config import (
     CloudExecutorSize,
     ResolvedSessionConfig,
@@ -61,6 +62,7 @@ class CloudSessionState(BaseSessionState):
     session_id: Optional[str] = None
     session_name: Optional[str] = None
     session_canonical_name: Optional[str] = None
+    cloud_catalog: BaseCatalog = None
 
     def __init__(
         self,
@@ -119,7 +121,16 @@ def execution(self) -> CloudExecution:
 
     @property
     def catalog(self):
-        pass
+        from fenic._backends.cloud.catalog import CloudCatalog
+        from fenic._backends.cloud.manager import CloudSessionManager
+
+        if self.cloud_catalog is None:
+            self.cloud_catalog = CloudCatalog(
+                ephemeral_catalog_id=self.ephemeral_catalog_id,
+                asyncio_loop=self.asyncio_loop,
+                cloud_session_manager=CloudSessionManager(),
+            )
+        return self.cloud_catalog
 
     # properties and methods referencing dynamic state managed by the CloudSessionManager
     @property
@@ -231,6 +242,7 @@ async def _entrypoint_get_or_create_session_engine(self):
         self.session_canonical_name = response.canonical_name
         self.engine_uri = response.uris.remote_actions_uri
         self.arrow_ipc_uri = response.uris.remote_results_uri_prefix
+        self.ephemeral_catalog_id = response.ephemeral_catalog_id
         logger.info(
             f"{'Found' if existing else 'Created'} Executor with session_id: {self.session_uuid}"
         )

tests/_backends/cloud/catalog/test_cloud_catalog.py

@@ -10,6 +10,9 @@
 
 pytest.importorskip("fenic_cloud")
 from fenic_cloud.hasura_client.generated_graphql_client import Client
+from fenic_cloud.hasura_client.generated_graphql_client.client import (
+    CatalogDispatchInput,
+)
 from fenic_cloud.hasura_client.generated_graphql_client.create_namespace import (
     CreateNamespace,
     CreateNamespaceInsertCatalogNamespaceOne,
@@ -103,6 +106,7 @@
 TEST_NEW_TABLE_NAME = "new_table"
 TEST_SAMPLE_LOCATION = "s3://test-bucket/test-path"
 TEST_NEW_CATALOG_NAME = "new_catalog"
+TEST_NONEXISTENT_IDENTIFIER = "nonexistent"
 
 
 @pytest.fixture(scope="function")
@@ -121,6 +125,55 @@ def schema(): # noqa: D103
         ]
     )
 
+def _load_table_side_effect(
+            dispatch: CatalogDispatchInput,
+            namespace: str,
+            name: str) -> LoadTable:
+    if TEST_NONEXISTENT_IDENTIFIER in name or name in [TEST_NEW_TABLE_NAME]:
+        raise Exception("Table not found")
+    return LoadTable(
+        simple_catalog=LoadTableSimpleCatalog(
+            load_table=LoadTableSimpleCatalogLoadTable(
+                created_at=datetime.now(),
+                updated_at=datetime.now(),
+                schema_=SimpleCatalogTableDetailsSchema(
+                    schema_id=1,
+                    identifier_field_ids=[1, 2, 3],
+                    fields=[
+                        SimpleCatalogSchemaDetailsFields(
+                            id=1,
+                            name="id",
+                            data_type="int64",
+                            arrow_data_type="int64",
+                            nullable=False,
+                            metadata=None,
+                        ),
+                        SimpleCatalogSchemaDetailsFields(
+                            id=2,
+                            name="name",
+                            data_type="string",
+                            arrow_data_type="string",
+                            nullable=False,
+                            metadata=None,
+                        ),
+                        SimpleCatalogSchemaDetailsFields(
+                            id=3,
+                            name="account_balance",
+                            data_type="Decimal128",
+                            arrow_data_type="Decimal128",
+                            nullable=False,
+                            metadata=None,
+                        ),
+                    ],
+                ),
+                name=TEST_TABLE_NAME_1,
+                location=None,
+                external=True,
+                file_format=None,
+                partition_field_names=None,
+            ),
+        ),
+    )
 
 @pytest.fixture
 def mock_user_client(schema):
@@ -181,49 +234,7 @@ def mock_user_client(schema):
         )
     )
 
-    user_client.load_table.return_value = LoadTable(
-        simple_catalog=LoadTableSimpleCatalog(
-            load_table=LoadTableSimpleCatalogLoadTable(
-                created_at=datetime.now(),
-                updated_at=datetime.now(),
-                schema_=SimpleCatalogTableDetailsSchema(
-                    schema_id=1,
-                    identifier_field_ids=[1, 2, 3],
-                    fields=[
-                        SimpleCatalogSchemaDetailsFields(
-                            id=1,
-                            name="id",
-                            data_type="int64",
-                            arrow_data_type="int64",
-                            nullable=False,
-                            metadata=None,
-                        ),
-                        SimpleCatalogSchemaDetailsFields(
-                            id=2,
-                            name="name",
-                            data_type="string",
-                            arrow_data_type="string",
-                            nullable=False,
-                            metadata=None,
-                        ),
-                        SimpleCatalogSchemaDetailsFields(
-                            id=3,
-                            name="account_balance",
-                            data_type="Decimal128",
-                            arrow_data_type="Decimal128",
-                            nullable=False,
-                            metadata=None,
-                        ),
-                    ],
-                ),
-                name=TEST_TABLE_NAME_1,
-                location=None,
-                external=True,
-                file_format=None,
-                partition_field_names=None,
-            ),
-        ),
-    )
+    user_client.load_table.side_effect = _load_table_side_effect
 
     user_client.create_namespace.return_value = CreateNamespace(
         insert_catalog_namespace_one=CreateNamespaceInsertCatalogNamespaceOne(
@@ -590,9 +601,12 @@ def _init_cloud_catalog(
 ) -> Any:
     os.environ["TYPEDEF_USER_ID"] = "mock_user_id"
     os.environ["TYPEDEF_USER_SECRET"] = "mock_user_secret"  # nosec B105
-    os.environ["HASURA_GRAPHQL_ADMIN_SECRET"] = "mock_admin_secret"  # nosec B105
     os.environ["REMOTE_SESSION_AUTH_PROVIDER_URI"] = "mock_auth_provider_uri"
-    cloud_catalog = CloudCatalog(session_state, cloud_session_manager)
+    cloud_catalog = CloudCatalog(
+        ephemeral_catalog_id=session_state.ephemeral_catalog_id,
+        asyncio_loop=session_state.asyncio_loop,
+        cloud_session_manager=cloud_session_manager,
+    )
     cloud_catalog.user_client = client
     cloud_catalog.user_id = TEST_DEFAULT_USER_ID
     cloud_catalog.organization_id = TEST_DEFAULT_ORGANIZATION_ID