feat: support multiple streams from source

Author: activescott

airbyte-integrations/connectors/destination-glide/destination_glide/destination.py

@@ -29,7 +29,7 @@
 logger.setLevel(LOG_LEVEL_DEFAULT)
 
 
-def mapJsonSchemaTypeToGlideType(json_type: str) -> str:
+def airbyteTypeToGlideType(json_type: str) -> str:
     jsonSchemaTypeToGlideType = {
         "string": "string",
         "number": "number",
@@ -75,37 +75,39 @@ def write(
         api_host = config.get('api_host', CONFIG_GLIDE_API_HOST_DEFAULT)
         api_path_root = config['api_path_root']
         api_key = config['api_key']
-        table_id = config['table_id']
         glide_api_version = config.get(
             'glide_api_version', CONFIG_GLIDE_API_VERSION_DEFAULT)
 
+        # configure the table based on the stream catalog:
         # choose a strategy based on config:
-        glide = GlideBigTableFactory.create(glide_api_version)
-        logger.debug(f"Using glide api strategy '{glide.__class__.__name__}' for glide_api_version '{glide_api_version}'.")  # nopep8
-        glide.init(api_host, api_key, api_path_root, table_id)
 
-        if len(configured_catalog.streams) > 1:
-            # TODO: Consider how we might want to support multiple streams. Review other airbyte destinations
-            raise Exception(f"The Glide destination expects only a single stream to be streamed into a Glide Table but found '{len(configured_catalog.streams)}' streams. Please select only one stream from the source.")  # nopep8
+        def create_table_client_for_stream(stream_name):
+            # TODO: sanitize stream_name chars and length for GBT name
+            glide = GlideBigTableFactory.create(glide_api_version)
+            logger.debug(f"Using glide api strategy '{glide.__class__.__name__}' for glide_api_version '{glide_api_version}'.")  # nopep8
+            glide.init(api_host, api_key, api_path_root, stream_name)
+            return glide
 
-        # configure the table based on the stream catalog:
+        table_clients = {}
         stream_names = {s.stream.name for s in configured_catalog.streams}
         for configured_stream in configured_catalog.streams:
             if configured_stream.destination_sync_mode != DestinationSyncMode.overwrite:
                 raise Exception(f'Only destination sync mode overwrite is supported, but received "{configured_stream.destination_sync_mode}".')  # nopep8 because https://github.com/hhatto/autopep8/issues/712
 
+            glide = create_table_client_for_stream(
+                configured_stream.stream.name)
             # upsert the GBT with schema to set_schema for dumping the data into it
             columns = []
             properties = configured_stream.stream.json_schema["properties"]
-            for prop_name in properties.keys():
-                prop = properties[prop_name]
+            for (prop_name, prop) in properties.items():
                 prop_type = prop["type"]
-                prop_format = prop["format"] if "format" in prop else ""
-                logger.debug(f"Found column/property '{prop_name}' with type '{prop_type}' and format '{prop_format}' in stream {configured_stream.stream.name}.")  # nopep8
+                logger.debug(f"Found column/property '{prop_name}' with type '{prop_type}' in stream {configured_stream.stream.name}.")  # nopep8
                 columns.append(
-                    Column(prop_name, mapJsonSchemaTypeToGlideType(prop_type)))
+                    Column(prop_name, airbyteTypeToGlideType(prop_type))
+                )
 
             glide.set_schema(columns)
+            table_clients[configured_stream.stream.name] = glide
 
         # stream the records into the GBT:
         buffers = defaultdict(list)
@@ -120,7 +122,6 @@ def write(
                         f"Stream {stream_name} was not present in configured streams, skipping")
                     continue
 
-                # TODO: check the columns match the columns that we saw in configured_catalog per https://docs.airbyte.com/understanding-airbyte/airbyte-protocol#destination
                 # add to buffer
                 record_data = message.record.data
                 record_id = str(uuid.uuid4())
@@ -133,24 +134,32 @@ def write(
                 # `Type.State` is a signal from the source that we should save the previous batch of `Type.RECORD` messages to the destination.
                 #   It is a checkpoint that enables partial success.
                 #   See https://docs.airbyte.com/understanding-airbyte/airbyte-protocol#state--checkpointing
-                logger.info(f"Writing buffered records to Glide API from {len(buffers.keys())} streams...")  # nopep8 because
+                logger.info(f"Writing buffered records to Glide API from {len(buffers.keys())} streams...")  # nopep8
                 for stream_name in buffers.keys():
                     stream_buffer = buffers[stream_name]
-                    logger.info(f"Saving buffered records to Glide API (stream: '{stream_name}' count: '{len(stream_buffer)}')...")  # nopep8 because https://github.com/hhatto/autopep8/issues/712
+                    logger.info(f"Saving buffered records to Glide API (stream: '{stream_name}', record count: '{len(stream_buffer)}')...")  # nopep8
                     DATA_INDEX = 2
                     data_rows = [row_tuple[DATA_INDEX]
                                  for row_tuple in stream_buffer]
                     if len(data_rows) > 0:
+                        if stream_name not in table_clients:
+                            raise Exception(
+                                f"Stream '{stream_name}' not found in table_clients")
+                        glide = table_clients[stream_name]
                         glide.add_rows(data_rows)
                     stream_buffer.clear()
                     logger.info(f"Saving buffered records to Glide API complete.")  # nopep8 because https://github.com/hhatto/autopep8/issues/712
 
+                # dump all buffers now as we just wrote them to the table:
+                buffers = defaultdict(list)
                 yield message
             else:
                 logger.warn(f"Ignoring unknown Airbyte input message type: {message.type}")  # nopep8 because https://github.com/hhatto/autopep8/issues/712
 
         # commit the stash to the table
-        glide.commit()
+        for stream_name, glide in table_clients.items():
+            glide.commit()
+            logger.info(f"Committed stream '{stream_name}' to Glide.")
 
         pass
 

airbyte-integrations/connectors/destination-glide/destination_glide/glide.py

@@ -64,16 +64,15 @@ def url(self, path: str) -> str:
     The protocol is to call `init`, `set_schema`, `add_rows` one or more times, and `commit` in that order.
     """
 
-    def init(self, api_host, api_key, api_path_root, table_id):
+    def init(self, api_host, api_key, api_path_root, table_name):
         """
         Sets the connection information for the table.
         """
         self.api_host = api_host
         self.api_key = api_key
         self.api_path_root = api_path_root
-        self.table_id = table_id
+        self.table_name = table_name
         # todo: validate args
-        pass
 
     @abstractmethod
     def set_schema(self, columns: List[Column]) -> None:
@@ -108,8 +107,7 @@ def create(cls, strategy: str) -> GlideBigTableBase:
         Creates a new instance of the default implementation for the GlideBigTable API client.
         """
         implementation_map = {
-            "tables": GlideBigTableRestStrategy(),
-            "mutations": GlideBigTableMutationsStrategy()
+            "tables": GlideBigTableRestStrategy()
         }
         if strategy not in implementation_map:
             raise ValueError(f"Strategy '{strategy}' not found. Expected one of '{implmap.keys()}'.")  # nopep8
@@ -135,13 +133,13 @@ def set_schema(self, columns: List[Column]) -> None:
         self.columns = columns
         # Create stash we can stash records into for later
         r = requests.post(
-            self.url(f"/stashes"),
+            self.url(f"stashes"),
             headers=self.headers(),
         )
         try:
             r.raise_for_status()
         except Exception as e:
-            raise Exception(f"failed to create stash") from e  # nopep8
+            raise Exception(f"failed to create stash. Response was '{r.text}'") from e  # nopep8
 
         result = r.json()
         self.stash_id = result["data"]["stashID"]
@@ -157,7 +155,7 @@ def _add_row_batch(self, rows: List[BigTableRow]) -> None:
         # TODO: add rows to stash/serial https://web.postman.co/workspace/glideapps-Workspace~46b48d24-5fc1-44b6-89aa-8d6751db0fc5/request/9026518-c282ef52-4909-4806-88bf-08510ee80770
         logger.debug(f"Adding rows batch with size {len(rows)}")
         r = requests.post(
-            self.url(f"/stashes/{self.stash_id}/{self.stash_serial}"),
+            self.url(f"stashes/{self.stash_id}/{self.stash_serial}"),
             headers=self.headers(),
             json={
                 "data": rows,
@@ -170,9 +168,9 @@ def _add_row_batch(self, rows: List[BigTableRow]) -> None:
         try:
             r.raise_for_status()
         except Exception as e:
-            raise Exception(f"failed to add rows batch for serial '{self.stash_serial}'") from e  # nopep8
+            raise Exception(f"failed to add rows batch for serial '{self.stash_serial}'. Response was '{r.text}'") from e  # nopep8
 
-        logger.info(f"Add rows batch for serial '{self.stash_serial}' succeeded.")  # nopep8
+        logger.info(f"Added {len(rows)} rows as batch for serial '{self.stash_serial}' successfully.")  # nopep8
         self.stash_serial += 1
 
     def add_rows(self, rows: Iterator[BigTableRow]) -> None:
@@ -184,14 +182,15 @@ def add_rows(self, rows: Iterator[BigTableRow]) -> None:
             batch = rows[i:i + min(BATCH_SIZE, len(rows) - i)]
             self._add_row_batch(batch)
 
-    def finalize_stash(self) -> None:
-        # overwrite the existing table with the right schema and rows:
-        r = requests.put(
-            self.url(f"/tables/{self.table_id}"),
+    def create_table_from_stash(self) -> None:
+        logger.info(f"Creating new table for table name '{self.table_name}'...") # nopep8
+        r = requests.post(
+            self.url(f"tables"),
             headers=self.headers(),
             json={
+                "name": self.table_name,
                 "schema": {
-                    "columns": self.columns,
+                    "columns": self.columns
                 },
                 "rows": {
                     "$stashID": self.stash_id
@@ -201,161 +200,56 @@ def finalize_stash(self) -> None:
         try:
             r.raise_for_status()
         except Exception as e:
-            raise Exception(f"failed to finalize stash") from e  # nopep8
-        logger.info(f"Successfully finalized record stash for table '{self.table_id}' (stash ID '{self.stash_id}')")
-
-    def commit(self) -> None:
-        self.raise_if_set_schema_not_called()
-        self.finalize_stash()
+            raise Exception(f"failed to create table '{self.table_name}'. Response was '{r.text}'.") from e  # nopep8
 
+        logger.info(f"Creating table '{self.table_name}' succeeded.")
 
-class GlideBigTableMutationsStrategy(GlideBigTableBase):
-    def __init__(self):
-        # TODO: hardcoded for now using old api
-        self.hardcoded_app_id = "Ix9CEuP6DiFugfjhSG5t"
-        self.hardcoded_column_lookup = {
-            '_airtable_id': {'type': "string", 'name': "Name"},
-            '_airtable_created_time': {'type': "date-time", 'name': "AwkFL"},
-            '_airtable_table_name': {'type': "string", 'name': "QF0zI"},
-            'id': {'type': "string", 'name': "tLPjZ"},
-            'name': {'type': "string", 'name': "1ZqF1"},
-            'host_id': {'type': "string", 'name': "B7fYe"},
-            'host_name': {'type': "string", 'name': "oyVzO"},
-            'neighbourhood_group': {'type': "string", 'name': "15J8U"},
-            'neighbourhood': {'type': "string", 'name': "Fy28U"},
-            'latitude': {'type': "number", 'name': "TLpMC"},
-            'longitude': {'type': "number", 'name': "oazQO"},
-            'room_type': {'type': "string", 'name': "TPJDZ"},
-            'price': {'type': "number", 'name': "7xzlG"},
-            'minimum_nights': {'type': "number", 'name': "usoY5"},
-            'number_of_reviews': {'type': "number", 'name': "XFXmR"},
-            'last_review': {'type': "date-time", 'name': "oseZl"},
-            'reviews_per_month': {'type': "number", 'name': "alw2R"},
-            'calculated_host_listings_count': {'type': "number", 'name': "hKws0"},
-            'availability_365': {'type': "number", 'name': "qZsgl"},
-            'number_of_reviews_ltm': {'type': "number", 'name': "rWisS"},
-            'license': {'type': "string", 'name': "7PVig"}
-        }
-
-    def headers(self) -> Dict[str, str]:
-        return {
-            "Content-Type": "application/json",
-            f"Authorization": f"Bearer {self.api_key}"
-        }
-
-    def url(self, path: str) -> str:
-        return f"{self.api_host}/{self.api_path_root}/{path}"
-
-    def set_schema(self, columns: List[Column]) -> None:
-        logger.debug(f"set_schema for table '{self.table_id}. Expecting columns: '{[c.id for c in columns]}'.")  # nopep8
-        self.delete_all()
-
-        for col in columns:
-            if col.id not in self.hardcoded_column_lookup:
-                logger.warning(
-                    f"Column '{col.id}' not found in hardcoded column lookup. Will be ignored.")
-
-    def rows(self) -> Iterator[BigTableRow]:
-        """
-        Gets the rows as of the Glide Big Table.
-        """
-        r = requests.post(
-            self.url("function/queryTables"),
+    def overwrite_table_from_stash(self, table_id) -> None:
+        # overwrite the specified table's schema and rows with the stash:
+        r = requests.put(
+            self.url(f"tables/{table_id}"),
             headers=self.headers(),
             json={
-                "appID": self.hardcoded_app_id,
-                "queries": [
-                    {
-                        "tableName": self.table_id,
-                        "utc": True
-                    }
-                ]
-            }
-        )
-        if r.status_code != 200:
-            logger.error(f"get rows request failed with status {r.status_code}: {r.text}.")  # nopep8 because https://github.com/hhatto/autopep8/issues/712
-            r.raise_for_status()
-
-        result = r.json()
-
-        # the result looks like an array of results; each result has a rows member that has an array or JSON rows:
-        for row in result:
-            for r in row['rows']:
-                yield r
-
-    def delete_all(self) -> None:
-        # TODO: perf: don't put these in a list
-        rows = list(self.rows())
-        logger.debug(f"Iterating over {len(rows)} rows to delete")
-
-        for row in rows:
-            # TODO: lame. batch these into one request with multiple mutations
-            r = requests.post(
-                self.url("function/mutateTables"),
-                headers=self.headers(),
-                json={
-                    "appID": self.hardcoded_app_id,
-                    "mutations": [
-                        {
-                            "kind": "delete-row",
-                            "tableName": self.table_id,
-                            "rowID": row['$rowID']
-                        }
-                    ]
-                }
-            )
-            if r.status_code != 200:
-                logger.error(f"delete request failed with status {r.status_code}: {r.text} trying to delete row id {row['$rowID']} with row: {row}")  # nopep8 because https://github.com/hhatto/autopep8/issues/712
-                r.raise_for_status()
-            else:
-                logger.debug(
-                    f"Deleted row successfully (rowID:'{row['$rowID']}'")
-
-    def add_rows_batch(self, rows: Iterator[BigTableRow]) -> None:
-        mutations = []
-        for row in rows:
-            # row is columnLabel -> value, but glide's mutate uses a column "name". We hard-code the lookup for our table here:
-            mutated_row = dict()
-            for k, v in row.items():
-                if k in self.hardcoded_column_lookup:
-                    col_info = self.hardcoded_column_lookup[k]
-                    mutated_row[col_info["name"]] = v
-                else:
-                    logger.error(
-                        f"Column {k} not found in column lookup. Ignoring column")
-
-            mutations.append(
-                {
-                    "kind": "add-row-to-table",
-                    "tableName": self.table_id,
-                    "columnValues": mutated_row
+                "schema": {
+                    "columns": self.columns,
+                },
+                "rows": {
+                    "$stashID": self.stash_id
                 }
-            )
-        r = requests.post(
-            self.url("function/mutateTables"),
-            headers=self.headers(),
-            json={
-                "appID": self.hardcoded_app_id,
-                "mutations": mutations
             }
         )
-        if r.status_code != 200:
-            logger.error(f"add rows failed with status {r.status_code}: {r.text}")  # nopep8 because https://github.com/hhatto/autopep8/issues/712
+        try:
             r.raise_for_status()
-
-    def add_rows(self, rows: Iterator[BigTableRow]) -> None:
-        BATCH_SIZE = 100
-
-        batch = []
-        for row in rows:
-            batch.append(row)
-            if len(batch) >= BATCH_SIZE:
-                self.add_rows_batch(batch)
-                batch = []
-
-        if len(batch) > 0:
-            self.add_rows_batch(batch)
+        except Exception as e:
+            raise Exception(f"failed to overwrite table '{table_id}'. Response was '{r.text}'") from e  # nopep8
 
     def commit(self) -> None:
-        logger.debug("commit table (noop).")
-        pass
+        self.raise_if_set_schema_not_called()
+        # first see if the table already exists
+        r = requests.get(
+            self.url(f"tables"),
+            headers=self.headers()
+        )
+        try:
+            r.raise_for_status()
+        except Exception as e:
+            raise Exception(f"Failed to get table list. Response was '{r.text}'.") from e # nopep8
+
+        found_table_id = None
+        # confirm if table exists:
+        body = r.json()
+        if "data" not in body:
+            raise Exception(f"get tables response did not include data in body. Status was: {r.status_code}: {r.text}.")  # nopep8
+
+        for table in body["data"]:
+            if table["name"] == self.table_name:
+                found_table_id = table["id"]
+                logger.info(f"Found existing table to reuse for table name '{self.table_name}' with ID '{found_table_id}'.")  # nopep8
+                break
+
+        if found_table_id != None:
+            self.overwrite_table_from_stash(found_table_id)
+        else:
+            self.create_table_from_stash()
+
+        logger.info(f"Successfully committed record stash for table '{self.table_name}' (stash ID '{self.stash_id}')")  # nopep8