relationalai 0.11.2__py3-none-any.whl → 0.11.4__py3-none-any.whl

relationalai/clients/snowflake.py

@@ -2048,7 +2048,7 @@ Otherwise, remove it from your '{profile}' configuration profile.
             assert len(parsed) == 4, f"Invalid source: {source}"
             db, schema, entity, identity = parsed
             assert db and schema and entity and identity, f"Invalid source: {source}"
-            source_types[identity] = cast(SourceInfo, {"type": None, "state": ""})
+            source_types[identity] = cast(SourceInfo, {"type": None, "state": "", "columns_hash": None})
             partitioned_sources[db][schema].append(entity)
 
         # TODO: Move to NA layer
@@ -2057,6 +2057,7 @@ Otherwise, remove it from your '{profile}' configuration profile.
                 f"""SELECT
                     inf.FQN,
                     inf.KIND,
+                    inf.COLUMNS_HASH,
                     IFF(DATEDIFF(second, ds.created_at::TIMESTAMP, inf.LAST_DDL::TIMESTAMP) > 0, 'STALE', 'CURRENT') AS STATE
                 FROM (
                     SELECT (SELECT {app_name}.api.normalize_fq_ids(ARRAY_CONSTRUCT(FQ_OBJECT_NAME))[0]:identifier::string) as FQ_OBJECT_NAME,
@@ -2068,26 +2069,45 @@ Otherwise, remove it from your '{profile}' configuration profile.
                         (SELECT {app_name}.api.normalize_fq_ids(
                             ARRAY_CONSTRUCT(
                                 CASE
-                                    WHEN TABLE_CATALOG = UPPER(TABLE_CATALOG) THEN TABLE_CATALOG
-                                    ELSE '"' || TABLE_CATALOG || '"'
+                                    WHEN t.TABLE_CATALOG = UPPER(t.TABLE_CATALOG) THEN t.TABLE_CATALOG
+                                    ELSE '"' || t.TABLE_CATALOG || '"'
                                 END || '.' ||
                                 CASE
-                                    WHEN TABLE_SCHEMA = UPPER(TABLE_SCHEMA) THEN TABLE_SCHEMA
-                                    ELSE '"' || TABLE_SCHEMA || '"'
+                                    WHEN t.TABLE_SCHEMA = UPPER(t.TABLE_SCHEMA) THEN t.TABLE_SCHEMA
+                                    ELSE '"' || t.TABLE_SCHEMA || '"'
                                 END || '.' ||
                                 CASE
-                                    WHEN TABLE_NAME = UPPER(TABLE_NAME) THEN TABLE_NAME
-                                    ELSE '"' || TABLE_NAME || '"'
+                                    WHEN t.TABLE_NAME = UPPER(t.TABLE_NAME) THEN t.TABLE_NAME
+                                    ELSE '"' || t.TABLE_NAME || '"'
                                 END
                             )
                         )[0]:identifier::string) as FQN,
                         CONVERT_TIMEZONE('UTC', LAST_DDL) AS LAST_DDL,
-                        TABLE_TYPE as KIND
-                    FROM {db}.INFORMATION_SCHEMA.TABLES
-                    WHERE TABLE_CATALOG = {IdentityParser.to_sql_value(db)} AND ({" OR ".join(
-                        f"(TABLE_SCHEMA = {IdentityParser.to_sql_value(schema)} AND TABLE_NAME IN ({','.join(f'{IdentityParser.to_sql_value(table)}' for table in tables)}))"
+                        TABLE_TYPE as KIND,
+                        SHA2(LISTAGG(
+                            COLUMN_NAME ||
+                            CASE
+                                WHEN c.NUMERIC_PRECISION IS NOT NULL AND c.NUMERIC_SCALE IS NOT NULL
+                                    THEN c.DATA_TYPE || '(' || c.NUMERIC_PRECISION || ',' || c.NUMERIC_SCALE || ')'
+                                WHEN c.DATETIME_PRECISION IS NOT NULL
+                                    THEN c.DATA_TYPE || '(0,' || c.DATETIME_PRECISION || ')'
+                                WHEN c.CHARACTER_MAXIMUM_LENGTH IS NOT NULL
+                                    THEN c.DATA_TYPE || '(' || c.CHARACTER_MAXIMUM_LENGTH || ')'
+                                ELSE c.DATA_TYPE
+                            END ||
+                            IS_NULLABLE,
+                            ','
+                        ) WITHIN GROUP (ORDER BY COLUMN_NAME), 256) as COLUMNS_HASH
+                    FROM {db}.INFORMATION_SCHEMA.TABLES t
+                    JOIN {db}.INFORMATION_SCHEMA.COLUMNS c
+                        ON t.TABLE_CATALOG = c.TABLE_CATALOG
+                        AND t.TABLE_SCHEMA = c.TABLE_SCHEMA
+                        AND t.TABLE_NAME = c.TABLE_NAME
+                    WHERE t.TABLE_CATALOG = {IdentityParser.to_sql_value(db)} AND ({" OR ".join(
+                        f"(t.TABLE_SCHEMA = {IdentityParser.to_sql_value(schema)} AND t.TABLE_NAME IN ({','.join(f'{IdentityParser.to_sql_value(table)}' for table in tables)}))"
                         for schema, tables in schemas.items()
                     )})
+                    GROUP BY t.TABLE_CATALOG, t.TABLE_SCHEMA, t.TABLE_NAME, t.LAST_DDL, t.TABLE_TYPE
                 ) inf on inf.FQN = ds.FQ_OBJECT_NAME
             """
                 for db, schemas in partitioned_sources.items()
@@ -2102,6 +2122,7 @@ Otherwise, remove it from your '{profile}' configuration profile.
             assert fqn, f"Error parsing returned FQN: {row_fqn}"
 
             source_types[fqn]["type"] = "TABLE" if row["KIND"] == "BASE TABLE" else row["KIND"]
+            source_types[fqn]["columns_hash"] = row["COLUMNS_HASH"]
             source_types[fqn]["state"] = row["STATE"]
 
         return source_types
@@ -3058,10 +3079,13 @@ class DirectAccessResources(Resources):
         return self._retrieve_service_endpoint()
 
     def _retrieve_service_endpoint(self, enforce_update=False) -> str:
+        account = self.config.get("account")
+        app_name = self.config.get("rai_app_name")
+        service_endpoint_key = f"{account}.{app_name}.service_endpoint"
         if self._service_endpoint and not enforce_update:
             return self._service_endpoint
-        if self._endpoint_info.get("service_endpoint", "") and not enforce_update:
-            self._service_endpoint = str(self._endpoint_info.get("service_endpoint", ""))
+        if self._endpoint_info.get(service_endpoint_key, "") and not enforce_update:
+            self._service_endpoint = str(self._endpoint_info.get(service_endpoint_key, ""))
             return self._service_endpoint
 
         is_snowflake_notebook = isinstance(runtime_env, SnowbookEnvironment)
@@ -3073,7 +3097,7 @@ class DirectAccessResources(Resources):
         else:
             self._service_endpoint = f"https://{result[0]['SERVICE_ENDPOINT']}"
 
-        self._endpoint_info.set("service_endpoint", self._service_endpoint)
+        self._endpoint_info.set(service_endpoint_key, self._service_endpoint)
         # save the endpoint to `ENDPOINT_FILE` to avoid calling the endpoint with every
         # pyrel execution
         try:
@@ -3116,7 +3140,12 @@ class DirectAccessResources(Resources):
             try:
                 response = _send_request()
                 if response.status_code != 200:
-                    message = response.json().get("message", "")
+                    try:
+                        message = response.json().get("message", "")
+                    except requests.exceptions.JSONDecodeError:
+                        raise ResponseStatusException(
+                            f"Failed to parse error response from endpoint {endpoint}.", response
+                        )
 
                     # fix engine on engine error and retry
                     if _is_engine_issue(message):

relationalai/clients/types.py

@@ -41,6 +41,7 @@ class EngineState(TypedDict):
 class SourceInfo(TypedDict):
     type: str|None
     state: str
+    columns_hash: str|None
     source: str