def test_example_F(test_cases_for_sql_data_connector_sqlite_execution_engine): random.seed(0) db = test_cases_for_sql_data_connector_sqlite_execution_engine config = yaml.load( """ name: my_sql_data_connector datasource_name: FAKE_Datasource_NAME data_assets: table_partitioned_by_foreign_key__F: splitter_method: _split_on_column_value splitter_kwargs: column_name: session_id """, ) config["execution_engine"] = db my_data_connector = ConfiguredAssetSqlDataConnector(**config) report = my_data_connector.self_check() print(json.dumps(report, indent=2)) assert report == { "class_name": "ConfiguredAssetSqlDataConnector", "data_asset_count": 1, "example_data_asset_names": ["table_partitioned_by_foreign_key__F"], "data_assets": { "table_partitioned_by_foreign_key__F": { "batch_definition_count": 49, # TODO Abe 20201029 : These values should be sorted "example_data_references": [ { "session_id": 3 }, { "session_id": 2 }, { "session_id": 4 }, ], } }, "unmatched_data_reference_count": 0, "example_unmatched_data_references": [], "example_data_reference": { "n_rows": 2, "batch_spec": { "table_name": "table_partitioned_by_foreign_key__F", "partition_definition": { "session_id": 2 }, "splitter_method": "_split_on_column_value", "splitter_kwargs": { "column_name": "session_id" }, }, }, }
def test_example_A(test_cases_for_sql_data_connector_sqlite_execution_engine): random.seed(0) db = test_cases_for_sql_data_connector_sqlite_execution_engine config = yaml.load( """ name: my_sql_data_connector datasource_name: FAKE_Datasource_NAME data_assets: table_partitioned_by_date_column__A: splitter_method: _split_on_column_value splitter_kwargs: column_name: date """, ) config["execution_engine"] = db my_data_connector = ConfiguredAssetSqlDataConnector(**config) report = my_data_connector.self_check() print(json.dumps(report, indent=2)) assert report == { "class_name": "ConfiguredAssetSqlDataConnector", "data_asset_count": 1, "example_data_asset_names": ["table_partitioned_by_date_column__A"], "data_assets": { "table_partitioned_by_date_column__A": { "batch_definition_count": 30, "example_data_references": [ { "date": "2020-01-01" }, { "date": "2020-01-02" }, { "date": "2020-01-03" }, ], } }, "unmatched_data_reference_count": 0, "example_unmatched_data_references": [], "example_data_reference": { "n_rows": 8, "batch_spec": { "table_name": "table_partitioned_by_date_column__A", "partition_definition": { "date": "2020-01-02" }, "splitter_method": "_split_on_column_value", "splitter_kwargs": { "column_name": "date" }, }, }, }
def test_behavior_with_whole_table_splitter( test_cases_for_sql_data_connector_sqlite_execution_engine, ): db = test_cases_for_sql_data_connector_sqlite_execution_engine config = yaml.load( """ name: my_sql_data_connector datasource_name: FAKE_Datasource_NAME assets: table_partitioned_by_date_column__A: splitter_method : "_split_on_whole_table" splitter_kwargs : {} """, ) config["execution_engine"] = db my_data_connector = ConfiguredAssetSqlDataConnector(**config) report_object = my_data_connector.self_check() print(json.dumps(report_object, indent=2)) batch_definition_list = ( my_data_connector.get_batch_definition_list_from_batch_request( batch_request=BatchRequest( datasource_name="FAKE_Datasource_NAME", data_connector_name="my_sql_data_connector", data_asset_name="table_partitioned_by_date_column__A", ) ) ) assert len(batch_definition_list) == 1 assert batch_definition_list[0]["batch_identifiers"] == {} batch_definition_list = ( my_data_connector.get_batch_definition_list_from_batch_request( batch_request=BatchRequest( datasource_name="FAKE_Datasource_NAME", data_connector_name="my_sql_data_connector", data_asset_name="table_partitioned_by_date_column__A", data_connector_query={}, ) ) ) assert len(batch_definition_list) == 1 assert batch_definition_list[0]["batch_identifiers"] == {} batch_definition_list = ( my_data_connector.get_batch_definition_list_from_batch_request( batch_request=BatchRequest( datasource_name="FAKE_Datasource_NAME", data_connector_name="my_sql_data_connector", data_asset_name="table_partitioned_by_date_column__A", data_connector_query={"batch_filter_parameters": {}}, ) ) ) assert len(batch_definition_list) == 1 assert batch_definition_list[0]["batch_identifiers"] == {}
def test_basic_self_check( test_cases_for_sql_data_connector_sqlite_execution_engine): random.seed(0) execution_engine = test_cases_for_sql_data_connector_sqlite_execution_engine config = yaml.load( """ name: my_sql_data_connector datasource_name: FAKE_Datasource_NAME assets: table_partitioned_by_date_column__A: #table_name: events # If table_name is omitted, then the table_name defaults to the asset name splitter_method: _split_on_column_value splitter_kwargs: column_name: date """, ) config["execution_engine"] = execution_engine my_data_connector = ConfiguredAssetSqlDataConnector(**config) report = my_data_connector.self_check() print(json.dumps(report, indent=2)) assert report == { "class_name": "ConfiguredAssetSqlDataConnector", "data_asset_count": 1, "example_data_asset_names": ["table_partitioned_by_date_column__A"], "data_assets": { "table_partitioned_by_date_column__A": { "batch_definition_count": 30, "example_data_references": [ { "date": "2020-01-01" }, { "date": "2020-01-02" }, { "date": "2020-01-03" }, ], } }, "unmatched_data_reference_count": 0, "example_unmatched_data_references": [], # FIXME: (Sam) example_data_reference removed temporarily in PR #2590: # "example_data_reference": { # "n_rows": 8, # "batch_spec": { # "table_name": "table_partitioned_by_date_column__A", # "data_asset_name": "table_partitioned_by_date_column__A", # "batch_identifiers": {"date": "2020-01-02"}, # "splitter_method": "_split_on_column_value", # "splitter_kwargs": {"column_name": "date"}, # }, # }, }
def test_example_G(test_cases_for_sql_data_connector_sqlite_execution_engine): random.seed(0) db = test_cases_for_sql_data_connector_sqlite_execution_engine config = yaml.load( """ name: my_sql_data_connector datasource_name: FAKE_Datasource_NAME assets: table_partitioned_by_multiple_columns__G: splitter_method: _split_on_multi_column_values splitter_kwargs: column_names: - y - m - d """, ) config["execution_engine"] = db my_data_connector = ConfiguredAssetSqlDataConnector(**config) report = my_data_connector.self_check() print(json.dumps(report, indent=2)) assert report == { "class_name": "ConfiguredAssetSqlDataConnector", "data_asset_count": 1, "example_data_asset_names": ["table_partitioned_by_multiple_columns__G"], "data_assets": { "table_partitioned_by_multiple_columns__G": { "batch_definition_count": 30, # TODO Abe 20201029 : These values should be sorted "example_data_references": [ {"y": 2020, "m": 1, "d": 1}, {"y": 2020, "m": 1, "d": 2}, {"y": 2020, "m": 1, "d": 3}, ], } }, "unmatched_data_reference_count": 0, "example_unmatched_data_references": [], # FIXME: (Sam) example_data_reference removed temporarily in PR #2590: # "example_data_reference": { # "n_rows": 8, # "batch_spec": { # "table_name": "table_partitioned_by_multiple_columns__G", # "data_asset_name": "table_partitioned_by_multiple_columns__G", # "batch_identifiers": { # "y": 2020, # "m": 1, # "d": 2, # }, # "splitter_method": "_split_on_multi_column_values", # "splitter_kwargs": {"column_names": ["y", "m", "d"]}, # }, # }, }
def test_more_complex_instantiation_of_ConfiguredAssetSqlDataConnector_include_schema_name_prefix_suffix( splitter_method_name_prefix, test_cases_for_sql_data_connector_sqlite_execution_engine, ): my_data_connector: ConfiguredAssetSqlDataConnector = ConfiguredAssetSqlDataConnector( name="my_sql_data_connector", datasource_name="my_test_datasource", execution_engine="test_cases_for_sql_data_connector_sqlite_execution_engine", assets={ "table_partitioned_by_date_column__A": { "splitter_method": f"{splitter_method_name_prefix}split_on_column_value", "splitter_kwargs": {"column_name": "date"}, "include_schema_name": True, "schema_name": "main", "data_asset_name_prefix": "taxi__", "data_asset_name_suffix": "__asset", }, }, ) assert ( "taxi__main.table_partitioned_by_date_column__A__asset" in my_data_connector.assets ) # schema_name provided, but include_schema_name is set to False with pytest.raises(ge_exceptions.DataConnectorError) as e: ConfiguredAssetSqlDataConnector( name="my_sql_data_connector", datasource_name="my_test_datasource", execution_engine="test_cases_for_sql_data_connector_sqlite_execution_engine", assets={ "table_partitioned_by_date_column__A": { "splitter_method": f"{splitter_method_name_prefix}split_on_column_value", "splitter_kwargs": {"column_name": "date"}, "include_schema_name": False, "schema_name": "main", "data_asset_name_prefix": "taxi__", "data_asset_name_suffix": "__asset", }, }, ) assert ( e.value.message == "ConfiguredAssetSqlDataConnector ran into an error while initializing Asset names. Schema main was specified, but 'include_schema_name' flag was set to False." )
def test_example_B(test_cases_for_sql_data_connector_sqlite_execution_engine): random.seed(0) db = test_cases_for_sql_data_connector_sqlite_execution_engine config = yaml.load(""" name: my_sql_data_connector datasource_name: FAKE_Datasource_NAME assets: table_partitioned_by_timestamp_column__B: splitter_method: _split_on_converted_datetime splitter_kwargs: column_name: timestamp """) config["execution_engine"] = db my_data_connector = ConfiguredAssetSqlDataConnector(**config) report = my_data_connector.self_check() print(json.dumps(report, indent=2)) assert report == { "class_name": "ConfiguredAssetSqlDataConnector", "data_asset_count": 1, "example_data_asset_names": ["table_partitioned_by_timestamp_column__B"], "data_assets": { "table_partitioned_by_timestamp_column__B": { "batch_definition_count": 30, "example_data_references": [ { "timestamp": "2020-01-01" }, { "timestamp": "2020-01-02" }, { "timestamp": "2020-01-03" }, ], } }, "unmatched_data_reference_count": 0, "example_unmatched_data_references": [], # FIXME: (Sam) example_data_reference removed temporarily in PR #2590: # "example_data_reference": { # "n_rows": 8, # "batch_spec": { # "table_name": "table_partitioned_by_timestamp_column__B", # "data_asset_name": "table_partitioned_by_timestamp_column__B", # "batch_identifiers": {"timestamp": "2020-01-02"}, # "splitter_method": "_split_on_converted_datetime", # "splitter_kwargs": {"column_name": "timestamp"}, # }, # }, }
def test_example_C( splitter_method_name_prefix, test_cases_for_sql_data_connector_sqlite_execution_engine, ): random.seed(0) db = test_cases_for_sql_data_connector_sqlite_execution_engine config = yaml.load( f""" name: my_sql_data_connector datasource_name: FAKE_Datasource_NAME assets: table_partitioned_by_regularly_spaced_incrementing_id_column__C: splitter_method: {splitter_method_name_prefix}split_on_divided_integer splitter_kwargs: column_name: id divisor: 10 """, ) config["execution_engine"] = db my_data_connector = ConfiguredAssetSqlDataConnector(**config) report = my_data_connector.self_check() print(json.dumps(report, indent=2)) assert report == { "class_name": "ConfiguredAssetSqlDataConnector", "data_asset_count": 1, "example_data_asset_names": [ "table_partitioned_by_regularly_spaced_incrementing_id_column__C" ], "data_assets": { "table_partitioned_by_regularly_spaced_incrementing_id_column__C": { "batch_definition_count": 12, "example_data_references": [ {"id": 0}, {"id": 1}, {"id": 2}, ], } }, "unmatched_data_reference_count": 0, "example_unmatched_data_references": [], # FIXME: (Sam) example_data_reference removed temporarily in PR #2590: # "example_data_reference": { # "n_rows": 10, # "batch_spec": { # "table_name": "table_partitioned_by_regularly_spaced_incrementing_id_column__C", # "data_asset_name": "table_partitioned_by_regularly_spaced_incrementing_id_column__C", # "batch_identifiers": {"id": 1}, # "splitter_method": "_split_on_divided_integer", # "splitter_kwargs": {"column_name": "id", "divisor": 10}, # }, # }, }
def test_get_batch_definition_list_from_batch_request( test_cases_for_sql_data_connector_sqlite_execution_engine, ): random.seed(0) db = test_cases_for_sql_data_connector_sqlite_execution_engine config = yaml.load( """ name: my_sql_data_connector datasource_name: FAKE_Datasource_NAME data_assets: table_partitioned_by_date_column__A: splitter_method: _split_on_column_value splitter_kwargs: column_name: date """, ) config["execution_engine"] = db my_data_connector = ConfiguredAssetSqlDataConnector(**config) my_data_connector._refresh_data_references_cache() batch_definition_list = ( my_data_connector.get_batch_definition_list_from_batch_request( batch_request=BatchRequest( datasource_name="FAKE_Datasource_NAME", data_connector_name="my_sql_data_connector", data_asset_name="table_partitioned_by_date_column__A", partition_request={ "partition_identifiers": { "date": "2020-01-01" } }, ))) assert len(batch_definition_list) == 1 batch_definition_list = ( my_data_connector.get_batch_definition_list_from_batch_request( batch_request=BatchRequest( datasource_name="FAKE_Datasource_NAME", data_connector_name="my_sql_data_connector", data_asset_name="table_partitioned_by_date_column__A", partition_request={"partition_identifiers": {}}, ))) assert len(batch_definition_list) == 30 # Note: Abe 20201109: It would be nice to put in safeguards for mistakes like this. # In this case, "date" should go inside "partition_identifiers". # Currently, the method ignores "date" entirely, and matches on too many partitions. # I don't think this is unique to ConfiguredAssetSqlDataConnector. # with pytest.raises(DataConnectorError) as e: # batch_definition_list = my_data_connector.get_batch_definition_list_from_batch_request( # batch_request=BatchRequest( # datasource_name="FAKE_Datasource_NAME", # data_connector_name="my_sql_data_connector", # data_asset_name="table_partitioned_by_date_column__A", # partition_request={ # "partition_identifiers" : {}, # "date" : "2020-01-01", # } # )) # assert "Unmatched key" in e.value.message batch_definition_list = ( my_data_connector.get_batch_definition_list_from_batch_request( batch_request=BatchRequest( datasource_name="FAKE_Datasource_NAME", data_connector_name="my_sql_data_connector", data_asset_name="table_partitioned_by_date_column__A", ))) assert len(batch_definition_list) == 30 with pytest.raises(TypeError): my_data_connector.get_batch_definition_list_from_batch_request( batch_request=BatchRequest( datasource_name="FAKE_Datasource_NAME", data_connector_name="my_sql_data_connector", )) with pytest.raises(TypeError): my_data_connector.get_batch_definition_list_from_batch_request( batch_request=BatchRequest(datasource_name="FAKE_Datasource_NAME", )) with pytest.raises(TypeError): my_data_connector.get_batch_definition_list_from_batch_request( batch_request=BatchRequest())
def test_example_E(test_cases_for_sql_data_connector_sqlite_execution_engine): random.seed(0) db = test_cases_for_sql_data_connector_sqlite_execution_engine config = yaml.load( """ name: my_sql_data_connector datasource_name: FAKE_Datasource_NAME data_assets: table_partitioned_by_incrementing_batch_id__E: splitter_method: _split_on_column_value splitter_kwargs: column_name: batch_id """, ) config["execution_engine"] = db my_data_connector = ConfiguredAssetSqlDataConnector(**config) report = my_data_connector.self_check() print(json.dumps(report, indent=2)) assert report == { "class_name": "ConfiguredAssetSqlDataConnector", "data_asset_count": 1, "example_data_asset_names": ["table_partitioned_by_incrementing_batch_id__E"], "data_assets": { "table_partitioned_by_incrementing_batch_id__E": { "batch_definition_count": 11, "example_data_references": [ { "batch_id": 0 }, { "batch_id": 1 }, { "batch_id": 2 }, ], } }, "unmatched_data_reference_count": 0, "example_unmatched_data_references": [], # FIXME: (Sam) example_data_reference removed temporarily in PR #2590: # "example_data_reference": { # "n_rows": 9, # "batch_spec": { # "table_name": "table_partitioned_by_incrementing_batch_id__E", # "data_asset_name": "table_partitioned_by_incrementing_batch_id__E", # "partition_definition": {"batch_id": 1}, # "splitter_method": "_split_on_column_value", # "splitter_kwargs": {"column_name": "batch_id"}, # }, # }, }
"class_name": "SqlAlchemyExecutionEngine", "connection_string": connection_string, }, ) # 2. Set sampler in data connector config data_connector_name: str = "test_data_connector" data_asset_name: str = table_name # Read from generated table name column_name: str = taxi_splitting_test_cases.test_column_name data_connector: ConfiguredAssetSqlDataConnector = ( ConfiguredAssetSqlDataConnector( name=data_connector_name, datasource_name=datasource_name, execution_engine=context.datasources[datasource_name]. execution_engine, assets={ data_asset_name: { "sampling_method": test_case.sampling_method_name, "sampling_kwargs": test_case.sampling_kwargs, } }, )) # 3. Check if resulting batches are as expected # using data_connector.get_batch_definition_list_from_batch_request() batch_request: BatchRequest = BatchRequest( datasource_name=datasource_name, data_connector_name=data_connector_name, data_asset_name=data_asset_name, ) batch_definition_list: List[ BatchDefinition] = data_connector.get_batch_definition_list_from_batch_request(