acryl-datahub 0.14.1.13rc2__py3-none-any.whl → 0.14.1.13rc4__py3-none-any.whl

{acryl_datahub-0.14.1.13rc2.dist-info → acryl_datahub-0.14.1.13rc4.dist-info}/RECORD

@@ -1,4 +1,4 @@
-datahub/__init__.py,sha256=e3nkmi7bdUTg9Q8NH4-2X1R6P-S6P3c6Kcfa42TRdEs,577
+datahub/__init__.py,sha256=WMIoTkprsOt2SjEA2gbvdoJ3VkTIeWCVPilQAqNdCQw,577
 datahub/__main__.py,sha256=pegIvQ9hzK7IhqVeUi1MeADSZ2QlP-D3K0OQdEg55RU,106
 datahub/entrypoints.py,sha256=3-qSfXAx3Z0FEkBV5tlO8fQr4xk4ySeDRMVTpS5Xd6A,7793
 datahub/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -96,7 +96,8 @@ datahub/configuration/datetimes.py,sha256=nayNc0mmlVKH6oVv9ud6C1dDUiZPGabW-YZxvr
 datahub/configuration/git.py,sha256=s55eUHxKqVZgtVsISaDyS-1F4iZBiybbjYsjbp5LU5o,6135
 datahub/configuration/import_resolver.py,sha256=b4Ie9L7knN1LALEVMxTcNFSklDD6CVE-4Ipy4ZYhNYA,369
 datahub/configuration/json_loader.py,sha256=vIDnjwXWi9yHDO8KW64EupOzOb_sspehGCD7xGHzg84,302
-datahub/configuration/kafka.py,sha256=GFHXnTyDAEHhpxnVUOepQ7aKYZOd9mTQyx1-MM5m88A,2126
+datahub/configuration/kafka.py,sha256=MlIwpd5FFyOyjdDXW_X9JTLNk7f988sPMgevkcZYVgI,2579
+datahub/configuration/kafka_consumer_config.py,sha256=DSwUU4HoqNyK4CNk9eIbX3eYsJMGQvORDiy1ZxkjlRc,1022
 datahub/configuration/pattern_utils.py,sha256=Q5IB9RfWOOo5FvRVBU7XkhiwHCxSQ1NTMfUlWtWI9qc,699
 datahub/configuration/pydantic_migration_helpers.py,sha256=4C_COAVZ5iJ8yxcWNgXZNWsY7ULogICNZ368oNF7zWg,1462
 datahub/configuration/source_common.py,sha256=68LZOuB23zSEcfgQJE1wZQnyYQHVVnEZK3Sniv_nEQs,2107
@@ -182,7 +183,7 @@ datahub/ingestion/sink/datahub_rest.py,sha256=17MjPW4F6LwLjZyATIWdDdIYvC74VH6g1j
 datahub/ingestion/sink/file.py,sha256=SxXJPJpkIGoaqRjCcSmj2ZE3xE4rLlBABBGwpTj5LWI,3271
 datahub/ingestion/sink/sink_registry.py,sha256=JRBWx8qEYg0ubSTyhqwgSWctgxwyp6fva9GoN2LwBao,490
 datahub/ingestion/source/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datahub/ingestion/source/confluent_schema_registry.py,sha256=K2LrEKCukw-kqa2AHsRE7MQYS4A416RY5YtpJt4esvg,18783
+datahub/ingestion/source/confluent_schema_registry.py,sha256=_h9D8bUXoaGcwgwB94dX6aTyLY5ve7XGdcVFSJHGSJc,18804
 datahub/ingestion/source/csv_enricher.py,sha256=xjCbcsSMM8l_ASCRAnNsUGKuYMrD1lec19Waixub1EM,29498
 datahub/ingestion/source/demo_data.py,sha256=yzA_R-wfSX2WPz0i5ukYlscpmpb0Pt8D7EkhtKfftvo,1286
 datahub/ingestion/source/elastic_search.py,sha256=qFUVNzynTVJTabASTjGMu8Qhf9UpNbEtSBFjaPQjBJE,22641
@@ -191,9 +192,6 @@ datahub/ingestion/source/file.py,sha256=pH-Qkjh5FQ2XvyYPE7Z8XEY4vUk_SUHxm8p8IxG1
 datahub/ingestion/source/ge_data_profiler.py,sha256=7oUvvADX4t1WiXwOruCURh3sNEY_7I41wbwXFvKaKWM,63587
 datahub/ingestion/source/ge_profiling_config.py,sha256=WusEMGFPu17y99jTT-1tTOiN87Q1sY7nyxXvXnPs1-E,10489
 datahub/ingestion/source/glue_profiling_config.py,sha256=vpMJH4Lf_qgR32BZy58suabri1yV5geaAPjzg2eORDc,2559
-datahub/ingestion/source/kafka.py,sha256=lGTkbwPvi8ysa0HaPaoWVHYblsPLB1js92cHbj2keKY,24876
-datahub/ingestion/source/kafka_connect.py,sha256=5KUlhn3876c41Z3kx5l4oJhbu0ekXZQRdxmu52vb_v8,55167
-datahub/ingestion/source/kafka_schema_registry_base.py,sha256=13XjSwqyVhH1CJUFHAbWdmmv_Rw0Ju_9HQdBmIzPNNA,566
 datahub/ingestion/source/ldap.py,sha256=Vnzg8tpwBYeyM-KBVVsUJvGZGBMJiCJ_i_FhxaFRQ9A,18627
 datahub/ingestion/source/metabase.py,sha256=oemiMdzjfr82Hx6rdwTNBzFM8962LDkosYh7SD_I5cY,31717
 datahub/ingestion/source/mlflow.py,sha256=SxCt4jtxQcpPWEI2rRNagCiE_6TWr2RroqmxRd_td1Y,11565
@@ -320,6 +318,10 @@ datahub/ingestion/source/iceberg/iceberg_profiler.py,sha256=hLT1Le_TEUoFXvsJSlrR
 datahub/ingestion/source/identity/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/ingestion/source/identity/azure_ad.py,sha256=GdmJFD4UMsb5353Z7phXRf-YsXR2woGLRJwBXUkgXq0,28809
 datahub/ingestion/source/identity/okta.py,sha256=PnRokWLG8wSoNZlXJiRZiW6APTEHO09q4n2j_l6m3V0,30756
+datahub/ingestion/source/kafka/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+datahub/ingestion/source/kafka/kafka.py,sha256=wC5tZ31fvXWLjse9qH3YONPhMxXIZ9QV9vbBcWDAP_w,25352
+datahub/ingestion/source/kafka/kafka_connect.py,sha256=5KUlhn3876c41Z3kx5l4oJhbu0ekXZQRdxmu52vb_v8,55167
+datahub/ingestion/source/kafka/kafka_schema_registry_base.py,sha256=13XjSwqyVhH1CJUFHAbWdmmv_Rw0Ju_9HQdBmIzPNNA,566
 datahub/ingestion/source/looker/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/ingestion/source/looker/lkml_patched.py,sha256=XShEU7Wbz0DubDhYMjKf9wjKZrBJa2XPg9MIjp8rPhk,733
 datahub/ingestion/source/looker/looker_common.py,sha256=cIrsc0nDEsODxx9tjvSm2A3-OlCkPKCu4W2L8PxzLWs,59926
@@ -447,7 +449,7 @@ datahub/ingestion/source/sql/mysql.py,sha256=nDWK4YbqomcJgnit9b8geUGrp_3eix4bt0_
 datahub/ingestion/source/sql/oracle.py,sha256=exJJvMSC42Oj0IiFtpAPFatWkRu9mVu8sd7ofOdanqU,22460
 datahub/ingestion/source/sql/postgres.py,sha256=uC1kYEI8VdxiZ1Y9IxMWzwmg11wtMqYN0e2fkok1rxo,11972
 datahub/ingestion/source/sql/presto.py,sha256=PB-CS5MX2dSRFRHjlxfkLHGXLZXFNCsVAAyRBtY6HMg,3611
-datahub/ingestion/source/sql/sql_common.py,sha256=Fv9ccVE-4qp9R4DdhQdy3er9u4ojHro73jogguSrpNU,51794
+datahub/ingestion/source/sql/sql_common.py,sha256=SZkPK7aUiTdMHf5alG5724i1pyAn1jFUjcuDkKUoeL0,51813
 datahub/ingestion/source/sql/sql_config.py,sha256=M-l_uXau0ODolLZHBzAXhy-Rq5yYxvJ6cLbCIea7Mww,9449
 datahub/ingestion/source/sql/sql_generic.py,sha256=9AERvkK8kdJUeDOzCYJDb93xdv6Z4DGho0NfeHj5Uyg,2740
 datahub/ingestion/source/sql/sql_generic_profiler.py,sha256=vWna43rv9ClfK9MBNTx8tdoUr35nD5d8ppgeamEEhSQ,12528
@@ -859,10 +861,10 @@ datahub/sql_parsing/_sqlglot_patch.py,sha256=iYJ8zOThHqqbamD5jdNr9iHTWD7ewNeHzPi
 datahub/sql_parsing/query_types.py,sha256=AxoWZiLuXM6iZ2AXdSlD82gqPNr1ZFh78wn2MSTu2wY,2479
 datahub/sql_parsing/schema_resolver.py,sha256=9wbJT80K4nsIHHOuQLso9QoFQAYwfSJZnqHwsaU3UTY,10197
 datahub/sql_parsing/sql_parsing_aggregator.py,sha256=YzYmWNQjXuQYcfh91ZnqjuzCDWdmZpE-ZX14Kz6XFMs,69938
-datahub/sql_parsing/sql_parsing_common.py,sha256=ptTwHXdAHuw0jbhIalEkB3mncDx5LsQeclRPecPzPFE,2339
+datahub/sql_parsing/sql_parsing_common.py,sha256=h_V_m54hJ9EUh5kczq7cYOIeNeo4bgf0Px0H-Nq-UIg,2602
 datahub/sql_parsing/sql_parsing_result_utils.py,sha256=prwWTj1EB2fRPv1eMB4EkpFNafIYAt-X8TIK0NWqank,796
-datahub/sql_parsing/sqlglot_lineage.py,sha256=jVMYZ5S36rHALiGmefQSGliRVcafLrodmfBUrTu_-9A,44787
-datahub/sql_parsing/sqlglot_utils.py,sha256=nVFG0UnHibApwNyJMEjdnJxgHreR9p9lN0Z3X6vhbnU,14448
+datahub/sql_parsing/sqlglot_lineage.py,sha256=K532_zvj7vEt2Ci2d3kxYZLZkq4S0ECBRZ2nHuY11a4,45927
+datahub/sql_parsing/sqlglot_utils.py,sha256=8MYzkyekhup3ihVStRPuwneWPNu17xhBg5SG8iVfFRY,14431
 datahub/sql_parsing/tool_meta_extractor.py,sha256=pE-pkRKBfNTXEJkaQM9NlG807mc-X6OtetgskJySCs8,2908
 datahub/telemetry/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/telemetry/stats.py,sha256=YltbtC3fe6rl1kcxn1A-mSnVpECTPm5k-brrUt7QxTI,967
@@ -888,7 +890,7 @@ datahub/utilities/dedup_list.py,sha256=dUSpe1AajfuwlHVJKNv-CzDXSCkaw0HgSMOsxqUkQ
 datahub/utilities/delayed_iter.py,sha256=XlsI0DCXkVVejFKOW_uMT0E8DTqqOHQN3Ooak4EcULE,645
 datahub/utilities/delta.py,sha256=hkpF8W7Lvg2gUJBQR3mmIzOxsRQ6i5cchRPFlAVoV10,1128
 datahub/utilities/docs_build.py,sha256=uFMK3z1d4BExpsrvguHunidbEDAzQ8hoOP7iQ0A_IVw,211
-datahub/utilities/file_backed_collections.py,sha256=5Og8AJV5WxbY7N8QOQWcRh8aNa7C3zrFKyAMaHddRmE,20232
+datahub/utilities/file_backed_collections.py,sha256=tS6hN0kxMaaTb8rB-vDqAd973mTKshERSC55JIe_3Cw,20557
 datahub/utilities/global_warning_util.py,sha256=adrEl3WhetQ-bymrPINjd976ZFndhbvk3QosUYGsos8,261
 datahub/utilities/hive_schema_to_avro.py,sha256=2-7NI9haCAYbyUmHTb-QPxjn4WbmnDDKAIGmHJ11-1E,11622
 datahub/utilities/is_pytest.py,sha256=2m9T4S9IIKhI5RfTqrB2ZmumzHocdxBHpM1HroWj2XQ,138
@@ -965,8 +967,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-0.14.1.13rc2.dist-info/METADATA,sha256=ocLjEMF3mpEq7e1fUg7cTaA_SkEUYNfDKL7DSKp3TEY,171138
-acryl_datahub-0.14.1.13rc2.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-acryl_datahub-0.14.1.13rc2.dist-info/entry_points.txt,sha256=tOZmkU-DL9O6-BRhkwS095b_y2RlSKQ7j0hvp57iw1Q,9420
-acryl_datahub-0.14.1.13rc2.dist-info/top_level.txt,sha256=iLjSrLK5ox1YVYcglRUkcvfZPvKlobBWx7CTUXx8_GI,25
-acryl_datahub-0.14.1.13rc2.dist-info/RECORD,,
+acryl_datahub-0.14.1.13rc4.dist-info/METADATA,sha256=pKvm7N_prnqdEylMVyCtB6aReDQm2T-0LmvkhiwUdrU,171138
+acryl_datahub-0.14.1.13rc4.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+acryl_datahub-0.14.1.13rc4.dist-info/entry_points.txt,sha256=VcQx0dnqaYLyeY_L5OaX7bLmmE-Il7TAXkxCKvEn2bA,9432
+acryl_datahub-0.14.1.13rc4.dist-info/top_level.txt,sha256=iLjSrLK5ox1YVYcglRUkcvfZPvKlobBWx7CTUXx8_GI,25
+acryl_datahub-0.14.1.13rc4.dist-info/RECORD,,

{acryl_datahub-0.14.1.13rc2.dist-info → acryl_datahub-0.14.1.13rc4.dist-info}/entry_points.txt

@@ -56,8 +56,8 @@ hive = datahub.ingestion.source.sql.hive:HiveSource
 hive-metastore = datahub.ingestion.source.sql.hive_metastore:HiveMetastoreSource
 iceberg = datahub.ingestion.source.iceberg.iceberg:IcebergSource
 json-schema = datahub.ingestion.source.schema.json_schema:JsonSchemaSource
-kafka = datahub.ingestion.source.kafka:KafkaSource
-kafka-connect = datahub.ingestion.source.kafka_connect:KafkaConnectSource
+kafka = datahub.ingestion.source.kafka.kafka:KafkaSource
+kafka-connect = datahub.ingestion.source.kafka.kafka_connect:KafkaConnectSource
 ldap = datahub.ingestion.source.ldap:LDAPSource
 looker = datahub.ingestion.source.looker.looker_source:LookerDashboardSource
 lookml = datahub.ingestion.source.looker.lookml_source:LookMLSource

datahub/__init__.py

@@ -3,7 +3,7 @@ import warnings
 
 # Published at https://pypi.org/project/acryl-datahub/.
 __package_name__ = "acryl-datahub"
-__version__ = "0.14.1.13rc2"
+__version__ = "0.14.1.13rc4"
 
 
 def is_dev_mode() -> bool:

datahub/configuration/kafka.py

@@ -1,6 +1,7 @@
 from pydantic import Field, validator
 
-from datahub.configuration.common import ConfigModel
+from datahub.configuration.common import ConfigModel, ConfigurationError
+from datahub.configuration.kafka_consumer_config import CallableConsumerConfig
 from datahub.configuration.validate_host_port import validate_host_port
 
 
@@ -36,6 +37,16 @@ class KafkaConsumerConnectionConfig(_KafkaConnectionConfig):
         description="Extra consumer config serialized as JSON. These options will be passed into Kafka's DeserializingConsumer. See https://docs.confluent.io/platform/current/clients/confluent-kafka-python/html/index.html#deserializingconsumer and https://github.com/edenhill/librdkafka/blob/master/CONFIGURATION.md .",
     )
 
+    @validator("consumer_config")
+    @classmethod
+    def resolve_callback(cls, value: dict) -> dict:
+        if CallableConsumerConfig.is_callable_config(value):
+            try:
+                value = CallableConsumerConfig(value).callable_config()
+            except Exception as e:
+                raise ConfigurationError(e)
+        return value
+
 
 class KafkaProducerConnectionConfig(_KafkaConnectionConfig):
     """Configuration class for holding connectivity information for Kafka producers"""

datahub/configuration/kafka_consumer_config.py

@@ -0,0 +1,35 @@
+import logging
+from typing import Any, Dict, Optional
+
+from datahub.ingestion.api.registry import import_path
+
+logger = logging.getLogger(__name__)
+
+
+class CallableConsumerConfig:
+    CALLBACK_ATTRIBUTE: str = "oauth_cb"
+
+    def __init__(self, config: Dict[str, Any]):
+        self._config = config
+
+        self._resolve_oauth_callback()
+
+    def callable_config(self) -> Dict[str, Any]:
+        return self._config
+
+    @staticmethod
+    def is_callable_config(config: Dict[str, Any]) -> bool:
+        return CallableConsumerConfig.CALLBACK_ATTRIBUTE in config
+
+    def get_call_back_attribute(self) -> Optional[str]:
+        return self._config.get(CallableConsumerConfig.CALLBACK_ATTRIBUTE)
+
+    def _resolve_oauth_callback(self) -> None:
+        if not self.get_call_back_attribute():
+            return
+
+        call_back = self.get_call_back_attribute()
+
+        assert call_back  # to silent lint
+        # Set the callback
+        self._config[CallableConsumerConfig.CALLBACK_ATTRIBUTE] = import_path(call_back)

datahub/ingestion/source/confluent_schema_registry.py

@@ -16,8 +16,10 @@ from confluent_kafka.schema_registry.schema_registry_client import (
 from datahub.ingestion.extractor import protobuf_util, schema_util
 from datahub.ingestion.extractor.json_schema_util import JsonSchemaTranslator
 from datahub.ingestion.extractor.protobuf_util import ProtobufSchema
-from datahub.ingestion.source.kafka import KafkaSourceConfig, KafkaSourceReport
-from datahub.ingestion.source.kafka_schema_registry_base import KafkaSchemaRegistryBase
+from datahub.ingestion.source.kafka.kafka import KafkaSourceConfig, KafkaSourceReport
+from datahub.ingestion.source.kafka.kafka_schema_registry_base import (
+    KafkaSchemaRegistryBase,
+)
 from datahub.metadata.com.linkedin.pegasus2avro.schema import (
     KafkaSchema,
     SchemaField,

datahub/ingestion/source/{kafka.py → kafka/kafka.py}

@@ -18,6 +18,7 @@ from confluent_kafka.schema_registry.schema_registry_client import SchemaRegistr
 
 from datahub.configuration.common import AllowDenyPattern
 from datahub.configuration.kafka import KafkaConsumerConnectionConfig
+from datahub.configuration.kafka_consumer_config import CallableConsumerConfig
 from datahub.configuration.source_common import (
     DatasetSourceConfigMixin,
     LowerCaseDatasetUrnConfigMixin,
@@ -49,7 +50,9 @@ from datahub.ingestion.api.source import (
 )
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.common.subtypes import DatasetSubTypes
-from datahub.ingestion.source.kafka_schema_registry_base import KafkaSchemaRegistryBase
+from datahub.ingestion.source.kafka.kafka_schema_registry_base import (
+    KafkaSchemaRegistryBase,
+)
 from datahub.ingestion.source.state.stale_entity_removal_handler import (
     StaleEntityRemovalHandler,
     StaleEntityRemovalSourceReport,
@@ -143,7 +146,7 @@ class KafkaSourceConfig(
 def get_kafka_consumer(
     connection: KafkaConsumerConnectionConfig,
 ) -> confluent_kafka.Consumer:
-    return confluent_kafka.Consumer(
+    consumer = confluent_kafka.Consumer(
         {
             "group.id": "test",
             "bootstrap.servers": connection.bootstrap,
@@ -151,6 +154,13 @@ def get_kafka_consumer(
         }
     )
 
+    if CallableConsumerConfig.is_callable_config(connection.consumer_config):
+        # As per documentation, we need to explicitly call the poll method to make sure OAuth callback gets executed
+        # https://docs.confluent.io/platform/current/clients/confluent-kafka-python/html/index.html#kafka-client-configuration
+        consumer.poll(timeout=30)
+
+    return consumer
+
 
 @dataclass
 class KafkaSourceReport(StaleEntityRemovalSourceReport):

datahub/ingestion/source/sql/sql_common.py

@@ -32,6 +32,7 @@ from datahub.emitter.mce_builder import (
     make_data_platform_urn,
     make_dataplatform_instance_urn,
     make_dataset_urn_with_platform_instance,
+    make_schema_field_urn,
     make_tag_urn,
 )
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
@@ -669,7 +670,7 @@ class SQLAlchemySource(StatefulIngestionSourceBase, TestableSource):
         )
 
         source_fields = [
-            f"urn:li:schemaField:({dataset_urn},{f})"
+            make_schema_field_urn(dataset_urn, f)
             for f in fk_dict["constrained_columns"]
         ]
         foreign_dataset = make_dataset_urn_with_platform_instance(
@@ -679,7 +680,7 @@ class SQLAlchemySource(StatefulIngestionSourceBase, TestableSource):
             env=self.config.env,
         )
         foreign_fields = [
-            f"urn:li:schemaField:({foreign_dataset},{f})"
+            make_schema_field_urn(foreign_dataset, f)
             for f in fk_dict["referred_columns"]
         ]
 

datahub/sql_parsing/sql_parsing_common.py

@@ -21,6 +21,9 @@ DIALECTS_WITH_CASE_INSENSITIVE_COLS = {
     # See more below:
     # https://documentation.sas.com/doc/en/pgmsascdc/9.4_3.5/acreldb/n0ejgx4895bofnn14rlguktfx5r3.htm
     "teradata",
+    # For SQL server, the default collation rules mean that all identifiers (schema, table, column names)
+    # are case preserving but case insensitive.
+    "mssql",
 }
 DIALECTS_WITH_DEFAULT_UPPERCASE_COLS = {
     # In some dialects, column identifiers are effectively case insensitive
@@ -28,6 +31,9 @@ DIALECTS_WITH_DEFAULT_UPPERCASE_COLS = {
     # automatically lowercase unquoted identifiers.
     "snowflake",
 }
+assert DIALECTS_WITH_DEFAULT_UPPERCASE_COLS.issubset(
+    DIALECTS_WITH_CASE_INSENSITIVE_COLS
+)
 
 
 class QueryType(enum.Enum):

datahub/sql_parsing/sqlglot_lineage.py

@@ -5,7 +5,7 @@ import functools
 import logging
 import traceback
 from collections import defaultdict
-from typing import Any, Dict, List, Optional, Set, Tuple, Union
+from typing import Any, Dict, List, Optional, Set, Tuple, TypeVar, Union
 
 import pydantic.dataclasses
 import sqlglot
@@ -873,6 +873,49 @@ def _translate_internal_column_lineage(
     )
 
 
+_StrOrNone = TypeVar("_StrOrNone", str, Optional[str])
+
+
+def _normalize_db_or_schema(
+    db_or_schema: _StrOrNone,
+    dialect: sqlglot.Dialect,
+) -> _StrOrNone:
+    if db_or_schema is None:
+        return None
+
+    # In snowflake, table identifiers must be uppercased to match sqlglot's behavior.
+    if is_dialect_instance(dialect, "snowflake"):
+        return db_or_schema.upper()
+
+    # In mssql, table identifiers must be lowercased.
+    elif is_dialect_instance(dialect, "mssql"):
+        return db_or_schema.lower()
+
+    return db_or_schema
+
+
+def _simplify_select_into(statement: sqlglot.exp.Expression) -> sqlglot.exp.Expression:
+    """
+    Check if the expression is a SELECT INTO statement. If so, converts it into a CTAS.
+    Other expressions are returned as-is.
+    """
+
+    if not (isinstance(statement, sqlglot.exp.Select) and statement.args.get("into")):
+        return statement
+
+    # Convert from SELECT <cols> INTO <out> <expr>
+    # to CREATE TABLE <out> AS SELECT <cols> <expr>
+    into_expr: sqlglot.exp.Into = statement.args["into"].pop()
+    into_table = into_expr.this
+
+    create = sqlglot.exp.Create(
+        this=into_table,
+        kind="TABLE",
+        expression=statement,
+    )
+    return create
+
+
 def _sqlglot_lineage_inner(
     sql: sqlglot.exp.ExpOrStr,
     schema_resolver: SchemaResolverInterface,
@@ -885,12 +928,9 @@ def _sqlglot_lineage_inner(
     else:
         dialect = get_dialect(default_dialect)
 
-    if is_dialect_instance(dialect, "snowflake"):
-        # in snowflake, table identifiers must be uppercased to match sqlglot's behavior.
-        if default_db:
-            default_db = default_db.upper()
-        if default_schema:
-            default_schema = default_schema.upper()
+    default_db = _normalize_db_or_schema(default_db, dialect)
+    default_schema = _normalize_db_or_schema(default_schema, dialect)
+
     if is_dialect_instance(dialect, "redshift") and not default_schema:
         # On Redshift, there's no "USE SCHEMA <schema>" command. The default schema
         # is public, and "current schema" is the one at the front of the search path.
@@ -918,6 +958,8 @@ def _sqlglot_lineage_inner(
     #     original_statement.sql(pretty=True, dialect=dialect),
     # )
 
+    statement = _simplify_select_into(statement)
+
     # Make sure the tables are resolved with the default db / schema.
     # This only works for Unionable statements. For other types of statements,
     # we have to do it manually afterwards, but that's slightly lower accuracy

datahub/sql_parsing/sqlglot_utils.py

@@ -61,7 +61,7 @@ def is_dialect_instance(
     else:
         platforms = list(platforms)
 
-    dialects = [sqlglot.Dialect.get_or_raise(platform) for platform in platforms]
+    dialects = [get_dialect(platform) for platform in platforms]
 
     if any(isinstance(dialect, dialect_class.__class__) for dialect_class in dialects):
         return True

datahub/utilities/file_backed_collections.py

@@ -228,6 +228,12 @@ class FileBackedDict(MutableMapping[str, _VT], Closeable, Generic[_VT]):
         else:
             self._conn = ConnectionWrapper()
 
+        if sqlite3.sqlite_version_info < (3, 24, 0):
+            # We use the ON CONFLICT clause to implement UPSERTs with sqlite.
+            # This was added in 3.24.0 from 2018-06-04.
+            # See https://www.sqlite.org/lang_conflict.html
+            raise RuntimeError("SQLite version 3.24.0 or later is required")
+
         # We keep a small cache in memory to avoid having to serialize/deserialize
         # data from the database too often. We use an OrderedDict to build
         # a poor-man's LRU cache.