relationalai 1.0.0a2__py3-none-any.whl → 1.0.0a4__py3-none-any.whl

v0/relationalai/tools/cli.py

@@ -27,7 +27,9 @@ from v0.relationalai.loaders.loader import Loader, rel_schema_to_type
 from ..clients.types import ImportSource, ImportSourceFile, ImportSourceTable
 from ..clients.client import ResourcesBase
 from ..tools import debugger as deb, qb_debugger as qb_deb
-from ..clients import config, snowflake
+from ..clients import config
+from v0.relationalai.tools.constants import RAI_APP_NAME
+from v0.relationalai.clients.resources.snowflake.cli_resources import CLIResources
 from v0.relationalai.tools.cli_helpers import (
     EMPTY_STRING_REGEX,
     ENGINE_NAME_ERROR,
@@ -337,7 +339,7 @@ def check_snowflake_connections_flow(cfg:config.Config):
     return True
 
 def role_flow(provider:ResourcesBase, cfg:config.Config):
-    roles = cast(snowflake.Resources, provider).list_roles()
+    roles = cast(CLIResources, provider).list_roles()
     result = controls.fuzzy_with_refetch(
             "Select a role:",
             "roles",
@@ -355,7 +357,7 @@ def warehouse_flow(provider:ResourcesBase, cfg:config.Config):
     result = controls.fuzzy_with_refetch(
             "Select a warehouse:",
             "warehouses",
-            lambda: [w["name"] for w in cast(snowflake.Resources, provider).list_warehouses()],
+            lambda: [w["name"] for w in cast(CLIResources, provider).list_warehouses()],
             default=cfg.get("warehouse", None),
         )
     if not result or isinstance(result, Exception):
@@ -367,12 +369,12 @@ def warehouse_flow(provider:ResourcesBase, cfg:config.Config):
 def rai_app_flow(provider:ResourcesBase, cfg:config.Config):
     auto_select = None
     if provider.config.get("platform") == "snowflake":
-        auto_select=snowflake.RAI_APP_NAME
+        auto_select=RAI_APP_NAME
 
     result = controls.fuzzy_with_refetch(
             "Select RelationalAI app name:",
             "apps",
-            lambda: [w["name"] for w in cast(snowflake.Resources, provider).list_apps()],
+            lambda: [w["name"] for w in cast(CLIResources, provider).list_apps()],
             default=cfg.get("rai_app_name", None),
             auto_select=auto_select
         )
@@ -454,7 +456,8 @@ def name_profile_flow(cfg: config.Config):
         )
         return name_profile_flow(cfg)
     config_store = ConfigStore()
-    if profile in config_store.get_profiles():
+    profiles = config_store.get_profiles() or {}
+    if profile in profiles:
         overwrite = controls.confirm(f"[yellow]Overwrite existing {profile} profile?")
         if overwrite:
             return profile
@@ -464,7 +467,8 @@ def name_profile_flow(cfg: config.Config):
 
 def save_flow(cfg:config.Config):
     config_store = ConfigStore()
-    if cfg.profile != PARTIAL_PROFILE_NAME and cfg.profile in config_store.get_profiles():
+    profiles = config_store.get_profiles() or {}
+    if cfg.profile != PARTIAL_PROFILE_NAME and cfg.profile in profiles:
         if not controls.confirm(f"Overwrite existing {cfg.profile} profile"):
             rich.print()
             profile_name = controls.text("Profile name:")
@@ -570,7 +574,7 @@ def profile_switch(profile:str|None=None):
         rich.print(f"\n[yellow]Error: {e}")
         pass
 
-    profiles = list(config_store.get_profiles().keys())
+    profiles = list((config_store.get_profiles() or {}).keys())
     divider()
     if not profile:
         if len(profiles) == 0:
@@ -672,7 +676,8 @@ def config_explain(profile:str|None=None, all_profiles:bool=False):
             ))
 
     if all_profiles:
-        for profile, props in config_store.get_profiles().items():
+        profiles = config_store.get_profiles() or {}
+        for profile, props in profiles.items():
             if profile == cfg.profile:
                 continue
             if len(props):
@@ -872,7 +877,6 @@ def engines_get(name:str|None=None):
 def create_engine_flow(cfg:config.Config, name=None, size=None, auto_suspend_mins=None):
     provider = get_resource_provider(None, cfg)
     engine = None
-    create_exception = None
     is_engine_present = False
     is_engine_suspended = False
     is_interactive = name is None or size is None
@@ -957,15 +961,7 @@ def create_engine_flow(cfg:config.Config, name=None, size=None, auto_suspend_min
         f"Engine '{name}' created!",
         failed_message="Error:"
     ):
-        try:
-            provider.create_engine(name, size, auto_suspend_mins)
-        except Exception as e:
-            create_exception = e
-            # We do not want to print the success context message if the engine creation fails
-            # Since we are raising here and passing a "fail_message", the spinner will print the actual error message
-            raise e
-    if isinstance(create_exception, Exception):
-        exit_with_error(f"[yellow]{create_exception}")
+        provider.create_engine(name, size, auto_suspend_mins)
     return name
 
 @cli.command(name="engines:create", help="Create a new engine")
@@ -1164,19 +1160,19 @@ def engines_alter_pool(size:str|None=None, min:int|None=None, max:int|None=None)
 def import_source_flow(provider: ResourcesBase) -> Sequence[ImportSource]:
     provider_type = type(provider)
 
-    if isinstance(provider, snowflake.Resources):
+    if isinstance(provider, CLIResources):
         return snowflake_import_source_flow(provider)
     else:
         # Lazy import for azure to avoid optional dependency issues
         try:
-            from v0.relationalai.clients import azure
-            if isinstance(provider, azure.Resources):
+            from v0.relationalai.clients.resources.azure.azure import Resources as AzureResources
+            if isinstance(provider, AzureResources):
                 return azure_import_source_flow(provider)
         except ImportError:
             pass
         raise Exception(f"No import source flow available for {provider_type.__module__}.{provider_type.__name__}")
 
-def snowflake_import_source_flow(provider: snowflake.Resources) -> Sequence[ImportSource]:
+def snowflake_import_source_flow(provider: CLIResources) -> Sequence[ImportSource]:
     with Spinner("Fetching databases", "Databases fetched"):
         try:
             dbs = provider.list_databases()
@@ -1227,7 +1223,7 @@ def azure_import_source_flow(provider: azure.Resources) -> Sequence[ImportSource
 
 def import_source_options_flow(provider: ResourcesBase, source: ImportSource, default_options:dict) -> dict:
     if isinstance(source, ImportSourceFile):
-        type: LoadType = default_options.get("type", None)
+        type: LoadType | None = default_options.get("type", None)
         if type is None or type == "auto":
             type = Loader.get_type_for(source)
         if type == "csv":
@@ -1329,7 +1325,7 @@ def parse_source(provider: ResourcesBase, raw: str) -> ImportSource:
 def imports_setup(engine_size:str|None=None, resume:bool=False, suspend:bool=False):
     divider(flush=True)
     ensure_config()
-    provider = get_resource_provider()
+    provider = cast(CLIResources, get_resource_provider())
     data = None
 
     if resume or suspend:
@@ -1394,7 +1390,7 @@ def imports_setup(engine_size:str|None=None, resume:bool=False, suspend:bool=Fal
 def imports_get(id:str|None=None):
     divider(flush=True)
     ensure_config()
-    provider = get_resource_provider()
+    provider = cast(CLIResources, get_resource_provider())
     import_streams = []
     import_response = []
     with Spinner("Fetching imports", "Imports fetched", "Error:"):
@@ -1567,7 +1563,7 @@ def imports_stream(
 ):
     divider(flush=True)
     ensure_config()
-    provider = get_resource_provider()
+    provider = cast(CLIResources, get_resource_provider())
     default_options = ImportOptionsType.reduce(options)
 
     # Resume or suspend import stream
@@ -1713,7 +1709,7 @@ def imports_snapshot(source:str|None, model:str|None, name:str|None, type:str|No
 def imports_delete(object, model, force):
     divider(flush=True)
     ensure_config()
-    provider = cast(snowflake.Resources, get_resource_provider())
+    provider = cast(CLIResources, get_resource_provider())
     if not model:
         with Spinner("Fetching models", "Models fetched"):
             try:
@@ -1768,7 +1764,7 @@ def imports_delete(object, model, force):
 def exports_list(model):
     divider(flush=True)
     ensure_config()
-    provider = cast(snowflake.Resources, get_resource_provider())
+    provider = cast(CLIResources, get_resource_provider())
     coming_soon()
     if not model:
         with Spinner("Fetching models", "Models fetched"):
@@ -1810,7 +1806,7 @@ def exports_list(model):
 def exports_delete(export, model):
     divider(flush=True)
     ensure_config()
-    provider = cast(snowflake.Resources, get_resource_provider())
+    provider = cast(CLIResources, get_resource_provider())
     coming_soon()
     if not model:
         with Spinner("Fetching models", "Models fetched"):

v0/relationalai/tools/cli_helpers.py

@@ -11,7 +11,6 @@ import functools
 import pytz
 
 from v0.relationalai.util.constants import TOP_LEVEL_PROFILE_NAME
-from .. import Resources
 from rich.table import Table
 from typing import Callable, Dict, Any
 from ..clients import config
@@ -39,7 +38,16 @@ def get_config(profile:str|Dict[str, Any]|None = None):
 @functools.cache
 def get_resource_provider(platform:str|None=None, _cfg:config.Config|None = None) -> ResourcesBase:
     cfg = _cfg or get_config()
-    provider = Resources(config=cfg)
+    platform = platform or cfg.get("platform", "snowflake")
+    if platform == "snowflake":
+        from v0.relationalai.clients.resources.snowflake.cli_resources import CLIResources
+        provider = CLIResources(config=cfg)
+    elif platform == "azure":
+        from v0.relationalai.clients.resources.azure.azure import Resources
+        provider = Resources(config=cfg)
+    else:
+        from .. import Resources
+        provider = Resources(config=cfg)
     return provider
 
 def unexpand_user_path(path):

v0/relationalai/util/otel_configuration.py

@@ -1,7 +1,8 @@
 from typing import cast
 
 from v0.relationalai.clients.config import Config
-from v0.relationalai.clients.snowflake import Resources, ResourcesBase
+from v0.relationalai.clients.resources.snowflake import Resources
+from v0.relationalai.clients.client import ResourcesBase
 from v0.relationalai.environments import runtime_env, SnowbookEnvironment
 
 def configure_otel(enable_otel_handler: bool, config: Config, resources: ResourcesBase):

v0/relationalai/util/otel_handler.py

@@ -17,7 +17,7 @@ from v0.relationalai.util.span_tracker import record_is_span, get_span_path_as_s
 from v0.relationalai.util.span_tracker import get_span_allowed_attributes_values, get_span_value_from_key
 from v0.relationalai.debugging import logger, Span, filter_span_attrs, otel_traceid_to_opentracing, get_current_span
 
-from v0.relationalai.clients.snowflake import Resources
+from v0.relationalai.clients.resources.snowflake import Resources
 
 MAX_PAYLOAD_SIZE = 25*1024 # 25KB
 MAX_ATTRIBUTE_LENGTH = 1000