ingestify 0.8.0__py3-none-any.whl → 0.9.0__py3-none-any.whl

ingestify/infra/store/dataset/sqlalchemy/tables.py

@@ -149,185 +149,202 @@ class IngestionJobStateString(TypeDecorator):
         return IngestionJobState[value]
 
 
-metadata = MetaData()
-
-dataset_table = Table(
-    "dataset",
-    metadata,
-    Column("bucket", String(255), default=None),
-    Column("dataset_id", String(255), primary_key=True),
-    Column("provider", String(255), index=True),
-    Column("dataset_type", String(255), index=True),
-    Column("state", DatasetStateString),
-    Column("name", String(255)),
-    Column(
-        "identifier",
-        # Use JSONB when available
-        JSON().with_variant(JSONB(), "postgresql"),
-    ),
-    Column("metadata", JSON),
-    Column("created_at", TZDateTime(6)),
-    Column("updated_at", TZDateTime(6)),
-    Column("last_modified_at", TZDateTime(6)),
-    # Required for performance querying when there are a lot of Datasets
-    # with the same provider and dataset_type
-    Index(
-        "idx_bucket_type_provider_last_modified",
-        "bucket",
-        "provider",
-        "dataset_type",
-        "last_modified_at",
-    ),
-)
-
-revision_table = Table(
-    "revision",
-    metadata,
-    Column(
-        "dataset_id", String(255), ForeignKey("dataset.dataset_id"), primary_key=True
-    ),
-    Column("revision_id", Integer, primary_key=True),
-    Column("description", String(255)),
-    Column("created_at", TZDateTime(6)),
-    Column("state", RevisionStateString, default=RevisionState.PENDING_VALIDATION),
-    Column("source", JSONType()),
-)
-
-file_table = Table(
-    "file",
-    metadata,
-    Column("dataset_id", String(255), primary_key=True),
-    Column("revision_id", Integer, primary_key=True),
-    Column("file_id", String(255), primary_key=True),
-    Column("created_at", TZDateTime(6)),
-    Column("modified_at", TZDateTime(6)),
-    Column("tag", String(255)),
-    Column("content_type", String(255)),
-    Column("size", BigInteger),
-    Column("data_feed_key", String(255)),
-    Column("data_spec_version", String(255)),
-    Column("data_serialization_format", String(255)),
-    Column("storage_compression_method", String(255)),
-    Column("storage_size", BigInteger),
-    Column("storage_path", PathString),
-    ForeignKeyConstraint(
-        ("dataset_id", "revision_id"),
-        [revision_table.c.dataset_id, revision_table.c.revision_id],
-        ondelete="CASCADE",
-    ),
-)
-
-ingestion_job_summary_table = Table(
-    "ingestion_job_summary",
-    metadata,
-    Column("ingestion_job_summary_id", String(255), primary_key=True),
-    Column("ingestion_job_id", String(255), index=True),
-    # From the IngestionPlan
-    Column("source_name", String(255)),
-    Column("provider", String(255)),
-    Column("dataset_type", String(255)),
-    Column(
-        "data_spec_versions",
-        JSONType(
-            serializer=lambda data_spec_versions: {
-                key: list(value) for key, value in data_spec_versions.items()
-            },
-            deserializer=lambda data_spec_versions: DataSpecVersionCollection.from_dict(
-                data_spec_versions
+def get_tables(table_prefix: str = ""):
+    """
+    Create all SQLAlchemy table definitions with an optional prefix.
+
+    Args:
+        table_prefix: Optional prefix for all table names (e.g., "prod_" would create "prod_dataset")
+
+    Returns:
+        A dictionary containing all table objects and metadata
+    """
+    metadata = MetaData()
+
+    dataset_table = Table(
+        f"{table_prefix}dataset",
+        metadata,
+        Column("bucket", String(255), default=None),
+        Column("dataset_id", String(255), primary_key=True),
+        Column("provider", String(255), index=True),
+        Column("dataset_type", String(255), index=True),
+        Column("state", DatasetStateString),
+        Column("name", String(255)),
+        Column(
+            "identifier",
+            # Use JSONB when available
+            JSON().with_variant(JSONB(), "postgresql"),
+        ),
+        Column("metadata", JSON),
+        Column("created_at", TZDateTime(6)),
+        Column("updated_at", TZDateTime(6)),
+        Column("last_modified_at", TZDateTime(6)),
+        # Required for performance querying when there are a lot of Datasets
+        # with the same provider and dataset_type
+        Index(
+            "idx_bucket_type_provider_last_modified",
+            "bucket",
+            "provider",
+            "dataset_type",
+            "last_modified_at",
+        ),
+    )
+
+    revision_table = Table(
+        f"{table_prefix}revision",
+        metadata,
+        Column(
+            "dataset_id",
+            String(255),
+            ForeignKey(f"{table_prefix}dataset.dataset_id"),
+            primary_key=True,
+        ),
+        Column("revision_id", Integer, primary_key=True),
+        Column("description", String(255)),
+        Column("created_at", TZDateTime(6)),
+        Column("state", RevisionStateString, default=RevisionState.PENDING_VALIDATION),
+        Column("source", JSONType()),
+    )
+
+    file_table = Table(
+        f"{table_prefix}file",
+        metadata,
+        Column("dataset_id", String(255), primary_key=True),
+        Column("revision_id", Integer, primary_key=True),
+        Column("file_id", String(255), primary_key=True),
+        Column("created_at", TZDateTime(6)),
+        Column("modified_at", TZDateTime(6)),
+        Column("tag", String(255)),
+        Column("content_type", String(255)),
+        Column("size", BigInteger),
+        Column("data_feed_key", String(255)),
+        Column("data_spec_version", String(255)),
+        Column("data_serialization_format", String(255)),
+        Column("storage_compression_method", String(255)),
+        Column("storage_size", BigInteger),
+        Column("storage_path", PathString),
+        ForeignKeyConstraint(
+            ("dataset_id", "revision_id"),
+            [revision_table.c.dataset_id, revision_table.c.revision_id],
+            ondelete="CASCADE",
+        ),
+    )
+
+    ingestion_job_summary_table = Table(
+        f"{table_prefix}ingestion_job_summary",
+        metadata,
+        Column("ingestion_job_summary_id", String(255), primary_key=True),
+        Column("ingestion_job_id", String(255), index=True),
+        # From the IngestionPlan
+        Column("source_name", String(255)),
+        Column("provider", String(255)),
+        Column("dataset_type", String(255)),
+        Column(
+            "data_spec_versions",
+            JSONType(
+                serializer=lambda data_spec_versions: {
+                    key: list(value) for key, value in data_spec_versions.items()
+                },
+                deserializer=lambda data_spec_versions: DataSpecVersionCollection.from_dict(
+                    data_spec_versions
+                ),
             ),
         ),
-    ),
-    Column(
-        "selector",
-        JSONType(
-            serializer=lambda selector: selector.filtered_attributes,
-            deserializer=lambda selector: Selector(**selector),
+        Column(
+            "selector",
+            JSONType(
+                serializer=lambda selector: selector.filtered_attributes,
+                deserializer=lambda selector: Selector(**selector),
+            ),
         ),
-    ),
-    Column("started_at", TZDateTime(6)),
-    Column("ended_at", TZDateTime(6)),
-    # Some task counters
-    Column("state", IngestionJobStateString),
-    Column("total_tasks", Integer),
-    Column("successful_tasks", Integer),
-    Column("ignored_successful_tasks", Integer),
-    Column("skipped_tasks", Integer),
-    Column("failed_tasks", Integer),
-    Column(
-        "timings",
-        JSONType(
-            serializer=lambda timings: [
-                # Timing is probably already a dictionary. Load it into Timing first, so it can be dumped
-                # in json mode
-                Timing.model_validate(timing).model_dump(mode="json")
-                for timing in timings
-            ],
-            deserializer=lambda timings: [
-                Timing.model_validate(timing) for timing in timings
-            ],
+        Column("started_at", TZDateTime(6)),
+        Column("ended_at", TZDateTime(6)),
+        # Some task counters
+        Column("state", IngestionJobStateString),
+        Column("total_tasks", Integer),
+        Column("successful_tasks", Integer),
+        Column("ignored_successful_tasks", Integer),
+        Column("skipped_tasks", Integer),
+        Column("failed_tasks", Integer),
+        Column(
+            "timings",
+            JSONType(
+                serializer=lambda timings: [
+                    # Timing is probably already a dictionary. Load it into Timing first, so it can be dumped
+                    # in json mode
+                    Timing.model_validate(timing).model_dump(mode="json")
+                    for timing in timings
+                ],
+                deserializer=lambda timings: [
+                    Timing.model_validate(timing) for timing in timings
+                ],
+            ),
         ),
-    ),
-    # Column(
-    #     "task_summaries",
-    #     JSONType(
-    #         serializer=lambda task_summaries: [
-    #             task_summary.model_dump(mode="json") for task_summary in task_summaries
-    #         ],
-    #         deserializer=lambda task_summaries: [
-    #             TaskSummary.model_validate(task_summary)
-    #             for task_summary in task_summaries
-    #         ],
-    #     ),
-    # ),
-)
-
-
-task_summary_table = Table(
-    "task_summary",
-    metadata,
-    Column(
-        "ingestion_job_summary_id",
-        String(255),
-        ForeignKey("ingestion_job_summary.ingestion_job_summary_id"),
-        primary_key=True,
-    ),
-    Column("task_id", String(255), primary_key=True),
-    Column("started_at", TZDateTime(6)),
-    Column("ended_at", TZDateTime(6)),
-    Column("operation", OperationString),
-    Column(
-        "dataset_identifier", JSONType(deserializer=lambda item: Identifier(**item))
-    ),
-    Column("persisted_file_count", Integer),
-    Column("bytes_retrieved", Integer),
-    Column("last_modified", TZDateTime(6)),
-    Column("state", TaskStateString),
-    Column(
-        "timings",
-        JSONType(
-            serializer=lambda timings: [
-                Timing.model_validate(timing).model_dump(mode="json")
-                for timing in timings
-            ],
-            deserializer=lambda timings: [
-                Timing.model_validate(timing) for timing in timings
-            ],
+    )
+
+    task_summary_table = Table(
+        f"{table_prefix}task_summary",
+        metadata,
+        Column(
+            "ingestion_job_summary_id",
+            String(255),
+            ForeignKey(f"{table_prefix}ingestion_job_summary.ingestion_job_summary_id"),
+            primary_key=True,
         ),
-    ),
-    # Column("description", String(255)),
-    # Column("created_at", TZDateTime(6)),
-    # Column("state", RevisionStateString, default=RevisionState.PENDING_VALIDATION),
-    # Column("source", JSONType()),
-)
-
-store_version_table = Table(
-    "store_version",
-    metadata,
-    Column("id", Integer, primary_key=True, default=1),
-    Column("ingestify_version", String(255), nullable=False),
-    Column("created_at", TZDateTime(6), nullable=False),
-    Column("updated_at", TZDateTime(6), nullable=False),
-)
+        Column("task_id", String(255), primary_key=True),
+        Column("started_at", TZDateTime(6)),
+        Column("ended_at", TZDateTime(6)),
+        Column("operation", OperationString),
+        Column(
+            "dataset_identifier", JSONType(deserializer=lambda item: Identifier(**item))
+        ),
+        Column("persisted_file_count", Integer),
+        Column("bytes_retrieved", Integer),
+        Column("last_modified", TZDateTime(6)),
+        Column("state", TaskStateString),
+        Column(
+            "timings",
+            JSONType(
+                serializer=lambda timings: [
+                    Timing.model_validate(timing).model_dump(mode="json")
+                    for timing in timings
+                ],
+                deserializer=lambda timings: [
+                    Timing.model_validate(timing) for timing in timings
+                ],
+            ),
+        ),
+    )
+
+    store_version_table = Table(
+        f"{table_prefix}store_version",
+        metadata,
+        Column("id", Integer, primary_key=True, default=1),
+        Column("ingestify_version", String(255), nullable=False),
+        Column("created_at", TZDateTime(6), nullable=False),
+        Column("updated_at", TZDateTime(6), nullable=False),
+    )
+
+    return {
+        "metadata": metadata,
+        "dataset_table": dataset_table,
+        "revision_table": revision_table,
+        "file_table": file_table,
+        "ingestion_job_summary_table": ingestion_job_summary_table,
+        "task_summary_table": task_summary_table,
+        "store_version_table": store_version_table,
+    }
+
+
+# Create default tables without prefix for backward compatibility
+_default_tables = get_tables("")
+metadata = _default_tables["metadata"]
+dataset_table = _default_tables["dataset_table"]
+revision_table = _default_tables["revision_table"]
+file_table = _default_tables["file_table"]
+ingestion_job_summary_table = _default_tables["ingestion_job_summary_table"]
+task_summary_table = _default_tables["task_summary_table"]
+store_version_table = _default_tables["store_version_table"]
 #
 #
 # mapper_registry = registry()

ingestify/main.py

@@ -82,14 +82,28 @@ def build_file_repository(file_url: str, identifier_transformer) -> FileReposito
 
 
 def get_dataset_store_by_urls(
-    metadata_url: str, file_url: str, bucket: str, dataset_types
+    metadata_url: str,
+    file_url: str,
+    bucket: str,
+    dataset_types,
+    metadata_options: dict = None,
 ) -> DatasetStore:
     """
     Initialize a DatasetStore by a DatasetRepository and a FileRepository
+
+    Args:
+        metadata_url: Database connection URL
+        file_url: File storage URL
+        bucket: Bucket name
+        dataset_types: Dataset type configurations
+        metadata_options: Optional dict with metadata store options (e.g., table_prefix)
     """
     if not bucket:
         raise Exception("Bucket is not specified")
 
+    if metadata_options is None:
+        metadata_options = {}
+
     identifier_transformer = IdentifierTransformer()
     for dataset_type in dataset_types:
         for id_key, id_config in dataset_type["identifier_keys"].items():
@@ -110,7 +124,12 @@ def get_dataset_store_by_urls(
     if metadata_url.startswith("postgres://"):
         metadata_url = metadata_url.replace("postgress://", "postgress+")
 
-    sqlalchemy_session_provider = SqlAlchemySessionProvider(metadata_url)
+    # Extract table_prefix from metadata_options
+    table_prefix = metadata_options.get("table_prefix", "")
+
+    sqlalchemy_session_provider = SqlAlchemySessionProvider(
+        metadata_url, table_prefix=table_prefix
+    )
 
     dataset_repository = SqlAlchemyDatasetRepository(sqlalchemy_session_provider)
 
@@ -124,11 +143,16 @@ def get_dataset_store_by_urls(
 def get_datastore(config_file, bucket: Optional[str] = None) -> DatasetStore:
     config = parse_config(config_file, default_value="")
 
+    # Extract metadata_options if provided
+    main_config = config["main"]
+    metadata_options = main_config.get("metadata_options", {})
+
     return get_dataset_store_by_urls(
-        metadata_url=config["main"]["metadata_url"],
-        file_url=config["main"]["file_url"],
-        bucket=bucket or config["main"].get("default_bucket"),
+        metadata_url=main_config["metadata_url"],
+        file_url=main_config["file_url"],
+        bucket=bucket or main_config.get("default_bucket"),
         dataset_types=config.get("dataset_types", []),
+        metadata_options=metadata_options,
     )
 
 
@@ -219,11 +243,16 @@ def get_engine(
             sources[name] = build_source(name=name, source_args=source_args)
 
     logger.info("Initializing IngestionEngine")
+
+    # Extract metadata_options if provided
+    metadata_options = config["main"].get("metadata_options", {})
+
     store = get_dataset_store_by_urls(
         metadata_url=config["main"]["metadata_url"],
         file_url=config["main"]["file_url"],
         bucket=bucket or config["main"].get("default_bucket"),
         dataset_types=config.get("dataset_types", []),
+        metadata_options=metadata_options,
     )
 
     # Setup an EventBus and wire some more components

ingestify/tests/conftest.py

@@ -0,0 +1,17 @@
+import tempfile
+
+import pytest
+import os
+
+
+@pytest.fixture(scope="function", autouse=True)
+def datastore_dir():
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        os.environ["TEST_DIR"] = tmpdirname
+        os.environ["INGESTIFY_RUN_EAGER"] = "true"
+        yield tmpdirname
+
+
+@pytest.fixture(scope="session")
+def config_file():
+    return os.path.abspath(os.path.dirname(__file__) + "/config.yaml")