Encoding: fix b/268011976

Author: mjquinn-google

client/README.md

@@ -298,6 +298,7 @@ More advanced macro replacement schemes are possible. Please see the
 an example.
 
 ### Extraction of Heredoc SQL Statements from KSH Inputs
+
 By default, during the preprocessing phase, any input paths ending in
 extension `.ksh` will be scanned for heredoc SQL statements. For example,
 given the following file named `foo.ksh`:
@@ -312,6 +313,7 @@ echo Trying another select.
 ```
 
 will be preprocessed to:
+
 ```sql
 SELECT 123, 'foo', 456 from bar;
 ```

client/bqms_run/encoding.py

@@ -0,0 +1,43 @@
+# Copyright 2022 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Encoding utilities."""
+
+import logging
+
+import icu
+
+logger = logging.getLogger(__name__)
+
+
+class EncodingDetector:
+    """
+    An encoding detector.
+    """
+
+    def detect(self, data: bytes) -> str:
+        """
+        Detect the encoding of the provided bytes, return the encoding name.
+        """
+        encoding = icu.CharsetDetector(data).detect().getName()
+        if not isinstance(encoding, str):
+            return "utf-8"
+        return encoding
+
+    def decode(self, data: bytes) -> str:
+        """
+        Detect the encoding of the provided bytes, then decode them to string.
+        """
+        encoding = self.detect(data)
+        logger.debug("Detected encoding: %s", encoding)
+        return data.decode(encoding)

client/bqms_run/hooks.py

@@ -24,7 +24,6 @@
 )
 from bqms_run.paths import Path
 
-
 logger = logging.getLogger(__name__)
 
 

client/bqms_run/ksh.py

@@ -18,7 +18,6 @@
 from enum import Enum
 from typing import List, Optional
 
-
 logger = logging.getLogger(__name__)
 
 
@@ -34,10 +33,7 @@ class ShellFragment:
     """
 
     def __init__(
-        self, text: str,
-        line: int,
-        fragment_type: ShellFragmentType,
-        quoted: bool
+        self, text: str, line: int, fragment_type: ShellFragmentType, quoted: bool
     ) -> None:
         super().__init__()
         self.text = text
@@ -91,15 +87,19 @@ def __init__(
         :param command_replace_to: the suffix of the heredoc command
         """
         super().__init__()
-        self.pattern = self._check(re.compile(
-            KshExtractor.to_pattern_for_command(command)))
+        self.pattern = self._check(
+            re.compile(KshExtractor.to_pattern_for_command(command))
+        )
         self.command_replace_from = command_replace_from
         self.command_replace_to = command_replace_to
 
     @staticmethod
     def to_pattern_for_command(command: str) -> str:
-        return KshExtractor.PATTERN_DEFAULT_PREFIX + command \
-               + KshExtractor.PATTERN_DEFAULT_SUFFIX
+        return (
+            KshExtractor.PATTERN_DEFAULT_PREFIX
+            + command
+            + KshExtractor.PATTERN_DEFAULT_SUFFIX
+        )
 
     @staticmethod
     def _check(pattern: re.Pattern) -> re.Pattern:  # type: ignore[type-arg]
@@ -126,9 +126,10 @@ def filter_heredoc_sql_texts(fragments: List[ShellFragment]) -> List[str]:
         return list(
             map(
                 lambda f: f.get_text(),
-                filter(lambda f:
-                       f.get_fragment_type() == ShellFragmentType.HEREDOC,
-                       fragments),
+                filter(
+                    lambda f: f.get_fragment_type() == ShellFragmentType.HEREDOC,
+                    fragments,
+                ),
             )
         )
 

client/bqms_run/main.py

@@ -29,6 +29,7 @@
 from marshmallow import ValidationError
 
 from bqms_run import workflow
+from bqms_run.encoding import EncodingDetector
 from bqms_run.gcp.bqms.object_name_mapping import ObjectNameMappingListSchema
 from bqms_run.gcp.bqms.request import build as build_bqms_request
 from bqms_run.gcp.bqms.source_env import SourceEnvSchema
@@ -78,8 +79,8 @@ def _parse_paths() -> Paths:
 
 def _read_config(config_path: Path) -> Dict[str, object]:
     logger.debug("Parsing config: %s.", config_path.as_uri())
-    with config_path.open(mode="r", encoding="utf-8") as config_file:
-        config_text = config_file.read()
+    with config_path.open(mode="rb") as config_file:
+        config_text = EncodingDetector().decode(config_file.read())
     config: Dict[str, object] = yaml.load(config_text, Loader=yaml.SafeLoader)
     return config
 
@@ -147,10 +148,8 @@ def _parse_object_name_mapping(
     object_name_mapping_path: Path,
 ) -> ObjectNameMappingList:
     logger.debug("Parsing object name mapping: %s.", object_name_mapping_path.as_uri())
-    with object_name_mapping_path.open(
-        mode="r", encoding="utf-8"
-    ) as macro_mapping_file:
-        object_name_mapping_text = macro_mapping_file.read()
+    with object_name_mapping_path.open(mode="rb") as macro_mapping_file:
+        object_name_mapping_text = EncodingDetector().decode(macro_mapping_file.read())
     object_name_mapping = json.loads(object_name_mapping_text)
     try:
         object_name_mapping_list = ObjectNameMappingListSchema.from_mapping(
@@ -172,8 +171,8 @@ def _parse_object_name_mapping(
 
 def _parse_macro_mapping(macro_mapping_path: Path) -> MacroExpanderRouter:
     logger.debug("Parsing macro mapping: %s.", macro_mapping_path.as_uri())
-    with macro_mapping_path.open(mode="r", encoding="utf-8") as macro_mapping_file:
-        macro_mapping_text = macro_mapping_file.read()
+    with macro_mapping_path.open(mode="rb") as macro_mapping_file:
+        macro_mapping_text = EncodingDetector().decode(macro_mapping_file.read())
     macro_mapping = yaml.load(macro_mapping_text, Loader=yaml.SafeLoader)
     try:
         validated_macro_mapping = MacroMapping.from_mapping(macro_mapping)

client/bqms_run/workflow.py

@@ -20,6 +20,7 @@
 from google.cloud.bigquery_migration_v2 import CreateMigrationWorkflowRequest
 from typing_extensions import ParamSpec
 
+from bqms_run.encoding import EncodingDetector
 from bqms_run.gcp.bqms.request import execute as execute_bqms_request
 from bqms_run.macros import MacroExpanderRouter
 from bqms_run.paths import Path, Paths, iterdirfiles
@@ -124,8 +125,8 @@ def _preprocess(source_file_path: Path) -> None:
                 target_file.write(source_bytes)
             return
 
-        with source_file_path.open(mode="r", encoding="utf-8") as source_file:
-            source_text = source_file.read()
+        with source_file_path.open(mode="rb") as source_file:
+            source_text = EncodingDetector().decode(data=source_file.read())
 
         preprocessed_text = preprocess_hook(relative_file_path, source_text)
         macro_expanded_text = (
@@ -147,8 +148,8 @@ def _postprocess(source_file_path: Path) -> None:
         target_file_path = paths.postprocessed_path / relative_file_path
         target_file_path.parent.mkdir(parents=True, exist_ok=True)
 
-        with source_file_path.open(mode="r", encoding="utf-8") as source_file:
-            source_text = source_file.read()
+        with source_file_path.open(mode="rb") as source_file:
+            source_text = EncodingDetector().decode(data=source_file.read())
 
         macro_unexpanded_text = (
             macro_expander_router.un_expand(relative_file_path, source_text)