search.py

#
#  Copyright (C) 2017-2025 Dremio Corporation
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#
import asyncio

from pydantic import (
    BaseModel,
    Field,
    ConfigDict,
    field_validator,
)
from typing import (
    Any,
    List,
    Union,
    Optional, Dict,
)
from dremioai.api.util import UStrEnum
from datetime import datetime
from enum import auto
from dremioai.config import settings
from dremioai.api.transport import DremioAsyncHttpClient as AsyncHttpClient
from dremioai.api.dremio.catalog import get_schemas, get_schema
from dremioai import log
import pandas as pd

logger = log.logger(__name__)


class QueryType(UStrEnum):
    UI_RUN = auto()
    UI_PREVIEW = auto()
    UI_INTERNAL_PREVIEW = auto()
    UI_INTERNAL_RUN = auto()
    UI_EXPORT = auto()
    ODBC = auto()
    JDBC = auto()
    REST = auto()
    ACCELERATOR_CREATE = auto()
    ACCELERATOR_DROP = auto()
    UNKNOWN = auto()
    PREPARE_INTERNAL = auto()
    ACCELERATOR_EXPLAIN = auto()
    UI_INITIAL_PREVIEW = auto()
    FLIGHT = auto()
    METADATA_REFRESH = auto()
    INTERNAL_ICEBERG_METADATA_DROP = auto()
    D2D = auto()
    ACCELERATOR_OPTIMIZE = auto()
    COPY_ERRORS_PLAN = auto()


class JobStatus(UStrEnum):
    NOT_SUBMITTED = auto()
    STARTING = auto()
    RUNNING = auto()
    COMPLETED = auto()
    CANCELED = auto()
    FAILED = auto()
    CANCELLATION_REQUESTED = auto()
    ENQUEUED = auto()
    PLANNING = auto()
    PENDING = auto()
    METADATA_RETRIEVAL = auto()
    QUEUED = auto()
    ENGINE_START = auto()
    EXECUTION_PLANNING = auto()
    INVALID_STATE = auto()


class Category(UStrEnum):
    JOB = auto()
    VIEW = auto()
    TABLE = auto()
    FOLDER = auto()
    UDF = auto()
    SPACE = auto()
    REFLECTION = auto()
    SCRIPT = auto()
    SOURCE = auto()


class UserOrRole(UStrEnum):
    UNSPECIFIED = auto()
    USER = auto()
    ROLE = auto()


class EnterpriseDatasetType(UStrEnum):
    TABLE = auto()
    VIEW = auto()


class EnterpriseSearchUserOrRoleObject(BaseModel):
    id: Optional[str] = None
    type: Optional[UserOrRole] = None
    username: Optional[str] = None
    rolename: Optional[str] = None


class EnterpriseDatasetObject(BaseModel):
    type: Optional[str] = Field(None, alias="datasetType")
    path: Optional[List[str]] = Field(None, alias="datasetPath")


class EnterpriseSearchJobObject(BaseModel):
    id: Optional[str] = None
    queried_ds: Optional[List[EnterpriseDatasetObject]] = Field(
        None, alias="queriedDatasets"
    )
    sql: Optional[str] = None
    job_type: Optional[QueryType] = Field(None, alias="jobType")
    user: Optional[EnterpriseSearchUserOrRoleObject] = None
    start_time: Optional[datetime] = Field(None, alias="startTime")
    finish_time: Optional[datetime] = Field(None, alias="finishTime")
    job_status: Optional[JobStatus] = Field(None, alias="jobStatus")
    error: Optional[str] = None


class EnterpriseSearchScriptObject(BaseModel):
    id: Optional[str] = None
    name: Optional[str] = None
    owner: Optional[EnterpriseSearchUserOrRoleObject] = None
    description: Optional[str] = None
    content: Optional[str] = None
    created_at: Optional[datetime] = Field(default=None, alias="createdAt")
    modified_at: Optional[datetime] = Field(default=None, alias="lastModifiedAt")


class EnterpriseSearchReflectionObject(BaseModel):
    id: Optional[str] = None
    name: Optional[str] = None
    type: Optional[str] = Field(default=None, alias="datasetType")
    path: Optional[List[str]] = Field(default=None, alias="datasetPath")
    branch: Optional[str] = Field(default=None, alias="datasetBranch")
    created_at: Optional[datetime] = Field(default=None, alias="createdAt")
    modified_at: Optional[datetime] = Field(default=None, alias="lastModifiedAt")


class EnterpriseSearchCatalogObject(BaseModel):
    path: Optional[List[str]] = None
    sub_paths: Optional[List[str]] = Field(default=None, alias="subPaths")
    type: Optional[str] = None
    branch: Optional[str] = None
    labels: Optional[List[str]] = None
    wiki: Optional[str] = None
    created_at: Optional[datetime] = Field(default=None, alias="createdAt")
    modified_at: Optional[datetime] = Field(default=None, alias="lastModifiedAt")
    func_sql: Optional[str] = Field(default=None, alias="functionSql")
    owner: Optional[EnterpriseSearchUserOrRoleObject] = None
    schema: Optional[Dict[str, Any]] = None

    async def populate_schemas(self):
        if self.path:
            try:
                data = await get_schema(dataset_path_or_id=self.path, include_tags=True, flatten=True)
                self.schema = data.get('schema') if data else None
            except Exception as e:
                logger.error("Schema not found for search",
                             path=self.path, reason=str(e))

    def as_df_dict(self):
        return {
            "path": self.path,
            "name": ".".join(f'"{p}"' for p in self.path),
            "type": self.type,
            "tags": ",".join(self.labels),
            "description": self.wiki,
            "schema": self.schema,
        }


class EnterpriseSearchResultsObject(BaseModel):
    category: Optional[Category] = None
    job: Optional[EnterpriseSearchJobObject] = Field(default=None, alias="jobObject")
    script: Optional[EnterpriseSearchScriptObject] = Field(
        default=None, alias="scriptObject"
    )
    reflection: Optional[EnterpriseSearchReflectionObject] = Field(
        default=None, alias="reflectionObject"
    )
    catalog: Optional[EnterpriseSearchCatalogObject] = Field(
        default=None, alias="catalogObject"
    )


class EnterpriseSearchResults(BaseModel):
    session_id: Optional[str] = Field(default=None, alias="sessionId")
    next_page_token: Optional[str] = Field(default=None, alias="nextPageToken")
    results: Optional[List[EnterpriseSearchResultsObject]] = Field(default_factory=list)
    error: Optional[str] = Field(default=None, alias="errorMessage")
    more: Optional[str] = Field(default=None, alias="moreInfo")


class Search(BaseModel):
    max_results: Optional[int] = Field(default=50, alias="maxResults")
    next_page_token: Optional[str] = Field(default=None, alias="pageToken")
    filter: Optional[Union[str, List[Category]]] = ""
    query: str = None

    @field_validator("filter", mode="after")
    @classmethod
    def validate_filter(cls, v: Union[str, List[Category]]) -> str:
        if isinstance(v, str) and v:
            v = f'category in ["{Category[v.upper()].name}"]'
        elif isinstance(v, list):
            v = ",".join([f'"{c.name}"' for c in v if isinstance(c, Category)])
            v = f"category in [{v}]"
        else:
            v = ""
        return v

    model_config: ConfigDict = ConfigDict(serialize_by_alias=True)


class EnterpriseSearchResultsWrapper(BaseModel):
    results: List[EnterpriseSearchResultsObject] = Field(default_factory=list)


async def get_search_results(
    search: str | Search, use_df: bool = False,
    remove_catalog_name: Optional[bool] = True
) -> EnterpriseSearchResultsWrapper | pd.DataFrame:
    if isinstance(search, str):
        search = Search(query=search)

    client = AsyncHttpClient()
    endpoint = (
        f"/v0/projects/{settings.instance().dremio.project_id}/search"
        if settings.instance().dremio.project_id
        else "/api/v3/search"
    )

    params = {"removeCatalogName": str(remove_catalog_name).lower()}

    result = []
    response = await client.post(
        endpoint,
        body=search.model_dump(exclude_none=True),
        deser=EnterpriseSearchResults,
        params=params,
    )
    while response.results and response.error is None and response.more is None:
        result.extend(response.results)
        if response.next_page_token is None:
            break
        search.next_page_token = response.next_page_token
        response = await client.post(
            endpoint,
            body=search.model_dump(exclude_none=True),
            deser=EnterpriseSearchResults,
            params=params,
        )
    result = [r for r in result if r.category in (Category.TABLE, Category.VIEW)]
    tasks = [r.catalog.populate_schemas() for r in result]
    await asyncio.gather(*tasks, return_exceptions=True)

    if use_df:
        data = [r.catalog.as_df_dict() for r in result]
        df = pd.DataFrame(data=data)
        if not df.empty:
            df["not_found"] = df.schema.isna()
        return df
    return EnterpriseSearchResultsWrapper(results=result)