diff --git a/application_sdk/transformers/atlas/__init__.py b/application_sdk/transformers/atlas/__init__.py index ca6653868..3348eb687 100644 --- a/application_sdk/transformers/atlas/__init__.py +++ b/application_sdk/transformers/atlas/__init__.py @@ -52,6 +52,7 @@ def __init__(self, connector_name: str, tenant_id: str, **kwargs: Any): Database, Function, Procedure, + Query, Schema, Table, TagAttachment, @@ -70,6 +71,7 @@ def __init__(self, connector_name: str, tenant_id: str, **kwargs: Any): "FUNCTION": Function, "TAG_REF": TagAttachment, "PROCEDURE": Procedure, + "QUERY": Query, } def transform_metadata( diff --git a/application_sdk/transformers/atlas/sql.py b/application_sdk/transformers/atlas/sql.py index 03f9c558b..60de940f9 100644 --- a/application_sdk/transformers/atlas/sql.py +++ b/application_sdk/transformers/atlas/sql.py @@ -1207,3 +1207,80 @@ def get_attributes(cls, obj: Dict[str, Any]) -> Dict[str, Any]: } except Exception as e: raise ValueError(f"Error creating TagAttachment Entity: {str(e)}") + + +class Query(assets.Query): + """Query entity transformer for Atlas. + + This class handles the transformation of query metadata into Atlas Query entities. + """ + + @classmethod + def get_attributes(cls, obj: Dict[str, Any]) -> Dict[str, Any]: + """Parse a dictionary into a Query entity. + + Args: + obj (Dict[str, Any]): Dictionary containing query metadata. + + Returns: + Dict[str, Any]: Dictionary with 'attributes', 'custom_attributes', and 'entity_class'. + + Raises: + ValueError: If required fields are missing or invalid. + """ + try: + # Required fields from entity spec + assert ( + obj.get("parentQualifiedName") is not None + ), "parentQualifiedName is required" + assert ( + obj.get("collectionQualifiedName") is not None + ), "collectionQualifiedName is required" + + attributes = {} + custom_attributes = {} + + # Map fields from input to Query entity attributes + attributes["longRawQuery"] = obj.get("querytxt") + attributes["rawQueryText"] = obj.get("querytxt") + attributes["defaultDatabaseQualifiedName"] = obj.get("dbname") + attributes["defaultSchemaQualifiedName"] = obj.get("schemaname") + attributes["parentQualifiedName"] = obj["parentQualifiedName"] + attributes["collectionQualifiedName"] = obj["collectionQualifiedName"] + attributes["isPrivate"] = obj.get("isPrivate") + attributes["isSqlSnippet"] = obj.get("isSqlSnippet") + attributes["isVisualQuery"] = obj.get("isVisualQuery") + attributes["visualBuilderSchemaBase64"] = obj.get( + "visualBuilderSchemaBase64" + ) + attributes["variablesSchemaBase64"] = obj.get("variablesSchemaBase64") + + # Deprecated field for backward compatibility + attributes["rawQuery"] = obj.get("rawQuery") + + # Place username and sessionid in custom_attributes + if obj.get("username"): + custom_attributes["username"] = obj["username"] + if obj.get("sessionid"): + custom_attributes["sessionid"] = obj["sessionid"] + if obj.get("starttime"): + custom_attributes["starttime"] = obj["starttime"] + if obj.get("endtime"): + custom_attributes["endtime"] = obj["endtime"] + if obj.get("session_starttime"): + custom_attributes["session_starttime"] = obj["session_starttime"] + if obj.get("session_endtime"): + custom_attributes["session_endtime"] = obj["session_endtime"] + + # Add any other fields not mapped above to custom_attributes + for k, v in obj.items(): + if k not in attributes and k not in custom_attributes: + custom_attributes[k] = v + + return { + "attributes": attributes, + "custom_attributes": custom_attributes, + "entity_class": Query, + } + except AssertionError as e: + raise ValueError(f"Error creating Query Entity: {str(e)}") diff --git a/application_sdk/transformers/query/templates/query.yaml b/application_sdk/transformers/query/templates/query.yaml new file mode 100644 index 000000000..49f1ed4f5 --- /dev/null +++ b/application_sdk/transformers/query/templates/query.yaml @@ -0,0 +1,161 @@ + # Query entity template for Atlas +# This template defines the structure for Query entities, including all relevant attributes. +# Follows the conventions of other templates in this directory. + +name: Query +superTypes: + - SQL +serviceType: atlan +description: Instance of a query in Atlan. +typeVersion: 1.2 + +attributeDefs: + - name: rawQuery + description: Deprecated. See 'longRawQuery' instead. + typeName: string + isOptional: true + cardinality: SINGLE + valuesMinCount: 0 + valuesMaxCount: 1 + isUnique: false + isIndexable: true + includeInNotification: true + - name: longRawQuery + description: Raw SQL query string. + typeName: string + isOptional: true + cardinality: SINGLE + valuesMinCount: 0 + valuesMaxCount: 1 + isUnique: false + isIndexable: false + includeInNotification: true + - name: rawQueryText + description: '' + typeName: string + isOptional: true + cardinality: SINGLE + isUnique: false + isIndexable: false + includeInNotification: true + indexTypeESConfig: + analyzer: truncate_analyzer + - name: defaultSchemaQualifiedName + description: Unique name of the default schema to use for this query. + typeName: string + indexType: STRING + isOptional: true + cardinality: SINGLE + valuesMinCount: 0 + valuesMaxCount: 1 + isUnique: false + skipScrubbing: true + isIndexable: true + includeInNotification: true + indexTypeESFields: + text: + type: text + analyzer: atlan_text_analyzer + - name: defaultDatabaseQualifiedName + description: Unique name of the default database to use for this query. + typeName: string + indexType: STRING + isOptional: true + cardinality: SINGLE + valuesMinCount: 0 + valuesMaxCount: 1 + isUnique: false + skipScrubbing: true + isIndexable: true + includeInNotification: true + indexTypeESFields: + text: + type: text + analyzer: atlan_text_analyzer + - name: variablesSchemaBase64 + description: Base64-encoded string of the variables to use in this query. + typeName: string + isOptional: true + cardinality: SINGLE + valuesMinCount: 0 + valuesMaxCount: 1 + isUnique: false + isIndexable: true + skipScrubbing: true + includeInNotification: true + - name: isPrivate + description: Whether this query is private (true) or shared (false). + typeName: boolean + isOptional: true + cardinality: SINGLE + valuesMinCount: 0 + valuesMaxCount: 1 + isUnique: false + skipScrubbing: true + isIndexable: true + includeInNotification: true + - name: isSqlSnippet + description: Whether this query is a SQL snippet (true) or not (false). + typeName: boolean + isOptional: true + cardinality: SINGLE + valuesMinCount: 0 + valuesMaxCount: 1 + isUnique: false + isIndexable: true + skipScrubbing: true + includeInNotification: true + - name: parentQualifiedName + description: Unique name of the parent collection or folder in which this query exists. + typeName: string + indexType: STRING + isOptional: false + cardinality: SINGLE + valuesMinCount: 1 + valuesMaxCount: 1 + isUnique: false + isIndexable: true + skipScrubbing: true + includeInNotification: true + indexTypeESFields: + text: + type: text + analyzer: atlan_text_analyzer + - name: collectionQualifiedName + description: Unique name of the collection in which this query exists. + typeName: string + indexType: STRING + isOptional: false + cardinality: SINGLE + valuesMinCount: 1 + valuesMaxCount: 1 + isUnique: false + skipScrubbing: true + isIndexable: true + includeInNotification: true + indexTypeESFields: + text: + type: text + analyzer: atlan_text_analyzer + - name: isVisualQuery + description: Whether this query is a visual query (true) or not (false). + typeName: boolean + isOptional: true + cardinality: SINGLE + valuesMinCount: 0 + valuesMaxCount: 1 + isUnique: false + isIndexable: true + skipScrubbing: true + includeInNotification: true + - name: visualBuilderSchemaBase64 + description: Base64-encoded string for the visual query builder. + typeName: string + isOptional: true + cardinality: SINGLE + valuesMinCount: 0 + valuesMaxCount: 1 + isUnique: false + isIndexable: false + skipScrubbing: true + includeInNotification: true