task: Add query transformer

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(

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)}")

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