Dynamic PlanetScale detection

tap_mysql/client.py

@@ -43,6 +43,21 @@ def patched_conform(
 class MySQLConnector(SQLConnector):
     """Connects to the MySQL SQL source."""
 
+    def __init__(
+        self,
+        is_vitess: bool = False,
+        *args,
+        **kwargs
+    ) -> None:
+        """Initialize the SQL connector.
+
+        Args:
+            is_vitess: Is this a vitess instance?
+            sqlalchemy_url: Optional URL for the connection.
+        """
+        self.is_vitess = is_vitess
+        super().__init__(*args, **kwargs)
+
     def create_engine(self) -> Engine:
         """Creates and returns a new engine. Do not call outside of _engine.
 
@@ -218,6 +233,9 @@ def discover_catalog_entry(
         Returns:
             `CatalogEntry` object for the given table or a view
         """
+        if self.is_vitess is False and is_view is False:
+            return super().discover_catalog_entry(engine, inspected, schema_name, table_name, is_view)
+        # For vitess views, we can't use DESCRIBE as it's not supported for views so we do the below
         # Initialize unique stream name
         unique_stream_id = self.get_fully_qualified_name(
             db_name=None,
@@ -337,6 +355,9 @@ def get_table_columns(
         Returns:
             An ordered list of column objects.
         """
+        if self.is_vitess:
+            return self.get_table_columns_vitess(full_table_name, column_names)
+        # If Vitess Instance then we can't use DESCRIBE as it's not supported for views so we do below
         if full_table_name not in self._table_cols_cache:
             _, schema_name, table_name = self.parse_full_table_name(full_table_name)
             with self._connect() as conn:
@@ -405,6 +426,7 @@ def get_records(self, context: dict | None) -> Iterable[dict[str, Any]]:
                 query = query.filter(replication_key_col >= start_val)
 
         with self.connector._connect() as conn:
-            conn.exec_driver_sql("set workload=olap")
+            if self.connector.is_vitess:
+                conn.exec_driver_sql("set workload=olap")
             for row in conn.execute(query):
                 yield dict(row)

tap_mysql/tap.py

@@ -301,6 +301,13 @@ def discover_streams(self) -> list[Stream]:
         Returns:
             List of discovered Stream objects.
         """
+        with self.connector._connect() as conn:
+            # Check if we are using PlanetScale, if so we need to let our connector know
+            # Ideally we'd just check to see if we're on Vitess, but I don't know how to do that quickly
+            output = conn.execute("select variable_value from performance_schema.global_variables where variable_name='version_comment' and variable_value like 'PlanetScale%'")
+            if len(output.fetchall()) > 0:
+                self.logger.info("Instance has been detected to be a Vitess instance. Using Vitess configuration.")
+                self.connector.is_vitess = True
         return [
             MySQLStream(self, catalog_entry, connector=self.connector)
             for catalog_entry in self.catalog_dict["streams"]