acryl-datahub 1.1.0.5rc4__py3-none-any.whl → 1.1.0.5rc5__py3-none-any.whl

{acryl_datahub-1.1.0.5rc4.dist-info → acryl_datahub-1.1.0.5rc5.dist-info}/RECORD

@@ -1,7 +1,7 @@
-acryl_datahub-1.1.0.5rc4.dist-info/licenses/LICENSE,sha256=9xNHpsD0uYF5ONzXsKDCuHHB-xbiCrSbueWXqrTNsxk,11365
+acryl_datahub-1.1.0.5rc5.dist-info/licenses/LICENSE,sha256=9xNHpsD0uYF5ONzXsKDCuHHB-xbiCrSbueWXqrTNsxk,11365
 datahub/__init__.py,sha256=aq_i5lVREmoLfYIqcx_pEQicO855YlhD19tWc1eZZNI,59
 datahub/__main__.py,sha256=pegIvQ9hzK7IhqVeUi1MeADSZ2QlP-D3K0OQdEg55RU,106
-datahub/_version.py,sha256=FxmAFDylE_4rYCNPJsg5swySRi1b4oeflnNCP_Du7ks,323
+datahub/_version.py,sha256=asGELDofJg3dNOgo2IEO-VgcFEFlz2Bkyep1JJAZd9U,323
 datahub/entrypoints.py,sha256=H-YFTvxTJOgpWsFBVlxyb1opjkq-hjTzNmjy5Fq3RHg,8992
 datahub/errors.py,sha256=p5rFAdAGVCk4Lqolol1YvthceadUSwpaCxLXRcyCCFQ,676
 datahub/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -67,7 +67,7 @@ datahub/api/graphql/assertion.py,sha256=o_q6SV7N1rJTVMNKSUBGJnZPk6TcVYoVShgDmPw6
 datahub/api/graphql/base.py,sha256=zk724_oYSJ0nK7X7Z80MijnA6ry9JqpxnBsJeYuONKA,1737
 datahub/api/graphql/operation.py,sha256=7E80HyE-5JLfLbFkQbgJeNwIaKngjBCrWES8eJO4OYc,5112
 datahub/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datahub/cli/check_cli.py,sha256=CsOF57GEJAxnRS7a2mrzNkbQSTxcvtr6nBcB9s5MBXA,15197
+datahub/cli/check_cli.py,sha256=1R_OwSx47GPtWMP0e9Ljef7mBr-xxDG1O8ls7YG8khQ,15767
 datahub/cli/cli_utils.py,sha256=2uvPv6WqxbRdH7UteHwhRash4E0ncU5P6XebrFLeECo,13584
 datahub/cli/config_utils.py,sha256=EeBGfhmf4AxYoTfnZ4GSiGIgpzJFkduNjN_FwmxZGhA,4889
 datahub/cli/container_cli.py,sha256=uDOwewGEPYHQt-ppYEb8ESXhZjPNIZG0Rt3cm2FzPqc,1569
@@ -141,7 +141,7 @@ datahub/ingestion/api/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3h
 datahub/ingestion/api/closeable.py,sha256=k12AT--s4GDtZ-po_rVm5QKgvGIDteeRPByZPIOfecA,599
 datahub/ingestion/api/committable.py,sha256=4S6GuBzvX2vb1A8P506NbspOKfZ1621sBG8t0lvRb8o,886
 datahub/ingestion/api/common.py,sha256=26lXJiM4YfdnVH1xfe2bpZNp2VKCdJcJ8ynK7rhh0FY,3029
-datahub/ingestion/api/decorators.py,sha256=qsI2HkyruhVR7VU-K0sCZmpAR9kg49CDdL9Tu24QpsA,4016
+datahub/ingestion/api/decorators.py,sha256=okTDYtOQygyC9-Mu3jiTHuW0Wwt9Td25utxc5yMMSDY,4420
 datahub/ingestion/api/global_context.py,sha256=OdSJg4a_RKE52nu8MSiEkK2UqRRDhDTyOleHEAzPKho,575
 datahub/ingestion/api/incremental_lineage_helper.py,sha256=7a6FTJ_uz4EEJS1vPtbYB2KvNlcZB3py28_FKxmRiSk,5993
 datahub/ingestion/api/incremental_properties_helper.py,sha256=KzdxdrQtaMV2XMHfPsCtRa7ffDGPA1w1hgPUjeenZBU,2514
@@ -179,7 +179,7 @@ datahub/ingestion/glossary/classifier.py,sha256=daLxnVv_JlfB_jBOxH5LrU_xQRndrsGo
 datahub/ingestion/glossary/classifier_registry.py,sha256=yFOYLQhDgCLqXYMG3L1BquXafeLcZDcmp8meyw6k9ts,307
 datahub/ingestion/glossary/datahub_classifier.py,sha256=O7wm6gQT1Jf2QSKdWjJQbS5oSzJwplXzfza26Gdq5Mg,7555
 datahub/ingestion/graph/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datahub/ingestion/graph/client.py,sha256=VnrUsC2jGN5W3nkurAxIR887J4O9mP_S0VjdheekNQA,72412
+datahub/ingestion/graph/client.py,sha256=kuxNcwW6O3o9i_AHoRLabeMSrYBMOSCExiZYnzOBYV0,74221
 datahub/ingestion/graph/config.py,sha256=rmkcqAL8fJoY9QyAeS0Xm8HvwHzV3pCjY-Om-50JJTI,1015
 datahub/ingestion/graph/connections.py,sha256=9462L0ZWGKURyypAln25eMPhK3pcufBar9tNDoqspXs,741
 datahub/ingestion/graph/entity_versioning.py,sha256=nrcNz0Qm6kpE6oTu_mrYUQDx14KPspBTc6R9SyFUY6c,6901
@@ -223,7 +223,7 @@ datahub/ingestion/source/openapi_parser.py,sha256=T87e2r-oPGgQl_FDMHnSGFZzApvWDC
 datahub/ingestion/source/preset.py,sha256=fncn-fgYcITsYEHVsvV6cGTQ9_xc_R06ejrw6ZbY3QA,3966
 datahub/ingestion/source/pulsar.py,sha256=u5F8QnCLJsht5-7XCiUTsnfhCPIpKVB_l32CgMCU-As,20187
 datahub/ingestion/source/redash.py,sha256=YxjSad-X_wPmxYH8dJmFz_VCFhiLTCTSlK99WdvcYiA,30653
-datahub/ingestion/source/salesforce.py,sha256=CQtDFv1OsbC1vyzNbKOc6GxhFQ5GdYj45hgAF0-oIcw,40487
+datahub/ingestion/source/salesforce.py,sha256=Pa_w1XszxFd8fyhpSWOfc2nOnevHwwstIvnRrQT4R9M,40584
 datahub/ingestion/source/source_registry.py,sha256=a2mLjJPLkSI-gYCTb_7U7Jo4D8jGknNQ_yScPIihXFk,1208
 datahub/ingestion/source/sql_queries.py,sha256=Ip7UZub7fgMh7P5jL_zJPY7lSkc9GGTy8GJ8lqZrcsE,9502
 datahub/ingestion/source/superset.py,sha256=4wEjhBj_Zf__1EhXDE9ltbaR9NTMZVdXFviHLf5VFL4,48387
@@ -237,7 +237,7 @@ datahub/ingestion/source/apply/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5
 datahub/ingestion/source/apply/datahub_apply.py,sha256=xTD-Iq3UHhxcz61RwNuI2kJjRrnQEfZFSgvS1X6loV4,7703
 datahub/ingestion/source/aws/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/ingestion/source/aws/aws_common.py,sha256=DN49pHY-AOVHqI3smF9cVPwvPmh5UqOF4ga-wga_gPc,17928
-datahub/ingestion/source/aws/glue.py,sha256=E_oCgf-taI2ZVXSn9LA9MoDk6JJBFEDBUbNIFkqUUSA,66547
+datahub/ingestion/source/aws/glue.py,sha256=Pa5QmQf8ZTQ9KAA6frp5LGrls3gFhRATEvhjaGTJnB8,66533
 datahub/ingestion/source/aws/s3_boto_utils.py,sha256=Y54jlLV5gLcuZ4Zs57kIW5dYHD89RSFfsVNlFbRnSkQ,3901
 datahub/ingestion/source/aws/s3_util.py,sha256=OFypcgmVC6jnZM90-gjcPpAMtTV1lbnreCaMhCzNlzs,2149
 datahub/ingestion/source/aws/sagemaker.py,sha256=Bl2tkBYnrindgx61VHYgNovUF_Kp_fXNcivQn28vC2w,5254
@@ -283,7 +283,7 @@ datahub/ingestion/source/common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm
 datahub/ingestion/source/common/data_platforms.py,sha256=HhuP3YIEi2WpyKDjUU8RiM0a2qjHWQcvc8kcqub0cVo,548
 datahub/ingestion/source/common/data_reader.py,sha256=XbSxiRTYrk6seOz0ZjVjzSpGvP8lEjmqXrNI4cdYYmQ,1819
 datahub/ingestion/source/common/gcp_credentials_config.py,sha256=_NapGkAqZMbXNClLlmOfANS7U6rChhdthRX9s9iUv9k,2411
-datahub/ingestion/source/common/subtypes.py,sha256=wkFfCwfjkvllpEYWjC6Bsv3htVYBGVCAkZPs3oFXhII,3306
+datahub/ingestion/source/common/subtypes.py,sha256=SvsaT2PM5VG7iH3c0iyvAK3oyWS_EpB3LhFOoX-e5Ws,4493
 datahub/ingestion/source/data_lake_common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/ingestion/source/data_lake_common/config.py,sha256=qUk83B01hjuBKHvVz8SmXnVCy5eFj-2-2QLEOrAdbgk,359
 datahub/ingestion/source/data_lake_common/data_lake_utils.py,sha256=55mK0nsehqGDTUOol9Oi3jZs8Pb04PIsHdC2WPP0dkg,6576
@@ -299,7 +299,7 @@ datahub/ingestion/source/datahub/report.py,sha256=VHBfCbwFRzdLdB7hQG9ST4EiZxl_vB
 datahub/ingestion/source/datahub/state.py,sha256=PZoT7sSK1wadVf5vN6phrgr7I6LL7ePP-EJjP1OO0bQ,3507
 datahub/ingestion/source/dbt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/ingestion/source/dbt/dbt_cloud.py,sha256=435blM28-RIE7ZfQReh69zBXppv8tjbku10oFcyJY8Q,17891
-datahub/ingestion/source/dbt/dbt_common.py,sha256=kcpIKAhrJHja8wabqtpzLY7JsRGDJUddoH4TWOqszq8,82062
+datahub/ingestion/source/dbt/dbt_common.py,sha256=5iqlHbH1eaRcysxvFmDfgNrfVpygnmmeHbLq_sLg2P4,82190
 datahub/ingestion/source/dbt/dbt_core.py,sha256=WVI2ZYXOMxgFzJnJqsqmEGS-5xdfiVIDsCb78lvSeQ0,24930
 datahub/ingestion/source/dbt/dbt_tests.py,sha256=pOZJaP4VsbaE5j4qVlE_E3ifno_KQpidfGTvOi5fr6I,9839
 datahub/ingestion/source/debug/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -487,7 +487,7 @@ datahub/ingestion/source/snowflake/snowflake_usage_v2.py,sha256=ySFm7WDk8FW9KjCn
 datahub/ingestion/source/snowflake/snowflake_utils.py,sha256=w-SftpjHSOf_6S2T2srHJyJta3MHS0usmn4Z9jgx4QE,13858
 datahub/ingestion/source/snowflake/snowflake_v2.py,sha256=eqOVRYmFiWQPAIDWUcOo5QZkG05xanlEbxz4aDTvSNM,34706
 datahub/ingestion/source/sql/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datahub/ingestion/source/sql/athena.py,sha256=U-btrqnbOgrBptl6b--QJrio_AN-c6pF0yL-wbUt4P0,23986
+datahub/ingestion/source/sql/athena.py,sha256=gF32rRkQ4TXQUKkyH3HPLINEuMDO0Tc1awGW8qEIUKw,24212
 datahub/ingestion/source/sql/clickhouse.py,sha256=_uf-6fQYxI1-kMw0dNUvJf1f75XX0Qps4HMVlgzdpmo,25670
 datahub/ingestion/source/sql/cockroachdb.py,sha256=WoOKCq7YjsuzSPm1SmKIYZ9CrvlSF8zWmP1fNHn4G3Q,1360
 datahub/ingestion/source/sql/druid.py,sha256=_tzgTa5jhPUXk6WCmS7p10feCwJm6yUFcOgMZA-OcE8,2922
@@ -499,7 +499,7 @@ datahub/ingestion/source/sql/mysql.py,sha256=34Vu3otULxUY0-JUEYdZw2aoyuTlc9KLcoJ
 datahub/ingestion/source/sql/oracle.py,sha256=ftnrk3iiEelyv9PBHPYbairuP1WgxZbi1gu6YdqY69E,29965
 datahub/ingestion/source/sql/postgres.py,sha256=t-28dTWLa3c_NgWlGS9ZVNFlZLxd6xTSguIKP4S4NHs,11887
 datahub/ingestion/source/sql/presto.py,sha256=tATa0M2q0PjUC_E9W_jSUsmKTP7cVJayLgrFMzG_eao,4223
-datahub/ingestion/source/sql/sql_common.py,sha256=ktCDlJeT1l_ZgwFwmq9SpssotKfyJ8xc4RFK6nmWb74,51648
+datahub/ingestion/source/sql/sql_common.py,sha256=wSCfcqEONichvl4D9_tQ1HQa8puYF-5BpNnES93f7-Y,51786
 datahub/ingestion/source/sql/sql_config.py,sha256=u3nGZYYl1WtaxfNsDU5bglgZ5Jq3Fxk9xei_CUIAXB0,8222
 datahub/ingestion/source/sql/sql_generic.py,sha256=9AERvkK8kdJUeDOzCYJDb93xdv6Z4DGho0NfeHj5Uyg,2740
 datahub/ingestion/source/sql/sql_generic_profiler.py,sha256=Zr39j4SI1fPTx1JdopVJyBslFnyp3lZCeb1th9eEB5c,11723
@@ -509,13 +509,13 @@ datahub/ingestion/source/sql/sql_utils.py,sha256=q-Bsk6WxlsRtrw9RXBxvqI3zuaMTC_F
 datahub/ingestion/source/sql/sqlalchemy_data_reader.py,sha256=FvHZ4JEK3aR2DYOBZiT_ZsAy12RjTu4t_KIR_92B11k,2644
 datahub/ingestion/source/sql/sqlalchemy_uri.py,sha256=u0ZvgdJjXZdo_vl7YIQfYuuWbGwpnH6OSozI2e8ZV4I,858
 datahub/ingestion/source/sql/sqlalchemy_uri_mapper.py,sha256=KOpbmDIE2h1hyYEsbVHJi2B7FlsyUMTXZx4diyzltQg,1826
-datahub/ingestion/source/sql/teradata.py,sha256=9WdrxDy02lRJi9IZgsAATFsmxcQnIw5Gr6yCqHJQy5k,33507
+datahub/ingestion/source/sql/teradata.py,sha256=yRU_F3dp2rH0nMC0iiL4X2DWXLyNn7kSWZ_ewSIgYEA,33537
 datahub/ingestion/source/sql/trino.py,sha256=zIfQ6GvW8Sbw4sxqsTcnibT51STka_nzNYvmld6HfHw,18947
 datahub/ingestion/source/sql/two_tier_sql_source.py,sha256=AB3Gtx4omAy_08zadHQpmUGmIGufkZ6o_ihWNnfvzYc,5783
-datahub/ingestion/source/sql/vertica.py,sha256=4keUJd6BMxaZZVYT13lT_8BWpUnZZeBn8Da46EltHy0,33441
+datahub/ingestion/source/sql/vertica.py,sha256=i1HLBW_cjAYbSCxB5MI-_7ZWD7Wmn4m6kRc_IZMD__o,33595
 datahub/ingestion/source/sql/mssql/__init__.py,sha256=1agpl8S_uDW40olkhCX_W19dbr5GO9qgjS3R7pLRZSk,87
 datahub/ingestion/source/sql/mssql/job_models.py,sha256=nAo3rciu-w2-dXCz6_ekDEbGMEjCMEfh8WvSfXoF2l0,9359
-datahub/ingestion/source/sql/mssql/source.py,sha256=krV80XcNFu-N73xoUA-W2zLiop3pJP7zkb3Dvj4PDKU,42783
+datahub/ingestion/source/sql/mssql/source.py,sha256=Uise_u6yXKU__9B_U3D3yObWNIVDzrz2AgEDZOlk6bQ,43101
 datahub/ingestion/source/sql/stored_procedures/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/ingestion/source/sql/stored_procedures/base.py,sha256=n0l5OaTuW-m3TRvkxs3TqvgMeWF6BagzW3tjyWUcC1A,8631
 datahub/ingestion/source/sql/stored_procedures/lineage.py,sha256=fryLhuAlsjr9SHIjHJ-PmtCMx89bjzWVnJZ3f1bwQVU,1905
@@ -535,7 +535,7 @@ datahub/ingestion/source/state_provider/datahub_ingestion_checkpointing_provider
 datahub/ingestion/source/state_provider/file_ingestion_checkpointing_provider.py,sha256=DziD57PbHn2Tcy51tYXCG-GQgyTGMUxnkuzVS_xihFY,4079
 datahub/ingestion/source/state_provider/state_provider_registry.py,sha256=SVq4mIyGNmLXE9OZx1taOiNPqDoQp03-Ot9rYnB5F3k,401
 datahub/ingestion/source/tableau/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datahub/ingestion/source/tableau/tableau.py,sha256=HhVcMBEWsPVWBVOwpSY4ShrwGnHL-O7qyYxl0U9tlgo,155161
+datahub/ingestion/source/tableau/tableau.py,sha256=ryLxfWwCRXhz7-PhPl1y3OwX08KOFT5agBc9XWZRrGQ,155291
 datahub/ingestion/source/tableau/tableau_common.py,sha256=MIkHKZg_v2IVCRk-YdPlLZl3m0LcWZm5Indwb3IV2ZQ,26931
 datahub/ingestion/source/tableau/tableau_constant.py,sha256=ZcAeHsQUXVVL26ORly0ByZk_GJAFbxaKuJAlX_sYMac,2686
 datahub/ingestion/source/tableau/tableau_server_wrapper.py,sha256=nSyx9RzC6TCQDm-cTVJ657qT8iDwzk_8JMKpohhmOc4,1046
@@ -1087,8 +1087,8 @@ datahub_provider/operators/datahub_assertion_operator.py,sha256=uvTQ-jk2F0sbqqxp
 datahub_provider/operators/datahub_assertion_sensor.py,sha256=lCBj_3x1cf5GMNpHdfkpHuyHfVxsm6ff5x2Z5iizcAo,140
 datahub_provider/operators/datahub_operation_operator.py,sha256=aevDp2FzX7FxGlXrR0khoHNbxbhKR2qPEX5e8O2Jyzw,174
 datahub_provider/operators/datahub_operation_sensor.py,sha256=8fcdVBCEPgqy1etTXgLoiHoJrRt_nzFZQMdSzHqSG7M,168
-acryl_datahub-1.1.0.5rc4.dist-info/METADATA,sha256=VA5JlGUBQDVDE7ZFv2nz76rrLvUFYpTcyKfcjjiidRk,182347
-acryl_datahub-1.1.0.5rc4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-acryl_datahub-1.1.0.5rc4.dist-info/entry_points.txt,sha256=bnGf6eX9UhiW8yVHtt6MJCVcmLErvrVQxTJAayA-PKc,9885
-acryl_datahub-1.1.0.5rc4.dist-info/top_level.txt,sha256=iLjSrLK5ox1YVYcglRUkcvfZPvKlobBWx7CTUXx8_GI,25
-acryl_datahub-1.1.0.5rc4.dist-info/RECORD,,
+acryl_datahub-1.1.0.5rc5.dist-info/METADATA,sha256=TVQ5LvGn85Nq01WoQbqeuRmjSXRjZZhtQiWGSpBgM8w,182347
+acryl_datahub-1.1.0.5rc5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+acryl_datahub-1.1.0.5rc5.dist-info/entry_points.txt,sha256=bnGf6eX9UhiW8yVHtt6MJCVcmLErvrVQxTJAayA-PKc,9885
+acryl_datahub-1.1.0.5rc5.dist-info/top_level.txt,sha256=iLjSrLK5ox1YVYcglRUkcvfZPvKlobBWx7CTUXx8_GI,25
+acryl_datahub-1.1.0.5rc5.dist-info/RECORD,,

datahub/_version.py

@@ -1,6 +1,6 @@
 # Published at https://pypi.org/project/acryl-datahub/.
 __package_name__ = "acryl-datahub"
-__version__ = "1.1.0.5rc4"
+__version__ = "1.1.0.5rc5"
 
 
 def is_dev_mode() -> bool:

datahub/cli/check_cli.py

@@ -482,26 +482,43 @@ def server_config() -> None:
 
 @check.command()
 @click.option(
-    "--urn", required=True, help="The urn or urn pattern (supports % for wildcard)"
+    "--urn", required=False, help="The urn or urn pattern (supports % for wildcard)"
 )
 @click.option("--aspect", default=None, help="Filter to a specific aspect name.")
 @click.option(
     "--start", type=int, default=None, help="Row number of sql store to restore from."
 )
 @click.option("--batch-size", type=int, default=None, help="How many rows to restore.")
+@click.option(
+    "--file",
+    required=False,
+    type=click.Path(exists=True, dir_okay=True, readable=True),
+    help="File absolute path containing URNs (one per line) to restore indices",
+)
 def restore_indices(
-    urn: str,
+    urn: Optional[str],
     aspect: Optional[str],
     start: Optional[int],
     batch_size: Optional[int],
+    file: Optional[str],
 ) -> None:
     """Resync metadata changes into the search and graph indices."""
+    if urn is None and file is None:
+        raise click.UsageError("Either --urn or --file must be provided")
     graph = get_default_graph(ClientMode.CLI)
 
-    result = graph.restore_indices(
+    graph.restore_indices(
         urn_pattern=urn,
         aspect=aspect,
         start=start,
         batch_size=batch_size,
+        file=file,
     )
-    click.echo(result)
+
+
+@check.command()
+def get_kafka_consumer_offsets() -> None:
+    """Get Kafka consumer offsets from the DataHub API."""
+    graph = get_default_graph(ClientMode.CLI)
+    result = graph.get_kafka_consumer_offsets()
+    pprint.pprint(result)

datahub/ingestion/api/decorators.py

@@ -1,12 +1,16 @@
+# So that SourceCapabilityModifier can be resolved at runtime
+from __future__ import annotations
+
 from dataclasses import dataclass
 from enum import Enum, auto
-from typing import Callable, Dict, Optional, Type
+from typing import Callable, Dict, List, Optional, Type
 
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.source import (
     Source,
     SourceCapability as SourceCapability,
 )
+from datahub.ingestion.source.common.subtypes import SourceCapabilityModifier
 
 
 def config_class(config_cls: Type) -> Callable[[Type], Type]:
@@ -88,10 +92,14 @@ class CapabilitySetting:
     capability: SourceCapability
     description: str
     supported: bool
+    subtype_modifier: Optional[List[SourceCapabilityModifier]] = None
 
 
 def capability(
-    capability_name: SourceCapability, description: str, supported: bool = True
+    capability_name: SourceCapability,
+    description: str,
+    supported: bool = True,
+    subtype_modifier: Optional[List[SourceCapabilityModifier]] = None,
 ) -> Callable[[Type], Type]:
     """
     A decorator to mark a source as having a certain capability
@@ -114,7 +122,10 @@ def capability(
                     cls.__capabilities.update(base_caps)
 
         cls.__capabilities[capability_name] = CapabilitySetting(
-            capability=capability_name, description=description, supported=supported
+            capability=capability_name,
+            description=description,
+            supported=supported,
+            subtype_modifier=subtype_modifier,
         )
         return cls
 

datahub/ingestion/graph/client.py

@@ -22,6 +22,7 @@ from typing import (
     Union,
 )
 
+import progressbar
 from avro.schema import RecordSchema
 from pydantic import BaseModel
 from requests.models import HTTPError
@@ -504,7 +505,7 @@ class DataHubGraph(DatahubRestEmitter, EntityVersioningAPI):
             "limit": limit,
             "filter": filter,
         }
-        end_point = f"{self.config.server}/aspects?action=getTimeseriesAspectValues"
+        end_point = f"{self._gms_server}/aspects?action=getTimeseriesAspectValues"
         resp: Dict = self._post_generic(end_point, query_body)
 
         values: Optional[List] = resp.get("value", {}).get("values")
@@ -524,7 +525,7 @@ class DataHubGraph(DatahubRestEmitter, EntityVersioningAPI):
     def get_entity_raw(
         self, entity_urn: str, aspects: Optional[List[str]] = None
     ) -> Dict:
-        endpoint: str = f"{self.config.server}/entitiesV2/{Urn.url_encode(entity_urn)}"
+        endpoint: str = f"{self._gms_server}/entitiesV2/{Urn.url_encode(entity_urn)}"
         if aspects is not None:
             assert aspects, "if provided, aspects must be a non-empty list"
             endpoint = f"{endpoint}?aspects=List(" + ",".join(aspects) + ")"
@@ -654,15 +655,15 @@ class DataHubGraph(DatahubRestEmitter, EntityVersioningAPI):
 
     @property
     def _search_endpoint(self):
-        return f"{self.config.server}/entities?action=search"
+        return f"{self._gms_server}/entities?action=search"
 
     @property
     def _relationships_endpoint(self):
-        return f"{self.config.server}/openapi/relationships/v1/"
+        return f"{self._gms_server}/openapi/relationships/v1/"
 
     @property
     def _aspect_count_endpoint(self):
-        return f"{self.config.server}/aspects?action=getCount"
+        return f"{self._gms_server}/aspects?action=getCount"
 
     def get_domain_urn_by_name(self, domain_name: str) -> Optional[str]:
         """Retrieve a domain urn based on its name. Returns None if there is no match found"""
@@ -1209,7 +1210,7 @@ class DataHubGraph(DatahubRestEmitter, EntityVersioningAPI):
         operation_name: Optional[str] = None,
         format_exception: bool = True,
     ) -> Dict:
-        url = f"{self.config.server}/api/graphql"
+        url = f"{self._gms_server}/api/graphql"
 
         body: Dict = {
             "query": query,
@@ -1434,40 +1435,82 @@ class DataHubGraph(DatahubRestEmitter, EntityVersioningAPI):
         related_aspects = response.get("relatedAspects", [])
         return reference_count, related_aspects
 
+    def get_kafka_consumer_offsets(
+        self,
+    ) -> dict:
+        """
+        Get Kafka consumer offsets from the DataHub API.
+
+        Args:
+            graph (DataHubGraph): The DataHub graph client
+
+        """
+        urls = {
+            "mcp": f"{self.config.server}/openapi/operations/kafka/mcp/consumer/offsets",
+            "mcl": f"{self.config.server}/openapi/operations/kafka/mcl/consumer/offsets",
+            "mcl-timeseries": f"{self.config.server}/openapi/operations/kafka/mcl-timeseries/consumer/offsets",
+        }
+
+        params = {"skipCache": "true", "detailed": "true"}
+        results = {}
+        for key, url in urls.items():
+            response = self._get_generic(url=url, params=params)
+            results[key] = response
+            if "errors" in response:
+                logger.error(f"Error: {response['errors']}")
+        return results
+
+    def _restore_index_call(self, payload_obj: dict) -> None:
+        result = self._post_generic(
+            f"{self._gms_server}/operations?action=restoreIndices", payload_obj
+        )
+        logger.debug(f"Restore indices result: {result}")
+
     def restore_indices(
         self,
-        urn_pattern: str,
+        urn_pattern: Optional[str] = None,
         aspect: Optional[str] = None,
         start: Optional[int] = None,
         batch_size: Optional[int] = None,
-    ) -> str:
+        file: Optional[str] = None,
+    ) -> None:
         """Restore the indices for a given urn or urn-like pattern.
 
         Args:
-            urn_pattern: The exact URN or a pattern (with % for wildcard) to match URNs.
+            urn_pattern: The exact URN or a pattern (with % for wildcard) to match URNs. If not provided, will restore indices from the file.
             aspect: Optional aspect string to restore indices for a specific aspect.
-            start: Optional integer to decide which row number of sql store to restore from. Default: 0.
-            batch_size: Optional integer to decide how many rows to restore. Default: 10.
+            start: Optional integer to decide which row number of sql store to restore from. Default: 0. Ignored in case file is provided.
+            batch_size: Optional integer to decide how many rows to restore. Default: 10. Ignored in case file is provided.
+            file: Optional file path to a file containing URNs to restore indices for.
 
         Returns:
             A string containing the result of the restore indices operation. This format is subject to change.
         """
-        if "%" in urn_pattern:
-            payload_obj: dict = {"urnLike": urn_pattern}
+        payload_obj = {}
+        if file is not None:
+            with open(file) as f:
+                for urn in progressbar.progressbar(f.readlines()):
+                    urn = urn.strip()
+                    if "%" in urn:
+                        payload_obj["urnLike"] = urn
+                    else:
+                        payload_obj["urn"] = urn
+                    if aspect is not None:
+                        payload_obj["aspect"] = aspect
+                    self._restore_index_call(payload_obj)
         else:
-            payload_obj = {"urn": urn_pattern}
-        if aspect is not None:
-            payload_obj["aspect"] = aspect
-        if start is not None:
-            payload_obj["start"] = start
-        if batch_size is not None:
-            payload_obj["batchSize"] = batch_size
-        raw_result = self._post_generic(
-            f"{self._gms_server}/operations?action=restoreIndices", payload_obj
-        )
-        result = raw_result["value"]
-        logger.debug(f"Restore indices result: {result}")
-        return result
+            if urn_pattern is not None:
+                if "%" in urn_pattern:
+                    payload_obj["urnLike"] = urn_pattern
+                else:
+                    payload_obj["urn"] = urn_pattern
+            if aspect is not None:
+                payload_obj["aspect"] = aspect
+            if start is not None:
+                payload_obj["start"] = start
+            if batch_size is not None:
+                payload_obj["batchSize"] = batch_size
+            self._restore_index_call(payload_obj)
 
     @functools.lru_cache
     def _make_schema_resolver(
@@ -1774,7 +1817,7 @@ class DataHubGraph(DatahubRestEmitter, EntityVersioningAPI):
             "Accept": "application/json",
             "Content-Type": "application/json",
         }
-        url = f"{self.config.server}/openapi/v2/entity/batch/{entity_name}"
+        url = f"{self._gms_server}/openapi/v2/entity/batch/{entity_name}"
         response = self._session.post(url, data=json.dumps(payload), headers=headers)
         response.raise_for_status()
 
@@ -1831,7 +1874,7 @@ class DataHubGraph(DatahubRestEmitter, EntityVersioningAPI):
             "Content-Type": "application/json",
         }
 
-        url = f"{self.config.server}/openapi/v3/entity/{entity_name}/batchGet"
+        url = f"{self._gms_server}/openapi/v3/entity/{entity_name}/batchGet"
         if with_system_metadata:
             url += "?systemMetadata=true"
 

datahub/ingestion/source/aws/glue.py

@@ -269,7 +269,7 @@ class GlueSourceReport(StaleEntityRemovalSourceReport):
 @capability(SourceCapability.DOMAINS, "Supported via the `domain` config field")
 @capability(
     SourceCapability.DELETION_DETECTION,
-    "Enabled by default when stateful ingestion is turned on.",
+    "Enabled by default via stateful ingestion.",
 )
 @capability(SourceCapability.LINEAGE_COARSE, "Enabled by default")
 @capability(

datahub/ingestion/source/common/subtypes.py

@@ -1,5 +1,10 @@
+import logging
+from typing import Any, Dict
+
 from datahub.utilities.str_enum import StrEnum
 
+logger = logging.getLogger(__name__)
+
 
 class DatasetSubTypes(StrEnum):
     # Generic SubTypes
@@ -76,6 +81,9 @@ class JobContainerSubTypes(StrEnum):
 
 
 class BIAssetSubTypes(StrEnum):
+    DASHBOARD = "Dashboard"
+    CHART = "Chart"
+
     # Generic SubTypes
     REPORT = "Report"
 
@@ -118,3 +126,36 @@ class MLAssetSubTypes(StrEnum):
     VERTEX_PIPELINE = "Pipeline Job"
     VERTEX_PIPELINE_TASK = "Pipeline Task"
     VERTEX_PIPELINE_TASK_RUN = "Pipeline Task Run"
+
+
+def create_source_capability_modifier_enum():
+    all_values: Dict[str, Any] = {}
+    source_enums = [
+        DatasetSubTypes,
+        DatasetContainerSubTypes,
+        BIContainerSubTypes,
+        FlowContainerSubTypes,
+        JobContainerSubTypes,
+        BIAssetSubTypes,
+        MLAssetSubTypes,
+    ]
+
+    for enum_class in source_enums:
+        for member in enum_class:  # type: ignore[var-annotated]
+            if member.name in all_values:
+                logger.error(
+                    f"Warning: {member.name} already exists with value {all_values[member.name]}, skipping {member.value}"
+                )
+                continue
+            all_values[member.name] = member.value
+
+    enum_code = "class SourceCapabilityModifier(StrEnum):\n"
+    for name, value in all_values.items():
+        enum_code += f'    {name} = "{value}"\n'
+
+    exec(enum_code, globals())
+    return globals()["SourceCapabilityModifier"]
+
+
+# This will have all values from the enums above
+SourceCapabilityModifier = create_source_capability_modifier_enum()

datahub/ingestion/source/dbt/dbt_common.py

@@ -355,7 +355,7 @@ class DBTCommonConfig(
     # override default value to True.
     incremental_lineage: bool = Field(
         default=True,
-        description="When enabled, emits incremental/patch lineage for non-dbt entities. When disabled, re-states lineage on each run.",
+        description="When enabled, emits incremental/patch lineage for non-dbt entities. When disabled, re-states lineage on each run. This would also require enabling 'incremental_lineage' in the counterpart warehouse ingestion (_e.g._ BigQuery, Redshift, etc).",
     )
 
     _remove_use_compiled_code = pydantic_removed_field("use_compiled_code")

datahub/ingestion/source/salesforce.py

@@ -33,7 +33,10 @@ from datahub.ingestion.api.decorators import (
 )
 from datahub.ingestion.api.source import MetadataWorkUnitProcessor, SourceReport
 from datahub.ingestion.api.workunit import MetadataWorkUnit
-from datahub.ingestion.source.common.subtypes import DatasetSubTypes
+from datahub.ingestion.source.common.subtypes import (
+    DatasetSubTypes,
+    SourceCapabilityModifier,
+)
 from datahub.ingestion.source.state.stale_entity_removal_handler import (
     StaleEntityRemovalHandler,
     StaleEntityRemovalSourceReport,
@@ -532,11 +535,11 @@ class SalesforceApi:
 @capability(
     capability_name=SourceCapability.DATA_PROFILING,
     description="Only table level profiling is supported via `profiling.enabled` config field",
+    subtype_modifier=[SourceCapabilityModifier.TABLE],
 )
 @capability(
     capability_name=SourceCapability.DELETION_DETECTION,
-    description="Not supported yet",
-    supported=False,
+    description="Enabled by default via stateful ingestion",
 )
 @capability(
     capability_name=SourceCapability.SCHEMA_METADATA,

datahub/ingestion/source/sql/athena.py

@@ -29,7 +29,10 @@ from datahub.ingestion.api.decorators import (
 from datahub.ingestion.api.source import StructuredLogLevel
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.aws.s3_util import make_s3_urn
-from datahub.ingestion.source.common.subtypes import DatasetContainerSubTypes
+from datahub.ingestion.source.common.subtypes import (
+    DatasetContainerSubTypes,
+    SourceCapabilityModifier,
+)
 from datahub.ingestion.source.ge_profiling_config import GEProfilingConfig
 from datahub.ingestion.source.sql.sql_common import (
     SQLAlchemySource,
@@ -321,9 +324,18 @@ class Partitionitem:
 @capability(
     SourceCapability.DATA_PROFILING,
     "Optionally enabled via configuration. Profiling uses sql queries on whole table which can be expensive operation.",
+    subtype_modifier=[SourceCapabilityModifier.TABLE],
+)
+@capability(
+    SourceCapability.LINEAGE_COARSE,
+    "Supported for S3 tables",
+    subtype_modifier=[SourceCapabilityModifier.TABLE],
+)
+@capability(
+    SourceCapability.LINEAGE_FINE,
+    "Supported for S3 tables",
+    subtype_modifier=[SourceCapabilityModifier.TABLE],
 )
-@capability(SourceCapability.LINEAGE_COARSE, "Supported for S3 tables")
-@capability(SourceCapability.LINEAGE_FINE, "Supported for S3 tables")
 @capability(SourceCapability.DESCRIPTIONS, "Enabled by default")
 class AthenaSource(SQLAlchemySource):
     """

datahub/ingestion/source/sql/mssql/source.py

@@ -27,6 +27,7 @@ from datahub.ingestion.api.decorators import (
 from datahub.ingestion.api.source import StructuredLogLevel
 from datahub.ingestion.api.source_helpers import auto_workunit
 from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.ingestion.source.common.subtypes import SourceCapabilityModifier
 from datahub.ingestion.source.sql.mssql.job_models import (
     JobStep,
     MSSQLDataFlow,
@@ -177,10 +178,18 @@ class SQLServerConfig(BasicSQLAlchemyConfig):
 @capability(
     SourceCapability.LINEAGE_COARSE,
     "Enabled by default to get lineage for stored procedures via `include_lineage` and for views via `include_view_lineage`",
+    subtype_modifier=[
+        SourceCapabilityModifier.STORED_PROCEDURE,
+        SourceCapabilityModifier.VIEW,
+    ],
 )
 @capability(
     SourceCapability.LINEAGE_FINE,
     "Enabled by default to get lineage for stored procedures via `include_lineage` and for views via `include_view_column_lineage`",
+    subtype_modifier=[
+        SourceCapabilityModifier.STORED_PROCEDURE,
+        SourceCapabilityModifier.VIEW,
+    ],
 )
 class SQLServerSource(SQLAlchemySource):
     """

datahub/ingestion/source/sql/sql_common.py

@@ -54,6 +54,7 @@ from datahub.ingestion.source.common.data_reader import DataReader
 from datahub.ingestion.source.common.subtypes import (
     DatasetContainerSubTypes,
     DatasetSubTypes,
+    SourceCapabilityModifier,
 )
 from datahub.ingestion.source.sql.sql_config import SQLCommonConfig
 from datahub.ingestion.source.sql.sql_report import SQLSourceReport
@@ -305,10 +306,12 @@ class ProfileMetadata:
 @capability(
     SourceCapability.LINEAGE_COARSE,
     "Enabled by default to get lineage for views via `include_view_lineage`",
+    subtype_modifier=[SourceCapabilityModifier.VIEW],
 )
 @capability(
     SourceCapability.LINEAGE_FINE,
     "Enabled by default to get lineage for views via `include_view_column_lineage`",
+    subtype_modifier=[SourceCapabilityModifier.VIEW],
 )
 @capability(SourceCapability.TEST_CONNECTION, "Enabled by default")
 @capability(

datahub/ingestion/source/sql/teradata.py

@@ -445,7 +445,10 @@ class TeradataConfig(BaseTeradataConfig, BaseTimeWindowConfig):
 @capability(SourceCapability.DOMAINS, "Enabled by default")
 @capability(SourceCapability.CONTAINERS, "Enabled by default")
 @capability(SourceCapability.PLATFORM_INSTANCE, "Enabled by default")
-@capability(SourceCapability.DELETION_DETECTION, "Optionally enabled via configuration")
+@capability(
+    SourceCapability.DELETION_DETECTION,
+    "Enabled by default when stateful ingestion is turned on",
+)
 @capability(SourceCapability.DATA_PROFILING, "Optionally enabled via configuration")
 @capability(SourceCapability.LINEAGE_COARSE, "Optionally enabled via configuration")
 @capability(SourceCapability.LINEAGE_FINE, "Optionally enabled via configuration")

datahub/ingestion/source/sql/vertica.py

@@ -25,7 +25,10 @@ from datahub.ingestion.api.decorators import (
 )
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.common.data_reader import DataReader
-from datahub.ingestion.source.common.subtypes import DatasetSubTypes
+from datahub.ingestion.source.common.subtypes import (
+    DatasetSubTypes,
+    SourceCapabilityModifier,
+)
 from datahub.ingestion.source.sql.sql_common import (
     SQLAlchemySource,
     SqlWorkUnit,
@@ -114,6 +117,10 @@ class VerticaConfig(BasicSQLAlchemyConfig):
 @capability(
     SourceCapability.LINEAGE_COARSE,
     "Enabled by default, can be disabled via configuration `include_view_lineage` and `include_projection_lineage`",
+    subtype_modifier=[
+        SourceCapabilityModifier.VIEW,
+        SourceCapabilityModifier.PROJECTIONS,
+    ],
 )
 @capability(
     SourceCapability.DELETION_DETECTION,