relationalai 0.12.12__py3-none-any.whl → 0.13.0__py3-none-any.whl

relationalai/__init__.py

@@ -7,8 +7,8 @@ import importlib.metadata
 
 from typing import cast
 
-from .clients import config as cfg
 from .clients.config import Config, save_config
+from . import clients
 from . import dsl
 from . import debugging
 from . import metamodel
@@ -19,11 +19,14 @@ from . import tools
 from .util.otel_configuration import configure_otel
 from snowflake.snowpark import Session
 from .errors import RAIException, handle_missing_integration
-from .environments import runtime_env, SessionEnvironment
-from relationalai.tools.constants import USE_DIRECT_ACCESS, Generation
+from .environments import runtime_env, SessionEnvironment, SnowbookEnvironment
+from relationalai.tools.constants import Generation
 
 import __main__
 
+# Define cfg after all imports since it depends on clients being imported
+cfg = clients.config
+
 __version__ = importlib.metadata.version(__package__ or __name__)
 
 def Model(
@@ -113,9 +116,9 @@ def Model(
 
     try:
         if platform == "azure":
-            import relationalai.clients.azure as azure
+            from relationalai.clients.resources.azure.azure import Graph
             from .util.otel_handler import disable_otel_handling, is_otel_initialized
-            model = azure.Graph(
+            model = Graph(
                 name,
                 profile=profile,
                 config=config,
@@ -127,8 +130,18 @@ def Model(
             if is_otel_initialized:
                 disable_otel_handling()
         elif platform == "snowflake":
-            from relationalai.clients import snowflake
-            model = snowflake.Graph(
+            try:
+                from relationalai.clients.resources.snowflake import Graph
+            except ImportError as e:
+                # Provide a helpful error message for Snowflake notebook environments
+                if isinstance(runtime_env, SnowbookEnvironment):
+                    raise ImportError(
+                        "Failed to import relationalai.clients.resources.snowflake. "
+                        "This may indicate that the relationalai.zip package structure is incomplete. "
+                        "Please ensure the zip file includes the full clients/resources/snowflake directory structure."
+                    ) from e
+                raise
+            model = Graph(
                 name,
                 profile=profile,
                 config=config,
@@ -163,18 +176,41 @@ def Resources(
     config = config or Config(profile)
     platform = config.get("platform", "snowflake")
     if platform == "azure":
-        from relationalai.clients.azure import Resources
+        from relationalai.clients.resources.azure.azure import Resources
         return Resources(config=config)
     elif platform == "snowflake":
-        from relationalai.clients.snowflake import Resources, DirectAccessResources
-        use_direct_access = config.get("use_direct_access", USE_DIRECT_ACCESS)
-        if use_direct_access:
-            return DirectAccessResources(config=config, connection=connection, reset_session=reset_session, generation=generation)
-        else:
-            return Resources(config=config, connection=connection, reset_session=reset_session, generation=generation)
+        try:
+            from relationalai.clients.resources.snowflake.resources_factory import create_resources_instance
+        except ImportError as e:
+            # Provide a helpful error message for Snowflake notebook environments
+            from relationalai.environments import runtime_env, SnowbookEnvironment
+            if isinstance(runtime_env, SnowbookEnvironment):
+                raise ImportError(
+                    "Failed to import relationalai.clients.resources.snowflake. "
+                    "This may indicate that the relationalai.zip package structure is incomplete. "
+                    "Please ensure the zip file includes the full clients/resources/snowflake directory structure."
+                ) from e
+            raise
+        return create_resources_instance(
+            config=config,
+            profile=profile,
+            connection=connection,
+            reset_session=reset_session,
+            generation=generation or Generation.V0,
+            dry_run=False,
+            language="rel",
+        )
     elif platform == "local":
-        from relationalai.clients.local import LocalResources
-        return LocalResources(config=config, connection=connection, reset_session=reset_session, generation=generation)
+        from relationalai.clients.resources.snowflake.resources_factory import create_resources_instance
+        return create_resources_instance(
+            config=config,
+            profile=profile,
+            connection=connection,
+            reset_session=reset_session,
+            generation=generation or Generation.V0,
+            dry_run=False,
+            language="rel",
+        )
     else:
         raise Exception(f"Unknown platform: {platform}")
 
@@ -187,15 +223,26 @@ def Provider(
     resources = Resources(profile, config, connection, generation=generation)
     platform = resources.config.get("platform", "snowflake")
     if platform == "azure":
-        import relationalai.clients.azure
-        resources = cast(relationalai.clients.azure.Resources, resources)
-        return relationalai.clients.azure.Provider(
+        from relationalai.clients.resources.azure.azure import Resources as AzureResources, Provider as AzureProvider
+        resources = cast(AzureResources, resources)
+        return AzureProvider(
             resources=resources
         )
     elif platform == "snowflake":
-        import relationalai.clients.snowflake
-        resources = cast(relationalai.clients.snowflake.Resources, resources)
-        return relationalai.clients.snowflake.Provider(
+        try:
+            from relationalai.clients.resources.snowflake import Resources as SnowflakeResources, Provider as SnowflakeProvider
+        except ImportError as e:
+            # Provide a helpful error message for Snowflake notebook environments
+            from relationalai.environments import runtime_env, SnowbookEnvironment
+            if isinstance(runtime_env, SnowbookEnvironment):
+                raise ImportError(
+                    "Failed to import relationalai.clients.resources.snowflake. "
+                    "This may indicate that the relationalai.zip package structure is incomplete. "
+                    "Please ensure the zip file includes the full clients/resources/snowflake directory structure."
+                ) from e
+            raise
+        resources = cast(SnowflakeResources, resources)
+        return SnowflakeProvider(
             resources=resources,
             generation=generation
         )

relationalai/clients/__init__.py

@@ -1,5 +1,18 @@
 # Azure import is dropped because we need to to do runtime import of azure module to support our package running on Snowflake Notebook
-from . import config, snowflake, client, local
+from . import config
+
+# Lazy imports for client and local to avoid circular import issues
+# These modules import from relationalai which imports from clients, creating a cycle
+def __getattr__(name: str):
+    if name == 'client':
+        from . import client
+        return client
+    elif name == 'local':
+        from . import local
+        return local
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
 
 # note: user must do `import relationalai.clients.azure` to get `azure` submodule
-__all__ = ['snowflake', 'config', 'client', 'local']
+# note: snowflake module is now at relationalai.clients.resources.snowflake
+# note: 'client' and 'local' are lazy-loaded via __getattr__, so they're not in __all__
+__all__ = ['config']

relationalai/clients/client.py

@@ -10,12 +10,12 @@ from dataclasses import dataclass
 from pandas import DataFrame
 import time
 
-from relationalai.clients.hash_util import database_name_from_sproc_name
-from relationalai.tools.constants import USE_GRAPH_INDEX, USE_PACKAGE_MANAGER
+from .hash_util import database_name_from_sproc_name
+from ..tools.constants import USE_GRAPH_INDEX, USE_PACKAGE_MANAGER
 
 
 from .types import AvailableModel, EngineState, Import, ImportSource, ImportsStatus, SourceMapEntry
-from ..clients.config import Config
+from .config import Config
 from ..compiler import Compiler
 from ..environments import runtime_env, NotebookRuntimeEnvironment
 from .. import dsl, debugging, metamodel as m
@@ -210,7 +210,7 @@ class ResourcesBase(ABC):
         return engine
 
     @abstractmethod
-    def auto_create_engine_async(self, name: str | None = None):
+    def auto_create_engine_async(self, name: str | None = None) -> str:
         pass
 
     _active_engine: EngineState|None = None

relationalai/clients/local.py

@@ -13,10 +13,10 @@ from typing import Any, Dict, List, Iterable, Literal, Optional, Tuple, Union
 from requests.adapters import HTTPAdapter
 from urllib3.util.retry import Retry
 
-from relationalai.clients.client import ResourcesBase, ProviderBase
-from relationalai.clients.config import Config
-from relationalai.clients.types import TransactionAsyncResponse
-from relationalai.clients.util import get_pyrel_version
+from .client import ResourcesBase, ProviderBase
+from .config import Config
+from .types import TransactionAsyncResponse
+from .util import get_pyrel_version
 from ..errors import ResponseStatusException
 from .. import debugging
 
@@ -202,7 +202,7 @@ class LocalResources(ResourcesBase):
     def alter_engine_pool(self, size: str | None = None, mins: int | None = None, maxs: int | None = None):
         raise NotImplementedError("alter_engine_pool not supported in local mode")
 
-    def auto_create_engine_async(self, name: str | None = None):
+    def auto_create_engine_async(self, name: str | None = None) -> str:
         raise NotImplementedError("auto_create_engine_async not supported in local mode")
 
     #--------------------------------------------------

relationalai/clients/resources/__init__.py

@@ -0,0 +1,8 @@
+"""
+Resources module for RelationalAI clients.
+"""
+
+# This package contains resource implementations for different platforms:
+# - snowflake: Snowflake-specific resources
+# - azure: Azure-specific resources (optional, may not be available in all environments)
+

relationalai/clients/{azure.py → resources/azure/azure.py}

@@ -9,18 +9,18 @@ from urllib.error import HTTPError
 
 from pandas import DataFrame
 
-from relationalai import debugging
-from relationalai.clients.util import poll_with_specified_overhead
-
-from ..errors import EngineNotFoundException, RAIException, AzureUnsupportedQueryTimeoutException
-from ..rel_utils import assert_no_problems
-from ..loaders.loader import emit_delete_import, import_file, list_available_resources
-from .config import Config
-from .types import EngineState, ImportSource, ImportSourceFile, TransactionAsyncResponse
-from .client import Client, ExportParams, ProviderBase, ResourcesBase
-from .. import dsl, rel, metamodel as m
+from .... import debugging
+from ...util import poll_with_specified_overhead
+
+from ....errors import EngineNotFoundException, RAIException, AzureUnsupportedQueryTimeoutException
+from ....rel_utils import assert_no_problems
+from ....loaders.loader import emit_delete_import, import_file, list_available_resources
+from ...config import Config
+from ...types import EngineState, ImportSource, ImportSourceFile, TransactionAsyncResponse
+from ...client import Client, ExportParams, ProviderBase, ResourcesBase
+from .... import dsl, rel, metamodel as m
 from railib import api
-from . import result_helpers
+from ... import result_helpers
 
 #--------------------------------------------------
 # Constants
@@ -105,7 +105,7 @@ class Resources(ResourcesBase):
     def resume_engine_async(self, name:str):
         return api.resume_engine(self._api_ctx(), name)
 
-    def auto_create_engine_async(self, name: str | None = None):
+    def auto_create_engine_async(self, name: str | None = None) -> str:
         raise Exception("Azure doesn't support auto_create_engine_async")
 
     def alter_engine_pool(self, size: str | None = None, mins: int | None = None, maxs: int | None = None):

relationalai/clients/resources/snowflake/__init__.py

@@ -0,0 +1,20 @@
+"""
+Snowflake resources module.
+"""
+# Import order matters - Resources must be imported first since other classes depend on it
+from .snowflake import Resources, Provider, Graph, SnowflakeClient, APP_NAME, PYREL_ROOT_DB, ExecContext, INTERNAL_ENGINE_SIZES, ENGINE_SIZES_AWS, ENGINE_SIZES_AZURE, PrimaryKey
+
+# These imports depend on Resources, so they come after
+from .cli_resources import CLIResources
+from .use_index_resources import UseIndexResources
+from .direct_access_resources import DirectAccessResources
+from .resources_factory import create_resources_instance
+
+__all__ = [
+    'Resources', 'DirectAccessResources', 'Provider', 'Graph', 'SnowflakeClient',
+    'APP_NAME', 'PYREL_ROOT_DB', 'CLIResources', 'UseIndexResources', 'ExecContext',
+    'INTERNAL_ENGINE_SIZES', 'ENGINE_SIZES_AWS', 'ENGINE_SIZES_AZURE', 'PrimaryKey',
+    'create_resources_instance',
+]
+
+

relationalai/clients/resources/snowflake/cli_resources.py

@@ -0,0 +1,87 @@
+"""
+CLI Resources - Resources class for CLI operations.
+Re-raises exceptions without transformation and provides CLI utility methods.
+"""
+from __future__ import annotations
+import json
+from typing import Any
+
+from ....tools.constants import RAI_APP_NAME
+
+# Import Resources from snowflake - this creates a dependency but no circular import
+# since snowflake.py doesn't import from this file
+from .snowflake import Resources, ExecContext
+
+
+class CLIResources(Resources):
+    """
+    Resources class for CLI operations.
+    Re-raises exceptions without transformation and provides CLI utility methods.
+    """
+
+    def _handle_standard_exec_errors(self, e: Exception, ctx: ExecContext) -> Any | None:
+        """For CLI resources, re-raise exceptions without transformation."""
+        raise e
+
+    def list_warehouses(self):
+        """List all warehouses in the Snowflake account."""
+        results = self._exec("SHOW WAREHOUSES")
+        if not results:
+            return []
+        return [{"name": name}
+                for (name, *rest) in results]
+
+    def list_compute_pools(self):
+        """List all compute pools in the Snowflake account."""
+        results = self._exec("SHOW COMPUTE POOLS")
+        if not results:
+            return []
+        return [{"name": name, "status": status, "min_nodes": min_nodes, "max_nodes": max_nodes, "instance_family": instance_family}
+                for (name, status, min_nodes, max_nodes, instance_family, *rest) in results]
+
+    def list_roles(self):
+        """List all available roles in the Snowflake account."""
+        results = self._exec("SELECT CURRENT_AVAILABLE_ROLES()")
+        if not results:
+            return []
+        # the response is a single row with a single column containing
+        # a stringified JSON array of role names:
+        row = results[0]
+        if not row:
+            return []
+        return [{"name": name} for name in json.loads(row[0])]
+
+    def list_apps(self):
+        """List all applications in the Snowflake account."""
+        all_apps = self._exec(f"SHOW APPLICATIONS LIKE '{RAI_APP_NAME}'")
+        if not all_apps:
+            all_apps = self._exec("SHOW APPLICATIONS")
+            if not all_apps:
+                return []
+        return [{"name": name}
+                for (time, name, *rest) in all_apps]
+
+    def list_databases(self):
+        """List all databases in the Snowflake account."""
+        results = self._exec("SHOW DATABASES")
+        if not results:
+            return []
+        return [{"name": name}
+                for (time, name, *rest) in results]
+
+    def list_sf_schemas(self, database: str):
+        """List all schemas in a given database."""
+        results = self._exec(f"SHOW SCHEMAS IN {database}")
+        if not results:
+            return []
+        return [{"name": name}
+                for (time, name, *rest) in results]
+
+    def list_tables(self, database: str, schema: str):
+        """List all tables and views in a given schema."""
+        results = self._exec(f"SHOW OBJECTS IN {database}.{schema}")
+        items = []
+        if results:
+            for (time, name, db_name, schema_name, kind, *rest) in results:
+                items.append({"name": name, "kind": kind.lower()})
+        return items