diff --git a/mitm_tooling/extraction/sql/db/__init__.py b/mitm_tooling/extraction/sql/db/__init__.py index ed6feed6c9c413c46cf47a3105abbd91b1bff0b7..376868031902a49ecb144eee84890968cb95429b 100644 --- a/mitm_tooling/extraction/sql/db/__init__.py +++ b/mitm_tooling/extraction/sql/db/__init__.py @@ -3,7 +3,7 @@ from mitm_tooling.utilities.sql_utils import create_sa_engine # noinspection PyUnresolvedReferences from .db_reflection import connect_and_reflect, derive_table_meta_info # noinspection PyUnresolvedReferences -from .db_probing import create_table_probe, initialize_db_probe, test_query +from .db_probing import create_table_probe, initialize_db_probe, test_query, create_db_probe # noinspection PyUnresolvedReferences from .db_schema_query import SyntacticColumnCondition, SemanticColumnCondition, SyntacticTableCondition, SemanticTableCondition, DBMetaQuery, resolve_db_meta_query, resolve_db_meta_selection # noinspection PyUnresolvedReferences diff --git a/mitm_tooling/extraction/sql/db/db_probing.py b/mitm_tooling/extraction/sql/db/db_probing.py index fdb39d985a0677fa42cefbf4e44269a23dcaffb6..925bc6a38dfa7289b37707f7dd1df1cdf4c48d8a 100644 --- a/mitm_tooling/extraction/sql/db/db_probing.py +++ b/mitm_tooling/extraction/sql/db/db_probing.py @@ -247,3 +247,11 @@ def create_table_probe(db_session: Session, table_meta: TableMetaInfo, sample_si sampled_values = df.to_dict(orient='list') # {str(c): vs for c, vs in df.to_dict(orient='list').items()} return TableProbe(table_meta=table_meta, row_count=row_count, sampled_values=sampled_values, inferred_types=inferred_types, sample_summaries=sample_summaries) + + +def create_db_probe(db_session: Session, db_meta: DBMetaInfo, sample_size: int = 100) -> DBProbe: + db_probe = initialize_db_probe(db_meta) + table_probes = ((tm.short_table_identifier, create_table_probe(db_session, tm, sample_size=sample_size)) for tm in + db_meta.tables.values()) + db_probe.update_probes(*table_probes) + return db_probe