From 706b635ff39431f9acd2f12ab9e9067c252883e9 Mon Sep 17 00:00:00 2001 From: Michael Granitzer Date: Mon, 24 Jun 2024 08:05:33 +0200 Subject: [PATCH] feat: added pyarrow producer / consumer pipeline with opensearch consumer --- docs/source/streams.md | 23 +++ owilix/_version.py | 4 +- owilix/cli.py | 12 +- owilix/cmd/local.py | 2 +- owilix/cmd/query.py | 224 +++++++---------------------- owilix/cmd/remote.py | 28 ++-- owilix/core/duckdb.py | 75 ++-------- owilix/core/metadata.py | 109 +++++++------- owilix/core/repository.py | 5 +- owilix/core/stream.py | 165 +++++++++++++++++++++ owilix/plugins/__init__.py | 0 owilix/plugins/search_consumers.py | 97 +++++++++++++ tests/do_oa.py | 116 +++++++++++++++ tests/test_cli.py | 174 ++++++++++++++++++++++ 14 files changed, 726 insertions(+), 308 deletions(-) create mode 100644 docs/source/streams.md create mode 100644 owilix/core/stream.py create mode 100644 owilix/plugins/__init__.py create mode 100644 owilix/plugins/search_consumers.py create mode 100644 tests/do_oa.py create mode 100644 tests/test_cli.py diff --git a/docs/source/streams.md b/docs/source/streams.md new file mode 100644 index 0000000..938e72e --- /dev/null +++ b/docs/source/streams.md @@ -0,0 +1,23 @@ +# PyArrow Streams + +`owilix` supports creating pyarrow streams from a select / where query and consuming it with different consumers, configured via the command line. +Streaming offers a flexible way to work with OWI data that are either streamed from local files or +remote files. + +Streaming supports the following consumers (which can be again provide the stream to exgternal proceses) + +## Streaming to Stdout + +A standard case is to reuse the stream in another data via stdin. This can be done by using the `owilix.core.stream.ConsumeToStdout` consumer and the following command as example + +```sh +owilix query stream --remote lrz:2023-12-3 select=url,title "where=url_suffix='at'" | nc -l 1234 +``` + +## Streaming to a host:port + +Another common use case is to stream the data to a network port. This can be done by using the `owilix.core.stream.ConsumeToSocket` consumer and the following command as example + +```sh +owilix query stream --local all:2023-12-3 select=url,title "where=url_suffix='at'" consumer="owilix.core.stream.ConsumeToSocket" host="localhost" port=1234 +``` diff --git a/owilix/_version.py b/owilix/_version.py index ca3a9a8..0de2ddb 100644 --- a/owilix/_version.py +++ b/owilix/_version.py @@ -1,3 +1,3 @@ # These version placeholders will be replaced later during substitution. -__version__ = "0.10.0" -__version_tuple__ = (0, 8, 0, "post", 1, "9d9a2b7") +__version__ = "0.10.0-post.2+06b6e1e" +__version_tuple__ = (0, 10, 0, "post", 2, "06b6e1e") diff --git a/owilix/cli.py b/owilix/cli.py index 68c5fc0..b178099 100644 --- a/owilix/cli.py +++ b/owilix/cli.py @@ -198,7 +198,9 @@ def logs(ctx, module): """ if module=="lexis": fn = ctx.obj['OWI'].get_lexis_log_filename() - if not os.path.exists(fn): ctx.obj["CONSOLE"].log(f"Lexis log {fn} not found") + if not os.path.exists(fn): + ctx.obj["CONSOLE"].log(f"Lexis log {fn} not found") + return with open(ctx.obj['OWI'].get_lexis_log_filename(), "r") as log: lines = log.readlines() ctx.obj["CONSOLE"].log(lines) @@ -217,10 +219,14 @@ def logs(ctx, module): ctx.obj["CONSOLE"].print("Unkown module {module}: options are lexis, events, errors") -def main(): - cmds = [local, clean, logs, config, remote, admin, query] +def register_commands(cli): + cmds = [local, clean, logs, config, remote, admin, query] for i in cmds: cli.add_command(i) + + +def main(): + register_commands(cli) cli(obj={}) diff --git a/owilix/cmd/local.py b/owilix/cmd/local.py index 4404025..af39f3f 100644 --- a/owilix/cmd/local.py +++ b/owilix/cmd/local.py @@ -74,7 +74,7 @@ def ls(self, specifier, *args, **kwargs): @LocalCommands.register -def free(self, specifier, *args, **kwargs): +def free(self, specifier): """ removes the datasets locally diff --git a/owilix/cmd/query.py b/owilix/cmd/query.py index 5416778..15fb00e 100644 --- a/owilix/cmd/query.py +++ b/owilix/cmd/query.py @@ -1,3 +1,4 @@ +import importlib import os import sys import uuid @@ -7,7 +8,8 @@ from typing import Dict, List, Optional from fsspec import AbstractFileSystem from owilix.cmd.base import BaseCommand, SubCommand, currentItemProgress, CommandResult, ask_yes_no -from owilix.core.duckdb import OWIlixSQLQuery, OWIDuckDBSelect, OWIDuckDBCopy, OWIDuckDBArrow +from owilix.core.duckdb import OWIlixSQLQuery, OWIDuckDBSelect, OWIDuckDBCopy +from owilix.core.stream import OWIDuckDBArrow, ConsumeToSocket, ConsumeToStdouts from owilix.core.metadata import Dataset, fill_metadata @@ -197,6 +199,7 @@ def slice(self, local_specifier, remote_specifier, _ds = _ds[0] _same = _ds.update_provenance(_datasets, files=files, select=select, where=where) + _ds.metadata.update(kwargs) # udpate metadata with additional kwargs if len(_same)>0 and not ignore_provenance: self.console.print(f"Dataset with id {_ds.internalID} /{_ds.title} already contains the same " f"files and query in its provenance list. Skipping the following datasets " @@ -259,14 +262,14 @@ def slice(self, local_specifier, remote_specifier, @QueryCommands.register def stream(self, local_specifier: str, remote_specifier: str, - select: str = "url,domain_label,title,plain_text", - where: Optional[str] = "", limit: Optional[int] = None, - files: str = "**/*.parquet", verbose: bool = False, - host:str = None, port:int = None, - pq_batch_size: int = 1, batch_size: int = 100, prefetch: int = 2): + select: str = "url,domain_label,title,plain_text", + where: Optional[str] = "", limit: Optional[int] = None, + files: str = "**/*.parquet", verbose: bool = False, + pq_batch_size: int = 1, batch_size: int = 100, prefetch: int = 2, queue_size:int =5, + consumer: str = "owilix.core.stream.ConsumeToStdouts", **kwargs): """ Executes the query over the datasets selected by specified local and remote specifiers - and applies select and where clause provided in kwargs. Provides a stream of arrow data strucures over stdoout + and applies select and where clause provided in kwargs. Provides a stream of arrow data structures over stdout to be consumed in a pipe. Args: @@ -276,178 +279,59 @@ def stream(self, local_specifier: str, remote_specifier: str, where (str, optional): WHERE clause to be applied in the SELECT statement. Defaults to an empty string. limit (int, optional): Limit on the number of rows to return. Defaults to None. files (str): Glob pattern for selecting files in both local and remote locations. Defaults to "**/*.parquet". - explain (bool): Whether to explain the query instead of executing it. Defaults to False. + verbose (bool): Whether to print detailed information to stderr. Defaults to False. pq_batch_size (int): Number of parquet files to consider in one batch. Defaults to 1. batch_size (int): Number of rows per query to be yielded back. Defaults to 100. + queue_size (int): Size of buffer queue in (multiplied by prefetch). Defaults to 5. prefetch (int): Number of batches to prefetch. Defaults to 2. - """ - # todo: print control messages to stderr, but make it configurable - import pyarrow as pa - import pyarrow.ipc as ipc - import queue - import threading - import sys - - def producer(generator, buffer_queue): - for batch in generator: - buffer_queue.put(batch) # Put the batch in the queue - buffer_queue.put(None) # Signal the end of the stream - - def stream_to_socket(buffer_queue, host, port): - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: - sock.connect((host, port)) - print(f"Connected to {host}:{port}") - - schema_written = False - writer = None - while True: - batch = buffer_queue.get() - if batch is None: - break # End of stream - - if not schema_written: - # Send schema as the first message - writer = ipc.RecordBatchStreamWriter(sock.makefile('wb'), batch.schema) - schema_written = True - - # Write the batch to the socket - writer.write_table(batch) - buffer_queue.task_done() - - # Close the writer after processing all batches - if writer: - writer.close() - print("Finished sending data.") - - def stream_to_output(buffer_queue, output): - schema_written = False - writer = None # Initialize the writer as None - while True: - batch = buffer_queue.get() - if batch is None: - break # End of stream - - if not schema_written: - # Initialize the writer with the schema from the first batch - writer = ipc.RecordBatchStreamWriter(output, batch.schema) - schema_written = True - - writer.write_table(batch) - buffer_queue.task_done() - - # Close the writer after processing all batches - if writer: - writer.close() - - def configure_kafka_producer(): - from confluent_kafka import Producer - import json - - # Kafka Configuration - KAFKA_BOOTSTRAP_SERVERS = 'localhost:9092' - KAFKA_TOPIC = 'pyarrow_stream' - - # Kafka Producer configuration - producer_config = { - 'bootstrap.servers': KAFKA_BOOTSTRAP_SERVERS, - 'client.id': 'pyarrow-stream-producer' - } - - # Create a Kafka producer - return Producer(producer_config) - - def stream_to_kafka(buffer_queue, topic): - while True: - batch = buffer_queue.get() - if batch is None: - break # End of stream - - # Serialize the batch to a bytes buffer - sink = pa.BufferOutputStream() - writer = ipc.RecordBatchStreamWriter(sink, batch.schema) - writer.write_table(batch) - writer.close() - - # Get the serialized data as bytes - data = sink.getvalue().to_pybytes() - - # Send serialized data to Kafka - producer.produce(topic, data) - producer.flush() # Ensure the message is sent - - buffer_queue.task_done() + consumer (str): The fully qualified class name of the consumer to use for processing batches. Defaults to "owilix.core.stream.ConsumeToStdouts". + **kwargs: Additional parameters to be passed to the consumer's constructor. + Raises: + ImportError: If the consumer class cannot be imported. + AttributeError: If the consumer class does not exist in the specified module. + + Examples: + - local stream to stdout: + owilix query stream --local all:2023-12-3 select=url,title "where=url_suffix='at'" + - local stream to be written to a network port: + owilix query stream --local all:2023-12-3 select=url,title "where=url_suffix='at'" consumer="owilix.core.stream.ConsumeToSocket" host="localhost" port=1234 + """ + # Load the consumer class dynamically + module_name, class_name = consumer.rsplit('.', 1) + try: + module = importlib.import_module(module_name) + ConsumerClass = getattr(module, class_name) + except ImportError as e: + print(f"Error: Failed to import module '{module_name}'.", file=sys.stderr) + raise e + except AttributeError as e: + print(f"Error: Module '{module_name}' does not have a class '{class_name}'.", file=sys.stderr) + raise e + + # Instantiate the consumer class with any additional keyword arguments + _, kwargs = self._cast_args(ConsumerClass.__init__, [], kwargs) + consumer_instance = ConsumerClass(**kwargs) + + # Retrieve the list of files and datasets to process all_files = self.get_all_files(files, local_specifier, remote_specifier, print_it=False) _datasets = all_files.keys() all_files = self.group_all_files_by_fs(all_files) + if verbose: print(f"Found '{sum([len(v) for v in all_files.values()])}' parquet files " - f"in {len(all_files.keys())} filesystems over {len(_datasets)}. " - f"Running queries against them and providing results as binary backpressure controlled output stream.", file=sys.stderr) + f"in {len(all_files.keys())} filesystems over {len(_datasets)} datasets. " + f"Running queries against them and providing results as binary backpressure-controlled output stream.", file=sys.stderr) + + # Prepare the SQL query sql = (OWIlixSQLQuery.from_templates("pq_select") .select(select) .where(where) - .limit(limit)) - - db = OWIDuckDBArrow(all_files, sql, explain=False, - pq_batch_size=pq_batch_size, - batch_size=batch_size, - prefetch=prefetch) - buffer_queue = queue.Queue(maxsize=5*prefetch) - - # Start the producer thread - producer_thread = threading.Thread(target=producer, args=(db.query_aggregator(), buffer_queue)) - producer_thread.start() - - if host and port: - import socket - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: - sock.connect((host, port)) - stream_to_output(buffer_queue, sock.makefile('wb')) - else: - # Stream the result to stdout (Linux pipe) - with open(sys.stdout.fileno(), 'wb', closefd=False) as output: - stream_to_output(buffer_queue, output) - - # Wait for the producer to finish - producer_thread.join() - buffer_queue.join() - -@QueryCommands.register -def consume(self): - """ - Consumes an arrow stream from stdin and prints the record batches to the console. Used for testing pyarrow stream command - :param self: - :return: - """ - import pyarrow as pa - import pyarrow.ipc as ipc - import sys - from rich.console import Console - from rich.table import Table - - def print_batch(batch): - # Create a rich Table to display the batch - table = Table(title="Arrow Stream Batch") - - # Add columns to the table - for column in batch.schema.names: - table.add_column(column) - - # Add rows to the table - for row in batch.to_pylist(): - table.add_row(*(str(row[col]) for col in batch.schema.names)) - - # Print the table to the console - self.console.print(table) - - # Open stdin as a file-like object for reading binary data - input_stream = sys.stdin.buffer - - # Create a RecordBatchStreamReader to read the PyArrow stream - reader = ipc.RecordBatchStreamReader(input_stream) - - # Process each record batch in the stream - for i in range(reader.num_record_batches): - batch = reader.get_batch(i) - print_batch(batch) \ No newline at end of file + .limit(str(limit))) + + # Initialize and start the streaming process with the consumer instance + db_instance = OWIDuckDBArrow(all_files, sql, explain=False, + pq_batch_size=pq_batch_size, + batch_size=batch_size, + prefetch=prefetch) + db_instance.stream(consumer_instance, queue_size) diff --git a/owilix/cmd/remote.py b/owilix/cmd/remote.py index b50a4be..e893330 100644 --- a/owilix/cmd/remote.py +++ b/owilix/cmd/remote.py @@ -85,22 +85,20 @@ def ls(self, specifier, *args, details=False, **kwargs): return CommandResult(success=True, object=datasets, msg=f"Showed {len(datasets)} datasets") @RemoteCommands.register -def pull(self, specifier, *args, **kwargs): +def pull(self, specifier, files = "**/*", overwrite = False): """"pulls the datasets with the given specifier to the local repository Args: overwrite(bool) : if True, files are overwritten if they exists (default False) files(str): a glob pattern to select files within the datasets to be pulled """ - file_select = kwargs.get("files", "**/*") - overwrite = kwargs.get("overwrite", False) self.console.print(f"Fetching datasets for specifier {specifier}") datasets = self.list_remote_dataset_by_specifier(specifier) self.show_datasets(datasets) if self.autoyes or ask_yes_no(self.console, f"Download these datasets:"): for d in datasets: self.console.print(f"Fetching files for {d.title}") - _remote_files = [rel_can_path(_p, d.path) for _p in self.owi.remote_data.files(d, file_select)] + _remote_files = [rel_can_path(_p, d.path) for _p in self.owi.remote_data.files(d, files)] _ds_local = self.owi.local.list(access=d.access, query={"internalID": d.internalID, "collectionName": d.collectionName}) if len(_ds_local)==0: @@ -109,7 +107,7 @@ def pull(self, specifier, *args, **kwargs): raise ValueError(f"Multiple datasets found for {d.internalID}. Dataset inconsistent") self.console.print(f"Found {len(_remote_files)} remote files. Syncing with local files") _local_files = [rel_can_path(_p, _ds_local[0].path) - for _p in self.owi.local.files(_ds_local[0], file_select)] + for _p in self.owi.local.files(_ds_local[0], files)] _missing_files = set(_remote_files) - set(_local_files) if not overwrite else set(_remote_files) if len(_missing_files)==0: self.console.print(f"Dataset {d.title} is already up to date. Missing files are {len(_missing_files)}.") @@ -138,7 +136,7 @@ def doctor(self, *args, **kwargs): @RemoteCommands.register -def push(self, specifier, *args, **kwargs): +def push(self, specifier, files = "**/*", mdupdate = True, overwrite = False, **kwargs): """" push local datasets fitting the specifier to the local repository. Pushing may change the internalID of the local dataset if the dataset does not exist on the server. @@ -154,9 +152,6 @@ def push(self, specifier, *args, **kwargs): dataCenter(str): specify the dataCenter the dataset is created in. HNote that only works if the dataset is not already in a different, known datacenter. overwrites the dataCenter specified in the selector. """ - file_select = kwargs.get("files", "**/*") - md_update = kwargs.get("mdupdate", True) - overwrite = kwargs.get("overwrite", False) dataCenter= kwargs.get("dataCenter", self.owi.parse_specifier(specifier).get("data_center",None)) self.console.print(f"Fetching datasets for specifier {specifier}") @@ -165,7 +160,7 @@ def push(self, specifier, *args, **kwargs): if self.autoyes or ask_yes_no(self.console, f"Upload these datasets (overwrite:{overwrite}):"): for d in datasets: self.console.print(f"Fetching files for {d.title}") - _local_files = [rel_can_path(_p, d.path) for _p in self.owi.local.files(d, file_select)] + _local_files = [rel_can_path(_p, d.path) for _p in self.owi.local.files(d, files)] _ds_remote = self.owi.remote_data.list(access=d.access, query={"internalID": d.internalID, "collectionName": d.collectionName}) if len(_ds_remote)==0: @@ -184,7 +179,7 @@ def push(self, specifier, *args, **kwargs): elif len(_ds_remote)>1: raise ValueError(f"Multiple datasets found for {d.internalID}. Dataset inconsistent") else: - if md_update: + if mdupdate: _ds_remote[0].metadata.update(cast_metadata({k:v for k,v in d.metadata.items() if k!="dataCenter"})) _ds_remote[0].reformat_metadata() self.owi.remote_data.update_metadata(_ds_remote[0]) @@ -195,7 +190,7 @@ def push(self, specifier, *args, **kwargs): self.console.print(f"Found {len(_local_files)} local files. " f"Syncing with remote files at {_ds_remote[0].dataCenter}") _remote_files = [rel_can_path(_p, _ds_remote[0].path) - for _p in self.owi.remote_data.files(_ds_remote[0], file_select)] + for _p in self.owi.remote_data.files(_ds_remote[0], files)] _missing_files = set(_local_files) - set(_remote_files) if not overwrite else set(_local_files) if len(_missing_files)==0: self.console.print(f"Dataset {d.title} is already up to date") @@ -219,14 +214,13 @@ def push(self, specifier, *args, **kwargs): return CommandResult(success=True, object=datasets, msg=f"Pushed {len(datasets)} datasets") @RemoteCommands.register -def diff(self, specifier, *args, **kwargs): +def diff(self, specifier, files=None, **kwargs): """ runs a dataset and/or file-level diff Args: files(str): a glob pattern to select files for. Use **/* for all files. If not set, diff will only be applied on the dataset level """ - file_select = kwargs.get("files", None) self.console.print(f"Fetching local and remote datasets for specifier {specifier}") _local_datasets = self.list_local_datasets_by_specifier(specifier) _remote_datasets = self.list_remote_dataset_by_specifier(specifier) @@ -241,7 +235,7 @@ def diff(self, specifier, *args, **kwargs): self.console.print("Dataets available only remotely:") self.show_datasets([_d for _d in _remote_datasets if _d.internalID not in _ids]) - if file_select is None: + if files is None: self.console.print(f"Diff done (for diff on a file level specify files=**/*") return @@ -257,8 +251,8 @@ def diff(self, specifier, *args, **kwargs): for i in _diff_print: i["key"] = i["key"] if i["local"]==i["remote"] else "[warning]"+str(i["key"])+"[/warning]" self.show_table(_diff_print, order="key,local,remote") # now check for file diffs - _local_files = [rel_can_path(_p, _localds.path) for _p in self.owi.local.files(_localds, file_select)] - _remote_files = [rel_can_path(_p, _remoteds.path) for _p in self.owi.remote_data.files(_remoteds, file_select)] + _local_files = [rel_can_path(_p, _localds.path) for _p in self.owi.local.files(_localds, files)] + _remote_files = [rel_can_path(_p, _remoteds.path) for _p in self.owi.remote_data.files(_remoteds, files)] self.console.print(f"Found {len(_local_files)} local files, {len(_remote_files)} remote files. ") _lonly = set(_local_files) - set(_remote_files) if len(_lonly)>1: diff --git a/owilix/core/duckdb.py b/owilix/core/duckdb.py index ce50817..00967e6 100644 --- a/owilix/core/duckdb.py +++ b/owilix/core/duckdb.py @@ -350,6 +350,7 @@ class OWIDuckDBSelect: self.as_dict = as_dict self.max_mem = max_mem self.retry_count = retry_count + self.logger = logging.getLogger("owilix") def retry_then_raise(self, connenction, create_owi_slice_query): @@ -396,11 +397,11 @@ class OWIDuckDBSelect: cursor = self.retry_then_raise(conn, _query.sql) columns = None - _logger.debug(f"Connection to filesystem '{fs.protocol}' for {len(pq_batch.files)} files opened.") + self.logger.debug(f"Connection to filesystem '{fs.protocol}' for {len(pq_batch.files)} files opened.") while True: - _logger.debug(f"Fetching {batch_size} rows from {fs.protocol} for {len(pq_batch.files)} files") + self.logger.debug(f"Fetching {batch_size} rows from {fs.protocol} for {len(pq_batch.files)} files") results = cursor.fetchmany(batch_size) - _logger.debug(f"Fetch of {batch_size} rows done from {fs.protocol} for {len(pq_batch.files)} files") + self.logger.debug(f"Fetch of {batch_size} rows done from {fs.protocol} for {len(pq_batch.files)} files") if not results: break if self.as_dict: @@ -410,12 +411,12 @@ class OWIDuckDBSelect: yield results except Exception as e: - _logger.exception(f"Error when executing {query.sql} on files {pq_batch}") + self.logger.exception(f"Error when executing {query.sql} on files {pq_batch}") raise e finally: if conn: conn.close() - _logger.debug(f"Connection to {fs.protocol} for {len(pq_batch.files)} files closed.") + self.logger.debug(f"Connection to {fs.protocol} for {len(pq_batch.files)} files closed.") tmp_dir.cleanup() def query_aggregator(self) -> Generator[List[tuple], None, None]: @@ -467,7 +468,7 @@ class OWIDuckDBSelect: for result_batch in future.result(): yield result_batch except Exception as e: - _logger.exception(f"Error processing task {task}: {e}") + self.logger.exception(f"Error processing task {task}: {e}") class OWIDuckDBCopy (OWIDuckDBSelect): @@ -543,7 +544,7 @@ class OWIDuckDBCopy (OWIDuckDBSelect): DROP TABLE IF EXISTS owi_slice; CREATE TABLE owi_slice AS """ + query.files([f[0] for f in pq_batch.files]).sql # could be also done in smaller batches - _logger.debug(f"Connection to filesystem '{fs.protocol}' for {len(pq_batch.files)} files opened.") + self.logger.debug(f"Connection to filesystem '{fs.protocol}' for {len(pq_batch.files)} files opened.") # Execute the query to create owi_slice self.retry_then_raise(conn, create_owi_slice_query) @@ -597,7 +598,7 @@ class OWIDuckDBCopy (OWIDuckDBSelect): conn.close() except Exception as e: - _logger.exception(f"Error when executing {query.sql} on files {pq_batch}") + self.logger.exception(f"Error when executing {query.sql} on files {pq_batch}") _error=str(e) finally: @@ -605,62 +606,6 @@ class OWIDuckDBCopy (OWIDuckDBSelect): yield [{"message": _error, "group":_sub_path, "num_files":len(pq_batch.files), "count":0, "success":0}] if conn: conn.close() - _logger.debug(f"Connection to {fs.protocol} for {len(pq_batch.files)} files closed.") + self.logger.debug(f"Connection to {fs.protocol} for {len(pq_batch.files)} files closed.") temp_dir.cleanup() -class OWIDuckDBArrow(OWIDuckDBSelect): - """ - A class to execute a DuckDB select, similar to OWIDuckDBSelect, but using Apache Arrow as return results - - Attributes: - see OWIDuckDBSelect - """ - - def run_query(self, fs: AbstractFileSystem, - pq_batch: ParquetBatch, - query: OWIlixSQLQuery, - batch_size: int = 0) -> Generator[List[tuple], None, None]: - """ - Run a SQL query on a batch of parquet files using DuckDB. - - Args: - fs (AbstractFileSystem): The filesystem containing the parquet files. - pq_batch (List[str]): A batch of parquet file paths to query. - query (OWIlixSQLQuery): The SQL query to execute. - batch_size (int, optional): The number of rows to fetch in each batch. Defaults to self.batch_size. - - Yields: - Generator[List[tuple], None, None]: A generator yielding batches of query results. - """ - if batch_size <= 0: - batch_size = self.batch_size - conn, tmp_dir = None, tempfile.TemporaryDirectory() - try: - temp_db_path = os.path.join(tmp_dir.name, 'temp_owi_select_duckdb.db') - conn = duckdb.connect(database=temp_db_path) - conn.execute(f"PRAGMA memory_limit='{self.max_mem}'") - conn.register_filesystem(fs) - _query = query.files([f[0] for f in pq_batch.files]) if len(pq_batch.files) > 0 else query - # Format the SQL query using the query_args from the ParquetBatch - if pq_batch.query_args: - _query = _query.format(**pq_batch.query_args) - - cursor = self.retry_then_raise(conn,_query.sql) - _logger.debug(f"Connection to filesystem '{fs.protocol}' for {len(pq_batch.files)} files opened.") - while True: - _logger.debug(f"Fetching {batch_size} rows from {fs.protocol} for {len(pq_batch.files)} files") - results = cursor.fetch_arrow_table(batch_size) - _logger.debug(f"Fetch of {batch_size} rows done from {fs.protocol} for {len(pq_batch.files)} files") - if not results or results.num_rows == 0: - break - yield results - - except Exception as e: - _logger.exception(f"Error when executing {query.sql} on files {pq_batch}") - raise e - finally: - if conn: - conn.close() - _logger.debug(f"Connection to {fs.protocol} for {len(pq_batch.files)} files closed.") - tmp_dir.cleanup() - diff --git a/owilix/core/metadata.py b/owilix/core/metadata.py index ae3f2b6..7ea46e5 100644 --- a/owilix/core/metadata.py +++ b/owilix/core/metadata.py @@ -1,8 +1,9 @@ -import datetime + import json import re import uuid from collections import defaultdict +from dateutil import parser from dataclasses import dataclass from datetime import datetime, timedelta import os @@ -417,24 +418,34 @@ class Dataset: """ Represents a dataset in the repository. Metadata are set in the .metadata property and can be accessed as attributes. - Workflow relevant properties are available as atttirbutes (e.g. path) - todo: integrate metadta checks and validation in dataset. + Workflow relevant properties are available as attributes (e.g. path). + TODO: integrate metadata checks and validation in dataset. """ def __init__(self, repository, path, **metadata): self.repository = repository self.path = path - self.metadata = cast_metadata(metadata) + self._metadata = cast_metadata(metadata) # Use a private attribute to store metadata self._change_log = None - if "internalID" not in metadata: + + if "internalID" not in self._metadata: raise ValueError("Dataset metadata must contain an internalID") - if "title" not in metadata: - self.metadata["title"] = "UNKNOWN TITLE" + if "title" not in self._metadata: + self._metadata["title"] = "UNKNOWN TITLE" + # Property to encapsulate metadata + @property + def metadata(self): + return self._metadata + + @metadata.setter + def metadata(self, value): + self._metadata = cast_metadata(value) + self.update_lastchanged() def __getattr__(self, name): - if name in self.metadata: - return self.metadata[name] + if name in self._metadata: + return self._metadata[name] else: raise AttributeError(f"Dataset has no attribute {name}") @@ -444,19 +455,19 @@ class Dataset: @property def startDate(self): try: - return datetime.strptime(self.metadata.get("startDate", None), '%Y-%m-%d') + return parser.parse(self._metadata.get("startDate", None)) except: return None @property def endDate(self): try: - return datetime.strptime(self.metadata.get("endDate", None), '%Y-%m-%d') + return parser.parse(self._metadata.get("endDate", None)) except: return None def get_changelog(self): - if self._change_log==None: + if self._change_log is None: _content = self.repository.readlines(self, "changelog.json") self._change_log = json.loads(_content) if _content else None self._change_log = self._change_log if self._change_log is not None else [] @@ -469,15 +480,16 @@ class Dataset: "msg": msg }) self._change_log = _change_log - if save: self.save_changelog() + if save: + self.save_changelog() def save_changelog(self): - if self._change_log!=None: + if self._change_log is not None: self.repository.writelines(self, "changelog.json", json.dumps(self._change_log)) return self._change_log def update_lastchanged(self): - self.metadata['lastChanged'] = self._get_now_formatted() + self._metadata['lastChanged'] = self._get_now_formatted() def _get_now_formatted(self): return datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f") @@ -487,27 +499,28 @@ class Dataset: Reformat the metadata to a dictionary """ self.metadata["title"] = reformat_metadata(self.metadata)["title"] + self.update_lastchanged() def consolidate_metadata(self, infer=False, **kwargs): """ Consolidate metadata by checking and updating the metadata fields. If infer is True, metadata is inferred from the files in the dataset and updated. - kwargs contains a set of metadata valus to be updated. Key is the metadata field name. - value is either the value to be set, None for deleting the metadata entry or (default) for getting the default value + kwargs contains a set of metadata values to be updated. Key is the metadata field name. + Value is either the value to be set, None for deleting the metadata entry, or (default) for getting the default value. """ - _defaults= fill_metadata({}, overwrite_inferred=False) - if infer: self.metadata.update(infer_metadata_from_files(self.repository.files_details(self))) - for k,v in _defaults.items(): + _defaults = fill_metadata({}, overwrite_inferred=False) + if infer: + self.metadata.update(infer_metadata_from_files(self.repository.files_details(self))) + for k, v in _defaults.items(): if k not in self.metadata: self.metadata[k] = v for key, value in kwargs.items(): if not value: - if key in self.metadata: del self.metadata[key] + if key in self.metadata: + del self.metadata[key] else: self.metadata[key] = value if value != "default" else _defaults[key] - self.metadata = cast_metadata(self.metadata) - - + self.metadata = cast_metadata(self.metadata) # This will trigger the setter and update_lastchanged def filter_by_date_range(self, start_date=None, end_date=None): start = self.startDate @@ -522,29 +535,30 @@ class Dataset: return True # No filtering @staticmethod - def filter_datasets(_datasets, day, duration, query)-> list: - - # query filter + def filter_datasets(_datasets, day, duration, query) -> list: + # Query filter if query is not None and query != {}: def matches_query(_d): _all = [] for key, value in query.items(): regex = key.endswith("*") key = key if not regex else key[:-1] - if not hasattr(_d, key): _all.append(False) + if not hasattr(_d, key): + _all.append(False) if not regex: _all.append(getattr(_d, key) == value) else: - _all.append(re.match(value, getattr(_d,key))) + _all.append(re.match(value, getattr(_d, key))) return all(_all) _datasets = [_d for _d in _datasets if matches_query(_d)] - if len(_datasets)==0: return [] - # time filter + if len(_datasets) == 0: + return [] + # Time filter if day == "latest": - day = max( _d.startDate for _d in _datasets) + day = max(_d.startDate for _d in _datasets) elif day is not None and not isinstance(day, datetime): day = datetime.strptime(day, '%Y-%m-%d') @@ -554,24 +568,23 @@ class Dataset: return [_d for _d in _datasets if _d.filter_by_date_range(startdate, day)] - @staticmethod def from_lexis_http(repository, path, **metadata): def convert_value(k, v): if isinstance(v, list) and len(v) == 1: return v[0] return v + lower_case = ["AlternateIdentifier", "CreationDate", "CustomMetadataSchema", "RelatedSoftware"] for l in lower_case: if l in metadata: - metadata[l[0].lower()+l[1]] = metadata.pop(l) - _md = extract_metadata_from_title(metadata["metadata"].get("title",["UKNOWN TITLE"])[0]) + metadata[l[0].lower() + l[1]] = metadata.pop(l) + _md = extract_metadata_from_title(metadata["metadata"].get("title", ["UKNOWN TITLE"])[0]) _md = _md | metadata["flags"] | metadata["location"] | {k: convert_value(k, v) for k, v in metadata["metadata"].items()} return Dataset(repository, path, **_md) - def update_provenance(self, datasets: List, files:str = None, select:str = None, - where:str=None): - _provenance_new = set([f"{create_provenance_url(d, files, select=select, where=where)}" for d in datasets]) + def update_provenance(self, datasets: List, files: str = None, select: str = None, where: str = None): + _provenance_new = set([f"{create_provenance_url(d, files, select=select, where=where)}" for d in datasets]) _provenance_old = set(self.metadata["provenance"]) _provenance = list(_provenance_old.union(_provenance_new)) self.metadata["provenance"] = _provenance @@ -581,19 +594,19 @@ class Dataset: return [d for d in datasets if d.internalID in _overlapping] @staticmethod - def merge_into_new(repository, datasets: List, files:str = None, select:str = None, - where:str=None, access:ACCESSTYPES="project", - collectionName:str="userslice", **kwargs): + def merge_into_new(repository, datasets: List, files: str = None, select: str = None, + where: str = None, access: str = "project", collectionName: str = "userslice", **kwargs): """ - create a new dataset by merging the provided ones and applying files, select and where filters (for updating the provenance) + Create a new dataset by merging the provided ones and applying files, select and where filters (for updating the provenance) """ _new_md = fill_metadata({}, None, False) - _new_md["provenance"] = [ f"{create_provenance_url(d, files, select=select, where=where)}" for d in datasets] + _new_md["provenance"] = [f"{create_provenance_url(d, files, select=select, where=where)}" for d in datasets] _new_md.update(kwargs) - _new_md["access"]=access - _new_md["collectionName"]=collectionName - if "description" not in kwargs: _new_md["description"] = (f"Merged dataset from {len(datasets)} datasets at " - f"{datetime.now()} using parameters files={files}, " - f"select={select}, where={where}") + _new_md["access"] = access + _new_md["collectionName"] = collectionName + if "description" not in kwargs: + _new_md["description"] = (f"Merged dataset from {len(datasets)} datasets at " + f"{datetime.now()} using parameters files={files}, " + f"select={select}, where={where}") return repository.create(**_new_md) diff --git a/owilix/core/repository.py b/owilix/core/repository.py index c9da7b5..2e5b3c8 100644 --- a/owilix/core/repository.py +++ b/owilix/core/repository.py @@ -609,7 +609,8 @@ class FileBasedRepository(AbstractRepository): if len(_md_errors)>0: logger.warning(f"Metadata inconsistencies for dataset {dataset.internalID}: "+",".join(_md_errors.values())) try: - dataset.update_lastchanged() + if not self.fs.exists(os.path.dirname(_p)): + self.fs.mkdirs(os.path.dirname(_p)) with self.fs.open(_p+".json", "w") as _fh: json.dump(dataset.metadata, _fh) except Exception as e: @@ -741,7 +742,7 @@ class IRODSRepository(FileBasedRepository): _md_errors = validate_metadata(dataset.metadata) if len(_md_errors)>0: logger.warning(f"Metadata inconsistencies for dataset {dataset.internalID}: "+",".join(_md_errors.values())) - dataset.update_lastchanged() + return update_metadata_for_irods_collection(self.session.collections.get(_p), dataset.metadata) def change_id(self, dataset, new_id): diff --git a/owilix/core/stream.py b/owilix/core/stream.py new file mode 100644 index 0000000..afd85ef --- /dev/null +++ b/owilix/core/stream.py @@ -0,0 +1,165 @@ +import os +import queue +import socket +import tempfile +import threading +from typing import Generator, List, Callable + +import duckdb +import pyarrow.ipc as ipc +from abc import ABC, abstractmethod +import pyarrow as pa +from fsspec import AbstractFileSystem + +from owilix.core.duckdb import OWIDuckDBSelect, ParquetBatch, OWIlixSQLQuery +import sys +import pyarrow.ipc as ipc + + + +class Consumer(ABC): + """An abstract class for consuming data from a stream of pyarrow table batches.""" + @abstractmethod + def consume(self, batch: pa.Table): + """Process a batch of data.""" + pass + + @abstractmethod + def close(self): + """Perform any cleanup necessary.""" + pass + +class ConsumeToSocket(Consumer): + """ + Class to consume data from a stream of pyarrow table batches and send it over a socket. + """ + def __init__(self, host: str, port: int): + self.host = host + self.port = port + self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + self.socket.connect((self.host, self.port)) + self.output = self.socket.makefile('wb') + self.writer = None + + def consume(self, batch: pa.Table): + if self.writer is None and isinstance(batch, pa.Table): + self.writer = ipc.RecordBatchStreamWriter(self.output, batch.schema) + + if self.writer and isinstance(batch, pa.Table): + self.writer.write_table(batch) + else: + print("Warning: The provided batch is not a valid pyarrow Table.", file=sys.stderr) + + def close(self): + if self.writer: + self.writer.close() + self.output.close() + self.socket.close() + + +class ConsumeToStdouts(Consumer): + def __init__(self, **kwargs): + self.output = open(sys.stdout.fileno(), 'wb', closefd=False) + self.writer = None + + def consume(self, batch: pa.Table): + if self.writer is None and isinstance(batch, pa.Table): + self.writer = ipc.RecordBatchStreamWriter(self.output, batch.schema) + + if self.writer and isinstance(batch, pa.Table): + self.writer.write_table(batch) + else: + print("Warning: The provided batch is not a valid pyarrow Table.", file=sys.stderr) + + def close(self): + if self.writer: + self.writer.close() + self.output.close() + + +class OWIDuckDBArrow(OWIDuckDBSelect): + """ + A class to execute a DuckDB select, similar to OWIDuckDBSelect, but using Apache Arrow as return results + + Attributes: + see OWIDuckDBSelect + """ + + def run_query(self, fs: AbstractFileSystem, + pq_batch: ParquetBatch, + query: OWIlixSQLQuery, + batch_size: int = 0) -> Generator[List[tuple], None, None]: + """ + Run a SQL query on a batch of parquet files using DuckDB. + + Args: + fs (AbstractFileSystem): The filesystem containing the parquet files. + pq_batch (List[str]): A batch of parquet file paths to query. + query (OWIlixSQLQuery): The SQL query to execute. + batch_size (int, optional): The number of rows to fetch in each batch. Defaults to self.batch_size. + + Yields: + Generator[List[tuple], None, None]: A generator yielding batches of query results. + """ + if batch_size <= 0: + batch_size = self.batch_size + conn, tmp_dir = None, tempfile.TemporaryDirectory() + try: + temp_db_path = os.path.join(tmp_dir.name, 'temp_owi_select_duckdb.db') + conn = duckdb.connect(database=temp_db_path) + conn.execute(f"PRAGMA memory_limit='{self.max_mem}'") + conn.register_filesystem(fs) + _query = query.files([f[0] for f in pq_batch.files]) if len(pq_batch.files) > 0 else query + # Format the SQL query using the query_args from the ParquetBatch + if pq_batch.query_args: + _query = _query.format(**pq_batch.query_args) + + cursor = self.retry_then_raise(conn,_query.sql) + self.logger.debug(f"Connection to filesystem '{fs.protocol}' for {len(pq_batch.files)} files opened.") + while True: + self.logger.debug(f"Fetching {batch_size} rows from {fs.protocol} for {len(pq_batch.files)} files") + results = cursor.fetch_arrow_table(batch_size) + self.logger.debug(f"Fetch of {batch_size} rows done from {fs.protocol} for {len(pq_batch.files)} files") + if not results or results.num_rows == 0: + break + yield results + + except Exception as e: + self.logger.exception(f"Error when executing {query.sql} on files {pq_batch}") + raise e + finally: + if conn: + conn.close() + self.logger.debug(f"Connection to {fs.protocol} for {len(pq_batch.files)} files closed.") + tmp_dir.cleanup() + + def producer(self, buffer_queue): + for batch in self.query_aggregator(): + buffer_queue.put(batch) # Put the batch in the queue + buffer_queue.put(None) # Signal the end of the stream + + def stream(self, consumer: Consumer, queue_size=5): + """ + starts streaming the query, buffer them via a queue of size queue_size*self.prefetch and consume them using the consumer + """ + buffer_queue = queue.Queue(maxsize=queue_size*self.prefetch) + # Start the producer thread + producer_thread = threading.Thread(target=self.producer, args=(buffer_queue,)) + producer_thread.start() + try: + # Consumer logic + while True: + batch = buffer_queue.get() + if batch is None: + break # End of stream + + consumer.consume(batch) + buffer_queue.task_done() + + # Wait for the producer to finish + producer_thread.join() + buffer_queue.join() + + finally: + # Perform any cleanup needed by the consumer + consumer.close() \ No newline at end of file diff --git a/owilix/plugins/__init__.py b/owilix/plugins/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/owilix/plugins/search_consumers.py b/owilix/plugins/search_consumers.py new file mode 100644 index 0000000..bd0a092 --- /dev/null +++ b/owilix/plugins/search_consumers.py @@ -0,0 +1,97 @@ +import os +import json +import logging +from typing import Any +from opensearchpy import OpenSearch, helpers +import pyarrow as pa + +from owilix.core.stream import Consumer + + +class OpenSearchConsumer(Consumer): + def __init__(self, index_name: str, **kwargs): + """ + Initializes the OpenSearch consumer. + + Args: + index_name (str): The name of the index where data will be inserted. + **kwargs: Additional configuration parameters for OpenSearch. + These can include 'host', 'port', 'timeout', etc. + """ + self.index_name = index_name + self.host = kwargs.get('host', 'localhost') + self.port = kwargs.get('port', 9200) + self.scheme = kwargs.get('scheme', 'http') + self.timeout = kwargs.get('timeout', 30) + + self.username = os.getenv('OWIXILX_OPENSEARCH_USERNAME', kwargs.get('username', None)) + self.password = os.getenv('OWIXLIS_OPENSEARCH_PASSWORD', kwargs.get('password', None)) + + self.client = self._create_opensearch_client() + + self.log_file = kwargs.get('log_file', 'opensearch_errors.log') + logging.basicConfig(filename=self.log_file, level=logging.ERROR) + + self._ensure_index_exists() + + def _create_opensearch_client(self) -> OpenSearch: + """ + Creates an OpenSearch client using provided credentials and configurations. + """ + auth = None + if self.username and self.password: + auth = (self.username, self.password) + + return OpenSearch( + hosts=[{'host': self.host, 'port': self.port}], + http_auth=auth, + use_ssl=(self.scheme == 'https'), + timeout=self.timeout + ) + + def _ensure_index_exists(self): + """ + Checks if the specified index exists in OpenSearch. Creates the index if it does not exist. + """ + if not self.client.indices.exists(index=self.index_name): + self.client.indices.create(index=self.index_name) + print(f"Index '{self.index_name}' created in OpenSearch.") + + def consume(self, batch: pa.Table): + """ + Consumes a batch of data and pushes it to OpenSearch. + + Args: + batch (pa.Table): The batch of data to be inserted. + """ + records = batch.to_pylist() + actions = [ + { + "_index": self.index_name, + "_source": record + } + for record in records + ] + + try: + helpers.bulk(self.client, actions) + except Exception as e: + logging.error(f"Failed to insert records into OpenSearch: {e}") + self._log_failed_records(actions) + + def _log_failed_records(self, actions: Any): + """ + Logs the records that failed to be inserted into OpenSearch to a log file. + + Args: + actions (Any): The list of records that failed to insert. + """ + with open(self.log_file, 'a') as log_file: + for action in actions: + log_file.write(json.dumps(action) + '\n') + + def close(self): + """ + Performs any necessary cleanup. For OpenSearch, no explicit cleanup is required. + """ + pass diff --git a/tests/do_oa.py b/tests/do_oa.py new file mode 100644 index 0000000..1697cef --- /dev/null +++ b/tests/do_oa.py @@ -0,0 +1,116 @@ +import inspect +from pydantic import BaseModel, create_model, ValidationError, TypeAdapter +from typing import Optional, Any, Tuple, Dict + + +def convert_args_kwargs_with_pydantic(func, args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> Tuple[ + Tuple[Any, ...], Dict[str, Any]]: + """ + Converts args and kwargs to match the types specified in the function's signature using pydantic. + + Parameters: + func (callable): The function whose signature will be used for type conversion. + args (tuple): The positional arguments to convert. + kwargs (dict): The keyword arguments to convert. + + Returns: + tuple: A tuple containing the converted args and kwargs. + """ + sig = inspect.signature(func) + parameters = list(sig.parameters.values()) + + converted_args = [] + converted_kwargs = {} + + # Convert positional arguments (*args) that match the function signature + for i, arg in enumerate(args): + if i < len(parameters): + param = parameters[i] + expected_type = param.annotation + + if expected_type == inspect.Parameter.empty: + converted_args.append(arg) + else: + try: + # Use TypeAdapter to convert the argument + type_adapter = TypeAdapter(expected_type) + converted_args.append(type_adapter.validate_python(arg)) + except (ValueError, TypeError) as e: + print(f"WARNING - Could not convert arg[{i}]='{arg}' to {expected_type}: {e}") + converted_args.append(arg) + + # Include remaining *args as-is + if len(args) > len(parameters): + converted_args.extend(args[len(parameters):]) + + # Convert keyword arguments (*kwargs) that match the function signature + for name, param in sig.parameters.items(): + if param.kind in (inspect.Parameter.KEYWORD_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD): + if name in kwargs: + expected_type = param.annotation + value = kwargs[name] + if expected_type != inspect.Parameter.empty: + try: + type_adapter = TypeAdapter(expected_type) + converted_kwargs[name] = type_adapter.validate_python(value) + except (ValueError, TypeError) as e: + print(f"WARNING - Could not convert kwarg '{name}'='{value}' to {expected_type}: {e}") + converted_kwargs[name] = value + else: + converted_kwargs[name] = value + + # Include any additional kwargs that weren't in the function signature + for k, v in kwargs.items(): + if k not in converted_kwargs: + converted_kwargs[k] = v + + # Ensure no positional argument conflicts with keyword arguments + for i, arg in enumerate(converted_args): + if i < len(parameters): + param_name = parameters[i].name + if param_name in converted_kwargs: + raise TypeError(f"Got multiple values for argument '{param_name}'") + + return tuple(converted_args), converted_kwargs + + +# Example function +def example_function( + local_specifier: str, + remote_specifier: str, + select: str = "url,domain_label,title,plain_text", + where: Optional[str] = "", + limit: Optional[int] = None, + files: str = "**/*.parquet", + explain: bool = False, + pq_batch_size: int = 1, + batch_size: int = 100, + prefetch: int = 2, + page_size: int = 10, + *args: Any, + **kwargs: Any +): + return locals() + + +# Example usage +args = ("example_local", "example_remote") +kwargs = { + "select": "url,domain_label,title", + "where": "url_suffix='at'", + "limit": "10", + "files": "**/*.parquet", + "explain": "false", + "pq_batch_size": "5", + "batch_size": "200", + "prefetch": "3", + "page_size": "15" +} + +converted_args, converted_kwargs = convert_args_kwargs_with_pydantic(example_function, args, kwargs) +print("Converted Args:", converted_args) +print("Converted Kwargs:", converted_kwargs) + +# Now you can call the function with the converted arguments +result = example_function(*converted_args, **converted_kwargs) +print("Function Result:", result) diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 0000000..2fe4eec --- /dev/null +++ b/tests/test_cli.py @@ -0,0 +1,174 @@ +import shutil +import pytest +from click.testing import CliRunner +import owilix.cli +import os, fsspec +import re + + +def parse_summary_output(output: str) -> dict: + """ + Parses the summary output to extract key-value pairs where values are numbers, including totals. + + Args: + output (str): The output string from the CLI command. + + Returns: + dict: A dictionary with keys as the parsed keys and values as numerical values. + """ + # Dictionary to store the extracted key-value pairs + result = {} + + # Define regex patterns to match: + # - General key-value pairs where the value is a number + pattern = re.compile(r'^(.*?)\s+(\d[\d.,]*)\s*(?:\w*|)(?:\s*\w*|)(?:\s*GiB|)$') + # - Specific pattern for the "Total of N shown." line + total_pattern = re.compile(r'^Total of (\d+) shown\.$') + + # Iterate over each line in the output + for line in output.splitlines(): + line = line.strip() # Remove leading and trailing whitespace + + # Check if the line matches the general key-value pattern + match = pattern.match(line) + if match: + key = match.group(1).strip() + # Remove commas and convert to appropriate numerical type + value_str = match.group(2).replace(',', '') + if '.' in value_str: + value = float(value_str) + else: + value = int(value_str) + result[key] = value + continue + + # Check if the line matches the total pattern + total_match = total_pattern.match(line) + if total_match: + total_value = int(total_match.group(1)) + result["Total"] = total_value + + return result + +_fs = fsspec.filesystem("file") + +class TestCLI: + @pytest.fixture(scope="class") + def runner(self): + """Fixture that provides a Click CLI runner.""" + return CliRunner() + + @pytest.fixture(scope="session") + def temp_dir(self, tmpdir_factory, request): + """Fixture to provide a unique temporary directory for the session and ensure cleanup.""" + temp_dir = tmpdir_factory.mktemp("cli_tests", numbered=True) + + def cleanup(): + """Cleanup function to remove the temporary directory after the test session.""" + shutil.rmtree(str(temp_dir)) + + request.addfinalizer(cleanup) + return temp_dir + + def setup_method(self): + """Method to setup each test by registering commands.""" + owilix.cli.register_commands(owilix.cli.cli) + + def test_help_command(self, runner): + """Test the help command to ensure it displays the correct output.""" + result = runner.invoke(owilix.cli.cli, ['--help'], prog_name='owilix.cli') + assert result.exit_code == 0 + assert "Main command line interface group for OWI management tools." in result.output + + def test_local_empty_ls(self, runner, temp_dir): + """Test the 'local' command with a subcommand and specifier.""" + result = runner.invoke(owilix.cli.cli, ['--target', str(temp_dir), 'local', 'ls', 'all'], prog_name='owilix.cli') + assert result.exit_code == 0 + assert len(_fs.ls(str(temp_dir))) == 1 + assert _fs.exists(os.path.join(str(temp_dir),".logs")) + assert _fs.exists(os.path.join(str(temp_dir), ".logs","events.json")) + assert result.output=='Fetching datasets for specifier all\nNo data available to display.\n' + + @pytest.mark.parametrize("specifier", ["all", "lrz:2023-10-31", "lrz:latest", "lrz:2023-10-31#7", "lrz:2023-10-31/collectionName=main;resourceType=owi"]) + def test_remote_ls_non_empty_lrz(self, runner, temp_dir, specifier): + """Test the 'local' command with a subcommand and specifier.""" + result = runner.invoke(owilix.cli.cli, ['--target', str(temp_dir), 'remote', 'ls', specifier], prog_name='owilix.cli') + assert result.exit_code == 0 + summary = parse_summary_output(result.output) + assert summary["Total Files"]>0 + assert summary["DataCenter lrz"] > 0 + assert summary["resourceType owi"] > 0 + assert summary["Total"] > 0 + + @pytest.mark.parametrize("specifier", ["all", "it4i:2023-12-03", "it4i:latest", "it4i:2023-12-31#7", + "it4i:2023-12-31#7/collectionName=main;resourceType=owi"]) + def test_remote_ls_non_empty_it4i(self, runner, temp_dir, specifier): + """Test the 'local' command with a subcommand and specifier.""" + result = runner.invoke(owilix.cli.cli, ['--target', str(temp_dir), 'remote', 'ls', specifier], + prog_name='owilix.cli') + assert result.exit_code == 0 + summary = parse_summary_output(result.output) + assert summary["Total Files"] > 0 + assert summary["DataCenter it4i"] > 0 + assert summary["resourceType owi"] > 0 + assert summary["Total"] > 0 + + def test_remote_pull_files(self, runner, temp_dir): + """Test the 'remote' command with a subcommand and specifier.""" + result = runner.invoke(owilix.cli.cli, ['--yes', '--target', str(temp_dir), 'remote', 'pull', + 'it4i:2023-12-03', 'files=**/language=slv/*'], prog_name='owilix.cli') + assert result.exit_code == 0 + assert "Fetching files for OWI-Open Web Index-main.owi@it4i-2023-12-3:2023-12-3" in result.output + assert "Found 2 remote files. Syncing with local files" in result.output # Adjust this check based on expected output + result = runner.invoke(owilix.cli.cli, ['--target', str(temp_dir), 'local', 'ls', 'all'], + prog_name='owilix.cli') + assert result.exit_code == 0 + assert len(_fs.ls(str(temp_dir))) == 2 + assert _fs.exists(os.path.join(str(temp_dir), "public")) + assert any([f for f in _fs.ls(os.path.join(str(temp_dir),"public","main")) if f.endswith("json")]) + + + def test_config_command(self, runner, temp_dir): + result = runner.invoke(owilix.cli.cli, ['config', 'set', 'showfields=url,domain_label', '--target', str(temp_dir)], prog_name='owilix.cli') + assert result.exit_code == 0 + assert "config" in result.output # Adjust this check based on expected output + # Verify that configuration was correctly set in the target directory + config_path = os.path.join(temp_dir, "owilix.cfg") + assert os.path.exists(config_path) + + def test_query_command(self, runner, temp_dir): + """Test the 'query' command with a subcommand and options.""" + result = runner.invoke(owilix.cli.cli, ['query', 'run', '--local', 'all', '--remote', 'lrz', '--target', str(temp_dir)], prog_name='owilix.cli') + assert result.exit_code == 0 + assert "query" in result.output # Adjust this check based on expected output + + def test_clean_command(self, runner, temp_dir): + """Test the 'clean' command.""" + # Pre-create a '.env' file in the temp_dir + env_path = os.path.join(temp_dir, '.env') + with open(env_path, 'w') as f: + f.write('test content') + + result = runner.invoke(owilix.cli.cli, ['clean', '--target', str(temp_dir)], prog_name='owilix.cli') + assert result.exit_code == 0 + assert "Cleaning" in result.output # Check for expected 'clean' command output + + # Verify that the '.env' file has been removed + assert not os.path.exists(env_path) + + def test_logs_command(self, runner, temp_dir): + """Test the 'logs' command with a module argument.""" + result = runner.invoke(owilix.cli.cli, ['--target', str(temp_dir), 'logs', 'lexis', ], prog_name='owilix.cli') + assert result.exit_code == 0 + result = runner.invoke(owilix.cli.cli, ['--target', str(temp_dir), 'logs', 'errors', ], prog_name='owilix.cli') + assert result.exit_code == 0 + assert "log" in result.output + result = runner.invoke(owilix.cli.cli, ['--target', str(temp_dir), 'logs', 'events', ], prog_name='owilix.cli') + assert result.exit_code == 0 + assert "[" in result.output and "]" in result.output + + def test_invalid_command(self, runner): + """Test an invalid command to ensure proper error handling.""" + result = runner.invoke(owilix.cli.cli, ['nonexistent'], prog_name='owilix.cli') + assert result.exit_code != 0 + assert "No such command" in result.output -- 2.51.2