acryl-datahub 1.0.0.4rc6__py3-none-any.whl → 1.0.0.4rc7__py3-none-any.whl

{acryl_datahub-1.0.0.4rc6.dist-info → acryl_datahub-1.0.0.4rc7.dist-info}/RECORD

@@ -1,7 +1,7 @@
-acryl_datahub-1.0.0.4rc6.dist-info/licenses/LICENSE,sha256=9xNHpsD0uYF5ONzXsKDCuHHB-xbiCrSbueWXqrTNsxk,11365
+acryl_datahub-1.0.0.4rc7.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=lFUyplnN9ea9C0jyvH1T6H0kVRxk5DIQVEdHH_34ut4,323
+datahub/_version.py,sha256=8uGokiOv-wIz3FMuxJJSoDgRcEr4zbywxDQ8l-_DLVs,323
 datahub/entrypoints.py,sha256=AQN5MzCe6q3LKI4SS6WmwN56kgjF6AC1ld7yELWVP2w,8953
 datahub/errors.py,sha256=p5rFAdAGVCk4Lqolol1YvthceadUSwpaCxLXRcyCCFQ,676
 datahub/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -266,8 +266,8 @@ datahub/ingestion/source/bigquery_v2/queries_extractor.py,sha256=_5cAXVU8b8T_nAP
 datahub/ingestion/source/bigquery_v2/usage.py,sha256=A9c-ofclaRk0NSnc4IRaqJYqMPv6ecCld_TPy3V2qFs,40748
 datahub/ingestion/source/cassandra/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/ingestion/source/cassandra/cassandra.py,sha256=lKvPP0Uahi9xw_yh9cArPPtwvAauXolaEk-6f-jhpz4,14558
-datahub/ingestion/source/cassandra/cassandra_api.py,sha256=UVGQTsk6O57Q6wrWo54bQPLtStTWhw_Fq6fgW3Bjgk8,12515
-datahub/ingestion/source/cassandra/cassandra_config.py,sha256=vIMUOzazWTGi03B51vI0-YMxaMJHUGmCxJJgd8pKhC8,3791
+datahub/ingestion/source/cassandra/cassandra_api.py,sha256=b7MApc3_tEfHoj-6ub6snkcv_DweL1wi_TGJjAA1-yU,13516
+datahub/ingestion/source/cassandra/cassandra_config.py,sha256=Ga9915cDZukR5-u2tMNx5Jkf8eza2oAE5YS_sQIVEVQ,4222
 datahub/ingestion/source/cassandra/cassandra_profiling.py,sha256=ZqsAY8NFsrrLqOduV7Aem2eJLtc2_OU9tW4tc_dh0V8,10984
 datahub/ingestion/source/cassandra/cassandra_utils.py,sha256=j-LidYkaCTmGnpUVNLsax_c3z32PsQbsbHeYojygd1s,5105
 datahub/ingestion/source/common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -482,7 +482,7 @@ datahub/ingestion/source/sql/mariadb.py,sha256=Hm102kmfs_1rd4lsTYhzVMZq5S3B6cyfv
 datahub/ingestion/source/sql/mysql.py,sha256=nDWK4YbqomcJgnit9b8geUGrp_3eix4bt0_k94o7g-0,3350
 datahub/ingestion/source/sql/oracle.py,sha256=ftnrk3iiEelyv9PBHPYbairuP1WgxZbi1gu6YdqY69E,29965
 datahub/ingestion/source/sql/postgres.py,sha256=uC1kYEI8VdxiZ1Y9IxMWzwmg11wtMqYN0e2fkok1rxo,11972
-datahub/ingestion/source/sql/presto.py,sha256=PB-CS5MX2dSRFRHjlxfkLHGXLZXFNCsVAAyRBtY6HMg,3611
+datahub/ingestion/source/sql/presto.py,sha256=tATa0M2q0PjUC_E9W_jSUsmKTP7cVJayLgrFMzG_eao,4223
 datahub/ingestion/source/sql/sql_common.py,sha256=jsweel_-vesNtcPonnfS11OUrlcZnS3wGt5r0dYTPnM,48637
 datahub/ingestion/source/sql/sql_config.py,sha256=u3nGZYYl1WtaxfNsDU5bglgZ5Jq3Fxk9xei_CUIAXB0,8222
 datahub/ingestion/source/sql/sql_generic.py,sha256=9AERvkK8kdJUeDOzCYJDb93xdv6Z4DGho0NfeHj5Uyg,2740
@@ -494,7 +494,7 @@ datahub/ingestion/source/sql/sqlalchemy_data_reader.py,sha256=FvHZ4JEK3aR2DYOBZi
 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/trino.py,sha256=gSLDyETKavSVR8l9wdebrfoc41cqAWz6ApqIicW0BF8,17892
+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=_9OgSgIgqBml0av063rb8nACiT3SAmzpw0ouyF91wv8,33382
 datahub/ingestion/source/sql/mssql/__init__.py,sha256=1agpl8S_uDW40olkhCX_W19dbr5GO9qgjS3R7pLRZSk,87
@@ -1054,8 +1054,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.0.0.4rc6.dist-info/METADATA,sha256=Ry6qydIyyjAptNjSQr46P_yzOTtYyW2WaIzk8SQI6Rw,179949
-acryl_datahub-1.0.0.4rc6.dist-info/WHEEL,sha256=DnLRTWE75wApRYVsjgc6wsVswC54sMSJhAEd4xhDpBk,91
-acryl_datahub-1.0.0.4rc6.dist-info/entry_points.txt,sha256=o3mDeJXSKhsy7XLkuogihraiabBdLn9HaizYXPrxmk0,9710
-acryl_datahub-1.0.0.4rc6.dist-info/top_level.txt,sha256=iLjSrLK5ox1YVYcglRUkcvfZPvKlobBWx7CTUXx8_GI,25
-acryl_datahub-1.0.0.4rc6.dist-info/RECORD,,
+acryl_datahub-1.0.0.4rc7.dist-info/METADATA,sha256=GIO2L9JPpY18rRVypRWukqtZ0AEjbEN6whMxK5MChZY,179949
+acryl_datahub-1.0.0.4rc7.dist-info/WHEEL,sha256=DnLRTWE75wApRYVsjgc6wsVswC54sMSJhAEd4xhDpBk,91
+acryl_datahub-1.0.0.4rc7.dist-info/entry_points.txt,sha256=o3mDeJXSKhsy7XLkuogihraiabBdLn9HaizYXPrxmk0,9710
+acryl_datahub-1.0.0.4rc7.dist-info/top_level.txt,sha256=iLjSrLK5ox1YVYcglRUkcvfZPvKlobBWx7CTUXx8_GI,25
+acryl_datahub-1.0.0.4rc7.dist-info/RECORD,,

datahub/_version.py

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

datahub/ingestion/source/cassandra/cassandra_api.py

@@ -1,3 +1,4 @@
+import ssl
 from dataclasses import dataclass, field
 from typing import Any, Dict, List, Optional
 
@@ -128,6 +129,23 @@ class CassandraAPI:
 
                 self._cassandra_session = cluster.connect()
                 return True
+
+            ssl_context = None
+            if self.config.ssl_ca_certs:
+                ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
+                ssl_context.load_verify_locations(self.config.ssl_ca_certs)
+                if self.config.ssl_certfile and self.config.ssl_keyfile:
+                    ssl_context.load_cert_chain(
+                        certfile=self.config.ssl_certfile,
+                        keyfile=self.config.ssl_keyfile,
+                    )
+                elif self.config.ssl_certfile or self.config.ssl_keyfile:
+                    # If one is provided, the other must be too.
+                    # This is a simplification; real-world scenarios might allow one without the other depending on setup.
+                    raise ValueError(
+                        "Both ssl_certfile and ssl_keyfile must be provided if one is specified."
+                    )
+
             if self.config.username and self.config.password:
                 auth_provider = PlainTextAuthProvider(
                     username=self.config.username, password=self.config.password
@@ -136,12 +154,14 @@ class CassandraAPI:
                     [self.config.contact_point],
                     port=self.config.port,
                     auth_provider=auth_provider,
+                    ssl_context=ssl_context,
                     load_balancing_policy=None,
                 )
             else:
                 cluster = Cluster(
                     [self.config.contact_point],
                     port=self.config.port,
+                    ssl_context=ssl_context,
                     load_balancing_policy=None,
                 )
 

datahub/ingestion/source/cassandra/cassandra_config.py

@@ -79,6 +79,21 @@ class CassandraSourceConfig(
         description="Configuration for cloud-based Cassandra, such as DataStax Astra DB.",
     )
 
+    ssl_ca_certs: Optional[str] = Field(
+        default=None,
+        description="Path to the CA certificate file for SSL connections.",
+    )
+
+    ssl_certfile: Optional[str] = Field(
+        default=None,
+        description="Path to the SSL certificate file for SSL connections.",
+    )
+
+    ssl_keyfile: Optional[str] = Field(
+        default=None,
+        description="Path to the SSL key file for SSL connections.",
+    )
+
     keyspace_pattern: AllowDenyPattern = Field(
         default=AllowDenyPattern.allow_all(),
         description="Regex patterns to filter keyspaces for ingestion.",

datahub/ingestion/source/sql/presto.py

@@ -1,10 +1,12 @@
+import functools
 from textwrap import dedent
-from typing import Optional
+from typing import Dict, Optional
 
 from pydantic.fields import Field
 from pyhive.sqlalchemy_presto import PrestoDialect
 from sqlalchemy import exc, sql
 from sqlalchemy.engine import reflection
+from sqlalchemy.engine.base import Engine
 
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.decorators import (
@@ -114,3 +116,18 @@ class PrestoSource(TrinoSource):
     def create(cls, config_dict, ctx):
         config = PrestoConfig.parse_obj(config_dict)
         return cls(config, ctx)
+
+
+# Unfortunately, the Presto dialect provide catalog_name as a column
+# therefore we we need some workaround to not fail.
+# This is a workaround to not fail which casuses to only get the table comment as property which is still better than to fail.
+@functools.lru_cache
+def gen_catalog_connector_dict(engine: Engine) -> Dict[str, str]:
+    query = dedent(
+        """
+        SELECT *
+        FROM "system"."metadata"."catalogs"
+        """
+    ).strip()
+    res = engine.execute(sql.text(query))
+    return {row.catalog_name: "" for row in res}

datahub/ingestion/source/sql/trino.py

@@ -134,19 +134,41 @@ def get_table_comment(self, connection, table_name: str, schema: str = None, **k
         ):
             properties_table = self._get_full_table(f"{table_name}$properties", schema)
             query = f"SELECT * FROM {properties_table}"
-            row = connection.execute(sql.text(query)).fetchone()
+            rows = connection.execute(sql.text(query)).fetchall()
 
             # Generate properties dictionary.
             properties = {}
-            if row:
+
+            if len(rows) == 0:
+                # No properties found, return empty dictionary
+                return {}
+
+            # Check if using the old format (key, value columns)
+            if (
+                connector_name == "iceberg"
+                and len(rows[0]) == 2
+                and "key" in rows[0]
+                and "value" in rows[0]
+            ):
+                #  https://trino.io/docs/current/connector/iceberg.html#properties-table
+                for row in rows:
+                    if row["value"] is not None:
+                        properties[row["key"]] = row["value"]
+                return {"text": properties.get("comment"), "properties": properties}
+            elif connector_name == "hive" and len(rows[0]) > 1 and len(rows) == 1:
+                # https://trino.io/docs/current/connector/hive.html#properties-table
+                row = rows[0]
                 for col_name, col_value in row.items():
                     if col_value is not None:
                         properties[col_name] = col_value
+                return {"text": properties.get("comment"), "properties": properties}
 
-            return {"text": properties.get("comment"), "properties": properties}
-        else:
-            return self.get_table_comment_default(connection, table_name, schema)
-    except Exception:
+            # If we can't get the properties we still fallback to the default
+        return self.get_table_comment_default(connection, table_name, schema)
+    except Exception as e:
+        logging.warning(
+            f"Failed to get table comment for {table_name} in {schema}: {e}"
+        )
         return {}