genelastic 0.6.1__py3-none-any.whl → 0.8.0__py3-none-any.whl

genelastic/import_data/cli_import.py

@@ -0,0 +1,379 @@
+# vi: se tw=80
+
+# Elasticsearch Python API:
+# https://www.elastic.co/guide/en/elasticsearch/client/python-api/current/overview.html
+# https://elasticsearch-py.readthedocs.io/en/latest/api.html
+
+import argparse
+import csv
+import datetime
+import hashlib
+import logging
+import sys
+import time
+from pathlib import Path
+
+import vcf
+
+from genelastic.common import (
+    AnalysisDocument,
+    BulkItems,
+    ElasticImportConn,
+    MetadataDocument,
+    ProcessDocument,
+    add_es_connection_args,
+    add_verbose_control_args,
+)
+
+from .bi_processes import BioInfoProcesses
+from .data_file import DataFile
+from .import_bundle_factory import make_import_bundle_from_files
+from .logger import configure_logging
+from .wet_processes import WetProcesses
+
+logger = logging.getLogger("genelastic")
+logging.getLogger("elastic_transport").setLevel(
+    logging.WARNING
+)  # Disable excessive logging
+logging.getLogger("urllib3").setLevel(
+    logging.WARNING
+)  # Disable excessive logging
+
+
+def read_args() -> argparse.Namespace:
+    """Read arguments from command line."""
+    parser = argparse.ArgumentParser(
+        description="Genetics data importer.",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+        allow_abbrev=False,
+    )
+    add_verbose_control_args(parser)
+    add_es_connection_args(parser)
+    parser.add_argument(
+        "-D",
+        "--dry-run",
+        dest="dryrun",
+        action="count",
+        default=0,
+        help=(
+            "Dry-run level. -D for data files loading (VCF, coverage, etc) "
+            "without connecting or importing to database. "
+            "-DD for metadata YAML files loading only (no loading of data files)."
+        ),
+    )
+    parser.add_argument(
+        "--log-file", dest="log_file", help="Path to a log file."
+    )
+    parser.add_argument(
+        "--no-list",
+        dest="no_list",
+        action="store_true",
+        help="Do not print list of files to be imported.",
+    )
+    parser.add_argument(
+        "--no-confirm",
+        dest="no_confirm",
+        action="store_true",
+        help="Do not ask confirmation before importing.",
+    )
+    parser.add_argument(
+        "files",
+        type=Path,
+        nargs="+",
+        default=None,
+        help="Data files that describe what to import.",
+    )
+    return parser.parse_args()
+
+
+def import_cov_file(
+    es_import_conn: ElasticImportConn | None,
+    file_index: str,
+    file: Path,
+    dryrun: int = 0,
+) -> None:
+    """Import a coverage file to the Elasticsearch database."""
+    # Set field types
+    if dryrun == 0 and es_import_conn:
+        es_import_conn.client.indices.put_mapping(
+            index=file_index,
+            body={
+                "properties": {
+                    "pos": {"type": "integer"},
+                    "depth": {"type": "byte"},
+                }
+            },
+        )
+
+    # Open file
+    if dryrun > 1:
+        logger.info(
+            "Would load and import Coverage file %s " "into index %s.",
+            file,
+            file_index,
+        )
+    else:
+        logger.info("Load Coverage file %s.", file)
+        if dryrun == 1:
+            logger.info(
+                "Would import Coverage file %s into index %s.", file, file_index
+            )
+        else:
+            logger.info(
+                "Import Coverage file %s into index %s.", file, file_index
+            )
+        with file.open(newline="", encoding="utf-8") as f:
+            # Read file as CSV
+            reader = csv.reader(f, delimiter="\t", quotechar='"')
+
+            # Loop on al lines
+            for row in reader:
+                # Build document
+                # Position starts at 0 inside coverage file
+                doc: MetadataDocument = {
+                    "type": "coverage",
+                    "chr": row[0],
+                    "pos": int(row[1]) + 1,
+                    "depth": int(row[2]),
+                }
+
+                # Insert document
+                if dryrun == 0 and es_import_conn:
+                    es_import_conn.client.index(index=file_index, document=doc)
+
+
+def import_analysis_metadata(  # noqa: PLR0913
+    es_import_conn: ElasticImportConn | None,
+    index_prefix: str,
+    file_index: str,
+    file: DataFile,
+    analysis_type: str,
+    dryrun: int = 0,
+) -> None:
+    """Import analysis metadata into a dedicated index."""
+    doc: AnalysisDocument = {
+        "path": str(file.path.resolve()),
+        "bundle_path": str(file.bundle_path.resolve())
+        if file.bundle_path
+        else None,
+        "metadata": file.metadata,
+        "file_index": file_index,
+        "type": analysis_type,
+    }
+
+    bulk_items: BulkItems = [
+        {"_index": f"{index_prefix}-analyses", "_source": doc}
+    ]
+
+    if dryrun == 0 and es_import_conn:
+        es_import_conn.import_items(
+            bulk_items,
+            start_time=time.perf_counter(),
+            total_items=len(bulk_items),
+        )
+
+
+def import_vcf_file(
+    es_import_conn: ElasticImportConn | None,
+    file_index: str,
+    file: DataFile,
+    dryrun: int = 0,
+) -> None:
+    """Import a VCF file to the Elasticsearch database."""
+    logger.info('Import VCF file "%s".', file)
+
+    if dryrun > 1:
+        logger.info(
+            "Would load and import VCF file %s " "into index %s.",
+            file.path,
+            file_index,
+        )
+    else:
+        logger.info("Load VCF file %s.", file.path)
+        if dryrun == 1:
+            logger.info(
+                "Would import VCF file %s into index %s.", file.path, file_index
+            )
+        else:
+            logger.info(
+                "Importing VCF file %s into index %s...", file.path, file_index
+            )
+
+        try:
+            vcf_reader = vcf.Reader(filename=str(file.path))
+            n = 0
+            start = time.perf_counter()
+            bulk_sz = 256  # Bulk size
+            bulk_items: BulkItems = []
+            for record in vcf_reader:
+                # Correct values
+                if not record.CHROM.startswith("chr"):
+                    if record.CHROM.lower().startswith("chr"):
+                        record.CHROM = "chr" + record.CHROM[3:]
+                    else:
+                        record.CHROM = "chr" + record.CHROM
+
+                # Build document
+                alt = [x if x is None else x.type for x in record.ALT]
+                doc: MetadataDocument = {
+                    "type": "vcf",
+                    "chr": record.CHROM,
+                    "pos": record.POS,
+                    "alt": alt,
+                    "info": record.INFO,
+                }
+
+                if dryrun == 0:
+                    # Append item to bulk
+                    bulk_items.append({"_index": file_index, "_source": doc})
+                    n += 1
+
+                    # Insert bulk of items
+                    if len(bulk_items) >= bulk_sz and es_import_conn:
+                        es_import_conn.import_items(
+                            bulk_items, start_time=start, total_items=n
+                        )
+                        bulk_items = []
+
+            # Insert remaining items
+            if dryrun == 0 and es_import_conn:
+                es_import_conn.import_items(
+                    bulk_items, start_time=start, total_items=n
+                )
+
+        except StopIteration:
+            logger.error("Skipping empty file : %s.", file.path)
+
+
+def import_processes(
+    es_import_conn: ElasticImportConn | None,
+    index: str,
+    processes: WetProcesses | BioInfoProcesses,
+    dryrun: int = 0,
+) -> None:
+    """Import processes into their own index."""
+    bulk_items: BulkItems = []
+
+    for proc_id in processes.get_process_ids():
+        process = processes[proc_id]
+        process_type = process.__class__.__name__
+        doc: ProcessDocument = process.data | {
+            "proc_id": proc_id,
+            "type": process_type,
+        }
+        bulk_items.append({"_index": index, "_source": doc})
+
+    if dryrun == 0 and es_import_conn:
+        es_import_conn.import_items(
+            bulk_items,
+            start_time=time.perf_counter(),
+            total_items=len(bulk_items),
+        )
+
+
+def generate_unique_index(index_prefix: str, filepath: Path) -> str:
+    """Generate a unique index with the following format:
+    <index_prefix>_<current_date>_<md5_hashed_filepath>
+    """
+    current_date = datetime.datetime.now(tz=datetime.UTC).strftime("%Y%m%d")
+    hashed_filepath = hashlib.md5(
+        str(filepath).encode("utf-8"), usedforsecurity=False
+    ).hexdigest()
+    return f"{index_prefix}-file-{current_date}-{hashed_filepath}"
+
+
+def main() -> None:  # noqa: C901
+    """Entry point of the import script."""
+    # Read command line arguments
+    args = read_args()
+
+    # Configure logging
+    configure_logging(args.verbose, log_file=args.log_file)
+    logger.debug("Arguments: %s", args)
+    logger.debug("LOGGERS: %s", logging.root.manager.loggerDict)
+
+    # Open connection to ES
+    if args.dryrun == 0:
+        addr = f"https://{args.es_host}:{args.es_port}"
+        logger.info("Trying to connect to Elasticsearch at %s...", addr)
+        es_import_conn = ElasticImportConn(
+            addr, args.es_cert_fp, basic_auth=(args.es_usr, args.es_pwd)
+        )
+    else:
+        es_import_conn = None
+
+    # Load YAML import bundle
+    import_bundle = make_import_bundle_from_files(args.files, check=True)
+    all_bundled_files = import_bundle.get_files()
+
+    # CHECK
+    for f in all_bundled_files:
+        if not f.exists():
+            msg = f"Path {f.path} does not point to a valid file."
+            raise RuntimeError(msg)
+
+    # LIST
+    if not args.no_list:
+        for f in all_bundled_files:
+            logger.info("Will import %s.", f.path)
+
+    # Ask confirmation for importing
+    if not args.no_confirm:
+        answer: str = "maybe"
+        while answer not in ["", "n", "y"]:
+            answer = input("Import (y/N)? ").lower()
+        if answer != "y":
+            logger.info("Import canceled.")
+            sys.exit(0)
+
+    # IMPORT
+    # Loop on file categories
+    for cat in import_bundle.analyses.get_all_categories():
+        # Import all files in this category.
+        for f in import_bundle.get_files(cat):
+            logger.info("Import %s files from %s.", cat, f.path)
+            # First, generate a unique index name for each file.
+            file_index = generate_unique_index(args.es_index_prefix, f.path)
+            # Then, import the analysis metadata into a dedicated index.
+            import_analysis_metadata(
+                es_import_conn,
+                args.es_index_prefix,
+                file_index,
+                f,
+                cat,
+                args.dryrun,
+            )
+            # Finally, import the file in its own index.
+            globals()[f"import_{cat}_file"](
+                es_import_conn=es_import_conn,
+                file_index=file_index,
+                file=f,
+                dryrun=args.dryrun,
+            )
+
+    # Import processes
+    logger.info("Importing wet processes.")
+    logger.info(
+        "Wet processes IDs = %s",
+        str(import_bundle.wet_processes.get_process_ids()),
+    )
+    import_processes(
+        es_import_conn,
+        f"{args.es_index_prefix}-wet_processes",
+        import_bundle.wet_processes,
+    )
+
+    logger.info("Importing bio info processes.")
+    logger.info(
+        "Bio info processes IDs = %s",
+        str(import_bundle.bi_processes.get_process_ids()),
+    )
+    import_processes(
+        es_import_conn,
+        f"{args.es_index_prefix}-bi_processes",
+        import_bundle.bi_processes,
+    )
+
+
+if __name__ == "__main__":
+    main()

genelastic/import_data/{info.py → cli_info.py}

@@ -1,71 +1,100 @@
-# pylint: disable=missing-module-docstring
 import argparse
 import logging
-import typing
 
-from genelastic.common import (ElasticQueryConn, add_verbose_control_args,
-                               add_es_connection_args, Bucket)
+from genelastic.common import (
+    Bucket,
+    ElasticQueryConn,
+    add_es_connection_args,
+    add_verbose_control_args,
+)
 
 from .logger import configure_logging
 
-logger = logging.getLogger('genelastic')
-logging.getLogger('elastic_transport').setLevel(logging.WARNING)  # Disable excessive logging
+logger = logging.getLogger("genelastic")
+logging.getLogger("elastic_transport").setLevel(
+    logging.WARNING
+)  # Disable excessive logging
 
 
 def read_args() -> argparse.Namespace:
     """Read arguments from command line."""
-    parser = argparse.ArgumentParser(description='ElasticSearch database info.',
-                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter,
-                                     allow_abbrev=False)
+    parser = argparse.ArgumentParser(
+        description="ElasticSearch database info.",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+        allow_abbrev=False,
+    )
     add_verbose_control_args(parser)
     add_es_connection_args(parser)
-    parser.add_argument("-y", "--list-bundles", action="store_true",
-                        help="List all imported YAML bundles.")
-    parser.add_argument("-f", "--list-data-files", action="store_true",
-                        help="List all imported data files.")
-    parser.add_argument("-w", "--list-wet-processes", action="store_true",
-                        help="List all imported wet processes.")
-    parser.add_argument("-b", "--list-bi-processes", action="store_true",
-                        help="List all imported bio info processes.")
-    parser.add_argument("-Y", "--list-data-files-per-bundle", action="store_true",
-                        help="For each imported YAML bundle, "
-                             "display some info and list its data files.")
+    parser.add_argument(
+        "-y",
+        "--list-bundles",
+        action="store_true",
+        help="List all imported YAML bundles.",
+    )
+    parser.add_argument(
+        "-f",
+        "--list-data-files",
+        action="store_true",
+        help="List all imported data files.",
+    )
+    parser.add_argument(
+        "-w",
+        "--list-wet-processes",
+        action="store_true",
+        help="List all imported wet processes.",
+    )
+    parser.add_argument(
+        "-b",
+        "--list-bi-processes",
+        action="store_true",
+        help="List all imported bio info processes.",
+    )
+    parser.add_argument(
+        "-Y",
+        "--list-data-files-per-bundle",
+        action="store_true",
+        help="For each imported YAML bundle, "
+        "display some info and list its data files.",
+    )
     return parser.parse_args()
 
 
 def list_bundles(es_query_conn: ElasticQueryConn, index: str) -> None:
     """List all imported YAML bundles."""
-
     query = {
         "size": 0,
         "aggs": {
             "get_bundle_paths": {
                 "composite": {
-                    "sources": {"bundle_path": {"terms": {"field": "bundle_path.keyword"}}},
+                    "sources": {
+                        "bundle_path": {
+                            "terms": {"field": "bundle_path.keyword"}
+                        }
+                    },
                     "size": 1000,
                 }
             }
-        }
+        },
     }
 
-    buckets: typing.List[Bucket] = es_query_conn.run_composite_aggregation(index, query)
+    buckets: list[Bucket] = es_query_conn.run_composite_aggregation(
+        index, query
+    )
 
-    print("Imported YAML files")
-    print("===================")
+    logger.info("Imported YAML files")
+    logger.info("===================")
 
     if len(buckets) == 0:
-        print("Empty response.", end="\n")
+        logger.info("Empty response.")
         return
 
     for bucket in buckets:
-        bundle_path = bucket['key']['bundle_path']
-        print(f'- {bundle_path}')
-    print()
+        bundle_path = bucket["key"]["bundle_path"]
+        logger.info("- %s", bundle_path)
 
 
 def list_data_files(es_query_conn: ElasticQueryConn, index: str) -> None:
     """List all imported data files."""
-
     query = {
         "size": 0,
         "aggs": {
@@ -75,22 +104,23 @@ def list_data_files(es_query_conn: ElasticQueryConn, index: str) -> None:
                     "size": 1000,
                 }
             }
-        }
+        },
     }
 
-    buckets: typing.List[Bucket] = es_query_conn.run_composite_aggregation(index, query)
+    buckets: list[Bucket] = es_query_conn.run_composite_aggregation(
+        index, query
+    )
 
-    print("Imported data files")
-    print("===================")
+    logger.info("Imported data files")
+    logger.info("===================")
 
     if len(buckets) == 0:
-        print("Empty response.", end="\n")
+        logger.info("Empty response.")
         return
 
     for bucket in buckets:
-        bundle_path = bucket['key']['path']
-        print(f'- {bundle_path}')
-    print()
+        bundle_path = bucket["key"]["path"]
+        logger.info("- %s", bundle_path)
 
 
 def list_processes(es_query_conn: ElasticQueryConn, index: str) -> None:
@@ -98,29 +128,30 @@ def list_processes(es_query_conn: ElasticQueryConn, index: str) -> None:
     process_ids = es_query_conn.get_field_values(index, "proc_id")
 
     if len(process_ids) == 0:
-        print("Empty response.", end="\n")
+        logger.info("Empty response.")
         return
 
     for process_id in process_ids:
-        print(f'- {process_id}')
-    print()
+        logger.info("- %s", process_id)
 
 
 def list_wet_processes(es_query_conn: ElasticQueryConn, index: str) -> None:
     """List all wet processes."""
-    print("Imported wet processes")
-    print("======================")
+    logger.info("Imported wet processes")
+    logger.info("======================")
     list_processes(es_query_conn, index)
 
 
 def list_bi_processes(es_query_conn: ElasticQueryConn, index: str) -> None:
     """List all bio info processes."""
-    print("Imported bi processes")
-    print("=====================")
+    logger.info("Imported bi processes")
+    logger.info("=====================")
     list_processes(es_query_conn, index)
 
 
-def list_data_files_per_bundle(es_query_conn: ElasticQueryConn, index: str) -> None:
+def list_data_files_per_bundle(
+    es_query_conn: ElasticQueryConn, index: str
+) -> None:
     """For each imported YAML bundle, display some info and list its data files."""
     query = {
         "size": 0,
@@ -130,50 +161,47 @@ def list_data_files_per_bundle(es_query_conn: ElasticQueryConn, index: str) -> N
                     "sources": [
                         {
                             "bundle_path": {
-                                "terms": {
-                                    "field": "bundle_path.keyword"
-                                }
+                                "terms": {"field": "bundle_path.keyword"}
                             }
                         }
                     ],
-                    "size": 100
+                    "size": 100,
                 },
-                "aggs": {
-                    "docs": {
-                        "top_hits": {
-                            "size": 100
-                        }
-                    }
-                }
+                "aggs": {"docs": {"top_hits": {"size": 100}}},
             }
-        }
+        },
     }
 
-    buckets: typing.List[Bucket] = es_query_conn.run_composite_aggregation(index, query)
+    buckets: list[Bucket] = es_query_conn.run_composite_aggregation(
+        index, query
+    )
 
-    print("Data files per YAML bundle")
-    print("==========================")
+    logger.info("Data files per YAML bundle")
+    logger.info("==========================")
 
     if len(buckets) == 0:
-        print("Empty response.", end="\n")
+        logger.info("Empty response.")
         return
 
     for bucket in buckets:
-
         documents = bucket["docs"]["hits"]["hits"]
         if len(documents) == 0:
             continue
 
-        print(f"- Bundle Path: {bucket['key']['bundle_path']}")
-        print(f"    -> Wet process: {documents[0]['_source']['metadata']['wet_process']}")
-        print(f"    -> Bio info process: {documents[0]['_source']['metadata']['bi_process']}")
-        print("    -> Data files:")
+        logger.info("- Bundle Path: %s", bucket["key"]["bundle_path"])
+        logger.info(
+            "    -> Wet process: %s",
+            documents[0]["_source"]["metadata"]["wet_process"],
+        )
+        logger.info(
+            "    -> Bio info process: %s",
+            documents[0]["_source"]["metadata"]["bi_process"],
+        )
+        logger.info("    -> Data files:")
 
         for doc in documents:
-            print(f"        - Index: {doc['_source']['file_index']}")
-            print(f"          Path: {doc['_source']['path']}")
-
-    print()
+            logger.info("        - Index: %s", doc["_source"]["file_index"])
+            logger.info("          Path: %s", doc["_source"]["path"])
 
 
 def main() -> None:
@@ -185,8 +213,9 @@ def main() -> None:
 
     addr = f"https://{args.es_host}:{args.es_port}"
     logger.info("Trying to connect to Elasticsearch at %s...", addr)
-    es_query_conn = ElasticQueryConn(addr, args.es_cert_fp,
-                                     basic_auth=(args.es_usr, args.es_pwd))
+    es_query_conn = ElasticQueryConn(
+        addr, args.es_cert_fp, basic_auth=(args.es_usr, args.es_pwd)
+    )
 
     analysis_index = f"{args.es_index_prefix}-analyses"
     wet_processes_index = f"{args.es_index_prefix}-wet_processes"
@@ -223,5 +252,5 @@ def main() -> None:
         list_data_files_per_bundle(es_query_conn, analysis_index)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()