diff --git a/mario/athena.py b/mario/athena.py new file mode 100644 index 0000000..b70ec36 --- /dev/null +++ b/mario/athena.py @@ -0,0 +1,193 @@ +from pyathena import connect +from mario.data_extractor import DataExtractor, Configuration +import logging +import datetime +import numbers +import re + +from mario.mapping import rewrite_csv_header_with_fieldmapping +from mario.options import CsvOptions +from mario.utils import gzip_file + +logger = logging.getLogger(__name__) + + +class AthenaConfiguration(Configuration): + """ + Extended configuration + """ + + def __init__(self): + super().__init__() + self.aws_s3_staging_dir = None + self.aws_region_name = None + self.aws_athena_workgroup = 'primary' + self.catalog = 'awsdatacatalog' + self.query_format = 'snake_case' + + +class AthenaDataExtractor(DataExtractor): + """ + Streaming extractor using Athena + PyAthena. + Extends StreamingDataExtractor; only get_connection() is Athena-specific. + """ + + def __init__(self, configuration: AthenaConfiguration, dataset_specification, metadata): + super().__init__(configuration, dataset_specification, metadata) + self.configuration = configuration + + def get_connection(self): + """ + PyAthena provides a DBAPI connection compatible with pandas.read_sql(). + """ + cfg = self.configuration + + if cfg.hook: + # If the user provided a hook, delegate to it + # See: https://airflow.apache.org/docs/apache-airflow-providers-amazon/stable/_api/airflow/providers/amazon/aws/hooks/athena_sql/index.html + return cfg.hook.get_conn() + + # Expect these in configuration: + return connect( + s3_staging_dir=cfg.aws_s3_staging_dir, + region_name=cfg.aws_region_name, + work_group=cfg.aws_athena_workgroup, + schema_name=cfg.schema, + 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 + """ + 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 + 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) + else: + raise NotImplementedError + + totals_df = read_sql(totals_query[0], self.get_connection(), params=totals_query[1]) + return totals_df.iat[0, 0] + + def run_query(self) -> str: + """ + Runs the SQL query in Athena and returns the + result path in S3 + :return: the S3 path where the result is stored + """ + import awswrangler as wr + import boto3 + + # Build SQL + self.__build_query__() + raw_sql, parameters = self._query + sql = interpolate_athena_sql(raw_sql, parameters) + cfg = self.configuration + + # 1. Run SQL via Wrangler + qid = wr.athena.start_query_execution( + sql=sql, + database=cfg.schema, + workgroup=cfg.aws_athena_workgroup, + s3_output=cfg.aws_s3_staging_dir, + ) + wr.athena.wait_query(qid) + + # 2. Get S3 CSV result path + client = boto3.client("athena") + meta = client.get_query_execution(QueryExecutionId=qid) + s3_uri = meta["QueryExecution"]["ResultConfiguration"]["OutputLocation"] + + return s3_uri + + def save_data_as_csv(self, file_path: str, **kwargs): + """ + Athena doesn't support streaming, but natively saves CSV files + in S3 as output, so we really don't need to do anything else + other than run the query and download the results from S3 + :param file_path: + :param kwargs: + :return: + """ + import awswrangler as wr + + # Parse options + options = CsvOptions(**kwargs) + + # Run query and get output location + s3_uri = self.run_query() + + # Download raw Athena CSV + wr.s3.download(path=s3_uri, local_file=file_path) + + # Rewrite header with FieldMapping + rewrite_csv_header_with_fieldmapping(file_path, self.mapping) + + if options.compress_using_gzip: + gz_path = gzip_file(file_path) + return gz_path + + return file_path + + +ATHENA_PARAM_PATTERN = re.compile(r'%\((?P[a-zA-Z0-9_]+)\)s') + + +def athena_escape_string(value: str) -> str: + """ + Escape a Python string for safe Athena SQL literal usage. + Athena uses standard SQL single-quoted strings; internal quotes doubled. + """ + return value.replace("'", "''") + + +def athena_literal(value): + """ + Convert a Python value into an Athena-safe SQL literal. + """ + # None -> NULL + if value is None: + return "NULL" + + # Boolean -> true/false + if isinstance(value, bool): + return "true" if value else "false" + + # Number -> raw + if isinstance(value, numbers.Number): + return str(value) + + # Date / datetime + if isinstance(value, (datetime.date, datetime.datetime)): + return f"'{value.isoformat()}'" + + # List/tuple -> comma-separated list of literals + if isinstance(value, (list, tuple)): + return ", ".join(athena_literal(v) for v in value) + + # Everything else -> string literal + return f"'{athena_escape_string(str(value))}'" + + +def interpolate_athena_sql(sql: str, params: dict): + """ + Replace all %(name)s placeholders using Athena-safe literals. + """ + + def repl(match): + name = match.group("name") + if name not in params: + 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 diff --git a/mario/data_extractor.py b/mario/data_extractor.py index 268d905..8410669 100644 --- a/mario/data_extractor.py +++ b/mario/data_extractor.py @@ -7,6 +7,7 @@ from mario.dataset_specification import DatasetSpecification from mario.metadata import Metadata from mario.options import CsvOptions, HyperOptions +from mario.mapping import FieldMapping logger = logging.getLogger(__name__) @@ -28,7 +29,8 @@ def __init__(self, schema: str = None, file_path: str = None, query_builder=None, - user: str = None + user: str = None, + query_format=None ): self.connection_string = connection_string self.hook = hook @@ -37,6 +39,7 @@ def __init__(self, self.file_path = file_path self.query_builder = query_builder self.user = user + self.query_format = query_format class DataExtractor: @@ -51,6 +54,10 @@ def __init__(self, self._data = None self._query = None self._total = 0 + self.mapping = FieldMapping( + query_format=configuration.query_format, + items=dataset_specification.items + ) def __load__(self): if self.configuration is not None: @@ -76,7 +83,9 @@ def __load_from_sql__(self): logger.info("Executing query") from sqlalchemy import create_engine engine = create_engine(self.configuration.connection_string) - self._data = pd.read_sql(sql=self._query[0], con=engine.connect(), params=self._query[1]) + df = pd.read_sql(sql=self._query[0], con=engine.connect(), params=self._query[1]) + df = self.mapping.df_to_logical(df) + self._data = df def __build_query__(self): logger.info("Building query") @@ -105,22 +114,12 @@ def __load_from_hyper__(self): table=table ) - def __get_column_name__(self, item: str): - """ Returns the column name for a metadata item""" - meta = self.metadata.get_metadata(item) - if meta.get_property('output_name') is not None: - return meta.get_property('output_name') - elif meta.get_property('physical_column_name') is not None: - return meta.get_property('physical_column_name') - else: - return meta.name - def __minimise_data__(self): """ Minimise data so we only keep the columns in the spec """ columns_to_keep = [] for item in self.dataset_specification.items: if self.metadata.get_metadata(item) and not self.metadata.get_metadata(item).get_property('formula'): - columns_to_keep.append(self.__get_column_name__(item)) + columns_to_keep.append(item) self._data = self._data[columns_to_keep] def __get_measure__(self, measure=None): @@ -171,7 +170,7 @@ def __drop_row_numbers__(self): if self._data is None: self.__load__() if 'row_number' in self._data.columns: - self._data = self._data .drop(columns=['row_number']) + self._data = self._data.drop(columns=['row_number']) def save_query(self, file_path: str, formatted: bool = False): """ @@ -229,6 +228,7 @@ class HyperFile(DataExtractor): Wrapper for a HyperFile as a data extractor - use when no data needs to be extracted, and we just want to treat a hyper as a hyper with no conversion to/from dataframe """ + def __init__(self, configuration: Configuration, dataset_specification: DatasetSpecification, @@ -289,7 +289,7 @@ def save_data_as_hyper(self, file_path: str, **kwargs): ) save_hyper_as_hyper(hyper_file=self.configuration.file_path, file_path=file_path, **kwargs) - def save_data_as_csv(self,file_path: str, **kwargs): + def save_data_as_csv(self, file_path: str, **kwargs): from mario.hyper_utils import save_hyper_as_csv options = CsvOptions(**kwargs) if options.minimise: @@ -309,6 +309,7 @@ class StreamingDataExtractor(DataExtractor): supporting streaming data from SQL to output formats without holding any data in memory using a data frame """ + def __init__(self, configuration: Configuration, dataset_specification: DatasetSpecification, @@ -396,6 +397,7 @@ def stream_sql_to_hyper(self, file_path: str, **kwargs): table_name = TableName(options.schema, options.table) row_counter = 0 for df in pd.read_sql(self._query[0], connection, chunksize=options.chunk_size): + df = self.mapping.df_to_logical(df) if options.validate or options.minimise or options.include_row_numbers: self._data = df if options.validate: @@ -410,6 +412,7 @@ def stream_sql_to_hyper(self, file_path: str, **kwargs): frame_to_hyper(df, database=file_path, table=table_name, table_mode='a') def stream_sql_query_to_csv(self, file_path, query, connection, row_counter=0, **kwargs) -> int: + from mario.query_builder import get_formatted_query options = CsvOptions(**kwargs) if options.compress_using_gzip: compression_options = dict(method='gzip') @@ -420,7 +423,8 @@ def stream_sql_query_to_csv(self, file_path, query, connection, row_counter=0, * mode = 'w' header = True - for df in pd.read_sql(sql=query[0], params=query[1], con=connection, chunksize=options.chunk_size): + for df in pd.read_sql(get_formatted_query(query[0], query[1]), connection, chunksize=options.chunk_size): + df = self.mapping.df_to_logical(df) if options.validate or options.minimise: self._data = df if options.validate: @@ -544,6 +548,7 @@ class PartitioningExtractor(StreamingDataExtractor): A data extractor that loads from SQL in batches using a specified constraint to partition by """ + def __init__(self, configuration: Configuration, dataset_specification: DatasetSpecification, @@ -670,5 +675,3 @@ def stream_sql_to_hyper(self, file_path: str, **kwargs): else: logger.info(f"Saving {options.chunk_size} rows to file") frame_to_hyper(df, database=file_path, table=table_name, table_mode='a') - - diff --git a/mario/dataset_splitter.py b/mario/dataset_splitter.py index 7ff0639..001171a 100644 --- a/mario/dataset_splitter.py +++ b/mario/dataset_splitter.py @@ -123,7 +123,7 @@ def process_batch(self, batch, column_name, file_handles, file_name): if value not in file_handles: os.makedirs(os.path.join(self.output_path, value), exist_ok=True) file_path = self.get_output_path(field_value=value, file=file_name) - file_handles[value] = open(file_path, 'w', newline='') + file_handles[value] = open(file_path, 'w', newline='', encoding='utf-8') writer = csv.DictWriter(file_handles[value], fieldnames=row.keys()) writer.writeheader() writer = csv.DictWriter(file_handles[value], fieldnames=row.keys()) @@ -142,7 +142,7 @@ def split_gzipped_csv(self, file_name: str, batch_size=10000): output_file_name = file_name.rstrip('.gz') - with gzip.open(file_path, 'rt', newline='') as infile: + with gzip.open(file_path, 'rt', newline='', encoding='utf-8') as infile: reader = csv.DictReader(infile) # Process the input file in batches @@ -172,7 +172,7 @@ def split_csv(self, file_name: str, batch_size=10000, compression=None): file_path = os.path.join(self.source_path, file_name) # Open the input CSV file for reading - with open(file_path, 'r') as infile: + with open(file_path, 'r', encoding='utf-8') as infile: reader = csv.DictReader(infile) # Process the input file in batches diff --git a/mario/mapping.py b/mario/mapping.py new file mode 100644 index 0000000..ca52174 --- /dev/null +++ b/mario/mapping.py @@ -0,0 +1,70 @@ +from typing import List + +import pandas as pd +import csv +import os +import tempfile + +from mario.utils import to_snake_case + + +class FieldMapping: + """ + Utility class for holding a set of logical-to-physical item mappings + """ + + def __init__(self, query_format, items: List[str]): + self._format = query_format + + self.as_physical = {} + self.as_logical = {} + + for item in items: + self.as_physical[item] = self.map_item(item) + self.as_logical[self.map_item(item)] = item + + def map_item(self, item): + if self._format is None: + return item + if self._format == 'snake_case': + return to_snake_case(item) + + def df_to_logical(self, df: pd.DataFrame) -> pd.DataFrame: + """ + Renames the columns in the dataframe from their physical to their logical names + :param df: pandas Dataframe + :return: pandas Dataframe + """ + return df.rename(columns=self.as_logical) + + +def rewrite_csv_header_with_fieldmapping(local_path, field_mapping): + """ + Replace the first-line header in a CSV using field_mapping, but stream + all remaining lines without loading the file into memory. + """ + # Create a temp file in the same directory (safer for atomic replace) + dir_name = os.path.dirname(local_path) + fd, temp_path = tempfile.mkstemp(dir=dir_name) + os.close(fd) + + with open(local_path, "r", encoding="utf-8", newline="") as src, \ + open(temp_path, "w", encoding="utf-8", newline="") as dst: + + reader = csv.reader(src) + writer = csv.writer(dst, quoting=csv.QUOTE_MINIMAL) + + # --- Read & rewrite only the first line --- + physical_header = next(reader) + logical_header = [ + field_mapping.as_logical.get(col, col) + for col in physical_header + ] + writer.writerow(logical_header) + + # --- Stream the rest unchanged --- + for row in reader: + writer.writerow(row) + + # Atomic replace of original file + os.replace(temp_path, local_path) diff --git a/mario/query_builder.py b/mario/query_builder.py index 4093b23..1dd3c03 100644 --- a/mario/query_builder.py +++ b/mario/query_builder.py @@ -7,6 +7,7 @@ from mario.data_extractor import Configuration from mario.dataset_specification import DatasetSpecification from mario.metadata import Metadata +from mario.mapping import FieldMapping logger = logging.getLogger(__name__) @@ -31,6 +32,11 @@ def __init__(self, self.configuration = configuration self.metadata = metadata self.dataset_specification = dataset_specification + items = dataset_specification.items + if dataset_specification.constraints: + for constraint in dataset_specification.constraints: + items.append(constraint.item) + self.mapping = FieldMapping(query_format=configuration.query_format, items=items) def create_query(self) -> [str, List[any]]: raise NotImplementedError @@ -57,7 +63,7 @@ def create_totals_query(self, measure=None) -> [str, List[any]]: if measure is None: _sql = f'SELECT COUNT(*) FROM "' + self.configuration.schema + '"."' + self.configuration.view + '"' else: - _sql = f'SELECT SUM("'+measure+'") FROM "' + self.configuration.schema + '"."' + self.configuration.view + '"' + _sql = f'SELECT SUM("'+self.mapping.as_physical[measure]+'") FROM "' + self.configuration.schema + '"."' + self.configuration.view + '"' _params = [] return [_sql, _params] @@ -127,12 +133,13 @@ def create_query(self) -> [str, List[any]]: # Don't include calculated fields meta = self.metadata.get_metadata(field) if not meta.get_property('formula'): - select_fields.append(field) + select_fields.append(self.mapping.as_physical[field]) group_fields = select_fields.copy() # remove measures from regular select measures = [] - for measure in self.dataset_specification.measures: + for measure_field in self.dataset_specification.measures: + measure = self.mapping.as_physical[measure_field] if measure in select_fields: select_fields.remove(measure) group_fields.remove(measure) @@ -162,7 +169,7 @@ def create_constraints(self, q): clauses = [] parameters = {} for constraint in self.dataset_specification.constraints: - column = constraint.item + column = self.mapping.as_physical[constraint.item] placeholders = [] for i in range(len(constraint.allowed_values)): parameter_name = column.replace(" ", "_") + str(i) diff --git a/mario/utils.py b/mario/utils.py index ffe3149..6d43f6f 100644 --- a/mario/utils.py +++ b/mario/utils.py @@ -7,3 +7,25 @@ def append_current_date_to_file_name(file_name: str) -> str: filename = os.path.splitext(os.path.basename(file_name))[0] extension = os.path.splitext(os.path.basename(file_name))[1] return filename + '_' + date_time + extension + + +def to_snake_case(name: str) -> str: + """ + Convert a name to lowercase_with_underscores. + Example: "Academic Year" -> "academic_year". + """ + import re + name = name.strip().lower() + name = re.sub(r"[^\w]+", "_", name) + name = re.sub(r"__+", "_", name).strip("_") + return name + + +def gzip_file(input_path, output_path=None): + import gzip + import shutil + if output_path is None: + output_path = input_path + ".gz" + with open(input_path, "rb") as f_in, gzip.open(output_path, "wb") as f_out: + shutil.copyfileobj(f_in, f_out) + return output_path diff --git a/mario/validation.py b/mario/validation.py index 5645d9b..8e85bfa 100644 --- a/mario/validation.py +++ b/mario/validation.py @@ -3,6 +3,7 @@ from mario.data_extractor import Configuration from mario.dataset_specification import DatasetSpecification +from mario.mapping import FieldMapping from mario.metadata import Metadata, Item logger = logging.getLogger(__name__) @@ -419,9 +420,24 @@ def __init__(self, configuration: Configuration ): super().__init__(dataset_specification, metadata) - self.connection = self.__get_connection__(configuration.connection_string) + if configuration.hook: + # If the user provided a hook, delegate to it + self.connection = configuration.hook.get_conn() + else: + self.connection = self.__get_connection__(configuration.connection_string) self.schema = configuration.schema self.view = configuration.view + self.mapping = FieldMapping(query_format=configuration.query_format, items=dataset_specification.items) + self.cached_column_values = {} + + def __get_column_name__(self, item: Item): + """ Returns the column name for a metadata item""" + if item.get_property('output_name') is not None: + return item.get_property('output_name') + elif item.get_property('physical_column_name') is not None: + return item.get_property('physical_column_name') + else: + return self.mapping.as_physical[item.name] def __get_connection__(self, connection_string: str): from sqlalchemy import create_engine @@ -472,9 +488,22 @@ def __get_minimum_maximum_values__(self, item:Item): def __get_column_values__(self, item: Item): import pandas as pd column = self.__get_column_name__(item) + + # Get cached values if present + if column in self.cached_column_values: + return self.cached_column_values[column] + + # If the column is not present, clear the cache - we don't want to + # cache all the columns as that uses more memory + self.cached_column_values = {} + + # Load unique values sql = f"SELECT DISTINCT \"{column}\" AS checkfield FROM {self.schema}.{self.view}" df = pd.read_sql(sql, self.connection) values = df['checkfield'].to_list() + + # Cache values so we don't repeat this query + self.cached_column_values[column] = values return values def __get_data_for_hierarchy__(self, name): @@ -495,7 +524,7 @@ def check_column_present(self, item: Item): sql = f"SELECT COLUMN_NAME as col FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_NAME='{self.view}' AND TABLE_SCHEMA='{self.schema}'" df = pd.read_sql(sql, self.connection) values = df['col'].to_list() - if item.name not in values: + if self.__get_column_name__(item) not in values: self.errors.append(f"Validation error: '{item.name}' in specification is missing from dataset") return False return True diff --git a/requirements.txt b/requirements.txt index 7755b73..6b37b67 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,4 +4,7 @@ tableauhyperapi~=0.0.18161 tableau-builder==0.18 pypika sqlalchemy -openpyxl \ No newline at end of file +openpyxl +pyathena +awswrangler +boto3 \ No newline at end of file diff --git a/setup.py b/setup.py index d6862b8..310aff7 100644 --- a/setup.py +++ b/setup.py @@ -24,6 +24,7 @@ ], extras_require={ 'Airflow': ['apache-airflow-providers-common-sql'], - 'Tableau': ['pantab', 'tableauhyperapi', 'tableau-builder==0.18'] + 'Tableau': ['pantab', 'tableauhyperapi', 'tableau-builder==0.18'], + 'Athena': ['pyathena', 'awswrangler', 'boto3'] } ) \ No newline at end of file diff --git a/test/test_athena_extractor.py b/test/test_athena_extractor.py new file mode 100644 index 0000000..248903b --- /dev/null +++ b/test/test_athena_extractor.py @@ -0,0 +1,263 @@ +from copy import copy + +import pytest + +from mario.athena import AthenaConfiguration, AthenaDataExtractor +from mario.dataset_specification import DatasetSpecification, Constraint +from mario.query_builder import SubsetQueryBuilder +from mario.metadata import Metadata, Item +import os +import shutil +import pandas as pd + +from mario.validation import SqlValidator + +AWS_PROFILE = os.environ.get('AWS_PROFILE') +AWS_ATHENA_RESULTS_DIR = os.environ.get('AWS_ATHENA_RESULTS_DIR') +AWS_REGION = os.environ.get('AWS_REGION') + + +class MockHook: + def __init__(self, extractor: AthenaDataExtractor): + self.extractor = extractor + + def get_conn(self): + return self.extractor.get_connection() + + +def get_test_conf(): + dataset = DatasetSpecification() + dataset.dimensions = [ + "Academic Year", + "Mode of study", + "Country of HE provider" + ] + dataset.measures = ['Number'] + dataset.name = 'student_open_data' + metadata = Metadata() + academic_year = Item() + academic_year.name = 'Academic Year' + mode_of_study = Item() + mode_of_study.name = 'Mode of study' + mode_of_study.set_property('domain', ['Full-time', 'Part-time']) + country_of_he_provider = Item() + country_of_he_provider.name = 'Country of HE provider' + country_of_he_provider.set_property('domain', ['England', 'Wales', 'Scotland', 'Northern Ireland']) + number_field = Item() + number_field.name = 'Number' + metadata.add_item(academic_year) + metadata.add_item(mode_of_study) + metadata.add_item(number_field) + metadata.add_item(country_of_he_provider) + + config = AthenaConfiguration() + config.aws_s3_staging_dir = AWS_ATHENA_RESULTS_DIR + config.aws_region_name = AWS_REGION + + return dataset, metadata, config + + +def test_athena_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_total(measure=dataset.measures[0]) + assert total == 28_733_910 + + +def test_athena_save_data_as_csv(): + if not AWS_PROFILE or not AWS_ATHENA_RESULTS_DIR or not AWS_REGION: + pytest.skip("Skipping Athena test as AWS not configured") + + shutil.rmtree('output/test_athena', ignore_errors=True) + os.makedirs('output/test_athena', exist_ok=True) + + 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 + ) + + extractor.save_data_as_csv( + file_path='output/test_athena/test.csv', + minimise=False, + compress_using_gzip=False, + do_not_modify_source=True + ) + + # Load and test + df = pd.read_csv('output/test_athena/test.csv') + for column in dataset.dimensions: + assert column in df.columns + assert 'Number' in df.columns + for dimension in dataset.dimensions: + assert dimension in df.columns + assert len(df.columns) == len(dataset.items) + + +def test_athena_validate(): + if not AWS_PROFILE or not AWS_ATHENA_RESULTS_DIR or not AWS_REGION: + pytest.skip("Skipping Athena test as AWS not configured") + + shutil.rmtree('output/test_athena', ignore_errors=True) + os.makedirs('output/test_athena', exist_ok=True) + + 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 + ) + mock_hook = MockHook(extractor) + cfg_validator = copy(cfg) + cfg_validator.hook = mock_hook + validator = SqlValidator( + dataset_specification=dataset, + configuration=cfg_validator, + metadata=metadata + ) + validator.validate_data(allow_nulls=False) + + +def test_athena_save_data_as_csv_with_gzip(): + if not AWS_PROFILE or not AWS_ATHENA_RESULTS_DIR or not AWS_REGION: + pytest.skip("Skipping Athena test as AWS not configured") + + shutil.rmtree('output/test_athena', ignore_errors=True) + os.makedirs('output/test_athena', exist_ok=True) + + 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 + ) + + extractor.save_data_as_csv( + file_path='output/test_athena/test.csv', + minimise=False, + compress_using_gzip=True, + do_not_modify_source=True + ) + + # Load and test + df = pd.read_csv('output/test_athena/test.csv.gz') + for column in dataset.dimensions: + assert column in df.columns + assert 'Number' in df.columns + for dimension in dataset.dimensions: + assert dimension in df.columns + assert len(df.columns) == len(dataset.items) + + +def test_athena_with_constraints(): + if not AWS_PROFILE or not AWS_ATHENA_RESULTS_DIR or not AWS_REGION: + pytest.skip("Skipping Athena test as AWS not configured") + + shutil.rmtree('output/test_athena_with_constraints', ignore_errors=True) + os.makedirs('output/test_athena_with_constraints', exist_ok=True) + + dataset, metadata, cfg = get_test_conf() + level_of_study = Item() + level_of_study.name = 'Level of study' + level_constraint = Constraint() + level_constraint.item = level_of_study.name + level_constraint.allowed_values = ['Postgraduate (research)', 'Postgraduate (taught)'] + metadata.add_item(level_of_study) + dataset.dimensions.append(level_of_study.name) + dataset.constraints.append(level_constraint) + cfg.query_builder = SubsetQueryBuilder + cfg.schema = 'demo' + cfg.view = 'student_open_data' + + extractor = AthenaDataExtractor( + configuration=cfg, + metadata=metadata, + dataset_specification=dataset + ) + + extractor.save_data_as_csv( + file_path='output/test_athena_with_constraints/test.csv', + minimise=False, + compress_using_gzip=False, + do_not_modify_source=True + ) + + # Load and test + df = pd.read_csv('output/test_athena_with_constraints/test.csv') + for column in dataset.dimensions: + assert column in df.columns + assert 'Number' in df.columns + for dimension in dataset.dimensions: + assert dimension in df.columns + assert len(df.columns) == len(dataset.items) + assert len(df['Level of study'].unique()) == 2 + + +def test_athena_with_constraints_and_apostrophes(): + if not AWS_PROFILE or not AWS_ATHENA_RESULTS_DIR or not AWS_REGION: + pytest.skip("Skipping Athena test as AWS not configured") + + shutil.rmtree('output/test_athena_with_constraints_and_apostrophes', ignore_errors=True) + os.makedirs('output/test_athena_with_constraints_and_apostrophes', exist_ok=True) + + dataset, metadata, cfg = get_test_conf() + provider = Item() + provider.name = 'HE Provider' + provider_constraint = Constraint() + provider_constraint.item = provider.name + provider_constraint.allowed_values = ["Queen's University Belfast", "City St George's, University of London"] + metadata.add_item(provider) + dataset.dimensions.append(provider.name) + dataset.constraints.append(provider_constraint) + cfg.query_builder = SubsetQueryBuilder + cfg.schema = 'demo' + cfg.view = 'student_open_data' + + extractor = AthenaDataExtractor( + configuration=cfg, + metadata=metadata, + dataset_specification=dataset + ) + + extractor.save_data_as_csv( + file_path='output/test_athena_with_constraints_and_apostrophes/test.csv', + minimise=False, + compress_using_gzip=False, + do_not_modify_source=True + ) + + # Load and test + df = pd.read_csv('output/test_athena_with_constraints_and_apostrophes/test.csv') + for column in dataset.dimensions: + assert column in df.columns + assert 'Number' in df.columns + for dimension in dataset.dimensions: + assert dimension in df.columns + assert len(df.columns) == len(dataset.items) + assert len(df['HE Provider'].unique()) == 2 + assert "Queen's University Belfast" in df['HE Provider'].unique() \ No newline at end of file diff --git a/test/test_dataset_splitter.py b/test/test_dataset_splitter.py index 3a4873e..170dec0 100644 --- a/test/test_dataset_splitter.py +++ b/test/test_dataset_splitter.py @@ -9,7 +9,7 @@ def count_rows_in_csv(file_path): - with open(file_path, 'r') as file: + with open(file_path, 'r', encoding='utf-8') as file: reader = csv.reader(file) row_count = sum(1 for row in reader) - 1 # Subtract 1 to exclude the header row return row_count