openmetadata-ingestion 1.6.3.2.dev0__py3-none-any.whl → 1.6.4.0__py3-none-any.whl

metadata/ingestion/source/database/snowflake/lineage.py

@@ -47,6 +47,8 @@ class SnowflakeLineageSource(
         AND (
             QUERY_TYPE IN ('MERGE', 'UPDATE','CREATE_TABLE_AS_SELECT')
             OR (QUERY_TYPE = 'INSERT' and query_text ILIKE '%%insert%%into%%select%%')
+            OR (QUERY_TYPE = 'ALTER' and query_text ILIKE '%%alter%%table%%swap%%')
+            OR (QUERY_TYPE = 'CREATE_TABLE' and query_text ILIKE '%%clone%%')
         )
     """
 
@@ -60,6 +62,7 @@ class SnowflakeLineageSource(
         start, _ = get_start_and_end(self.source_config.queryLogDuration)
         query = self.stored_procedure_query.format(
             start_date=start,
+            account_usage=self.service_connection.accountUsageSchema,
         )
         queries_dict = self.procedure_queries_dict(
             query=query,

metadata/ingestion/source/database/snowflake/metadata.py

@@ -418,6 +418,7 @@ class SnowflakeSource(
                     SNOWFLAKE_FETCH_ALL_TAGS.format(
                         database_name=self.context.get().database,
                         schema_name=schema_name,
+                        account_usage=self.service_connection.accountUsageSchema,
                     )
                 )
 
@@ -431,6 +432,7 @@ class SnowflakeSource(
                         SNOWFLAKE_FETCH_ALL_TAGS.format(
                             database_name=f'"{self.context.get().database}"',
                             schema_name=f'"{self.context.get().database_schema}"',
+                            account_usage=self.service_connection.accountUsageSchema,
                         )
                     )
                 except Exception as inner_exc:
@@ -635,6 +637,7 @@ class SnowflakeSource(
             query.format(
                 database_name=self.context.get().database,
                 schema_name=self.context.get().database_schema,
+                account_usage=self.service_connection.accountUsageSchema,
             )
         ).all()
         for row in results:

metadata/ingestion/source/database/snowflake/queries.py

@@ -25,7 +25,7 @@ SNOWFLAKE_SQL_STATEMENT = textwrap.dedent(
       start_time "start_time",
       end_time "end_time",
       total_elapsed_time "duration"
-    from snowflake.account_usage.query_history
+    from {account_usage}.query_history
     WHERE query_text NOT LIKE '/* {{"app": "OpenMetadata", %%}} */%%'
     AND query_text NOT LIKE '/* {{"app": "dbt", %%}} */%%'
     AND start_time between to_timestamp_ltz('{start_time}') and to_timestamp_ltz('{end_time}')
@@ -39,7 +39,7 @@ SNOWFLAKE_SESSION_TAG_QUERY = 'ALTER SESSION SET QUERY_TAG="{query_tag}"'
 SNOWFLAKE_FETCH_ALL_TAGS = textwrap.dedent(
     """
     select TAG_NAME, TAG_VALUE, OBJECT_DATABASE, OBJECT_SCHEMA, OBJECT_NAME, COLUMN_NAME
-    from snowflake.account_usage.tag_references
+    from {account_usage}.tag_references
     where OBJECT_DATABASE = '{database_name}'
       and OBJECT_SCHEMA = '{schema_name}'
 """
@@ -234,11 +234,11 @@ SHOW EXTERNAL TABLES IN DATABASE "{database_name}"
 """
 
 SNOWFLAKE_TEST_FETCH_TAG = """
-select TAG_NAME from snowflake.account_usage.tag_references limit 1
+select TAG_NAME from {account_usage}.tag_references limit 1
 """
 
 SNOWFLAKE_TEST_GET_QUERIES = """
-SELECT query_text from snowflake.account_usage.query_history limit 1
+SELECT query_text from {account_usage}.query_history limit 1
 """
 
 SNOWFLAKE_TEST_GET_TABLES = """
@@ -296,7 +296,7 @@ SELECT
   ARGUMENT_SIGNATURE AS signature,
   COMMENT as comment,
   'StoredProcedure' as procedure_type
-FROM SNOWFLAKE.ACCOUNT_USAGE.PROCEDURES
+FROM {account_usage}.PROCEDURES
 WHERE PROCEDURE_CATALOG = '{database_name}'
   AND PROCEDURE_SCHEMA = '{schema_name}'
   AND DELETED IS NULL
@@ -313,7 +313,7 @@ SELECT
   ARGUMENT_SIGNATURE AS signature,
   COMMENT as comment,
   'UDF' as procedure_type
-FROM SNOWFLAKE.ACCOUNT_USAGE.FUNCTIONS
+FROM {account_usage}.FUNCTIONS
 WHERE FUNCTION_CATALOG = '{database_name}'
   AND FUNCTION_SCHEMA = '{schema_name}'
   AND DELETED IS NULL
@@ -336,7 +336,7 @@ WITH SP_HISTORY AS (
       SESSION_ID,
       START_TIME,
       END_TIME
-    FROM SNOWFLAKE.ACCOUNT_USAGE.QUERY_HISTORY SP
+    FROM {account_usage}.QUERY_HISTORY SP
     WHERE QUERY_TYPE = 'CALL'
       AND START_TIME >= '{start_date}'
       AND QUERY_TEXT <> ''
@@ -353,7 +353,7 @@ Q_HISTORY AS (
       USER_NAME,
       SCHEMA_NAME,
       DATABASE_NAME
-    FROM SNOWFLAKE.ACCOUNT_USAGE.QUERY_HISTORY SP
+    FROM {account_usage}.QUERY_HISTORY SP
     WHERE QUERY_TYPE <> 'CALL'
       AND QUERY_TEXT NOT LIKE '/* {{"app": "OpenMetadata", %%}} */%%'
       AND QUERY_TEXT NOT LIKE '/* {{"app": "dbt", %%}} */%%'

metadata/ingestion/source/database/snowflake/query_parser.py

@@ -60,6 +60,7 @@ class SnowflakeQueryParserSource(QueryParserSource, ABC):
             end_time=end_time,
             result_limit=self.config.sourceConfig.config.resultLimit,
             filters=self.get_filters(),
+            account_usage=self.service_connection.accountUsageSchema,
         )
 
     def check_life_cycle_query(

metadata/ingestion/source/database/sqlite/service_spec.py

@@ -1,8 +1,4 @@
-from metadata.ingestion.source.database.sqlite.lineage import SqliteLineageSource
 from metadata.ingestion.source.database.sqlite.metadata import SqliteSource
 from metadata.utils.service_spec.default import DefaultDatabaseSpec
 
-ServiceSpec = DefaultDatabaseSpec(
-    metadata_source_class=SqliteSource,
-    lineage_source_class=SqliteLineageSource,
-)
+ServiceSpec = DefaultDatabaseSpec(metadata_source_class=SqliteSource)

metadata/ingestion/source/database/teradata/service_spec.py

@@ -1,8 +1,4 @@
-from metadata.ingestion.source.database.teradata.lineage import TeradataLineageSource
 from metadata.ingestion.source.database.teradata.metadata import TeradataSource
 from metadata.utils.service_spec.default import DefaultDatabaseSpec
 
-ServiceSpec = DefaultDatabaseSpec(
-    metadata_source_class=TeradataSource,
-    lineage_source_class=TeradataLineageSource,
-)
+ServiceSpec = DefaultDatabaseSpec(metadata_source_class=TeradataSource)

metadata/ingestion/source/database/trino/connection.py

@@ -39,6 +39,7 @@ from metadata.ingestion.connections.builders import (
     create_generic_db_connection,
     get_connection_args_common,
     init_empty_connection_arguments,
+    init_empty_connection_options,
 )
 from metadata.ingestion.connections.secrets import connection_with_options_secrets
 from metadata.ingestion.connections.test_connections import (
@@ -135,6 +136,10 @@ def get_connection(connection: TrinoConnection) -> Engine:
     # here we are creating a copy of connection, because we need to dynamically
     # add auth params to connectionArguments, which we do no intend to store
     # in original connection object and in OpenMetadata database
+    from trino.sqlalchemy.dialect import TrinoDialect
+
+    TrinoDialect.is_disconnect = _is_disconnect
+
     connection_copy = deepcopy(connection)
     if connection_copy.verify:
         connection_copy.connectionArguments = (
@@ -183,3 +188,11 @@ def test_connection(
         queries=queries,
         timeout_seconds=timeout_seconds,
     )
+
+
+# pylint: disable=unused-argument
+def _is_disconnect(self, e, connection, cursor):
+    """is_disconnect method for the Databricks dialect"""
+    if "JWT expired" in str(e):
+        return True
+    return False

metadata/ingestion/source/database/unitycatalog/lineage.py

@@ -11,6 +11,7 @@
 """
 Databricks Unity Catalog Lineage Source Module
 """
+import traceback
 from typing import Iterable, Optional
 
 from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
@@ -27,6 +28,7 @@ from metadata.generated.schema.type.entityLineage import (
     EntitiesEdge,
     LineageDetails,
 )
+from metadata.generated.schema.type.entityLineage import Source as LineageSource
 from metadata.generated.schema.type.entityReference import EntityReference
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException, Source
@@ -109,9 +111,59 @@ class UnitycatalogLineageSource(Source):
                     )
                 )
         if col_lineage:
-            return LineageDetails(columnsLineage=col_lineage)
+            return LineageDetails(
+                columnsLineage=col_lineage, source=LineageSource.QueryLineage
+            )
         return None
 
+    def _handle_upstream_table(
+        self,
+        table_streams: LineageTableStreams,
+        table: Table,
+        databricks_table_fqn: str,
+    ) -> Iterable[Either[AddLineageRequest]]:
+        for upstream_table in table_streams.upstream_tables:
+            try:
+                if not upstream_table.name:
+                    continue
+                from_entity_fqn = fqn.build(
+                    metadata=self.metadata,
+                    entity_type=Table,
+                    database_name=upstream_table.catalog_name,
+                    schema_name=upstream_table.schema_name,
+                    table_name=upstream_table.name,
+                    service_name=self.config.serviceName,
+                )
+
+                from_entity = self.metadata.get_by_name(
+                    entity=Table, fqn=from_entity_fqn
+                )
+                if from_entity:
+                    lineage_details = self._get_lineage_details(
+                        from_table=from_entity,
+                        to_table=table,
+                        databricks_table_fqn=databricks_table_fqn,
+                    )
+                    yield Either(
+                        left=None,
+                        right=AddLineageRequest(
+                            edge=EntitiesEdge(
+                                toEntity=EntityReference(id=table.id, type="table"),
+                                fromEntity=EntityReference(
+                                    id=from_entity.id, type="table"
+                                ),
+                                lineageDetails=lineage_details,
+                            )
+                        ),
+                    )
+            except Exception:
+                logger.debug(
+                    "Error while processing lineage for "
+                    f"{upstream_table.catalog_name}.{upstream_table.schema_name}.{upstream_table.name}"
+                    f" -> {databricks_table_fqn}"
+                )
+                logger.debug(traceback.format_exc())
+
     def _iter(self, *_, **__) -> Iterable[Either[AddLineageRequest]]:
         """
         Based on the query logs, prepare the lineage
@@ -128,37 +180,9 @@ class UnitycatalogLineageSource(Source):
                 table_streams: LineageTableStreams = self.client.get_table_lineage(
                     databricks_table_fqn
                 )
-                for upstream_table in table_streams.upstream_tables:
-                    from_entity_fqn = fqn.build(
-                        metadata=self.metadata,
-                        entity_type=Table,
-                        database_name=upstream_table.catalog_name,
-                        schema_name=upstream_table.schema_name,
-                        table_name=upstream_table.name,
-                        service_name=self.config.serviceName,
-                    )
-
-                    from_entity = self.metadata.get_by_name(
-                        entity=Table, fqn=from_entity_fqn
-                    )
-                    if from_entity:
-                        lineage_details = self._get_lineage_details(
-                            from_table=from_entity,
-                            to_table=table,
-                            databricks_table_fqn=databricks_table_fqn,
-                        )
-                        yield Either(
-                            left=None,
-                            right=AddLineageRequest(
-                                edge=EntitiesEdge(
-                                    toEntity=EntityReference(id=table.id, type="table"),
-                                    fromEntity=EntityReference(
-                                        id=from_entity.id, type="table"
-                                    ),
-                                    lineageDetails=lineage_details,
-                                )
-                            ),
-                        )
+                yield from self._handle_upstream_table(
+                    table_streams, table, databricks_table_fqn
+                )
 
     def test_connection(self) -> None:
         test_connection_fn = get_test_connection_fn(self.service_connection)

metadata/ingestion/source/database/unitycatalog/query_parser.py

@@ -44,6 +44,13 @@ class UnityCatalogQueryParserSource(
 
     filters: str
 
+    def _init_super(
+        self,
+        config: WorkflowSource,
+        metadata: OpenMetadata,
+    ):
+        super().__init__(config, metadata, False)
+
     # pylint: disable=super-init-not-called
     def __init__(self, config: WorkflowSource, metadata: OpenMetadata):
         self._init_super(config=config, metadata=metadata)

metadata/ingestion/source/database/unitycatalog/service_spec.py

@@ -11,6 +11,9 @@ from metadata.ingestion.source.database.unitycatalog.usage import (
 from metadata.profiler.interface.sqlalchemy.unity_catalog.profiler_interface import (
     UnityCatalogProfilerInterface,
 )
+from metadata.profiler.interface.sqlalchemy.unity_catalog.sampler_interface import (
+    UnityCatalogSamplerInterface,
+)
 from metadata.utils.service_spec.default import DefaultDatabaseSpec
 
 ServiceSpec = DefaultDatabaseSpec(
@@ -19,4 +22,5 @@ ServiceSpec = DefaultDatabaseSpec(
     usage_source_class=UnitycatalogUsageSource,
     profiler_class=UnityCatalogProfilerInterface,
     test_suite_class=UnityCatalogTestSuiteInterface,
+    sampler_class=UnityCatalogSamplerInterface,
 )

metadata/ingestion/source/database/unitycatalog/usage.py

@@ -11,17 +11,22 @@
 """
 unity catalog usage module
 """
+import traceback
+from datetime import datetime
+from typing import Iterable
 
-from metadata.ingestion.source.database.databricks.usage import DatabricksUsageSource
+from metadata.generated.schema.type.basic import DateTime
+from metadata.generated.schema.type.tableQuery import TableQueries, TableQuery
 from metadata.ingestion.source.database.unitycatalog.query_parser import (
     UnityCatalogQueryParserSource,
 )
+from metadata.ingestion.source.database.usage_source import UsageSource
 from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
 
 
-class UnitycatalogUsageSource(UnityCatalogQueryParserSource, DatabricksUsageSource):
+class UnitycatalogUsageSource(UnityCatalogQueryParserSource, UsageSource):
     """
     UnityCatalog Usage Source
 
@@ -29,3 +34,37 @@ class UnitycatalogUsageSource(UnityCatalogQueryParserSource, DatabricksUsageSour
     DatabricksUsageSource as both the sources would call
     the same API for fetching Usage Queries
     """
+
+    def yield_table_queries(self) -> Iterable[TableQuery]:
+        """
+        Method to yield TableQueries
+        """
+        queries = []
+        data = self.client.list_query_history(
+            start_date=self.start,
+            end_date=self.end,
+        )
+        for row in data or []:
+            try:
+                if self.client.is_query_valid(row):
+                    queries.append(
+                        TableQuery(
+                            dialect=self.dialect.value,
+                            query=row.get("query_text"),
+                            userName=row.get("user_name"),
+                            startTime=str(row.get("query_start_time_ms")),
+                            endTime=str(row.get("execution_end_time_ms")),
+                            analysisDate=DateTime(datetime.now()),
+                            serviceName=self.config.serviceName,
+                            duration=row.get("duration")
+                            if row.get("duration")
+                            else None,
+                        )
+                    )
+            except Exception as err:
+                logger.debug(traceback.format_exc())
+                logger.warning(
+                    f"Failed to process query {row.get('query_text')} due to: {err}"
+                )
+
+        yield TableQueries(queries=queries)

metadata/ingestion/source/database/usage_source.py

@@ -153,7 +153,7 @@ class UsageSource(QueryParserSource, ABC):
                 if query:
                     logger.debug(
                         (
-                            f"###### USAGE QUERY #######\n{mask_query(query, self.dialect.value)}"
+                            f"###### USAGE QUERY #######\n{mask_query(query, self.dialect.value) or query}"
                             "\n##########################"
                         )
                     )

metadata/ingestion/source/pipeline/nifi/connection.py

@@ -17,8 +17,10 @@ from typing import Optional
 from metadata.generated.schema.entity.automations.workflow import (
     Workflow as AutomationWorkflow,
 )
+from metadata.generated.schema.entity.services.connections.pipeline.nifi.basicAuth import (
+    NifiBasicAuth,
+)
 from metadata.generated.schema.entity.services.connections.pipeline.nifiConnection import (
-    BasicAuthentication,
     NifiConnection,
 )
 from metadata.generated.schema.entity.services.connections.testConnectionResult import (
@@ -34,7 +36,7 @@ def get_connection(connection: NifiConnection) -> NifiClient:
     """
     Create connection
     """
-    if isinstance(connection.nifiConfig, BasicAuthentication):
+    if isinstance(connection.nifiConfig, NifiBasicAuth):
         return NifiClient(
             host_port=connection.hostPort,
             username=connection.nifiConfig.username,

metadata/profiler/interface/sqlalchemy/unity_catalog/sampler_interface.py

@@ -0,0 +1,29 @@
+#  Copyright 2021 Collate
+#  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.
+
+"""
+Interfaces with database for all database engine
+supporting sqlalchemy abstraction layer
+"""
+from metadata.ingestion.source.database.databricks.connection import (
+    get_connection as databricks_get_connection,
+)
+from metadata.sampler.sqlalchemy.sampler import SQASampler
+
+
+class UnityCatalogSamplerInterface(SQASampler):
+    def get_client(self):
+        """client is the session for SQA"""
+        self.connection = databricks_get_connection(self.service_connection_config)
+        self.client = super().get_client()
+        self.set_catalog(self.client)
+
+        return self.client

metadata/profiler/source/fetcher/config.py

@@ -42,3 +42,7 @@ class EntityFilterConfigInterface(Protocol):
     @property
     def useFqnForFiltering(self) -> Optional[bool]:
         ...
+
+    @property
+    def includeViews(self) -> Optional[bool]:
+        ...

metadata/profiler/source/fetcher/fetcher_strategy.py

@@ -17,6 +17,7 @@ from abc import ABC, abstractmethod
 from typing import Iterable, Iterator, Optional, cast
 
 from metadata.generated.schema.entity.data.database import Database
+from metadata.generated.schema.entity.data.table import TableType
 from metadata.generated.schema.entity.services.ingestionPipelines.status import (
     StackTraceError,
 )
@@ -115,13 +116,13 @@ class DatabaseFetcherStrategy(FetcherStrategy):
         super().__init__(config, metadata, global_profiler_config, status)
         self.source_config = cast(
             EntityFilterConfigInterface, self.source_config
-        )  # Satisfy typchecker
+        )  # Satisfy typechecker
 
     def _filter_databases(self, databases: Iterable[Database]) -> Iterable[Database]:
         """Filter databases based on the filter pattern
 
         Args:
-            database (Database): Database to filter
+            databases (Database): Database to filter
 
         Returns:
             bool
@@ -192,6 +193,21 @@ class DatabaseFetcherStrategy(FetcherStrategy):
 
         return False
 
+    def _filter_views(self, table: Table) -> bool:
+        """Filter the tables based on include views configuration"""
+        # If we include views, nothing to filter
+        if self.source_config.includeViews:
+            return False
+
+        # Otherwise, filter out views
+        if table.tableType == TableType.View:
+            self.status.filter(
+                table.name.root, f"We are not including views {table.name.root}"
+            )
+            return True
+
+        return False
+
     def _filter_column_metrics_computation(self):
         """Filter"""
 
@@ -242,6 +258,7 @@ class DatabaseFetcherStrategy(FetcherStrategy):
                 not self.source_config.classificationFilterPattern
                 or not self.filter_classifications(table)
             )
+            and not self._filter_views(table)
         ]
 
         return tables

metadata/sampler/sampler_interface.py

@@ -76,9 +76,6 @@ class SamplerInterface(ABC):
         self._columns: Optional[List[SQALikeColumn]] = None
         self.sample_config = sample_config
 
-        if not self.sample_config.profileSample:
-            self.sample_config.profileSample = 100
-
         self.entity = entity
         self.include_columns = include_columns
         self.exclude_columns = exclude_columns

metadata/sampler/sqlalchemy/bigquery/sampler.py

@@ -54,7 +54,6 @@ class BigQuerySampler(SQASampler):
         sample_query: Optional[str] = None,
         storage_config: DataStorageConfig = None,
         sample_data_count: Optional[int] = SAMPLE_DATA_DEFAULT_COUNT,
-        table_type: TableType = None,
         **kwargs,
     ):
         super().__init__(
@@ -68,7 +67,7 @@ class BigQuerySampler(SQASampler):
             sample_data_count=sample_data_count,
             **kwargs,
         )
-        self.raw_dataset_type: TableType = table_type
+        self.raw_dataset_type: Optional[TableType] = entity.tableType
 
     def set_tablesample(self, selectable: SqaTable):
         """Set the TABLESAMPLE clause for BigQuery

metadata/sampler/sqlalchemy/sampler.py

@@ -12,6 +12,7 @@
 Helper module to handle data sampling
 for the profiler
 """
+import hashlib
 import traceback
 from typing import List, Optional, Union, cast
 
@@ -32,6 +33,7 @@ from metadata.profiler.orm.functions.modulo import ModuloFn
 from metadata.profiler.orm.functions.random_num import RandomNumFn
 from metadata.profiler.processor.handle_partition import build_partition_predicate
 from metadata.sampler.sampler_interface import SamplerInterface
+from metadata.utils.constants import UTF_8
 from metadata.utils.helpers import is_safe_sql_query
 from metadata.utils.logger import profiler_interface_registry_logger
 
@@ -109,17 +111,28 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
             query = self.get_partitioned_query(query)
         return query
 
+    def get_sampler_table_name(self) -> str:
+        """Get the base name of the SQA table for sampling.
+        We use MD5 as a hashing algorithm to generate a unique name for the table
+        keeping its length controlled. Otherwise, we ended up having issues
+        with names getting truncated when we add the suffixes to the identifiers
+        such as _sample, or _rnd.
+        """
+        encoded_name = self.raw_dataset.__tablename__.encode(UTF_8)
+        hash_object = hashlib.md5(encoded_name)
+        return hash_object.hexdigest()
+
     def get_sample_query(self, *, column=None) -> Query:
         """get query for sample data"""
         if self.sample_config.profileSampleType == ProfileSampleType.PERCENTAGE:
             rnd = self._base_sample_query(
                 column,
                 (ModuloFn(RandomNumFn(), 100)).label(RANDOM_LABEL),
-            ).cte(f"{self.raw_dataset.__tablename__}_rnd")
+            ).cte(f"{self.get_sampler_table_name()}_rnd")
             session_query = self.client.query(rnd)
             return session_query.where(
                 rnd.c.random <= self.sample_config.profileSample
-            ).cte(f"{self.raw_dataset.__tablename__}_sample")
+            ).cte(f"{self.get_sampler_table_name()}_sample")
 
         table_query = self.client.query(self.raw_dataset)
         session_query = self._base_sample_query(
@@ -129,7 +142,7 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
         return (
             session_query.order_by(RANDOM_LABEL)
             .limit(self.sample_config.profileSample)
-            .cte(f"{self.raw_dataset.__tablename__}_rnd")
+            .cte(f"{self.get_sampler_table_name()}_rnd")
         )
 
     def get_dataset(self, column=None, **__) -> Union[DeclarativeMeta, AliasedClass]:
@@ -143,7 +156,7 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
         if not self.sample_config.profileSample:
             if self.partition_details:
                 partitioned = self._partitioned_table()
-                return partitioned.cte(f"{self.raw_dataset.__tablename__}_partitioned")
+                return partitioned.cte(f"{self.get_sampler_table_name()}_partitioned")
 
             return self.raw_dataset
 
@@ -162,23 +175,23 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
             return self._fetch_sample_data_from_user_query()
 
         # Add new RandomNumFn column
-        rnd = self.get_sample_query()
+        ds = self.get_dataset()
         if not columns:
-            sqa_columns = [col for col in inspect(rnd).c if col.name != RANDOM_LABEL]
+            sqa_columns = [col for col in inspect(ds).c if col.name != RANDOM_LABEL]
         else:
             # we can't directly use columns as it is bound to self.raw_dataset and not the rnd table.
             # If we use it, it will result in a cross join between self.raw_dataset and rnd table
             names = [col.name for col in columns]
             sqa_columns = [
                 col
-                for col in inspect(rnd).c
+                for col in inspect(ds).c
                 if col.name != RANDOM_LABEL and col.name in names
             ]
 
         try:
             sqa_sample = (
                 self.client.query(*sqa_columns)
-                .select_from(rnd)
+                .select_from(ds)
                 .limit(self.sample_limit)
                 .all()
             )
@@ -224,7 +237,7 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
         stmt = stmt.columns(*list(inspect(self.raw_dataset).c))
 
         return self.client.query(stmt.subquery()).cte(
-            f"{self.raw_dataset.__tablename__}_user_sampled"
+            f"{self.get_sampler_table_name()}_user_sampled"
         )
 
     def _partitioned_table(self) -> Query: