Firestore Transaction - Transaction Failed: TypeError: Transaction.set(...).then Is Not A Function

I have written a transaction to check if a user exists in Firestore before attempting to create it.

  setUserData(uid, email, displayName, photoURL, firstName, lastName, referralId) {
    const userRef = this.afs.firestore.collection('users').doc(`${uid}`);
    return this.afs.firestore.runTransaction(async (transaction: any) => {
      const doc = await transaction.get(userRef);
      let userData = {};
      if (!doc.exists) {
        userData = {
          uid: uid,
          email: email,
          displayName: displayName,
          photoURL: photoURL,
          emailVerified: true,
          firstName: firstName,
          lastName: lastName,
          referralId: referralId
        };
        transaction.set(userRef, { userData }, { merge: true }).then(() => {
          this.referralService.addUserToWaitlist(referralId);
        });
      }
    }).then(() => {
      if (!environment.production) {
        console.log(
          'Transaction successfully committed.'
        );
      }
    }).catch((error) => {
      if (!environment.production) {
        console.log('Transaction failed: ', error);
      }
    });
  }

However, I keep getting the following error:

Transaction failed:  TypeError: transaction.set(...).then is not a function

Does transaction not have the equivalent of then?

That's not how writing documents works in a transaction. Since a transaction is an all-or-nothing operation, you can't wait for a set() to complete before moving on to another one, so it wouldn't be helpful for it to return a promise. With transactions, you have to set() all your documents before the end of your function, then at the end, the transaction will attempt to write them all. If they cannot all be written atomically, your transaction function will be run again.

Also notice that the plain JavaScript API (not Angular) for transaction.set() is declared to return the same Transaction object rather than a promise.