From 402abb52c1db23959dd973fb0a9b163802ddcb6c Mon Sep 17 00:00:00 2001 From: scottwilson Date: Fri, 19 Jun 2026 11:29:42 +0100 Subject: [PATCH 1/8] feat: added get_row_count and ensured implementation of get_total and get_row_count are the same across Extracto classes --- mario/athena.py | 28 +++++++----- mario/data_extractor.py | 87 ++++++++++++++++++++++++++----------- test/test_data_extractor.py | 18 +++++--- 3 files changed, 92 insertions(+), 41 deletions(-) diff --git a/mario/athena.py b/mario/athena.py index b70ec36..9f95f00 100644 --- a/mario/athena.py +++ b/mario/athena.py @@ -56,18 +56,11 @@ def get_connection(self): catalog_name=cfg.catalog ) - def get_total(self, measure=None): - """ - For totals when streaming data we need to run a totals SQL query separate - from the main query and use the results of this - :return: the total value of the query - NOTE this is a direct copy from StreamingDataExtractor - """ + def __get_total__(self, measure): from pandas import read_sql - logger.info("Building totals query") - measure = self.__get_measure__(measure) if self.configuration.query_builder is not None: from mario.query_builder import QueryBuilder + logger.info("Building totals query") query_builder: QueryBuilder = self.configuration.query_builder( configuration=self.configuration, metadata=self.metadata, @@ -79,6 +72,21 @@ def get_total(self, measure=None): totals_df = read_sql(totals_query[0], self.get_connection(), params=totals_query[1]) return totals_df.iat[0, 0] + def get_row_count(self): + return self.__get_total__(None) + + def get_total(self, measure=None): + """ + For totals when streaming data we need to run a totals SQL query separate + from the main query and use the results of this + :return: the total value of the query + NOTE this is a direct copy from StreamingDataExtractor + """ + measure = self.__get_measure__(measure) + if measure is None: + return self.get_row_count() + return self.__get_total__(measure) + def run_query(self) -> str: """ Runs the SQL query in Athena and returns the @@ -190,4 +198,4 @@ def repl(match): raise KeyError(f"Missing SQL parameter: {name}") return athena_literal(params[name]) - return ATHENA_PARAM_PATTERN.sub(repl, sql) \ No newline at end of file + return ATHENA_PARAM_PATTERN.sub(repl, sql) diff --git a/mario/data_extractor.py b/mario/data_extractor.py index 8410669..d814515 100644 --- a/mario/data_extractor.py +++ b/mario/data_extractor.py @@ -131,11 +131,28 @@ def __get_measure__(self, measure=None): raise ValueError(f"Measure {measure} does not exist in dataset specification") return measure - def get_total(self, measure=None): + def get_row_count(self): + """ + :return: the number of rows in the data + """ df = self.get_data_frame() + return len(df) + + def get_total(self, measure=None): + """ + Returns the total from the dataset, using either the supplied measure, or + the default measure from the dataset specification if no measure is supplied. + If an invalid measure is supplied, a ValueError will be raised. + If no measure is supplied, and no measure exists in the dataset, the row_count() + is returned instead. + :param measure: the measure to use + :return: the sum of the data with the measure + """ measure = self.__get_measure__(measure) if measure is None: - return len(df) + return self.get_row_count() + + df = self.get_data_frame() self._total = df[measure].sum() return self._total @@ -256,22 +273,30 @@ def __minimise_data__(self): table=table ) - def get_total(self, measure=None): - from mario.hyper_utils import get_row_count, get_default_table_and_schema, get_total + def get_row_count(self): + from mario.hyper_utils import get_row_count as __get_row_count__ + from mario.hyper_utils import get_default_table_and_schema schema, table = get_default_table_and_schema(self.configuration.file_path) + return __get_row_count__ ( + hyper_file_path=self.configuration.file_path, + schema=schema, + table=table + ) + + def get_total(self, measure=None): + measure = self.__get_measure__(measure) if measure is None: - return get_row_count( - hyper_file_path=self.configuration.file_path, - schema=schema, - table=table - ) - else: - return get_total( - hyper_file_path=self.configuration.file_path, - schema=schema, - table=table, - measure=measure - ) + return self.get_row_count() + + from mario.hyper_utils import get_total as __get_total__ + from mario.hyper_utils import get_default_table_and_schema + schema, table = get_default_table_and_schema(self.configuration.file_path) + return __get_total__( + hyper_file_path=self.configuration.file_path, + schema=schema, + table=table, + measure=measure + ) def save_data_as_hyper(self, file_path: str, **kwargs): from mario.hyper_utils import save_hyper_as_hyper, add_row_numbers_to_hyper, get_default_table_and_schema @@ -360,26 +385,33 @@ def save_data_as_hyper(self, file_path: str, **kwargs): **kwargs ) - def get_total(self, measure=None): - """ - For totals when streaming data we need to run a totals SQL query separate - from the main query and use the results of this - :return: the total value of the query - """ - logger.info("Building totals query") - measure = self.__get_measure__(measure) + def __get_total__(self, measure): if self.configuration.query_builder is not None: from mario.query_builder import QueryBuilder + logger.info("Building totals query") query_builder: QueryBuilder = self.configuration.query_builder( configuration=self.configuration, metadata=self.metadata, dataset_specification=self.dataset_specification) totals_query = query_builder.create_totals_query(measure=measure) + totals_df = pd.read_sql(totals_query[0], self.get_connection(), params=totals_query[1]) + return totals_df.iat[0, 0] else: raise NotImplementedError - totals_df = pd.read_sql(totals_query[0], self.get_connection(), params=totals_query[1]) - return totals_df.iat[0, 0] + def get_row_count(self): + return self.__get_total__(None) + + def get_total(self, measure=None): + """ + For totals when streaming data we need to run a totals SQL query separate + from the main query and use the results of this + :return: the total value of the query + """ + measure = self.__get_measure__(measure) + if measure is None: + return self.get_row_count() + return self.__get_total__(measure) def stream_sql_to_hyper(self, file_path: str, **kwargs): """ @@ -609,6 +641,9 @@ def __load_from_sql__(self): def get_data_frame(self, minimise=True, include_row_numbers=False) -> DataFrame: raise NotImplementedError() + def get_row_count(self): + raise NotImplementedError() + def get_total(self, measure=None): raise NotImplementedError() diff --git a/test/test_data_extractor.py b/test/test_data_extractor.py index cee500b..2bdfb6a 100644 --- a/test/test_data_extractor.py +++ b/test/test_data_extractor.py @@ -518,7 +518,7 @@ def test_partitioning_extractor_streaming(): def test_hyper_total_measure(): dataset = dataset_from_json(os.path.join('test', 'dataset.json')) - dataset.measures = [] + dataset.measures = ['Sales'] metadata = metadata_from_json(os.path.join('test', 'metadata.json')) configuration = Configuration( file_path=os.path.join('test', 'orders.hyper') @@ -528,13 +528,19 @@ def test_hyper_total_measure(): metadata=metadata, configuration=configuration ) - assert extractor.get_total() == 10194 + + assert extractor.get_row_count() == 10194 assert extractor.get_total(measure='Sales') == 2326534.354299952 + # Defaults to "Sales" as measure is in dataset.measures() + assert extractor.get_total() == 2326534.354299952 + # Raise ValueError if measure not in dataset.measures() + with pytest.raises(ValueError): + extractor.get_total(measure='Stuff') def test_hyper_to_csv(): dataset = dataset_from_json(os.path.join('test', 'dataset.json')) - dataset.measures = [] + dataset.measures = ['Sales'] metadata = metadata_from_json(os.path.join('test', 'metadata.json')) configuration = Configuration( file_path=os.path.join('test', 'orders.hyper') @@ -558,9 +564,10 @@ def test_hyper_to_csv(): df = pd.read_csv(output_file) assert round(df['Sales'].sum(), 4) == 2326534.3543 + def test_hyper_to_csv_without_copy_to_tmp(): dataset = dataset_from_json(os.path.join('test', 'dataset.json')) - dataset.measures = [] + dataset.measures = ['Sales'] metadata = metadata_from_json(os.path.join('test', 'metadata.json')) # Copy the source data to avoid overwriting during other pytest runs shutil.copyfile( @@ -590,9 +597,10 @@ def test_hyper_to_csv_without_copy_to_tmp(): df = pd.read_csv(output_file) assert round(df['Sales'].sum(), 4) == 2326534.3543 + def test_hyper_to_csv_without_using_pantab(): dataset = dataset_from_json(os.path.join('test', 'dataset.json')) - dataset.measures = [] + dataset.measures = ['Sales'] metadata = metadata_from_json(os.path.join('test', 'metadata.json')) # Copy the source data to avoid overwriting during other pytest runs shutil.copyfile( From 3f2ebb3da0056909b0dc5364acc861aa6940c3f2 Mon Sep 17 00:00:00 2001 From: scottwilson Date: Fri, 19 Jun 2026 11:50:24 +0100 Subject: [PATCH 2/8] test: updated tests to fit new contract whereby asking for a measure not in dataset.measures raises an error rather than silently giving you the row count. --- test/test_data_extractor.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/test/test_data_extractor.py b/test/test_data_extractor.py index 2bdfb6a..7e938a6 100644 --- a/test/test_data_extractor.py +++ b/test/test_data_extractor.py @@ -558,7 +558,8 @@ def test_hyper_to_csv(): minimise=False, compress_using_gzip=False ) - assert extractor.get_total() == 10194 + assert extractor.get_row_count() == 10194 + assert round(extractor.get_total(), 2) == 2326534.35 assert round(extractor.get_total(measure='Sales'), 2) == 2326534.35 df = pd.read_csv(output_file) @@ -591,7 +592,7 @@ def test_hyper_to_csv_without_copy_to_tmp(): compress_using_gzip=False, do_not_modify_source=False ) - assert extractor.get_total() == 10194 + assert extractor.get_row_count() == 10194 assert round(extractor.get_total(measure='Sales'), 2) == 2326534.35 df = pd.read_csv(output_file) @@ -625,12 +626,13 @@ def test_hyper_to_csv_without_using_pantab(): do_not_modify_source=False, use_pantab=False ) - assert extractor.get_total() == 10194 + assert extractor.get_row_count() == 10194 assert round(extractor.get_total(measure='Sales'), 2) == 2326534.35 df = pd.read_csv(output_file) assert round(df['Sales'].sum(), 4) == 2326534.3543 + def test_partitioning_extractor_partition_sql_no_data_in_partition(): # Skip this test if we don't have a connection string if not os.environ.get('CONNECTION_STRING'): From 6b2bb7a053e3b068d98959469a969d74466e46c9 Mon Sep 17 00:00:00 2001 From: scottwilson Date: Fri, 19 Jun 2026 11:51:35 +0100 Subject: [PATCH 3/8] fix: removed duplicate logic for missing measures, instead when we get a request for a query with measure None we just return Count(*). Substituting measures happens in the Extractor layer now. --- mario/query_builder.py | 25 +++++++++++-------------- 1 file changed, 11 insertions(+), 14 deletions(-) diff --git a/mario/query_builder.py b/mario/query_builder.py index 780e45c..4fae70f 100644 --- a/mario/query_builder.py +++ b/mario/query_builder.py @@ -101,23 +101,23 @@ def contains_subject(self) -> bool: return True return False + def __get_measure_label__(self, measure:str): + if measure in ['FPE', 'FTE'] and not self.contains_subject(): + return 'Count' + return measure + def create_totals_query(self, measure=None) -> [str, List[any]]: """ Constructs a 'total only' query by combining the selections and constraints. :return: an array containing the query prepared statement, and the parameters - TODO make this more generic """ measures = [] - measure = self.mapping.as_physical[measure] if measure is None: - for measure in self.dataset_specification.measures: - # Decide column name based on whether 'subject' fields are present - if measure in ['FPE', 'FTE'] and not self.contains_subject(): - measures.append(fn.Sum(self.table[measure], 'Count')) - else: - measures.append(fn.Sum(self.table[measure], measure)) + measures.append(fn.Count('*', 'Count')) else: - measures.append(fn.Sum(self.table[measure], measure)) + measure_label = self.__get_measure_label__(measure) + measure = self.mapping.as_physical[measure] + measures.append(fn.Sum(self.table[measure], measure_label)) q = Query().from_(self.table).select(*measures) @@ -141,15 +141,12 @@ def create_query(self) -> [str, List[any]]: # remove measures from regular select measures = [] for measure_field in self.dataset_specification.measures: + measure_label = self.__get_measure_label__(measure_field) measure = self.mapping.as_physical[measure_field] if measure in select_fields: select_fields.remove(measure) group_fields.remove(measure) - # Decide column name based on whether 'subject' fields are present - if measure in ['FPE', 'FTE'] and not self.contains_subject(): - measures.append(fn.Sum(self.table[measure], 'Count')) - else: - measures.append(fn.Sum(self.table[measure], measure)) + measures.append(fn.Sum(self.table[measure], measure_label)) select_fields.extend(measures) q = Query(). \ From b021310318130f424c44f2bbeafcee266cfb62c7 Mon Sep 17 00:00:00 2001 From: scottwilson Date: Fri, 19 Jun 2026 12:16:31 +0100 Subject: [PATCH 4/8] test: Added a consistency test across implementations --- test/test_data_extractor.py | 67 ++++++++++++++++++++++++++++++++++++- 1 file changed, 66 insertions(+), 1 deletion(-) diff --git a/test/test_data_extractor.py b/test/test_data_extractor.py index 7e938a6..87fa1c4 100644 --- a/test/test_data_extractor.py +++ b/test/test_data_extractor.py @@ -805,4 +805,69 @@ def test_partitioning_extractor_with_row_numbers_apostrophes(): ) columns = pd.read_csv(csv_file).columns # Check row_number omitted in CSV output - assert 'row_number' not in columns \ No newline at end of file + assert 'row_number' not in columns + + +def test_consistency_across_extractors(tmp_path): + """ + Ensure get_row_count and get_total are consistent across: + - DataExtractor (CSV) + - HyperFile (Hyper) + """ + + # --- Setup dataset / metadata --- + dataset = dataset_from_json(os.path.join('test', 'dataset.json')) + metadata = metadata_from_json(os.path.join('test', 'metadata.json')) + + csv_path = os.path.join('test', 'orders.csv') + + # --- 1. DataExtractor (baseline) --- + base_config = Configuration(file_path=csv_path) + base_extractor = DataExtractor( + dataset_specification=dataset, + metadata=metadata, + configuration=base_config + ) + + base_extractor.validate_data() + + base_row_count = base_extractor.get_row_count() + base_total_default = base_extractor.get_total() + + # choose explicit measures + measures = dataset.measures.copy() + measures.remove('Profit Ratio') # This is a calculation! + + base_totals = { + m: base_extractor.get_total(measure=m) + for m in measures + } + + # --- 2. Convert to Hyper --- + hyper_path = tmp_path / "test.hyper" + + base_extractor.save_data_as_hyper(file_path=str(hyper_path)) + + # --- 3. HyperFile --- + hyper_config = Configuration(file_path=str(hyper_path)) + + hyper_extractor = HyperFile( + dataset_specification=dataset, + metadata=metadata, + configuration=hyper_config + ) + + hyper_row_count = hyper_extractor.get_row_count() + hyper_total_default = hyper_extractor.get_total() + + hyper_totals = { + m: hyper_extractor.get_total(measure=m) + for m in measures + } + + # --- Assertions: STRICT consistency --- + assert hyper_row_count == base_row_count + assert round(hyper_total_default, 6) == round(base_total_default, 6) + + for m in measures: + assert round(hyper_totals[m], 6) == round(base_totals[m], 6) \ No newline at end of file From 53dfe57930c7b07e5fae6d96adf93e40ff655aa2 Mon Sep 17 00:00:00 2001 From: scottwilson Date: Fri, 19 Jun 2026 12:36:34 +0100 Subject: [PATCH 5/8] fix: added a deprecation warning when calling get_total with no measure or None --- mario/athena.py | 4 +++- mario/data_extractor.py | 11 +++++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/mario/athena.py b/mario/athena.py index 9f95f00..b4d6d05 100644 --- a/mario/athena.py +++ b/mario/athena.py @@ -1,5 +1,5 @@ from pyathena import connect -from mario.data_extractor import DataExtractor, Configuration +from mario.data_extractor import DataExtractor, Configuration, GET_TOTAL_WITHOUT_MEASURE import logging import datetime import numbers @@ -82,6 +82,8 @@ def get_total(self, measure=None): :return: the total value of the query NOTE this is a direct copy from StreamingDataExtractor """ + if measure is None: + logger.warning(GET_TOTAL_WITHOUT_MEASURE) measure = self.__get_measure__(measure) if measure is None: return self.get_row_count() diff --git a/mario/data_extractor.py b/mario/data_extractor.py index d814515..8dc65b1 100644 --- a/mario/data_extractor.py +++ b/mario/data_extractor.py @@ -12,6 +12,11 @@ logger = logging.getLogger(__name__) +GET_TOTAL_WITHOUT_MEASURE = 'Deprecation warning: Calling get_total() without any measure specified, or with None, ' \ + 'is going to be deprecated in future. Instead you should either call get_row_count(), ' \ + 'or get_total(measure) where the measure exists in the dataset.' + + class Configuration: """ Configuration for a data extractor. This can include: @@ -148,6 +153,8 @@ def get_total(self, measure=None): :param measure: the measure to use :return: the sum of the data with the measure """ + if measure is None: + logger.warning(GET_TOTAL_WITHOUT_MEASURE) measure = self.__get_measure__(measure) if measure is None: return self.get_row_count() @@ -284,6 +291,8 @@ def get_row_count(self): ) def get_total(self, measure=None): + if measure is None: + logger.warning(GET_TOTAL_WITHOUT_MEASURE) measure = self.__get_measure__(measure) if measure is None: return self.get_row_count() @@ -408,6 +417,8 @@ def get_total(self, measure=None): from the main query and use the results of this :return: the total value of the query """ + if measure is None: + logger.warning(GET_TOTAL_WITHOUT_MEASURE) measure = self.__get_measure__(measure) if measure is None: return self.get_row_count() From 6339c9696c90abb0de5e0e3019303b2b99635034 Mon Sep 17 00:00:00 2001 From: scottwilson Date: Tue, 23 Jun 2026 09:44:36 +0100 Subject: [PATCH 6/8] refactor: delegate to __get_total__ for consistency --- mario/data_extractor.py | 26 +++++++++++++++----------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/mario/data_extractor.py b/mario/data_extractor.py index 8dc65b1..64340fc 100644 --- a/mario/data_extractor.py +++ b/mario/data_extractor.py @@ -143,6 +143,11 @@ def get_row_count(self): df = self.get_data_frame() return len(df) + def __get_total__(self, measure): + df = self.get_data_frame() + self._total = df[measure].sum() + return self._total + def get_total(self, measure=None): """ Returns the total from the dataset, using either the supplied measure, or @@ -158,10 +163,7 @@ def get_total(self, measure=None): measure = self.__get_measure__(measure) if measure is None: return self.get_row_count() - - df = self.get_data_frame() - self._total = df[measure].sum() - return self._total + return self.__get_total__(measure) def validate_data(self, allow_nulls=True): from mario.validation import DataFrameValidator @@ -290,13 +292,7 @@ def get_row_count(self): table=table ) - def get_total(self, measure=None): - if measure is None: - logger.warning(GET_TOTAL_WITHOUT_MEASURE) - measure = self.__get_measure__(measure) - if measure is None: - return self.get_row_count() - + def __get_total__(self, measure): from mario.hyper_utils import get_total as __get_total__ from mario.hyper_utils import get_default_table_and_schema schema, table = get_default_table_and_schema(self.configuration.file_path) @@ -307,6 +303,14 @@ def get_total(self, measure=None): measure=measure ) + def get_total(self, measure=None): + if measure is None: + logger.warning(GET_TOTAL_WITHOUT_MEASURE) + measure = self.__get_measure__(measure) + if measure is None: + return self.get_row_count() + return self.__get_total__(measure) + def save_data_as_hyper(self, file_path: str, **kwargs): from mario.hyper_utils import save_hyper_as_hyper, add_row_numbers_to_hyper, get_default_table_and_schema options = HyperOptions(**kwargs) From 0c10cebd1791b37f83314fa84e6f344f6611112d Mon Sep 17 00:00:00 2001 From: scottwilson Date: Tue, 23 Jun 2026 09:44:52 +0100 Subject: [PATCH 7/8] test: added test for row count in Athena --- test/test_athena_extractor.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/test/test_athena_extractor.py b/test/test_athena_extractor.py index a54db1b..eeaae45 100644 --- a/test/test_athena_extractor.py +++ b/test/test_athena_extractor.py @@ -106,6 +106,26 @@ def test_athena_count(): total = extractor.get_total(measure=dataset.measures[0]) assert total == 28_733_910 + +def test_athena_row_count(): + if not AWS_PROFILE or not AWS_ATHENA_RESULTS_DIR or not AWS_REGION: + pytest.skip("Skipping Athena test as AWS not configured") + + dataset, metadata, cfg = get_test_conf() + cfg.query_builder = SubsetQueryBuilder + cfg.schema = 'demo' + cfg.view = 'student_open_data' + + extractor = AthenaDataExtractor( + configuration=cfg, + metadata=metadata, + dataset_specification=dataset + ) + + total = extractor.get_row_count() + assert total == 28_733_910 + + def test_athena_count_with_space(): if not AWS_PROFILE or not AWS_ATHENA_RESULTS_DIR or not AWS_REGION: pytest.skip("Skipping Athena test as AWS not configured") From 77729ae2913b5a8e05adefd2bb9b189f88625180 Mon Sep 17 00:00:00 2001 From: scottwilson Date: Tue, 23 Jun 2026 09:45:13 +0100 Subject: [PATCH 8/8] test: added totals and row counts tests for streaming data extractor --- test/test_data_extractor.py | 124 ++++++++++++++++++++++++++++++++++++ 1 file changed, 124 insertions(+) diff --git a/test/test_data_extractor.py b/test/test_data_extractor.py index 87fa1c4..fc3c01b 100644 --- a/test/test_data_extractor.py +++ b/test/test_data_extractor.py @@ -808,6 +808,130 @@ def test_partitioning_extractor_with_row_numbers_apostrophes(): assert 'row_number' not in columns +def test_streaming_extractor_get_row_count(): + # Skip if no DB configured + if not os.environ.get('CONNECTION_STRING'): + pytest.skip("Skipping SQL test as no database configured") + + dataset = dataset_from_json(os.path.join('test', 'dataset.json')) + metadata = metadata_from_json(os.path.join('test', 'metadata.json')) + + configuration = Configuration( + connection_string=os.environ.get('CONNECTION_STRING'), + schema="dev", + view="superstore", + query_builder=ViewBasedQueryBuilder + ) + + extractor = StreamingDataExtractor( + dataset_specification=dataset, + metadata=metadata, + configuration=configuration + ) + + # Trigger streaming so internal totals/row_count are computed + with tempfile.NamedTemporaryFile(suffix=".csv") as file: + extractor.stream_sql_to_csv(file_path=file.name, chunk_size=1000) + + # Validate row count + assert extractor.get_row_count() == 10194 + + +def test_streaming_extractor_get_total(): + # Skip if no DB configured + if not os.environ.get('CONNECTION_STRING'): + pytest.skip("Skipping SQL test as no database configured") + + dataset = dataset_from_json(os.path.join('test', 'dataset.json')) + metadata = metadata_from_json(os.path.join('test', 'metadata.json')) + + configuration = Configuration( + connection_string=os.environ.get('CONNECTION_STRING'), + schema="dev", + view="superstore", + query_builder=ViewBasedQueryBuilder + ) + + extractor = StreamingDataExtractor( + dataset_specification=dataset, + metadata=metadata, + configuration=configuration + ) + + with tempfile.NamedTemporaryFile(suffix=".csv") as file: + extractor.stream_sql_to_csv(file_path=file.name, chunk_size=1000) + + # Default measure (dataset.measures[0]) + total_default = extractor.get_total() + + # Explicit measure + total_sales = extractor.get_total(measure="Sales") + + assert round(total_default, 6) == round(2326534.3543, 6) + assert round(total_sales, 6) == round(2326534.3543, 6) + + +def test_streaming_extractor_get_total_multiple_measures(): + if not os.environ.get('CONNECTION_STRING'): + pytest.skip("Skipping SQL test as no database configured") + + dataset = dataset_from_json(os.path.join('test', 'dataset.json')) + dataset.measures = ['Sales', 'Profit'] + + metadata = metadata_from_json(os.path.join('test', 'metadata.json')) + + configuration = Configuration( + connection_string=os.environ.get('CONNECTION_STRING'), + schema="dev", + view="superstore", + query_builder=ViewBasedQueryBuilder + ) + + extractor = StreamingDataExtractor( + dataset_specification=dataset, + metadata=metadata, + configuration=configuration + ) + + with tempfile.NamedTemporaryFile(suffix=".csv") as file: + extractor.stream_sql_to_csv(file_path=file.name, chunk_size=1000) + + assert round(extractor.get_total(measure="Sales"), 4) == 2326534.3543 + assert round(extractor.get_total(measure="Profit"), 4) == 292296.8146 + + +def test_streaming_extractor_get_total_no_measures(): + if not os.environ.get('CONNECTION_STRING'): + pytest.skip("Skipping SQL test as no database configured") + + dataset = dataset_from_json(os.path.join('test', 'dataset.json')) + dataset.measures = [] # no measures + + metadata = metadata_from_json(os.path.join('test', 'metadata.json')) + + configuration = Configuration( + connection_string=os.environ.get('CONNECTION_STRING'), + schema="dev", + view="superstore", + query_builder=ViewBasedQueryBuilder + ) + + extractor = StreamingDataExtractor( + dataset_specification=dataset, + metadata=metadata, + configuration=configuration + ) + + with tempfile.NamedTemporaryFile(suffix=".csv") as file: + extractor.stream_sql_to_csv(file_path=file.name, chunk_size=1000) + + # Falls back to row count + assert extractor.get_total() == 10194 + + with pytest.raises(ValueError): + extractor.get_total(measure="Sales") + + def test_consistency_across_extractors(tmp_path): """ Ensure get_row_count and get_total are consistent across: